Add calibration framework and perfect model experiment

.buildkite/pipeline.yml

@@ -39,6 +39,9 @@ steps:
       - "julia --project=perf -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=perf -e 'using Pkg; Pkg.status()'"
 
+      - echo "--- Instantiate calibration/test"
+      - "julia --project=calibration/test -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
+
       - echo "--- Download artifacts"
       - "julia --project=examples artifacts/download_artifacts.jl"
 
@@ -470,6 +473,14 @@ steps:
         agents:
           slurm_mem: 20GB
 
+  - group: "Calibration"
+    steps:
+      - label: "Calibration interface unit tests"
+        command: julia --project=calibration/test calibration/test/interface.jl
+      - label: "end to end test"
+        command: julia --project=calibration/test calibration/test/e2e_test.jl
+        artifact_paths: "calibration_end_to_end_test/*"
+
   - group: "Diagnostic EDMFX"
     steps:
 

.gitignore

@@ -28,6 +28,7 @@ docs/site/
 # environment.
 /Manifest.toml
 test/Manifest.toml
+calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Manifest.toml
 
 # ignore vscode artifacts
 *.vscode
@@ -43,6 +44,8 @@ test/Manifest.toml
 *.hdf5
 *.h5
 *.g
+*.lock
+*.loc
 
 # misc
 .DS_Store

Artifacts.toml

@@ -4,3 +4,10 @@ git-tree-sha1 = "9da4af348af2606764e5e2c94d9439714221dff1"
     [[aerosol2005.download]]
     sha256 = "952d723b8462439c266dc0df79d0d4e71f774559edb941e03b7979e8ad743f56"
     url = "https://caltech.box.com/shared/static/chmel1vdthfvfac0yl61ayw2jqnjzr2c.gz"
+
+[atmos_held_suarez_obs]
+git-tree-sha1 = "d93ff2958e12d6bba6e8343bfa73554390df52e4"
+
+    [[atmos_held_suarez_obs.download]]
+    sha256 = "f5d63df2bab849632bd2fcd5a4e3165465d1898e1a5fbee6e5fe07f5b3e57fc3"
+    url = "https://caltech.box.com/shared/static/sbn6afsgn2xzxi5n0ffs928otstshjvi.gz"

calibration/README.md

