Add external forcing types to gcm-driven scm calibration, allowing for both shallow and deep convection. Includes option for defining stretched calibration grid.

Add cal grid to edmf_ensemble_stats

gcm driven calibration updates: restart, plotting support for variable grids, add plot prior and normalize var scripts, add microphysics cal option

Add NN mixing length closure, add load pretrain weights logic in prior

Add clippings for NN inputs, dev calibration

Add scm_runner tools to run SCMs (in parallel) across cases for a given parameter set

scm runner updates

limit nn mixing length by smag length + (1/z), and add new dz input variable. Remove problematic deep convective cases from library. Update minibatcher to specify cases for each epoch. Add serialize_std_model with default in run_cal to allow different priors for NN weights and biases. Increase request limits for top-level slurm sbatch script, preventing minor memory leak from killing calibration. Increase noise by an order of mag. Update, improve restart calibration + allow to work with NNs. Limit number of processes per ensemble member by 5 to avoid queue waits .

Add back deep deep conv cases, clip ql, qi above 0 in les cases. Fix get_optimal_particle nearest neighbor mean getter. Simplify nn prior creation at beginning of run_calibration. Request less time on nodes for scm runs.

Limit num_cpu to 5 in restart script. Increase noise, use batch size = 20. And cfsite info to plot_ensemble plots. Add radiation metrics to runner model config. Add nn_helper functions. Add precal prior toml.

Increase noise and batch size (30). Increase prior sigma for NN (0.05 for weights). Lower prior std for mixing_length_diss_coeff. Turn off accelerator.

Make shallow library default in get_les_metadata. Make leaky relu NN default. Change default noise with leaky relu (std_weight = 0.1, std_bias = 0.00001). Increase default runtime for restart for ens members to 180 mins. Increase resource requests for scm_runner.

Make leaky relu NN default (and write to scratch). Pass args to scm_runner in sbatch script. Add CLI option for scm_runner, more flexible parsing of cfsite (for different AMIP experiments).

Switch between NN and standard priors in restart calibration script

Add batch size logic in edmf_ensemble_stats

Add forcing parameters as additional tomls

Author: costachris

calibration/experiments/gcm_driven_scm/Project.toml

@@ -11,6 +11,7 @@ ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -19,5 +20,5 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 
 [compat]
-ClimaCalibrate = "=0.0.3"
-EnsembleKalmanProcesses = "2"
+EnsembleKalmanProcesses = "=2.1.2"
+ClimaCalibrate = "=0.0.13"

calibration/experiments/gcm_driven_scm/README.md

@@ -1,10 +1,33 @@
 # Overview of Calibration Pipeline for GCM-Driven Single Column Model (EDMF)
 
-This setup provide tools for calibrating both prognostic and diagnostic EDMF variants to LES profiles, given the same forcings and boundary conditions. The gcm-driven EDMF setup is employed in single-column mode, which uses both interactive radiation and surface fluxes. Forcing profiles include resolved eddy advection, horizontal advection, subsidence, and GCM state relaxation. The setup is run to the top of the atmosphere to compute radiation, but calibrations statistics are computed only on the lower 4km (`z_max`), where LES output is available.
+## Pipeline Components
+
+### Configuration Files
+- `experiment_config.yml` - Configuration of calibration settings
+  - Defines spatiotemporal calibration window
+  - Specifies required pipeline file paths
+  - Controls batch processing parameters
+
+- `model_config_**.yml` - Configuration for ClimaAtmos single column model
+  - Defines model-specific parameters
+  - Allows for multiple model configuration variants
+
+## Best Practices
+- Ensure `batch_size` matches available LES configurations
+- Verify normalization factors for each variable
+- Monitor ensemble convergence using provided plotting tools
+
+This setup provide tools for calibrating both prognostic and diagnostic EDMF variants to LES profiles, given the same forcings and boundary conditions. The gcm-driven EDMF setup is employed in single-column mode, which uses both interactive radiation and surface fluxes. Forcing profiles include resolved eddy advection, horizontal advection, subsidence, and GCM state relaxation. The setup is run to the top of the atmosphere to compute radiation, but calibration statistics are only computed on the calibration grid `z_cal_grid`.
 
 LES profiles are available for different geolocations ("cfsites"), spanning seasons, forcing host models, and climates (AMIP, AMIP4K). A given LES simulation is referred to as a "configuration". Calibrations employ batching by default and stack multiple configurations (a number equal to the `batch_size`) in a given iteration. The observation vector for a single configuration is formed by concatenating profiles across calibration variables, where each geophysical variable is normalized to have approximately unit variance and zero mean. These variable-by-variable normalization factors are precomputed (`norm_factors_dict`) and applied to all observations. Following this operation, the spatiotemporal calibration window is applied and temporal means are computed to form the observation vector `y`. Because variables are normalized to have 0 mean and unit variance, a constant diagonal noise matrix is used (configurable as `const_noise`).
 
 
