Skip to content

Commit c0eb34c

Browse files
committed
Add subseasonal calibration pipeline
1 parent 85214b9 commit c0eb34c

File tree

15 files changed

+796
-56
lines changed

15 files changed

+796
-56
lines changed

.buildkite/longruns/pipeline.yml

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -177,21 +177,15 @@ steps:
177177
- group: "Calibration experiments"
178178

179179
steps:
180-
- label: "Perfect model calibration test"
181-
key: "amip_pm_calibration"
180+
- label: "Subseasonal calibration test"
181+
key: "subseasonal_calibration"
182182
command:
183-
- "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl"
184-
artifact_paths: "experiments/calibration/output/*"
185-
env:
186-
CLIMACOMMS_DEVICE: "CUDA"
187-
CLIMACOMMS_CONTEXT: "SINGLETON"
183+
- julia --color=yes --project=experiments/ClimaEarth experiments/calibration/subseasonal/generate_observations.jl
184+
- julia --color=yes --project=experiments/ClimaEarth experiments/calibration/subseasonal/run_calibration.jl
185+
artifact_paths: "experiments/calibration/subseasonal/output/*"
188186
agents:
189187
queue: clima
190-
slurm_mem: 96GB
191-
slurm_ntasks: 3
192-
slurm_gpus_per_task: 1
193-
slurm_cpus_per_task: 4
194-
slurm_time: 05:00:00
188+
slurm_time: 24:00:00
195189

196190
- wait
197191

.buildkite/pipeline.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,7 @@ steps:
373373
- label: "Perfect model calibration test"
374374
key: "amip_pm_calibration"
375375
command:
376-
- "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl"
376+
- "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/test/run_calibration.jl"
377377
artifact_paths: "experiments/calibration/output/*"
378378
env:
379379
CLIMACOMMS_DEVICE: "CUDA"

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ docs/src/generated/
3030
# Experiments
3131
!experiments/ClimaEarth/**/Manifest.toml
3232
!experiments/ClimaCore/**/Manifest.toml
33+
experiments/calibration/coarse_amip/output*/
3334

3435
# Output
3536
output/

config/subseasonal_configs/wxquest_diagedmf.yml

Lines changed: 10 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,34 +2,29 @@ FLOAT_TYPE: "Float32"
22
albedo_model: "CouplerAlbedo"
33
atmos_config_file: "config/atmos_configs/climaatmos_wx_diagedmf.yml"
44
coupler_toml: ["toml/amip.toml"]
5-
mode_name: "subseasonal"
6-
era5_initial_condition_dir: "/net/sampo/data1/wxquest_data/initial_conditions"
75
initial_condition: "WeatherModel"
86
checkpoint_dt: "7days"
97
dt: "30secs"
108
dt_cpl: "30secs"
119
energy_check: false
1210
h_elem: 16
13-
14-
### integrated land ###
15-
# land_model: "integrated"
16-
17-
### bucket land ###
11+
mode_name: "subseasonal"
12+
era5_initial_condition_dir: "/net/sampo/data1/wxquest_data/initial_conditions"
1813
land_model: "bucket"
1914
bucket_albedo_type: "map_temporal"
20-
bucket_initial_condition: "/net/sampo/data1/wxquest_data/initial_conditions/era5_bucket_processed_20250907_0000.nc"
21-
2215
land_spun_up_ic: false
2316
land_temperature_anomaly: "nothing"
24-
use_land_diagnostics: true
2517
radiation_reset_rng_seed: true
26-
start_date: "20250907"
18+
start_date: "20180901"
2719
surface_setup: "PrescribedSurface"
2820
topo_smoothing: true
2921
topography: "Earth"
30-
t_end: "7days"
3122
netcdf_output_at_levels: true
23+
output_default_diagnostics: false
24+
use_coupler_diagnostics: false
25+
use_land_diagnostics: false
26+
strict_params: false
3227
extra_atmos_diagnostics:
33-
- short_name: [tas, mslp, pr, ua, va, rhoa, pfull, hur, hus, clw, cli, clivi, clwvi, cl, arup, tke]
34-
period: 1hours
35-
reduction_time: average
28+
- short_name: [hfls, hfss, rsus, rlus]
29+
period: 1months
30+
reduction_time: average

