WeatherQuest Calibration pipeline with weekly/monthly support + CERES, ERA5 Pressure Levels

config/atmos_configs/climaatmos_wx_diagedmf.yml

@@ -26,7 +26,8 @@ h_elem: 15
 z_max: 70000.0
 z_elem: 70
 dz_bottom: 30.0
-edmfx_entr_model: "Generalized"
+edmfx_entr_model: "PiGroups"
+# edmfx_entr_model: "Generalized"
 edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
 edmfx_sgs_mass_flux: true
@@ -40,7 +41,7 @@ toml: [toml/diagnostic_edmfx_era5_ic.toml]
 ### 1M Microphysics ###
 # precip_model: "1M"
 # moist: "nonequil"
-# toml: [toml/diagnostic_edmfx_1M.toml]
+# # toml: [toml/diagnostic_edmfx_1M.toml]
 
 cloud_model: "quadrature"
 use_itime: true

config/subseasonal_configs/wxquest_diagedmf.yml

@@ -4,7 +4,7 @@ atmos_config_file: "config/atmos_configs/climaatmos_wx_diagedmf.yml"
 coupler_toml: ["toml/wxquest_diagedmf.toml"]
 strict_params: false
 mode_name: "subseasonal"
-era5_initial_condition_dir: "/net/sampo/data1/wxquest_data/initial_conditions/initial_conditions_v1_dev"
+era5_initial_condition_dir: "/glade/campaign/univ/ucit0011/cchristo/wxquest_data/initial_conditions/initial_conditions_v1_dev"
 initial_condition: "WeatherModel"
 checkpoint_dt: "40days"
 dt: "120secs"
@@ -13,14 +13,20 @@ dt_cloud_fraction: "1hours"
 dt_rad: "1hours"
 dt_seaice: "120secs"
 energy_check: false
-h_elem: 32
+h_elem: 12
 land_fraction_source: "era5"
 binary_area_fraction: false
 
 ### integrated land ###
 land_model: "integrated"
 lai_source: "modis_monthly_climatology"
 
+### gravity waves 
+# orographic_gravity_wave: "raw_topo"
+non_orographic_gravity_wave: true
+
+
+
 ### bucket land ###
 # land_model: "bucket"
 # bucket_albedo_type: "map_temporal"
@@ -30,18 +36,17 @@ lai_source: "modis_monthly_climatology"
 land_spun_up_ic: false
 land_temperature_anomaly: "nothing"
 radiation_reset_rng_seed: true
-start_date: "20100101"
+start_date: "20230115"
 surface_setup: "PrescribedSurface"
 topo_smoothing: true
 topography: "Earth"
-t_end: "40days"
+# t_end: "8days"
+t_end: "1days"
 netcdf_output_at_levels: true
 output_default_diagnostics: false
 use_coupler_diagnostics: false
 use_land_diagnostics: false
-strict_params: false
-insolation: "timevarying"
 extra_atmos_diagnostics:
-  - short_name: [zg, orog, tas, ta, mslp, pr, ua, va, ts, rhoa, pfull, hur, hus, clw, cl, hfss, hfls, evspsbl, hussfc, rsds, rlds, rlus, rsus, rld, wa, waup]
-    period: 24hours
-    reduction_time: average
+  - short_name: [tas, pr, mslp, rsut, rlut]
+    period: 1days
+    reduction_time: average

config/subseasonal_configs/wxquest_diagedmf_weekly_calibration.yml