+### Observation Map
+1. **Time-mean**: Time-mean of profiles taken between [`y_t_start_sec`, `y_t_end_sec`] for `y` and [`g_t_start_sec`, `g_t_end_sec`] for `G`.
+2. **Interpolation**: Case-specific (i.e., "shallow", "deep") interpolation to the calibration grid, defined with stretch-grid parameters in `z_cal_grid`.
+3. **Normalization**: Variable-specific normalization using the mean and standard deviation defined in `norm_factors_by_var`. Optionally, take log of variables using `log_vars` before normalization.
+4. **Concatenation**: Stack across cases in a batch, forming `y`, `G`.
+
 ## Getting Started
 
 ### Define calibration and model configurations:
@@ -20,6 +43,11 @@ LES profiles are available for different geolocations ("cfsites"), spanning seas
 ### Analyze output with:
 - `julia --project plot_ensemble.jl` - plots vertical profiles of all ensemble members in a given iteration, given path to calibration output
 - `julia --project edmf_ensemble_stats.jl` - computes and plots metrics offline [i.e., root mean squared error (RMSE)] as a function of iteration, given path to calibration output.
-- `julia --project plot_eki.jl` - plot eki metrics [loss, var-weighted loss] and `y`, `g` vectors vs iteration, display best particles
+- `julia --project plot_eki.jl` - plot eki metrics [loss, variance-weighted loss] and `y`, `G` vectors vs iteration, display best particles
 
+## Troubleshooting
 
+- **Memory Issues**: If you encounter out of memory errors, increase the memory allocation in both `run_calibration.sbatch` and the `experiment_config.yml` file. This is particularly important when working with larger batch sizes. Example error message:
+  ```
+  srun: error: hpc-92-10: task 9: Out Of Memory
+  ```

calibration/experiments/gcm_driven_scm/edmf_ensemble_stats.jl

@@ -1,12 +1,8 @@
 #!/usr/bin/env julia
 
-import ClimaComms
-@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
-
 using ArgParse
 using Distributed
-addprocs()
-
+addprocs(1)
 
 @everywhere begin
     using EnsembleKalmanProcesses: TOMLInterface
@@ -60,6 +56,13 @@ function parse_args()
     return parse_with_settings(s)
 end
 
+@everywhere function validate_ensemble_member(iteration_dir, batch_size)
+    config_dirs =
+        filter(x -> isdir(joinpath(iteration_dir, x)), readdir(iteration_dir))
+    num_configs = count(x -> startswith(x, "config_"), config_dirs)
+    return num_configs == batch_size
+end
+
 function main()
     args = parse_args()
 
@@ -87,6 +90,7 @@ function main()
     cal_vars = config_dict["y_var_names"]
     const_noise_by_var = config_dict["const_noise_by_var"]
     n_iterations = config_dict["n_iterations"]
+    batch_size = config_dict["batch_size"]
     model_config_dict =
         YAML.load_file(joinpath(output_dir, "configs", "model_config.yml"))
 
@@ -95,9 +99,6 @@ function main()
     end
 
     ref_paths, _ = get_les_calibration_library()
