From b1185c07320fd0d33474fc549400fbee22c8e51c Mon Sep 17 00:00:00 2001
From: amylu00 <alu3@caltech.edu>
Date: Thu, 19 Sep 2024 09:36:31 -0700
Subject: [PATCH] UKI to AIDA ice nucleation calibrations

---
 .../ice_nucleation_2024/AIDA_calibrations.jl  |  2 +-
 papers/ice_nucleation_2024/calibration.jl     | 52 ++++++++++++++++++-
 papers/ice_nucleation_2024/perfect_calib.jl   |  2 +-
 test/ice_nucleation_calibration.jl            |  2 +-
 4 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/papers/ice_nucleation_2024/AIDA_calibrations.jl b/papers/ice_nucleation_2024/AIDA_calibrations.jl
index e35f84bc0..d23119d18 100644
--- a/papers/ice_nucleation_2024/AIDA_calibrations.jl
+++ b/papers/ice_nucleation_2024/AIDA_calibrations.jl
@@ -85,7 +85,7 @@ 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(FT, "ABHOM", params, IC, frozen_frac_moving_mean, Γ)
+    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]]
diff --git a/papers/ice_nucleation_2024/calibration.jl b/papers/ice_nucleation_2024/calibration.jl
index 83f9c4a75..3239dbeff 100644
--- a/papers/ice_nucleation_2024/calibration.jl
+++ b/papers/ice_nucleation_2024/calibration.jl
@@ -247,7 +247,7 @@ function create_prior(FT, IN_mode, ; perfect_model = false)
     return prior
 end
 
-function calibrate_J_parameters(FT, IN_mode, params, IC, y_truth, Γ,; perfect_model = false)
+function calibrate_J_parameters_EKI(FT, IN_mode, params, IC, y_truth, Γ,; perfect_model = false)
     # Random number generator
     rng_seed = 24
     rng = Random.seed!(Random.GLOBAL_RNG, rng_seed)
@@ -295,6 +295,56 @@ function calibrate_J_parameters(FT, IN_mode, params, IC, y_truth, Γ,; perfect_m
     return [m_coeff_ekp, c_coeff_ekp, ϕ_n_values]
 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
+    α_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,
+    )
+
+    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)
+        # 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
+    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?
+
+    return [UKI_results, u_stored, g_stored, n_iter, N_iterations, err]
+end
+
 function ensemble_means(ϕ_n_values, N_iterations, N_ensemble)
     iterations = collect(1:N_iterations)
     m_mean = zeros(length(iterations))
diff --git a/papers/ice_nucleation_2024/perfect_calib.jl b/papers/ice_nucleation_2024/perfect_calib.jl
index c3105f57d..d581909d6 100644
--- a/papers/ice_nucleation_2024/perfect_calib.jl
+++ b/papers/ice_nucleation_2024/perfect_calib.jl
@@ -17,7 +17,7 @@ for IN_mode in IN_mode_list
     y_truth = pseudo_data[1]
     coeff_true = pseudo_data[3]
 
-    output = calibrate_J_parameters(
+    output = calibrate_J_parameters_EKI(
         FT,
         IN_mode,
         params,
diff --git a/test/ice_nucleation_calibration.jl b/test/ice_nucleation_calibration.jl
index a87204ce9..85ab87fa6 100644
--- a/test/ice_nucleation_calibration.jl
+++ b/test/ice_nucleation_calibration.jl
@@ -14,7 +14,7 @@ function test_J_calibration(FT, IN_mode)
     y_truth = pseudo_data[1]
     coeff_true = pseudo_data[3]
 
-    output = calibrate_J_parameters(
+    output = calibrate_J_parameters_EKI(
         FT,
         IN_mode,
         params,