Add subseasonal calibration pipeline

.buildkite/longruns/pipeline.yml

@@ -183,7 +183,7 @@ steps:
       - label: "Perfect model calibration test"
         key: "amip_pm_calibration"
         command:
-          - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl"
+          - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/coarse_amip/run_calibration.jl"
         artifact_paths: "experiments/calibration/output/*"
         env:
           CLIMACOMMS_DEVICE: "CUDA"

.gitignore

@@ -30,6 +30,7 @@ docs/src/generated/
 # Experiments
 !experiments/ClimaEarth/**/Manifest.toml
 !experiments/ClimaCore/**/Manifest.toml
+experiments/calibration/coarse_amip/output*/
 
 # Output
 output/

config/atmos_configs/climaatmos_wx_progedmf.yml

@@ -51,4 +51,4 @@ moist: equil
 
 cloud_model: "quadrature_sgs"
 use_itime: true
-deep_atmosphere: false
+deep_atmosphere: false

config/subseasonal_configs/wxquest_diagedmf.yml

@@ -1,36 +1,33 @@
 FLOAT_TYPE: "Float32"
 albedo_model: "CouplerAlbedo"
 atmos_config_file: "config/atmos_configs/climaatmos_wx_diagedmf.yml"
-coupler_toml: ["toml/amip.toml"]
-mode_name: "subseasonal"
-era5_initial_condition_dir: "/net/sampo/data1/wxquest_data/initial_conditions"
+coupler_toml: ["toml/amip.toml", "toml/wxquest_diagedmf.toml"]
 initial_condition: "WeatherModel"
-checkpoint_dt: "7days"
+checkpoint_dt: "93days"
 co2_model: "fixed"
-dt: "30secs"
-dt_cpl: "30secs"
+coupler_output_dir: "output_quick"
+dt: "90secs"
+dt_cpl: "90secs"
 energy_check: false
 h_elem: 16
-
-### integrated land ###
-# land_model: "integrated"
-
-### bucket land ###
+mode_name: "subseasonal"
+era5_initial_condition_dir: /net/sampo/data1/wxquest_data/initial_conditions
 land_model: "bucket"
 bucket_albedo_type: "map_temporal"
-bucket_initial_condition: "/net/sampo/data1/wxquest_data/initial_conditions/era5_bucket_processed_20250907_0000.nc"
-
 land_spun_up_ic: false
 land_temperature_anomaly: "nothing"
-use_land_diagnostics: true
 radiation_reset_rng_seed: true
-start_date: "20250907"
+start_date: "20180901"
 surface_setup: "PrescribedSurface"
 topo_smoothing: true
 topography: "Earth"
-t_end: "7days"
+t_end: "700days"  
 netcdf_output_at_levels: true
+output_default_diagnostics: false
+use_coupler_diagnostics: false
+use_land_diagnostics: true
+strict_params: false
 extra_atmos_diagnostics:
-  - short_name: [tas, mslp, pr, ua, va, rhoa, pfull, hur, hus, clw, cli, clivi, clwvi, cl, arup, tke]
-    period: 1hours
-    reduction_time: average
+  - short_name: [hfls, hfss, mslp, tas, pr, ta, pfull, hur, rsut, rsds, rsus, rsuscs, rsutcs, rlutcs, rlut, rlds, rlus, rlucs, ts, evspsbl, hfes]
+    period: 1months
+    reduction_time: average

experiments/calibration/README.md

@@ -1,18 +1,10 @@
 # ClimaCoupler Calibration Experiments
 
-This folder contains a trivial perfect-model calibration of the atmosphere coupled with the bucket model.
-The calibration uses 30-day and lat/lon averages of top-of-atmosphere shortwave 
-radiation to calibrate the `total_solar_irradiance` parameter in a perfect model setting. 
-The current run script uses the `ClimaCalibrate.SlurmManager` to add Slurm workers
- which run each ensemble member in parallel.
-
-To run this calibration on a Slurm cluster, ensure that `run_calibration.sh` is 
-configured for your cluster and run `sbatch run_calibration.sh`. The output will 
-be generated in `experiments/calibration/output`.
+This folder contains coupled calibration experiment scripts.
+The `coarse_amip` calibration uses the October monthly surface fluxes to calibrate the inverse entrainment timescale.
 
 Components:
-- run_calibration.sh: SBATCH script used to instantiate the project and run the calibration on a Slurm cluster.
-- run_calibration.jl: Julia script for the overall calibration and postprocessing. Contains the expriment configuration, such as ensemble size and number of iterations.
+- generate_observations.jl: Script to generate observations for the calibration pipeline
+- run_calibration.jl: Script for the overall calibration and postprocessing. Contains the expriment configuration, such as ensemble size and number of iterations.
 - model_interface.jl: Contains `forward_model`, the function that gets run during calibration. This basically just uses the `setup_run` function.
-- model_config.yml: Contains the configuration for the coupler
-- 
+- observation_map.jl: Contains the loss function and generalized plotting utilities.

experiments/calibration/api.jl