-    comms_ctx = ClimaComms.SingletonCommsContext()
-    atmos_config = CA.AtmosConfig(model_config_dict; comms_ctx)
-    zc_model = get_z_grid(atmos_config, z_max = z_max)
 
     @everywhere function calculate_statistics(y_var)
         non_nan_values = y_var[.!isnan.(y_var)]
@@ -124,9 +125,9 @@ function main()
         cal_vars,
         const_noise_by_var,
         ref_paths,
-        zc_model,
         reduction,
         ensemble_size,
+        batch_size,
     )
         println("Processing Iteration: $iteration")
         stats_df = DataFrame(
@@ -141,13 +142,25 @@ function main()
             rmse_std = Union{Missing, Float64}[],
         )
         config_indices = get_batch_indicies_in_iteration(iteration, output_dir)
+        iteration_dir =
+            joinpath(output_dir, "iteration_$(lpad(iteration, 3, '0'))")
+
+        valid_ensemble_members = filter(
+            config_i -> validate_ensemble_member(
+                joinpath(iteration_dir, "member_$(lpad(config_i, 3, '0'))"),
+                batch_size,
+            ),
+            config_indices,
+        )
+
         for var_name in var_names
             means = Float64[]
             maxs = Float64[]
             mins = Float64[]
             sum_squared_errors = zeros(Float64, ensemble_size)
-            for config_i in config_indices
-                data = ensemble_data(
+
+            for config_i in valid_ensemble_members
+                data, zc_model = ensemble_data(
                     process_profile_variable,
                     iteration,
                     config_i,
@@ -157,6 +170,7 @@ function main()
                     output_dir = output_dir,
                     z_max = z_max,
                     n_vert_levels = n_vert_levels,
+                    return_z_interp = true,
                 )
                 for i in 1:size(data, 2)
                     y_var = data[:, i]
@@ -166,25 +180,32 @@ function main()
                     push!(mins, col_min)
                 end
                 if in(var_name, cal_vars)
+                    ref_path = ref_paths[config_i]
+                    cfsite_number, _, _, _ = parse_les_path(ref_path)
+                    forcing_type = get_cfsite_type(cfsite_number)
+
+                    ti = config_dict["y_t_start_sec"]
+                    ti = isa(ti, AbstractFloat) ? ti : ti[forcing_type]
+                    tf = config_dict["y_t_end_sec"]
+                    tf = isa(tf, AbstractFloat) ? tf : tf[forcing_type]
+
                     y_true, Σ_obs, norm_vec_obs = get_obs(
-                        ref_paths[config_i],
+                        ref_path,
                         [var_name],
                         zc_model;
-                        ti = config_dict["y_t_start_sec"],
-                        tf = config_dict["y_t_end_sec"],
+                        ti = ti,
+                        tf = tf,
                         Σ_const = const_noise_by_var,
                         z_score_norm = false,
                     )
                     sum_squared_errors +=
                         compute_ensemble_squared_error(data, y_true)
                 end
             end
+
             if in(var_name, cal_vars)
-                # Compute RMSE per ensemble member
                 rmse_per_member = sqrt.(sum_squared_errors / n_vert_levels)
-                # Filter out NaNs (failed simulations)
                 valid_rmse = rmse_per_member[.!isnan.(rmse_per_member)]
-                non_nan_simulation_count = length(valid_rmse)
                 mean_rmse = mean(valid_rmse)
                 min_rmse = minimum(valid_rmse)
                 max_rmse = maximum(valid_rmse)
@@ -226,9 +247,9 @@ function main()
                 cal_vars,
                 const_noise_by_var,
                 ref_paths,
-                zc_model,
                 reduction,
                 ensemble_size,
+                batch_size,
             ),
             iterations_list,
         )

calibration/experiments/gcm_driven_scm/experiment_config.yml

@@ -1,25 +1,61 @@
-prior_path: prior_prognostic_pi_entr.toml
+prior_path: prior_prognostic_pi_entr_smooth_entr_detr_impl_0M_v1.toml
 ensemble_size: 100