@@ -0,0 +1,55 @@
+FLOAT_TYPE: "Float32"
+albedo_model: "CouplerAlbedo"
+atmos_config_file: "config/atmos_configs/climaatmos_wx_diagedmf.yml"
+coupler_toml: ["toml/wxquest_diagedmf.toml"]
+strict_params: false
+mode_name: "subseasonal"
+era5_initial_condition_dir: "/glade/campaign/univ/ucit0011/cchristo/wxquest_data/initial_conditions/initial_conditions_v1_dev"
+initial_condition: "WeatherModel"
+checkpoint_dt: "100days"
+dt: "150secs"
+dt_cpl: "150secs"
+dt_cloud_fraction: "1hours"
+dt_rad: "1hours"
+dt_seaice: "150secs"
+ode_algo: "ARS232"
+energy_check: false
+h_elem: 12
+z_elem: 63
+z_max: 60000.0
+land_fraction_source: "era5"
+binary_area_fraction: false
+
+### integrated land ###
+land_model: "integrated"
+lai_source: "modis_monthly_climatology"
+
+### gravity waves 
+# orographic_gravity_wave: "raw_topo"
+non_orographic_gravity_wave: true
+
+
+
+### bucket land ###
+# land_model: "bucket"
+# bucket_albedo_type: "map_temporal"
+# bucket_albedo_type: "era5"
+# bucket_initial_condition: "/net/sampo/data1/cchristo/WeatherQuest_data/archive_dev4/era5_bucket_processed_20250831_0000.nc"
+
+land_spun_up_ic: false
+land_temperature_anomaly: "nothing"
+radiation_reset_rng_seed: true
+start_date: "20230115"
+surface_setup: "PrescribedSurface"
+topo_smoothing: true
+topography: "Earth"
+# t_end: "8days"
+t_end: "1days"
+netcdf_output_at_levels: true
+output_default_diagnostics: false
+use_coupler_diagnostics: false
+use_land_diagnostics: false
+extra_atmos_diagnostics:
+  - short_name: [tas, pr, mslp, rsut, rlut]
+    period: 1days
+    reduction_time: average

experiments/calibration/README.md

