Merge pull request #84 from CliMA/ne/model_interface

Simplify model interface and remove experiment ID

.github/workflows/ci.yml

@@ -30,7 +30,16 @@ jobs:
           - x64
     steps:
       - uses: actions/checkout@v4
-      - uses: julia-actions/cache@v1
+      - uses: actions/cache@v1
+        env:
+          cache-name: cache-artifacts
+        with:
+          path: ~/.julia/artifacts
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-test-${{ env.cache-name }}-
+            ${{ runner.os }}-test-
+            ${{ runner.os }}-
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}

docs/src/api.md

@@ -3,9 +3,8 @@
 ## Model Interface
 
 ```@docs
-ClimaCalibrate.get_config
+ClimaCalibrate.set_up_forward_model
 ClimaCalibrate.run_forward_model
-ClimaCalibrate.get_forward_model
 ClimaCalibrate.observation_map
 ```
 

docs/src/atmos_setup_guide.md

@@ -16,9 +16,9 @@ To calibrate parameters, you need:
 - Truth and noise data
 - Observation map script with a function `observation_map(iteration)`
 
-These components are detailed in the guide below. Examples of all of these can also be found in `experiments/sphere_held_suarez_rhoe_equilmoist`
+These components are detailed in the guide below. Examples of all of these can also be found in ClimaAtmos in `calibration/experiments/sphere_held_suarez_rhoe_equilmoist`
 
-First, create a folder for your experiment with a descriptive experiment id in the `experiments/` folder. All of the components described below will be stored in this folder.
+First, create a folder for your experiment with a descriptive name in the `calibration/experiments/` folder. All of the components described below will be stored in this folder.
 
 ## Atmos Configuration File
 
@@ -96,8 +96,7 @@ Constraint constructors:
 
 The observation map is applied to process model output diagnostics into the exact observable used to fit to observations. In a perfect model setting it is used also to generate the observation.
 
-Your observation map file must export a function `observation_map(::Val{:<experiment_id>}, iteration)`, this function is specific to each experiment, so it is dispatched on the `experiment_id`.
-These requirements arise from the update step, which runs the function with your given experiment ID.
+These requirements arise from the update step, which runs the `observation_map` function.
 This function must load in model diagnostics for each ensemble member in the iteration and construct an array `arr = Array{Float64}(undef, dims..., ensemble_size)` such that
 `arr[:, i]` will return the i-th ensemble member's observation map output. Note this floating point precision is required for the EKI update step.
 
@@ -113,7 +112,7 @@ First, construct the array to store the ensemble's observations. Then, for each
 Pseudocode for `observation_map(iteration)`:
 
 ```julia
-function observation_map(::Val{:sphere_held_suarez_rhoe_equilmoist}, iteration)
+function observation_map(iteration)
     # Get Configuration
     config_file = joinpath("calibration", "sphere_held_suarez_rhoe_equilmoist")
     config = ExperimentConfig(config_file)

docs/src/emulate_sample.md

@@ -16,7 +16,8 @@ import ClimaCalibrate as CAL
 ```
 
 Next, load in the data, EKP object, and prior distribution. These values are taken
-from the perfect model experiment with experiment ID `sphere_held_suarez_rhoe_equilmoist`.
+from the Held-Suarez perfect model experiment in ClimaAtmos.
+
 ```julia
 asset_path = joinpath(
     pkgdir(CAL),

docs/src/precompilation.md

@@ -7,7 +7,7 @@ For ClimaCalibrate, this is useful under certain conditions:
 2. **The model runtime is short compared to the compile time.** If the model runtime is an order of magnitude or more than the compilation time, any benefit from reduced compilation time will be trivial.
 
 # How do I precompile my configuration?
-The easiest way is by copying and pasting the code snippet below into `src/ClimaCalibrate.jl` and replacing the `job_id` with your experiment ID.
+The easiest way is by copying and pasting the code snippet below into `src/ClimaCalibrate.jl`.
 This will precompile the model step and all callbacks for the given configuration.
 ```julia
 using PrecompileTools
@@ -16,10 +16,8 @@ import ClimaAtmos as CA
 import YAML
 
 @setup_workload begin
-    config_file = "your experiment dir"
-    ExperimentConfig(config_file; output_dir = "precompilation")
+    config_file = Dict("FLOAT_TYPE" => "Float64")
     @compile_workload begin
-        initialize(job_id)
         config = CA.AtmosConfig(config_dict)
         simulation = CA.get_simulation(config)
         (; integrator) = simulation

experiments/surface_fluxes_perfect_model/generate_data.jl

@@ -1,5 +1,4 @@
 # generate_data: generate true y, noise and x_inputs
-experiment_id = "surface_fluxes_perfect_model"
 
 import SurfaceFluxes as SF
 import SurfaceFluxes.Parameters as SFPP
@@ -10,8 +9,9 @@ import SurfaceFluxes.Parameters: SurfaceFluxesParameters
 using ClimaCalibrate
 
 pkg_dir = pkgdir(ClimaCalibrate)
