Remove experiment ID

Update tests Fixes after rebase test fixes fix default output w
CliMA · May 16, 2024 · 8e0db64 · 8e0db64
1 parent b8a2a3a
commit 8e0db64
Show file tree

Hide file tree

Showing 16 changed files with 68 additions and 168 deletions.
diff --git a/docs/src/api.md b/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
 ```
 

diff --git a/docs/src/atmos_setup_guide.md b/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)

diff --git a/docs/src/emulate_sample.md b/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),

diff --git a/docs/src/precompilation.md b/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

diff --git a/experiments/surface_fluxes_perfect_model/generate_data.jl b/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

diff --git a/experiments/surface_fluxes_perfect_model/model_interface.jl b/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)

diff --git a/experiments/surface_fluxes_perfect_model/observation_map.jl b/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"

diff --git a/experiments/surface_fluxes_perfect_model/postprocessing.jl b/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"])
 

diff --git a/src/backends.jl b/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

diff --git a/src/ekp_interface.jl b/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,

diff --git a/src/model_interface.jl b/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
diff --git a/src/slurm.jl b/src/slurm.jl
@@ -2,7 +2,7 @@
 kwargs(; kwargs...) = Dict{Symbol, Any}(kwargs...)
 
 """
-generate_sbatch_script
+    generate_sbatch_script
 
 
 """
@@ -47,11 +47,7 @@ function generate_sbatch_script(
     model_interface = "$model_interface"; include(model_interface)
 
     experiment_dir = "$experiment_dir"
-    experiment_config = CAL.ExperimentConfig(experiment_dir)
-    experiment_id = experiment_config.id
-    physical_model = CAL.get_forward_model(Val(Symbol(experiment_id)))
-    CAL.run_forward_model(physical_model, CAL.get_config(physical_model, member, iteration, experiment_dir))
-    @info "Forward Model Run Completed" experiment_id physical_model iteration member'
+    CAL.run_forward_model(CAL.set_up_forward_model(member, iteration, experiment_dir))'
     """
     return sbatch_contents
 end
@@ -157,7 +153,8 @@ If verbose, includes the ensemble member's output.
 """
 function log_member_error(output_dir, iteration, member, verbose = false)
     member_log = path_to_model_log(output_dir, iteration, member)
-    warn_str = "Ensemble member $member raised an error. See model log at $(abspath(member_log)) for stacktrace"
+    warn_str = """Ensemble member $member raised an error. See model log at \
+    $(abspath(member_log)) for stacktrace"""
     if verbose
         stacktrace = replace(readchomp(member_log), "\\n" => "\n")
         warn_str = warn_str * ": \n$stacktrace"
@@ -167,9 +164,11 @@ end
 
 function report_iteration_status(statuses, output_dir, iter)
     all(job_completed.(statuses)) || error("Some jobs are not complete")
+
     if all(job_failed, statuses)
         error(
-            "Full ensemble for iteration $iter has failed. See model logs in $(abspath(path_to_iteration(output_dir, iter))) for details.",
+            """Full ensemble for iteration $iter has failed. See model logs in
+            $(abspath(path_to_iteration(output_dir, iter)))""",
         )
     elseif any(job_failed, statuses)
         @warn "Failed ensemble members: $(findall(job_failed, statuses))"
@@ -246,5 +245,4 @@ function format_slurm_time(minutes::Int)
         )
     end
 end
-
 format_slurm_time(str::AbstractString) = str
diff --git a/test/ekp_interface.jl b/test/ekp_interface.jl
@@ -17,7 +17,6 @@ n_iterations = 1
 ensemble_size = 10
 
 config = CAL.ExperimentConfig(
-    "test",
     n_iterations,
     ensemble_size,
     observations,