experiments/calibration/README.md

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,13 @@
11
# ClimaCoupler Calibration Experiments
22

3-
This folder contains a trivial perfect-model calibration of the atmosphere coupled with the bucket model.
4-
The calibration uses 30-day and lat/lon averages of top-of-atmosphere shortwave
5-
radiation to calibrate the `total_solar_irradiance` parameter in a perfect model setting.
6-
The current run script uses the `ClimaCalibrate.SlurmManager` to add Slurm workers
7-
which run each ensemble member in parallel.
83

9-
To run this calibration on a Slurm cluster, ensure that `run_calibration.sh` is
10-
configured for your cluster and run `sbatch run_calibration.sh`. The output will
11-
be generated in `experiments/calibration/output`.
4+
This folder contains pipelines to reproduce coupled calibration experiments.
5+
Each pipeline has its own subfolder:
126

13-
Components:
14-
- run_calibration.sh: SBATCH script used to instantiate the project and run the calibration on a Slurm cluster.
15-
- run_calibration.jl: Julia script for the overall calibration and postprocessing. Contains the expriment configuration, such as ensemble size and number of iterations.
16-
- model_interface.jl: Contains `forward_model`, the function that gets run during calibration. This basically just uses the `setup_run` function.
17-
- model_config.yml: Contains the configuration for the coupler
18-
-
7+
- perfect_model: A trivial perfect-model calibration of the atmosphere coupled
8+
with the bucket model. The calibration uses 30-day and lat/lon averages of
9+
top-of-atmosphere shortwave radiation to calibrate the `total_solar_irradiance`
10+
parameter in a perfect model setting. The current run script uses the
11+
`ClimaCalibrate.SlurmManager` to add Slurm workers which run each ensemble
12+
member in parallel.
13+
- subseasonal: Calibrates the inverse entrainment timescale to ERA5 October monthly surface fluxes and surface temperature from 2018 to 2024.