@@ -0,0 +1,23 @@
+# ClimaAtmos Calibration Experiments
+
+This directory contains the model interface (`model_interface.jl`) for ClimaCalibrate and folders (within `experiments/`) for reproducing calibration experiments.
+
+## Current Experiments
+
+### sphere_held_suarez_rhoe_equilmoist
+
+A perfect-model calibration, using ClimaAtmos v0.24.0.
+
+- Configuration: 0-moment microphysics scheme, spherical with Held-Suarez forcing.
+- Observational data: zonal and 60-day time average air temperature at 242m altitude. Provided by [ClimaArtifacts](https://github.com/CliMA/ClimaArtifacts/tree/main/atmos_held_suarez_obs).
+- Parameter being calibrated: `equator_pole_temperature_gradient_wet`
+
+For more details, see the experiment directory.
+
+To reproduce the results, on the Caltech central cluster run
+
+```sbatch calibration/experiments/sphere_held_suarez_rhoe_equilmoist/pipeline.sbatch```
+
+## New Experiments
+
+To set up your own experiment, please follow the [setup guide](https://clima.github.io/ClimaCalibrate.jl/dev/atmos_setup_guide/).

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Artifacts.toml

@@ -0,0 +1,6 @@
+[atmos_held_suarez_obs]
+git-tree-sha1 = "d93ff2958e12d6bba6e8343bfa73554390df52e4"
+
+    [[atmos_held_suarez_obs.download]]
+    sha256 = "f5d63df2bab849632bd2fcd5a4e3165465d1898e1a5fbee6e5fe07f5b3e57fc3"
+    url = "https://caltech.box.com/shared/static/sbn6afsgn2xzxi5n0ffs928otstshjvi.gz"

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Project.toml

@@ -0,0 +1,13 @@
+[deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ClimaAnalysis = "29b5916a-a76c-4e73-9657-3c8fd22e65e6"
+ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
+ClimaCalibrate = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
+ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
+
+[compat]
+ClimaAtmos = "=0.24.0"

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/model_config.yml

@@ -0,0 +1,15 @@
+dt: 500secs
+t_end: 560days
+ode_algo: ARS343
+moist: equil 
+forcing: held_suarez
+precip_model: 0M
+job_id: sphere_held_suarez_rhoe_equilmoist
+output_dir: output/sphere_held_suarez_rhoe_equilmoist
+restart_file: /groups/esm/ClimaArtifacts/artifacts/atmos_held_suarez_obs/day200.0.hdf5
+output_default_diagnostics: false
+diagnostics:
+  - reduction_time: average
+    short_name: ta
+    period: 60days
+    writer: nc

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/observation_map.jl

@@ -0,0 +1,34 @@
+import EnsembleKalmanProcesses: TOMLInterface
+import ClimaCalibrate: observation_map, ExperimentConfig
+import ClimaAnalysis: SimDir, get, slice, average_lat, average_lon
+
+function observation_map(iteration)
+    ensemble_size = 10
+    output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
+    single_member_dims = (1,)
+    G_ensemble = Array{Float64}(undef, single_member_dims..., ensemble_size)
+
+    for m in 1:ensemble_size
+        member_path =
+            TOMLInterface.path_to_ensemble_member(output_dir, iteration, m)
+        simdir_path = joinpath(member_path, "output_active")
+        if isdir(simdir_path)
+            simdir = SimDir(simdir_path)
+            G_ensemble[:, m] .= process_member_data(simdir)
+        else
+            G_ensemble[:, m] .= NaN
+        end
+    end
+    return G_ensemble
+end
+
+const meters = 1.0
+const days = 86400.0
+# Cut off first 120 day to get equilibrium, take second level slice
+function process_member_data(simdir::SimDir)
+    isempty(simdir.vars) && return NaN
+    ta = get(simdir; short_name = "ta", reduction = "average", period = "60d")
+    zonal_avg_temp_observations =
+        slice(average_lat(average_lon(ta)), z = 242meters)
+    return slice(zonal_avg_temp_observations, time = 240days).data
+end

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/pipeline.sbatch

@@ -0,0 +1,32 @@
+#!/bin/bash
+#SBATCH --time=3:00:00
+#SBATCH --output=sphere_held_suarez_rhoe_equilmoist_calibration.txt
+
+julia --project=calibration/experiments/sphere_held_suarez_rhoe_equilmoist -e '
+import Pkg; Pkg.instantiate(;verbose=true)
+import JLD2
+using ClimaUtilities.ClimaArtifacts
+import ClimaCalibrate: calibrate, ExperimentConfig, CaltechHPC, get_prior, kwargs
+
+experiment_dir = dirname(Base.active_project())
+include(joinpath(experiment_dir, "observation_map.jl"))
+artifact_path = @clima_artifact("atmos_held_suarez_obs")
+
+# Load data and configurations
+n_iterations = 2
+ensemble_size = 10
+observations = JLD2.load_object(joinpath(artifact_path, "obs_mean.jld2"))
+noise = JLD2.load_object(joinpath(artifact_path, "obs_noise_cov.jld2"))
+prior = get_prior(joinpath(experiment_dir, "prior.toml"))
+output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
+
+experiment_config = ExperimentConfig(; n_iterations, ensemble_size, 
+                                    observations, noise, output_dir, prior)
+slurm_kwargs = kwargs(time = 50, ntasks = 16)
+eki = calibrate(CaltechHPC, experiment_config; slurm_kwargs, verbose=true)
+
+include(joinpath(experiment_dir, "postprocessing.jl"))
+prior = get_prior(joinpath(experiment_dir, "prior.toml"))
+convergence_plot(eki, prior)
+scatter_plot(eki)
+'

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/postprocessing.jl

@@ -0,0 +1,105 @@
+import EnsembleKalmanProcesses as EKP
+using EnsembleKalmanProcesses.ParameterDistributions
+using EnsembleKalmanProcesses.TOMLInterface
+using Distributions
+import Statistics: mean
+import CairoMakie: Makie
+import ClimaCalibrate
+
+function convergence_plot(
+    eki::EKP.EnsembleKalmanProcess,
+    prior::ParameterDistribution;
+    theta_star = 65.0,
+    output = joinpath("output", "sphere_held_suarez_rhoe_equilmoist"),
+)
+    u_vec = EKP.get_u(eki)
+    error_vec = Float64[]
+    spread_vec = Float64[]
+    for ensemble in u_vec
+        ensemble_error = 0.0
+        ensemble_spread = 0.0
+        ensemble_mean = mean(ensemble)
+        for i in ensemble
+            ensemble_error += abs(i - theta_star)^2
+            ensemble_spread += abs(i - ensemble_mean)^2
+        end
+        ensemble_error /= length(ensemble)
+        ensemble_spread /= length(ensemble)
+
+        push!(error_vec, ensemble_error)
+        push!(spread_vec, ensemble_spread)
+    end
+
+    phi_vec = transform_unconstrained_to_constrained(prior, u_vec)
+    u_series = eachcol(reduce(vcat, u_vec))
+    phi_series = eachcol(reduce(vcat, phi_vec))
+
+    f = Makie.Figure(title = "Convergence Plot", resolution = (800, 800))
+
+    ax = Makie.Axis(
+        f[1, 1],
+        xlabel = "Iteration",
+        ylabel = "Error",
+        xticks = 0:50,
+    )
+    Makie.lines!(ax, 0.0:(length(error_vec) - 1), error_vec)
+
+    ax = Makie.Axis(
+        f[1, 2],
+        xlabel = "Iteration",
+        ylabel = "Spread",
+        xticks = 0:50,
+    )
+    Makie.lines!(ax, 0.0:(length(spread_vec) - 1), spread_vec)
+
+    ax = Makie.Axis(
+        f[2, 1],
+        xlabel = "Iteration",
+        ylabel = "Unconstrained Parameters",
+        xticks = 0:50,
+    )
+    Makie.lines!.(ax, tuple(0.0:(length(u_series[1]) - 1)), u_series)
+
+    ax = Makie.Axis(
+        f[2, 2],
+        xlabel = "Iteration",
+        ylabel = "Constrained Parameters",
+        xticks = 0:50,
+    )
+    Makie.lines!.(ax, tuple(0.0:(length(phi_series[1]) - 1)), phi_series)
+    Makie.hlines!(ax, [theta_star], color = :red, linestyle = :dash)
+    Makie.save(joinpath(output, "convergence.png"), f)
+end
+
+function scatter_plot(
+    eki,
+    output = joinpath("output", "sphere_held_suarez_rhoe_equilmoist"),
+)
+    # Define figure with explicit size
+    f = Makie.Figure(resolution = (800, 600))
+    a = Makie.Axis(
+        f[1, 1],
+        title = "60-Day Zonal Avg Temp at 242m versus Unconstrained Equator-Pole Temp Gradient",
+        ylabel = "Parameter Value",
+        xlabel = "Temperature (K)",
+    )
+
+    g = vec.(EKP.get_g(eki; return_array = true))
+    u = vec.(EKP.get_u(eki; return_array = true))
+
+    for (gg, uu) in zip(g, u)
+        Makie.scatter!(a, gg, uu)
+    end
+
+    # Save the figure
+    Makie.save(joinpath(output, "scatter.png"), f)
+end
+
+# Uncomment for easy plotting
+# import JLD2
+# iteration = 2
+# output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
+# eki_filepath = joinpath(ClimaCalibrate.path_to_iteration(output_dir, iteration), "eki_file.jld2")
+# eki = JLD2.load_object(eki_filepath)
+# prior_path = joinpath("calibration", "experiments", "sphere_held_suarez_rhoe_equilmoist", "prior.toml")
+# prior = ClimaCalibrate.get_prior(prior_path)

calibration/experiments/sphere_held_suarez_rhoe_equilmoist/prior.toml

@@ -0,0 +1,5 @@
+["equator_pole_temperature_gradient_wet"]
+prior = "Parameterized(Normal(4.779568,0.31223328))"
+constraint = "[bounded_below(0)]"
+# EKP.ParameterDistributions.constrained_gaussian("name", 125, 40, 0, Inf)
+type = "float"

calibration/model_interface.jl

@@ -0,0 +1,54 @@
+import ClimaAtmos as CA
+import YAML
+import ClimaComms
+@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+using ClimaUtilities.ClimaArtifacts
+import ClimaCalibrate:
+    set_up_forward_model, run_forward_model, path_to_ensemble_member
+
+"""
+    set_up_forward_model(member, iteration, experiment_dir::AbstractString)
+    set_up_forward_model(member, iteration, config_dict::AbstractDict)
+
+Return an AtmosConfig object for the given member and iteration.
+
+Turns off default diagnostics and sets the TOML parameter file to the member's path.
+This assumes that the  config dictionary has an `output_dir` key.
+"""
+function set_up_forward_model(member, iteration, experiment_dir::AbstractString)
+    config_dict = YAML.load_file(joinpath(experiment_dir, "model_config.yml"))
+    set_up_forward_model(member, iteration, config_dict::AbstractDict)
+end
+
+function set_up_forward_model(member, iteration, config_dict::AbstractDict)
+    output_dir = config_dict["output_dir"]
+    member_path = path_to_ensemble_member(output_dir, iteration, member)
+    config_dict["output_dir"] = member_path
+    parameter_path = joinpath(member_path, "parameters.toml")
+    if haskey(config_dict, "toml")
+        push!(config_dict["toml"], parameter_path)
+    else
+        config_dict["toml"] = [parameter_path]
+    end
+
+    # Turn off default diagnostics
+    config_dict["output_default_diagnostics"] = false
+    return CA.AtmosConfig(config_dict)
+end
+
+"""
+    run_forward_model(atmos_config::CA.AtmosConfig)
+
+Run the atmosphere model with the given an AtmosConfig object.
+Currently only has basic error handling.
+"""
+function run_forward_model(atmos_config::CA.AtmosConfig)
+    simulation = CA.get_simulation(atmos_config)
+    sol_res = CA.solve_atmos!(simulation)
+    if sol_res.ret_code == :simulation_crashed
+        !isnothing(sol_res.sol) && sol_res.sol .= eltype(sol_res.sol)(NaN)
+        error(
+            "The ClimaAtmos simulation has crashed. See the stack trace for details.",
+        )
+    end
+end