Merge branch 'main' into ro/p3testing

CliMA · Aug 19, 2024 · 9e7ce7c · 9e7ce7c
2 parents 5def380 + 1c9b899
commit 9e7ce7c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/src/CalibrateCMP/OptimizationUtils.jl b/src/CalibrateCMP/OptimizationUtils.jl
@@ -47,11 +47,17 @@ end
 
 # EKP
 function calibrate(::EKPStyle, priors, config, ref_stats_list::Vector{ReferenceStatistics}; verbose::Int = 3)
-    @assert config["process"]["EKP_method"] == "EKI" || config["process"]["EKP_method"] == "UKI"
+    @assert config["process"]["EKP_method"] == "EKI" ||
+            config["process"]["EKP_method"] == "UKI" ||
+            config["process"]["EKP_method"] == "ETKI"
 
     u_names = collect(keys(config["prior"]["parameters"]))
     n_iter = config["process"]["n_iter"]
-    n_ensemble = config["process"]["EKP_method"] == "EKI" ? config["process"]["n_ens"] : 2 * length(u_names) + 1
+    if config["process"]["EKP_method"] == "UKI"
+        n_ensemble = 2 * length(u_names) + 1
+    else
+        n_ensemble = config["process"]["n_ens"]
+    end
     param_ensemble = [initial_parameter_ensemble(priors, n_ensemble)]
     n_cases = length(ref_stats_list)
     batch_size = config["process"]["batch_size"]
@@ -104,7 +110,7 @@ function calibrate(::EKPStyle, priors, config, ref_stats_list::Vector{ReferenceS
 
             ekpobj = generate_ekp(param_ensemble[end], RS, priors, config["process"])
             ϕ_n = get_ϕ_ensemble(param_ensemble[end], priors)
-            G_n = [run_dyn_model(ϕ_n[:, i], u_names, config, RS = RS) for i in 1:n_ensemble]
+            G_n = run_ensembles(ϕ_n, u_names, config, n_ensemble, RS = RS)
             G_ens = hcat(G_n...)
             if config["process"]["augmented"]
                 G_ens = vcat(G_ens, param_ensemble[end])
@@ -142,6 +148,14 @@ function generate_ekp(
             Inversion(),
             scheduler = DefaultScheduler(process_settings["Δt"]),
         )
+    elseif process_settings["EKP_method"] == "ETKI"
+        ekpobj = EnsembleKalmanProcess(
+            param_ensemble,
+            y,
+            Γ,
+            TransformInversion(),
+            scheduler = DefaultScheduler(process_settings["Δt"]),
+        )
     elseif process_settings["EKP_method"] == "UKI"
         process = Unscented(
             mean(param_ensemble, dims = 2)[:],
@@ -173,3 +187,19 @@ function get_results(param_ensemble::Vector{Matrix{FT}}, priors::ParameterDistri
     _cov_optim = _cov_optim ./ tr(_cov_optim) .* length(_ϕ_optim)
     return (ϕ_optim = _ϕ_optim, cov_optim = _cov_optim)
 end
+
+function run_ensembles(
+    u::Array{FT, 2},
+    u_names::Array{String, 1},
+    config::Dict,
+    n_ensemble::Int;
+    RS::Union{Nothing, ReferenceStatistics} = nothing,
+    case_numbers::Vector{Int} = Int[],
+)
+    G_n = zeros(length(RS.y), n_ensemble)
+    Threads.@threads for i in 1:n_ensemble
+        # run the model with the current parameters, i.e., map θ to G(θ)
+        G_n[:, i] = run_dyn_model(u[:, i], u_names, config, RS = RS, case_numbers = case_numbers)
+    end
+    return [G_n[:, i] for i in 1:n_ensemble]
+end
diff --git a/test/experiments/calibrations/config.jl b/test/experiments/calibrations/config.jl
@@ -41,7 +41,7 @@ function get_process_config()
     config["n_ens"] = 15
     # Define EKP time step
     config["Δt"] = 1.0
-    config["EKP_method"] = "UKI"
+    config["EKP_method"] = "ETKI"
     # Define whether state vector is augmented by parameters for Bayesian regularization
     config["augmented"] = false
     # Define Bayesian regularization degree scale (prior cov * scale)