-n_iterations: 12
-batch_size: 2 # number of cases per iteration
-model_config : model_config_prognostic.yml # options {model_config_prognostic.yml, model_config_diagnostic.yml}
-output_dir : output/exp_1 # output dir
-y_var_names: [thetaa, hus, clw] # calibration variables clw
+n_iterations: 8
+batch_size: 5 # number of cases per iteration
+# model_config : model_config_prognostic.yml # options {model_config_prognostic.yml, model_config_diagnostic.yml}
+model_config : model_config_prognostic_impl.yml
+output_dir : /central/scratch/cchristo/debug/exp16
+
+# Slurm resource configuration
+slurm_time: "02:00:00"
+slurm_mem_per_cpu: "25G"
+slurm_cpus_per_task: 1
+
+y_var_names: [thetaa, hus, clw]  # calibration variables clw clw]
 log_vars: ["clw"] # take log(var) when forming y, g
-z_max : 4000 # spatial subsetting: use statistics from [0, z_max] (in [m]) for calibration
-dims_per_var : 29 # num dimensions per variable (num cells in vertical profile below z_max)
+# log_vars: []
+
+nice_loc_ug: 0.01
+nice_loc_gg: 0.5
+
+z_max: null
+z_cal_grid: # calibration grid (stretch-grid parameters). In general, `z_elem` should be the same for all types
+  shallow:
+    z_max: 4000.0
+    z_elem: 30
+    dz_bottom: 30
+  deep:
+    z_max: 15000.0
+    z_elem: 30
+    dz_bottom: 30
+dims_per_var : 30 # num dimensions per variable (num cells in vertical profile below z_max)
 # eki_timestep: 0.1 # timestep of eki, if using default
-y_t_start_sec : 475200.0 # start time of LES averaging window [s] : 5.5 days
-y_t_end_sec : 518400.0 # end time of LES averaging window [s] : 6 days (LES length = 6 days)
-g_t_start_sec : 216000.0 # start time of SCM averaging window [s] : 2.5 days
-g_t_end_sec : 259200.0 # end time of SCM averaging window [s] : 3 days (SCM length = 3 days)
+
+y_t_start_sec: # start time of LES averaging window [s]
+  shallow: 475200.0 # 5.5 days
+  deep: 302400.0 # 3.5 days
+y_t_end_sec: # end time of LES averaging window [s]
+  shallow: 518400.0 # 6 days (LES length = 6 days)
+  deep: 345600.0 # 4 days (LES length = 4 days)
+g_t_start_sec: 216000.0 # start time of SCM averaging window [s] : 2.5 days
+g_t_end_sec: 259200.0 # end time of SCM averaging window [s] : 3 days (SCM length = 3 days)
 
 norm_factors_by_var:
-  thetaa: [298.828, 8.617]
-  hus: [0.00676, 0.00423]
-  clw: [-9.808, 3.116] # log norm factors
+  thetaa: [301.218, 15.235]
+  hus: [0.00672, 0.00477]
+  clw: [-9.579, 3.164] # log norm factors
+  # cli: [-11.697, 1.304] # log norm factors
 
 const_noise_by_var:
-  thetaa: 0.00005
-  hus: 0.00005
-  clw: 0.00005
+  thetaa: 0.0016
+  hus: 0.0016
+  clw: 0.0045
+  # clw: 0.0016
+  # cli: 0.01
+
+pretrained_nn_path: "/home/cchristo/ml_mixing_length/nn_666p_leaky_relu.jld2"
+
+# Config files for deep and shallow cases
+forcing_toml_files:
+  shallow: "scm_tomls/gcmdriven_relaxation_shallow_forcing.toml"
+  deep: "scm_tomls/gcmdriven_relaxation_deep_forcing.toml"

calibration/experiments/gcm_driven_scm/get_les_metadata.jl

@@ -3,10 +3,22 @@ using Glob
 """
 """
 