experiments/calibration/api.jl

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
import Dates
2+
3+
"""
4+
struct CalibrateConfig{SPINUP <: Dates.Period, EXTEND <: Dates.Period}
5+
short_names::Vector{String}
6+
minibatch_size::Int64
7+
n_iterations::Int64
8+
sample_date_ranges::Vector{NTuple{2, DATE}}
9+
extend::EXTEND
10+
spinup::SPINUP
11+
nelements::Tuple{Int64, Int64}
12+
output_dir::String
13+
rng_seed::Int64
14+
end
15+
16+
A configuration struct for keeping track of multiple fields that are of interest
17+
to a user running calibration, or that are needed in multiple places (e.g., for
18+
ensemble members and generating observations).
19+
"""
20+
struct CalibrateConfig{SPINUP <: Dates.Period, EXTEND <: Dates.Period}
21+
"Configuration file to use for ClimaCoupler simulation"
22+
config_file::String
23+
24+
"The short names of the observations used for calibration. The short names
25+
should match the same names used for the diagnostics."
26+
short_names::Vector{String}
27+
28+
"The size of the minibatch for each iteration"
29+
minibatch_size::Int64
30+
31+
"The number of iterations to run the calibration for"
32+
n_iterations::Int64
33+
34+
"The date ranges of the samples for calibration and used to determine the
35+
start and end dates of a simulation for each iteration of calibration"
36+
sample_date_ranges::Vector{NTuple{2, Dates.DateTime}}
37+
38+
"The amount of time to run a simulation after the last date of the
39+
minibatch"
40+
extend::EXTEND
41+
42+
"The amount of time to run a simulation before the first date of the
43+
minibatch"
44+
spinup::SPINUP
45+
46+
"The directory to store the iterations and members of the calibration."
47+
output_dir::String
48+
49+
"An integer value for ensuring calibrations are the same between multiple
50+
calibrations with the same settings"
51+
rng_seed::Int64
52+
end
53+
54+
"""
55+
CalibrateConfig(;
56+
config_file,
57+
short_names,
58+
sample_date_ranges,
59+
extend,
60+
spinup = Dates.Month(3),
61+
minibatch_size,
62+
n_iterations,
63+
output_dir = "calibration/weatherquest",,
64+
rng_seed = 42,
65+
)
66+
67+
Initializes a CalibrateConfig, which is of interest to a user running
68+
calibration or contains values needed in multiple places during calibration.
69+
70+
Keyword arguments
71+
=====================
72+
73+
- `config_file`: Configuration file to use for ClimaCoupler simulation.
74+
75+
- `short_names`: Short names of the observations. The currently supported short
76+
names are `pr`, `tas`, and `mslp`.
77+
78+
- `minibatch_size`: The size of the minibatch for each iteration.
79+
80+
- `n_iterations`: The number of iterations to run the calibration for.
81+
82+
- `sample_date_ranges`: The date ranges for each sample. The dates should be the
83+
same as found in the time series data of the observations.
84+
85+
- `extend`: The amount of time to run the simulation after the end date
86+
determined by `sample_date_ranges`. For seasonal averages, `extend` should be
87+
`Dates.Month(3)` and for monthly averages, `extend` should be
88+
`Dates.Month(1)`.
89+
90+
- `spinup`: The amount of time to run the simulation before the start date
91+
determined by `sample_date_ranges`.
92+
93+
- `nelements`: The resolution of the model. This is also used to determine the
94+
mask of the observations.
95+
96+
- `output_dir`: The location to save the calibration at.
97+
98+
- `rng_seed`: An integer to ensure that calibration runs with the same settings
99+
are the same.
100+
"""
101+
function CalibrateConfig(;
102+
config_file,
103+
short_names,
104+
minibatch_size,
105+
n_iterations,
106+
sample_date_ranges,
107+
extend,
108+
spinup = Dates.Month(3),
109+
output_dir = "calibration/weatherquest",
110+
rng_seed = 42,
111+
)
112+
isempty(short_names) && error("Cannot run calibration with no short names")
113+
isempty(sample_date_ranges) &&
114+
error("Cannot run calibration with no date ranges for the samples")
115+
116+
sample_date_ranges = [
117+
(Dates.DateTime(date_pair[1]), Dates.DateTime(date_pair[2])) for
118+
date_pair in sample_date_ranges
119+
]
120+
121+
for (start_date, stop_date) in sample_date_ranges
122+
start_date <= stop_date || error(
123+
"The start date ($start_date) should be before the stop date ($stop_date)",
124+
)
125+
end
126+
issorted(sample_date_ranges) ||
127+
error("The samples in $sample_date_ranges should be sorted")
128+
129+
minibatch_size > 0 || error("The minibatch size ($minibatch_size) should be positive")
130+
n_iterations > 0 || error("The number of iterations ($n_iterations) should be positive")
131+
132+
num_samples = length(sample_date_ranges)
133+
minibatch_size > num_samples && error(
134+
"The minibatch size is $minibatch_size, but the number of samples is $num_samples",
135+
)
136+
137+
remaining = num_samples % minibatch_size
138+
remaining == 0 || @warn(
139+
"Number of samples is not divisible by the minibatch size; the last $remaining samples may be missing when running the calibration"
140+
)
141+
142+
return CalibrateConfig(
143+
config_file,
144+
short_names,
145+
minibatch_size,
146+
n_iterations,
147+
sample_date_ranges,
148+
extend,
149+
spinup,
150+
output_dir,
151+
rng_seed,
152+
)
153+
154+
end

experiments/calibration/run_calibration.sh

Lines changed: 0 additions & 13 deletions
This file was deleted.

0 commit comments

Comments
 (0)