@@ -17,4 +17,4 @@ To run a pipeline on a Slurm cluster, ensure that the given runscript
 cluster and run:
 ```
 julia --project=experiments/ClimaEarth experiments/calibration/<pipeline>/run_calibration.jl
-```
+```

experiments/calibration/subseasonal/generate_observations.jl

@@ -7,6 +7,10 @@ import EnsembleKalmanProcesses as EKP
 import JLD2
 import ClimaDiagnostics
 import ClimaCore
+
+# setup_run.jl provides CoupledSimulation (needed by resampled_lonlat fallback)
+include(joinpath(pkgdir(ClimaCoupler), "experiments/ClimaEarth/setup_run.jl"))
+
 # Access CalibrateConfig
 include(joinpath(@__DIR__, "run_calibration.jl"))
 

experiments/calibration/subseasonal/model_interface.jl

@@ -1,20 +1,29 @@
+#=
+Defines ClimaCalibrate.forward_model which runs the coupled simulation
+for a single ensemble member.
+
+Prerequisites: This file expects the following to be set up before inclusion:
+  - CALIBRATE_CONFIG must be defined (from run_calibration.jl)
+  - setup_run.jl must be included (provides setup_and_run function)
+
+Used by: subseasonal, subseasonal_weekly (and potentially other pipelines)
+=#
+
 ENV["CLIMACOMMS_DEVICE"] = "CUDA"
 ENV["CLIMACOMMS_CONTEXT"] = "SINGLETON"
 import ClimaCoupler
 import ClimaCalibrate
-import CUDA
-import EnsembleKalmanProcesses as EKP
-include(joinpath(pkgdir(ClimaCoupler), "experiments", "ClimaEarth", "setup_run.jl"))
-include(
-    joinpath(
-        pkgdir(ClimaCoupler),
-        "experiments",
-        "calibration",
-        "subseasonal",
-        "run_calibration.jl",
-    ),
-)
+using Dates: Date, Second
 using Pkg
+
+include(joinpath(pkgdir(ClimaCoupler), "experiments", "ClimaEarth", "setup_run.jl"))
+
+# Include run_calibration.jl only if CALIBRATE_CONFIG is not already defined
+# This allows other pipelines to include their own run_calibration.jl first
+if !@isdefined(CALIBRATE_CONFIG)
+    include(joinpath(@__DIR__, "run_calibration.jl"))
+end
+
 function ClimaCalibrate.forward_model(iter, member)
     Pkg.status()
     config_dict = ClimaCoupler.Input.get_coupler_config_dict(CALIBRATE_CONFIG.config_file)

experiments/calibration/subseasonal/observation_map.jl

@@ -84,44 +84,6 @@ function largest_period(sample_date_range)
     return period
 end
 
-function get_var(short_name, simdir)
-    if short_name == "tas - ta"
-        tas = get(simdir; short_name = "tas")
-        ta = get(simdir; short_name = "ta")
-        # TODO: figure out why this doesn't work
-        # pfull = get(simdir; short_name = "pfull")
-        # ta_900 = ClimaAnalysis.Atmos.to_pressure_coordinates(ta, pfull; target_pressure=[900])
-        ta_900hpa = slice(ta; z = 1000)
-        var = tas - ta_900hpa
-    else
-        # TODO: support multiple periods/reductions of same variable
-        var = get(simdir; short_name)
-    end
-
-    var.attributes["short_name"] = short_name
-    return var
-end
-
-"""
-    preprocess_var(var::ClimaAnalysis.OutputVar, reference_date)
-
-Preprocess `var` before flattening for G ensemble matrix.
-
-For "pr", weekly sums are computed. For "tas" and "mslp", weekly means are
-computed from daily means. The daily means are computing starting from
-`reference_date`.
-
-This function assumes that the data is monthly.
-"""
-function preprocess_var(var, sample_date_range)
-    period = largest_period(sample_date_range)
-    var = ClimaAnalysis.Var._shift_by(var, date -> date - period)
-    var = set_units(var, var_units[short_name(var)])
-    # TODO: Match dates instead of just windowing
-    var = window(var, "time"; left = sample_date_range[1], right = sample_date_range[2])
-    return var
-end
-
 """
     ClimaCalibrate.analyze_iteration(ekp,
                                      g_ensemble,

experiments/calibration/subseasonal/observation_utils.jl

@@ -3,17 +3,33 @@ using Statistics
 import Dates
 using ClimaAnalysis
 
-include(joinpath(pkgdir(ClimaCoupler), "experiments/ClimaEarth/setup_run.jl"))
+# Shared variable units (used by both subseasonal and subseasonal_weekly)
 
 var_units = Dict(
+    # Surface temperature/pressure
     "pr" => "kg m^-2 s^-1",
     "mslp" => "Pa",
     "tas" => "K",
     "tas - ta" => "K",
+    "ts" => "K",
+    "sp" => "Pa",
+    # Surface fluxes
     "hfls" => "W m^-2",
     "hfss" => "W m^-2",
-    "rsus" => "W m^-2",
-    "rlus" => "W m^-2",
+    # Radiation (TOA)
+    "rsdt" => "W m^-2",    # incoming solar
+    "rsut" => "W m^-2",    # TOA outgoing SW
+    "rlut" => "W m^-2",    # TOA outgoing LW
+    "rsutcs" => "W m^-2",  # TOA outgoing SW clear-sky
+    "rlutcs" => "W m^-2",  # TOA outgoing LW clear-sky
+    # Radiation (surface)
+    "rsds" => "W m^-2",    # surface downwelling SW
+    "rsus" => "W m^-2",    # surface upwelling SW
+    "rlds" => "W m^-2",    # surface downwelling LW
+    "rlus" => "W m^-2",    # surface upwelling LW
+    "rsdscs" => "W m^-2",  # surface downwelling SW clear-sky
+    "rsuscs" => "W m^-2",  # surface upwelling SW clear-sky
+    "rldscs" => "W m^-2",  # surface downwelling LW clear-sky
 )
 
 """
@@ -26,3 +42,40 @@ function remove_global_mean(var)
     mean_data = mean_var.data
     return ClimaAnalysis.replace(val -> val - mean_data[1], var)
 end
+
+"""
+    get_var(short_name, simdir)
+
+Get an `OutputVar` from `simdir` for the given `short_name`.
+Handles composite variables like "tas - ta" (surface-minus-boundary-layer temperature).
+"""
+function get_var(short_name, simdir)
+    if short_name == "tas - ta"
+        tas = get(simdir; short_name = "tas")
+        ta = get(simdir; short_name = "ta")
+        ta_900hpa = slice(ta; z = 1000)
+        var = tas - ta_900hpa
+    else
+        var = get(simdir; short_name)
+    end
+
+    var.attributes["short_name"] = short_name
+    return var
+end
+
+"""
+    preprocess_var(var::ClimaAnalysis.OutputVar, sample_date_range)
+
+Preprocess simulation output `var` before flattening for the G ensemble matrix.
+
+Shifts the time axis by the largest period in the date range, sets units, and
+windows to the date range. Calls `largest_period` (defined per-pipeline) to
+determine the shift.
+"""
+function preprocess_var(var, sample_date_range)
+    period = largest_period(sample_date_range)
+    var = ClimaAnalysis.Var._shift_by(var, date -> date - period)
+    var = set_units(var, var_units[short_name(var)])
+    var = window(var, "time"; left = sample_date_range[1], right = sample_date_range[2])
+    return var
+end

experiments/calibration/subseasonal_weekly/README.md

@@ -0,0 +1,83 @@
+# Subseasonal Calibration
+
+Calibrate ClimaCoupler parameters using EnsembleKalmanProcesses against ERA5 observations.
+
+## Quick Start
+
+
+### Option A: TransformUnscented (recommended, simpler)
+
+```bash
+cd /glade/u/home/cchristo/clima/copies2/ClimaCoupler.jl
+
+# 1. Generate observations (only needed once, or when changing obs settings)
+julia --project=experiments/ClimaEarth experiments/calibration/subseasonal_weekly/generate_observations.jl
+
+# 2. Run calibration (from tmux on login node)
+julia --project=experiments/ClimaEarth experiments/calibration/subseasonal_weekly/run_calibration.jl
+```
+
+### Option B: TransformInversion (more robust, needs precompute)
+
+```bash
+cd /glade/u/home/cchristo/clima/copies2/ClimaCoupler.jl
+
+# 1. Generate observations (only needed once)
+julia --project=experiments/ClimaEarth experiments/calibration/subseasonal_weekly/generate_observations.jl
+
+# 2. Run full calibration (from tmux - handles precompute automatically)
+./experiments/calibration/subseasonal_weekly/run_full_calibration.sh
+```
+
+## Switching Between Process Types
+
+Edit `run_calibration.jl` and comment/uncomment the appropriate OPTION block (~line 116-156).
+
+## Changing Parameters
+
+Edit `calibration_setup.jl` - this is the single source of truth for:
+- `CALIBRATION_PRIORS` - list of parameters to calibrate
+- `CALIBRATION_ENSEMBLE_SIZE` - ensemble size (for TransformInversion only)
+- `NORMALIZE_VARIABLES` - if true, normalize variables to zero mean and unit variance
+- `CALIBRATION_NOISE_SCALAR` - noise variance for observation covariance
+
+### Variable Normalization
+
+When `NORMALIZE_VARIABLES = true` (default), each variable is normalized to zero mean and unit variance:
+1. `generate_observations.jl` computes latitude-weighted mean and std from ERA5 data
+2. Stats are saved to `norm_stats.jld2` 
+3. Observations are normalized before creating the observation vector
+4. `observation_map.jl` loads the same stats and applies them to model output
+
+This ensures variables with different magnitudes (e.g., temperature in K vs radiation in W/m²) contribute equally to the calibration objective.
+
+## Data Source
+
+Uses weekly-averaged ERA5 data from:
+`/glade/campaign/univ/ucit0011/cchristo/wxquest_data/daily_weekly_stats/weekly`
+
+# Analyzing Calibration 
+    `julia --project=experiments/ClimaEarth experiments/calibration/subseasonal_weekly/analyze_calibration.jl`
+
+## Key Files
+
+| File | Purpose |
+|------|---------|
+| `calibration_setup.jl` | defines parameters and settings to calibrate |
+| `run_calibration.jl` | Main calibration script |
+| `generate_observations.jl` | Loads weekly ERA5 data, creates observation vector |
+| `observation_map.jl` | Maps model output to observation space |
+| `run_full_calibration.sh` | Orchestrator for TransformInversion |
+
+## Output
+
+Results go to the `output_dir` specified in `run_calibration.jl` (currently `/glade/derecho/scratch/cchristo/calibration/expXX`).
+Logs are saved to `calibration_YYYYMMDD_HHMMSS.log` in the output directory.
+
+Plot results:
+```bash
+julia --project=experiments/ClimaEarth experiments/calibration/subseasonal_weekly/quick_plot_cal.jl
+```
+
+## Troubleshooting
+**Bus error**: JLD2 mmap issue on Derecho Lustre - use JLD2 IOStream mode.