+# cfSite numbers
+CFSITE_TYPES = Dict(
+    "shallow" => (collect(4:15)..., collect(17:23)...),
+    "deep" =>
+        (collect(30:33)..., collect(66:70)..., 82, 92, 94, 96, 99, 100),
+)
+
 function get_les_calibration_library()
     les_library = get_shallow_LES_library()
-    # AMIP4K data: July, NE Pacific
-    cfsite_numbers = (17, 23)
+    # AMIP data: July, NE Pacific
+    # cfsite_numbers = (17, 18, 22, 23, 30, 94)
+    # cfsite_numbers = (17, 22, 23, 30, 33, 94)
+    cfsite_numbers = (17, 21, 23, 30, 33)# 94)
+    # cfsite_numbers = (30, 33,)# 94)
+
+    # cfsite_numbers = (17, 30,)# 94)
     les_kwargs = (forcing_model = "HadGEM2-A", month = 7, experiment = "amip")
     ref_paths = [
         get_stats_path(get_cfsite_les_dir(cfsite_number; les_kwargs...)) for
@@ -15,6 +27,20 @@ function get_les_calibration_library()
     return (ref_paths, cfsite_numbers)
 end
 
+function get_cfsite_type(i, cfsite_numbers)
+    return get_cfsite_type(cfsite_numbers[i])
+end
+
+function get_cfsite_type(cfsite_number::Int)
+    if cfsite_number in CFSITE_TYPES["shallow"]
+        return "shallow"
+    elseif cfsite_number in CFSITE_TYPES["deep"]
+        return "deep"
+    else
+        @error "cfSite number $(cfsite_number) not found in available sites."
+    end
+end
+
 """
     get_LES_library
 
@@ -25,7 +51,18 @@ and experiments.
 """
 function get_LES_library()
     LES_library = get_shallow_LES_library()
-    deep_sites = (collect(30:33)..., collect(66:70)..., 82, 92, 94, 96, 99, 100)
+    deep_sites = deepcopy(CFSITE_TYPES["deep"])
+
+
+    # remove <0 ql/cli cases 
+    # sites_07 = deepcopy(setdiff(deep_sites, [92, 99, 100]))
+    # append!(LES_library["HadGEM2-A"]["07"]["cfsite_numbers"], sites_07)
+    # sites_01 = deepcopy(setdiff(deep_sites, [99,]))
+    # append!(LES_library["HadGEM2-A"]["01"]["cfsite_numbers"], sites_01)
+    # sites_04 = deepcopy(setdiff(deep_sites, [32, 92, 94, 96, 99, 100]))
+    # append!(LES_library["HadGEM2-A"]["04"]["cfsite_numbers"], sites_04)
+    # sites_10 = deepcopy(setdiff(deep_sites, [92, 94, 99, 100]))
+    # append!(LES_library["HadGEM2-A"]["10"]["cfsite_numbers"], sites_10)
 
     append!(LES_library["HadGEM2-A"]["07"]["cfsite_numbers"], deep_sites)
     append!(LES_library["HadGEM2-A"]["01"]["cfsite_numbers"], deep_sites)
@@ -34,6 +71,7 @@ function get_LES_library()
     sites_10 = deepcopy(setdiff(deep_sites, [94, 100]))
     append!(LES_library["HadGEM2-A"]["10"]["cfsite_numbers"], sites_10)
 
+
     LES_library_full = deepcopy(LES_library)
     for model in keys(LES_library_full)
         for month in keys(LES_library_full[model])
@@ -103,8 +141,7 @@ function get_shallow_LES_library()
         "CNRM-CM5" => Dict(),
         "CNRM-CM6-1" => Dict(),
     )
-    Shen_et_al_sites = collect(4:15)
-    append!(Shen_et_al_sites, collect(17:23))
+    Shen_et_al_sites = collect(deepcopy(CFSITE_TYPES["shallow"]))
 
     # HadGEM2-A model (76 AMIP-AMIP4K pairs)
     LES_library["HadGEM2-A"]["10"] = Dict()