-experiment_path = "$pkg_dir/experiments/$experiment_id"
-data_path = "$experiment_path/data"
+experiment_path =
+    joinpath(pkg_dir, "experiments", "surface_fluxes_perfect_model")
+data_path = joinpath(experiment_path, "data")
 include(joinpath(experiment_path, "model_interface.jl"))
 
 FT = Float32

experiments/surface_fluxes_perfect_model/model_interface.jl

@@ -1,11 +1,6 @@
 import EnsembleKalmanProcesses as EKP
 using ClimaCalibrate
-import ClimaCalibrate:
-    AbstractPhysicalModel,
-    get_config,
-    run_forward_model,
-    get_forward_model,
-    ExperimentConfig
+import ClimaCalibrate: set_up_forward_model, run_forward_model, ExperimentConfig
 import YAML
 
 """
@@ -32,7 +27,6 @@ We need to follow the following steps for the calibration:
 4. define the prior distributions for θ (this is subjective and can be based on expert knowledge or previous studies)
 
 """
-struct SurfaceFluxModel <: AbstractPhysicalModel end
 
 experiment_dir = joinpath(
     pkgdir(ClimaCalibrate),
@@ -42,34 +36,23 @@ experiment_dir = joinpath(
 include(joinpath(experiment_dir, "sf_model.jl"))
 include(joinpath(experiment_dir, "observation_map.jl"))
 
-function get_forward_model(::Val{:surface_fluxes_perfect_model})
-    return SurfaceFluxModel()
-end
-
-function get_config(
-    model::SurfaceFluxModel,
-    member,
-    iteration,
-    experiment_dir::AbstractString,
-)
-    return get_config(
-        model,
+function set_up_forward_model(member, iteration, experiment_dir::AbstractString)
+    return set_up_forward_model(
         member,
         iteration,
         ExperimentConfig(experiment_dir),
     )
 end
 
 """
-    get_config(member, iteration, experiment_dir::AbstractString)
-    get_config(member, iteration, experiment_config::ExperimentConfig)
+    set_up_forward_model(member, iteration, experiment_dir::AbstractString)
+    set_up_forward_model(member, iteration, experiment_config::ExperimentConfig)
 
 Returns an config dictionary object for the given member and iteration.
 Given an experiment dir, it will load the ExperimentConfig
 This assumes that the config dictionary has the `output_dir` key.
 """
-function get_config(
-    ::SurfaceFluxModel,
+function set_up_forward_model(
     member,
     iteration,
     experiment_config::ExperimentConfig,
@@ -103,7 +86,7 @@ end
 Runs the model with the given an AbstractDict object.
 """
 
-function run_forward_model(::SurfaceFluxModel, config::AbstractDict)
+function run_forward_model(config::AbstractDict)
     x_inputs = load_profiles(config["x_data_file"])
     FT = typeof(x_inputs.profiles_int[1].T)
     obtain_ustar(FT, x_inputs, config)

experiments/surface_fluxes_perfect_model/observation_map.jl

@@ -16,7 +16,7 @@ experiment_dir = joinpath(
 Returns the observation map (from the raw model output to the observable y),
 as specified by process_member_data, for the given iteration.
 """
-function observation_map(::Val{:surface_fluxes_perfect_model}, iteration)
+function observation_map(iteration)
     config = ExperimentConfig(experiment_dir)
     (; output_dir, ensemble_size) = config
     model_output = "model_ustar_array.jld2"

experiments/surface_fluxes_perfect_model/postprocessing.jl

@@ -14,7 +14,6 @@ using ClimaCalibrate
 experiment_dir = dirname(Base.active_project())
 experiment_config = ClimaCalibrate.ExperimentConfig(experiment_dir)
 output_dir = experiment_config.output_dir
-experiment_id = experiment_config.id
 N_iter = experiment_config.n_iterations
 N_mem = experiment_config.ensemble_size
 
@@ -103,9 +102,7 @@ end
 
 
 pkg_dir = pkgdir(ClimaCalibrate)
-model_config = YAML.load_file(
-    joinpath(pkg_dir, "experiments", experiment_id, "model_config.yml"),
-)
+model_config = YAML.load_file(joinpath(experiment_dir, "model_config.yml"))
 
 eki_path = joinpath(
     ClimaCalibrate.path_to_iteration(output_dir, N_iter),
@@ -133,10 +130,7 @@ include(joinpath(experiment_dir, "model_interface.jl"))
 f = Makie.Figure()
 ax = Makie.Axis(f[1, 1], xlabel = "Iteration", ylabel = "Model Ustar")
 ustar_obs = JLD2.load_object(
-    joinpath(
-        pkg_dir,
-        "experiments/$experiment_id/data/synthetic_ustar_array_noisy.jld2",
-    ),
+    joinpath(pkg_dir, "$experiment_dir/data/synthetic_ustar_array_noisy.jld2"),
 )
 x_inputs = load_profiles(model_config["x_data_file"])
 

src/backends.jl

@@ -47,19 +47,15 @@ calibrate(b::Type{JuliaBackend}, experiment_dir::AbstractString) =
 
 function calibrate(::Type{JuliaBackend}, config::ExperimentConfig)
     initialize(config)
-    (; n_iterations, id, ensemble_size) = config
+    (; n_iterations, ensemble_size) = config
     eki = nothing
-    physical_model = get_forward_model(Val(Symbol(id)))
     for i in 0:(n_iterations - 1)
         @info "Running iteration $i"
         for m in 1:ensemble_size
-            run_forward_model(
-                physical_model,
-                get_config(physical_model, m, i, config),
-            )
+            run_forward_model(set_up_forward_model(m, i, config))
             @info "Completed member $m"
         end
-        G_ensemble = observation_map(Val(Symbol(id)), i)
+        G_ensemble = observation_map(i)
         save_G_ensemble(config, i, G_ensemble)
         eki = update_ensemble(config, i)
     end
@@ -150,7 +146,7 @@ function calibrate(
         )
         report_iteration_status(statuses, output_dir, iter)
         @info "Completed iteration $iter, updating ensemble"
-        G_ensemble = observation_map(Val(Symbol(config.id)), iter)
+        G_ensemble = observation_map(iter)
         save_G_ensemble(config, iter, G_ensemble)
         eki = update_ensemble(config, iter)
     end

src/ekp_interface.jl

@@ -16,7 +16,6 @@ ExperimentConfig holds the configuration for a calibration experiment.
 This can be constructed from a YAML configuration file or directly using individual parameters.
 """
 struct ExperimentConfig
-    id::AbstractString
     n_iterations::Integer
     ensemble_size::Integer
     observations::Any
@@ -40,10 +39,7 @@ function ExperimentConfig(filepath::AbstractString; kwargs...)
         error("Invalid experiment configuration filepath: `$filepath`")
     end
 
-    experiment_id =
-        get(config_dict, "experiment_id", last(splitdir(experiment_dir)))
-    default_output =
-        haskey(ENV, "CI") ? experiment_id : joinpath("output", experiment_id)
+    default_output = joinpath(experiment_dir, "output")
     output_dir = get(config_dict, "output_dir", default_output)
 
     n_iterations = config_dict["n_iterations"]
@@ -65,7 +61,6 @@ function ExperimentConfig(filepath::AbstractString; kwargs...)
     prior = get_prior(prior_path)
 
     return ExperimentConfig(
-        experiment_id,
         n_iterations,
         ensemble_size,
         observations,

src/model_interface.jl

@@ -1,66 +1,37 @@
 import EnsembleKalmanProcesses as EKP
 import YAML
 
-"""
-    AbstractPhysicalModel
 
-Abstract type to define the interface for physical models.
 """
-abstract type AbstractPhysicalModel end
+    set_up_forward_model(member, iteration, experiment_dir::AbstractString)
+    set_up_forward_model(member, iteration, experiment_config::ExperimentConfig)
 
-"""
-    get_config(physical_model::AbstractPhysicalModel, member, iteration, experiment_path::AbstractString)
-    get_config(physical_model::AbstractPhysicalModel, member, iteration, experiment_config)
+Set up and configure a single member's forward model. Used in conjunction with `run_forward_model`.
 
-Fetch the model information for a specific ensemble member and iteration based on a provided path.
 This function must be overriden by a component's model interface and 
 should set things like the parameter path and other member-specific settings.
 """
-function get_config(
-    physical_model::AbstractPhysicalModel,
-    member,
-    iteration,
-    experiment_path::AbstractString,
-)
-    experiment_config = ExperimentConfig(experiment_path)
-    return get_config(physical_model, member, iteration, experiment_config)
-end
+set_up_forward_model(member, iteration, experiment_dir::AbstractString) =
+    set_up_forward_model(member, iteration, ExperimentConfig(experiment_dir))
 
-get_config(
-    physical_model::AbstractPhysicalModel,
-    member,
-    iteration,
-    experiment_config,
-) = error("get_config not implemented for $physical_model")
+set_up_forward_model(member, iteration, experiment_config::ExperimentConfig) =
+    error("set_up_forward_model not implemented")
 
 """
-    run_forward_model(physical_model::AbstractPhysicalModel, config)
+    run_forward_model(config)
 
 Executes the forward model simulation with the given configuration.
-The `config` should be obtained from `get_config`. 
+`config` should be obtained from `set_up_forward_model`. 
 This function should be overridden with model-specific implementation details.
 """
-run_forward_model(physical_model::AbstractPhysicalModel, model_config) =
-    error("run_forward_model not implemented for $physical_model")
-
-"""
-    get_forward_model(experiment_id::Val)
-
-Retrieves a custom physical model struct for the specified experiment ID.
-Throws an error if the experiment ID is unrecognized.
-"""
-function get_forward_model(experiment_id::Val)
-    error("get_forward_model not implemented for $experiment_id")
-end
+run_forward_model(model_config) = error("run_forward_model not implemented")
 
 """
     observation_map(val:Val, iteration)
 
 Runs the observation map for the specified iteration.
 This function must be implemented for each calibration experiment.
 """
-function observation_map(val::Val, iteration)
-    error(
-        "observation_map not implemented for experiment $val at iteration $iteration",
-    )
+function observation_map(iteration)
+    error("observation_map not implemented")
 end