@@ -0,0 +1,154 @@
+import Dates
+
+"""
+    struct CalibrateConfig{SPINUP <: Dates.Period, EXTEND <: Dates.Period}
+        short_names::Vector{String}
+        minibatch_size::Int64
+        n_iterations::Int64
+        sample_date_ranges::Vector{NTuple{2, DATE}}
+        extend::EXTEND
+        spinup::SPINUP
+        nelements::Tuple{Int64, Int64}
+        output_dir::String
+        rng_seed::Int64
+    end
+
+A configuration struct for keeping track of multiple fields that are of interest
+to a user running calibration, or that are needed in multiple places (e.g., for
+ensemble members and generating observations).
+"""
+struct CalibrateConfig{SPINUP <: Dates.Period, EXTEND <: Dates.Period}
+    "Configuration file to use for ClimaCoupler simulation"
+    config_file::String
+
+    "The short names of the observations used for calibration. The short names
+    should match the same names used for the diagnostics."
+    short_names::Vector{String}
+
+    "The size of the minibatch for each iteration"
+    minibatch_size::Int64
+
+    "The number of iterations to run the calibration for"
+    n_iterations::Int64
+
+    "The date ranges of the samples for calibration and used to determine the
+    start and end dates of a simulation for each iteration of calibration"
+    sample_date_ranges::Vector{NTuple{2, Dates.DateTime}}
+
+    "The amount of time to run a simulation after the last date of the
+    minibatch"
+    extend::EXTEND
+
+    "The amount of time to run a simulation before the first date of the
+    minibatch"
+    spinup::SPINUP
+
+    "The directory to store the iterations and members of the calibration."
+    output_dir::String
+
+    "An integer value for ensuring calibrations are the same between multiple
+    calibrations with the same settings"
+    rng_seed::Int64
+end
+
+"""
+    CalibrateConfig(;
+        config_file,
+        short_names,
+        sample_date_ranges,
+        extend,
+        spinup = Dates.Month(3),
+        minibatch_size,
+        n_iterations,
+        output_dir = "calibration/weatherquest",,
+        rng_seed = 42,
+    )
+
+Initializes a CalibrateConfig, which is of interest to a user running
+calibration or contains values needed in multiple places during calibration.
+
+Keyword arguments
+=====================
+
+- `config_file`: Configuration file to use for ClimaCoupler simulation.
+
+- `short_names`: Short names of the observations. The currently supported short
+  names are `pr`, `tas`, and `mslp`.
+
+- `minibatch_size`: The size of the minibatch for each iteration.
+
+- `n_iterations`: The number of iterations to run the calibration for.
+
+- `sample_date_ranges`: The date ranges for each sample. The dates should be the
+  same as found in the time series data of the observations.
+
+- `extend`: The amount of time to run the simulation after the end date
+  determined by `sample_date_ranges`. For seasonal averages, `extend` should be
+  `Dates.Month(3)` and for monthly averages, `extend` should be
+  `Dates.Month(1)`.
+
+- `spinup`: The amount of time to run the simulation before the start date
+  determined by `sample_date_ranges`.
+
+- `nelements`: The resolution of the model. This is also used to determine the
+  mask of the observations.
+
+- `output_dir`: The location to save the calibration at.
+
+- `rng_seed`: An integer to ensure that calibration runs with the same settings
+  are the same.
+"""
+function CalibrateConfig(;
+    config_file,
+    short_names,
+    minibatch_size,
+    n_iterations,
+    sample_date_ranges,
+    extend,
+    spinup = Dates.Month(3),
+    output_dir = "calibration/weatherquest",
+    rng_seed = 42,
+)
+    isempty(short_names) && error("Cannot run calibration with no short names")
+    isempty(sample_date_ranges) &&
+        error("Cannot run calibration with no date ranges for the samples")
+
+    sample_date_ranges = [
+        (Dates.DateTime(date_pair[1]), Dates.DateTime(date_pair[2])) for
+        date_pair in sample_date_ranges
+    ]
+
+    for (start_date, stop_date) in sample_date_ranges
+        start_date <= stop_date || error(
+            "The start date ($start_date) should be before the stop date ($stop_date)",
+        )
+    end
+    issorted(sample_date_ranges) ||
+        error("The samples in $sample_date_ranges should be sorted")
+
+    minibatch_size > 0 || error("The minibatch size ($minibatch_size) should be positive")
+    n_iterations > 0 || error("The number of iterations ($n_iterations) should be positive")
+
+    num_samples = length(sample_date_ranges)
+    minibatch_size > num_samples && error(
+        "The minibatch size is $minibatch_size, but the number of samples is $num_samples",
+    )
+
+    remaining = num_samples % minibatch_size
+    remaining == 0 || @warn(
+        "Number of samples is not divisible by the minibatch size; the last $remaining samples may be missing when running the calibration"
+    )
+
+    return CalibrateConfig(
+        config_file,
+        short_names,
+        minibatch_size,
+        n_iterations,
+        sample_date_ranges,
+        extend,
+        spinup,
+        output_dir,
+        rng_seed,
+    )
+
+end