added UKI to AIDA_calibrations.jl

CliMA · Sep 20, 2024 · a342934 · a342934
1 parent b1185c0
commit a342934
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 38 deletions.
diff --git a/papers/Project.toml b/papers/Project.toml
@@ -9,4 +9,5 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
diff --git a/papers/ice_nucleation_2024/AIDA_calibrations.jl b/papers/ice_nucleation_2024/AIDA_calibrations.jl
@@ -85,11 +85,12 @@ for (exp_index, data_file_name) in enumerate(data_file_names)
     MK.save("$plot_name"*"_ICNC.svg", AIDA_data_fig)
 
     ### Calibration.
-    output = calibrate_J_parameters_EKI(FT, "ABHOM", params, IC, frozen_frac_moving_mean, Γ)
-    n_iterations = size(output[3])[1]
-    n_ensembles = size(output[3][1])[2]
-    calibrated_parameters = [output[1], output[2]]
-    calibrated_ensemble_means = ensemble_means(output[3], n_iterations, n_ensembles)
+    # EKI_output = calibrate_J_parameters_EKI(FT, "ABHOM", params, IC, frozen_frac_moving_mean, Γ)
+    UKI_output = calibrate_J_parameters_UKI(FT, "ABHOM", params, IC, frozen_frac_moving_mean, Γ)
+    n_iterations = size(EKI_output[3])[1]
+    n_ensembles = size(EKI_output[3][1])[2]
+    calibrated_parameters = [EKI_output[1], EKI_output[2]]
+    calibrated_ensemble_means = ensemble_means(EKI_output[3], n_iterations, n_ensembles)
     # Calibrated parcel
     parcel = run_calibrated_model(FT, "ABHOM", calibrated_parameters, params, IC)
     #parcel_default = run_calibrated_model(FT, "ABHOM", [FT(255.927125), FT(-68.553283)], params, IC)

diff --git a/papers/ice_nucleation_2024/calibration.jl b/papers/ice_nucleation_2024/calibration.jl
@@ -9,6 +9,7 @@ import ClimaParams as CP
 import CloudMicrophysics as CM
 import CloudMicrophysics.Parameters as CMP
 import Thermodynamics as TD
+using StatsBase
 
 #! format: off
 # definition of the ODE problem for parcel model
@@ -297,52 +298,38 @@ end
 
 function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, Γ,; perfect_model = false)
     # Random number generator
-    rng_seed = 24
-    rng = Random.seed!(Random.GLOBAL_RNG, rng_seed)
-
     prior = create_prior(FT, IN_mode, perfect_model = perfect_model)
     N_ensemble = 25      # runs N_ensemble trials per iteration
-    N_iterations = 100   # number of iterations the inverse problem goes through
+    N_iterations = 25   # number of iterations the inverse problem goes through
     α_reg = 1.0
     update_freq = 1
 
-    # Generate initial ensemble and set up EKI
-    initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ensemble)
-    process = EKP.Unscented(mean(prior), cov(prior); α_reg = α_reg, update_freq = update_freq) # TODO - replace with my variables
-    UKI_obj = EKP.EnsembleKalmanProcess(
-        initial_ensemble,
-        y_truth,
-        Γ,
-        process;
-        rng = rng,
-    )
+    @info("", y_truth)
+    truth = EKP.Observations.Observation(y_truth, Γ, "y_truth")
+
+    # Generate initial ensemble and set up UKI
+    process = EKP.Unscented(mean(prior), cov(prior); α_reg = α_reg, update_freq = update_freq)
+    UKI_obj = EKP.EnsembleKalmanProcess(truth, process)
 
-    err = []
-    final_iter = [N_iterations]  # TODO - think this is supposed to be final_it
     for n in 1:N_iterations
         # Return transformed parameters in physical/constrained space
         ϕ_n = EKP.get_ϕ_final(prior, UKI_obj)
+        @info("",ϕ_n)
         # Evaluate forward map
-        G_ens = hcat(
-            [
-                run_model(params, ϕ_n[:, i], IN_mode, FT, IC) for
-                i in 1:N_ensemble
-            ]...,
-        )
-        terminate = EKP.EnsembleKalmanProcesses.update_ensemble!(UKI_obj, G_ens)
-        push!(err, get_error(UKI_obj)[end])
-        if !isnothing(terminate)
-            final_iter[1] = n - 1
-            break
-        end
+        G_n = [
+            run_model(params, ϕ_n[:, i], IN_mode, FT, IC) for
+            i in 1:N_ensemble
+        ]
+        # Reformat into `d x N_ens` matrix
+        G_ens = hcat(G_n...)
+        # Update ensemble
+        EKP.EnsembleKalmanProcesses.update_ensemble!(UKI_obj, G_ens)
     end
-    n_iter = final_iter[1]
 
-    UKI_results = get_u_mean_final(UKI_obj)
-    u_stored = get_u(ukiobj, return_array = false)  # calibrated parameters
-    g_stored = get_g(ukiobj, return_array = false)  # calibrated parcel outputs?
+    UKI_mean = get_u_mean_final(UKI_obj)
+    UKI_cov = get_u_cov_final(UKI_obj)
 
-    return [UKI_results, u_stored, g_stored, n_iter, N_iterations, err]
+    return [UKI_mean, UKI_cov]
 end
 
 function ensemble_means(ϕ_n_values, N_iterations, N_ensemble)