Add UKI to AIDA calibrations

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"

papers/ice_nucleation_2024/calibration.jl

@@ -9,14 +9,15 @@ import ClimaParams as CP
 import CloudMicrophysics as CM
 import CloudMicrophysics.Parameters as CMP
 import Thermodynamics as TD
+import StatsBase as SB
 
 #! format: off
 # definition of the ODE problem for parcel model
 include(joinpath(pkgdir(CM), "parcel", "Parcel.jl"))
 include(joinpath(pkgdir(CM), "papers", "ice_nucleation_2024", "calibration_setup.jl"))
 
 # Define model which wraps around parcel and overwrites calibrated parameters
-function run_model(p, coefficients, IN_mode, FT, IC)
+function run_model(p, coefficients, IN_mode, FT, IC, end_sim)
     # grabbing parameters
     m_calibrated, c_calibrated = coefficients
     (; const_dt, w, t_max, aerosol_act, aerosol, r_nuc) = p
@@ -100,7 +101,7 @@ function run_model(p, coefficients, IN_mode, FT, IC)
 
     # solve ODE
     local sol = run_parcel(IC, FT(0), t_max, params)
-    return sol[9, :] ./ (IC[7] + IC[8] + IC[9])  # frozen fraction
+    return sol[9, end - end_sim:end] ./ (IC[7] + IC[8] + IC[9])
 end
 
 function run_calibrated_model(FT, IN_mode, coefficients, p, IC)
@@ -247,52 +248,121 @@ 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, end_sim, Γ,; 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_ensemble = 25       # runs N_ensemble trials per iteration
+    N_iterations = 50     # number of iterations the inverse problem goes through
 
     # Generate initial ensemble and set up EKI
     initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ensemble)
     EKI_obj = EKP.EnsembleKalmanProcess(
         initial_ensemble,
-        y_truth,
+        y_truth[end - end_sim:end],
         Γ,
         EKP.Inversion();
         rng = rng,
+        verbose = true,
+        localization_method = EKP.Localizers.NoLocalization(), # no localization
     )
 
     # Carry out the EKI calibration
     # ϕ_n_values[iteration] stores ensembles of calibrated coeffs in that iteration
     global ϕ_n_values = []
+    final_iter = N_iterations
     for n in 1:N_iterations
         ϕ_n = EKP.get_ϕ_final(prior, EKI_obj)
         G_ens = hcat(
             [
-                run_model(params, ϕ_n[:, i], IN_mode, FT, IC) for
+                run_model(params, ϕ_n[:, i], IN_mode, FT, IC, end_sim) for
                 i in 1:N_ensemble
             ]...,
         )
-        EKP.update_ensemble!(EKI_obj, G_ens)
+        # Update ensemble
+        terminated = EKP.update_ensemble!(EKI_obj, G_ens)
+        # if termination flagged, can stop at earlier iteration
+        if !isnothing(terminated)
+            final_iter = n - 1
+            break
+        end
 
         ϕ_n_values = vcat(ϕ_n_values, [ϕ_n])
     end
 
     # Mean coefficients of all ensembles in the final iteration
     m_coeff_ekp = round(
-        Distributions.mean(ϕ_n_values[N_iterations][1, 1:N_ensemble]),
+        Distributions.mean(ϕ_n_values[final_iter][1, 1:N_ensemble]),
         digits = 6,
     )
     c_coeff_ekp = round(
-        Distributions.mean(ϕ_n_values[N_iterations][2, 1:N_ensemble]),
+        Distributions.mean(ϕ_n_values[final_iter][2, 1:N_ensemble]),
         digits = 6,
     )
 
-    return [m_coeff_ekp, c_coeff_ekp, ϕ_n_values]
+    calibrated_coeffs = [m_coeff_ekp, c_coeff_ekp]
+
+    return [calibrated_coeffs, ϕ_n_values]
+end
+
+function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, end_sim, Γ,; perfect_model = false)
+    prior = create_prior(FT, IN_mode, perfect_model = perfect_model)
+    N_iterations = 25
+    α_reg = 1.0
+    update_freq = 1
+
+    # truth = EKP.Observations.Observation(y_truth, Γ, "y_truth")
+    truth = EKP.Observation(
+        Dict("samples" => vec(SB.mean(y_truth[end - end_sim: end], dims = 2)), "covariances" => Γ, "names" => "y_truth")
+    )
+
+    # Generate initial ensemble and set up UKI
+    process = EKP.Unscented(
+        SB.mean(prior),
+        SB.cov(prior);
+        α_reg = α_reg,
+        update_freq = update_freq,
+        impose_prior = false,
+    )
+    UKI_obj = EKP.EnsembleKalmanProcess(truth, process; verbose = true)
+
+    err = []
+    final_iter =[N_iterations]
+    for n in 1:N_iterations
+        # Return transformed parameters in physical/constrained space
+        ϕ_n = EKP.get_ϕ_final(prior, UKI_obj)
+        # Evaluate forward map
+        G_n = [
+            run_model(params, ϕ_n[:, i], IN_mode, FT, IC, end_sim) for
+            i in 1:size(ϕ_n)[2]  #i in 1:N_ensemble
+        ]
+        # Reformat into `d x N_ens` matrix
+        G_ens = hcat(G_n...)
+        # Update ensemble
+        terminate = EKP.EnsembleKalmanProcesses.update_ensemble!(UKI_obj, G_ens)
+        push!(err, EKP.get_error(UKI_obj)[end])
+        println(
+            "Iteration: " *
+            string(n) *
+            ", Error: " *
+            string(err[n]) *
+            " norm(Cov):" *
+            string(Distributions.norm(EKP.get_process(UKI_obj).uu_cov[n]))
+        )
+        if !isnothing(terminate)
+            final_iter[1] = n - 1
+            break
+        end
+    end
+
+    UKI_mean_u_space = EKP.get_u_mean_final(UKI_obj)
+    UKI_mean = EKP.transform_unconstrained_to_constrained(prior, UKI_mean_u_space)
+
+    ϕ_n = EKP.get_ϕ_final(prior, UKI_obj)
+
+    return [UKI_mean, ϕ_n, final_iter]
 end
 
 function ensemble_means(ϕ_n_values, N_iterations, N_ensemble)

papers/ice_nucleation_2024/calibration_setup.jl

@@ -129,11 +129,11 @@ function perf_model_pseudo_data(FT, IN_mode, params, IC)
     elseif IN_mode == "ABHOM"
         coeff_true = [FT(255.927125), FT(-68.553283)]
     end
-
-    G_truth = run_model(params, coeff_true, IN_mode, FT, IC)
+    sol_ICNC = run_calibrated_model(FT, IN_mode, coeff_true, params, IC)[9, :]
+    G_truth = sol_ICNC ./ (IC[7] + IC[8] + IC[9])
     dim_output = length(G_truth)
 
-    Γ = 0.03 * LinearAlgebra.I * (maximum(G_truth) - minimum(G_truth))
+    Γ = 0.001 * LinearAlgebra.I * (maximum(G_truth) - minimum(G_truth))
     noise_dist = Distributions.MvNormal(zeros(dim_output), Γ)
 
     y_truth = zeros(length(G_truth), n_samples) # where noisy ICNC will be stored

papers/ice_nucleation_2024/perfect_calib.jl

@@ -16,30 +16,32 @@ for IN_mode in IN_mode_list
     Γ = pseudo_data[2]
     y_truth = pseudo_data[1]
     coeff_true = pseudo_data[3]
+    end_sim = 25
 
-    output = calibrate_J_parameters(
+    output = calibrate_J_parameters_EKI(
         FT,
         IN_mode,
         params,
         IC,
         y_truth,
+        end_sim,
         Γ,
         perfect_model = true,
     )
 
-    iterations = collect(1:size(output[3])[1])
-    calibrated_parameters = [output[1], output[2]]
+    iterations = collect(1:size(output[2])[1])
+    calibrated_parameters = output[1]
     m_mean = []
     m_mean = ensemble_means(
-        output[3],
-        size(output[3])[1],
-        size(output[3][1])[2],
+        output[2],
+        size(output[2])[1],
+        size(output[2][1])[2],
     )[1]
     c_mean = []
     c_mean = ensemble_means(
-        output[3],
-        size(output[3])[1],
-        size(output[3][1])[2],
+        output[2],
+        size(output[2])[1],
+        size(output[2][1])[2],
     )[2]
 
     # Plotting calibrated parameters per iteration

test/Project.toml

@@ -31,7 +31,6 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 
 [compat]
 ClimaParams = "0.10.6"
-EnsembleKalmanProcesses = "1.1.5"
 KernelAbstractions = "0.9"
 MLJFlux = "0.4"
 Optim = "<1.9.3"

test/ice_nucleation_calibration.jl

@@ -13,34 +13,75 @@ function test_J_calibration(FT, IN_mode)
     Γ = pseudo_data[2]
     y_truth = pseudo_data[1]
     coeff_true = pseudo_data[3]
+    end_sim = 25
 
-    output = calibrate_J_parameters(
+    EKI_output = calibrate_J_parameters_EKI(
         FT,
         IN_mode,
         params,
         IC,
         y_truth,
+        end_sim,
         Γ,
         perfect_model = true,
     )
-    calibrated_parameters = [output[1], output[2]]
-    calibrated_soln =
-        run_calibrated_model(FT, IN_mode, calibrated_parameters, params, IC)
+    UKI_output = calibrate_J_parameters_UKI(
+        FT,
+        IN_mode,
+        params,
+        IC,
+        y_truth,
+        end_sim,
+        Γ,
+        perfect_model = true,
+    )
+
+    EKI_calibrated_parameters = EKI_output[1]
+    UKI_calibrated_parameters = UKI_output[1]
+    EKI_calibrated_soln =
+        run_calibrated_model(FT, IN_mode, EKI_calibrated_parameters, params, IC)
+    UKI_calibrated_soln =
+        run_calibrated_model(FT, IN_mode, UKI_calibrated_parameters, params, IC)
     true_soln = run_calibrated_model(FT, IN_mode, coeff_true, params, IC)
 
-    # test that coeffs are close to "true" values and that end ICNC are similar
-    if IN_mode == "ABDINM"
-        TT.@test calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
-        TT.@test calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(1.5)
-        TT.@test calibrated_soln[9, end] ≈ true_soln[9, end] rtol = FT(0.3)
-    elseif IN_mode == "ABIFM"
-        TT.@test calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
-        TT.@test calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
-        TT.@test calibrated_soln[9, end] ≈ true_soln[9, end] rtol = FT(0.3)
-    elseif IN_mode == "ABHOM"
-        TT.@test calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
-        TT.@test calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
-        TT.@test calibrated_soln[9, end] ≈ true_soln[9, end] rtol = FT(0.3)
+    TT.@testset "EKI Perfect Model Calibrations on AIDA" begin
+        # test that coeffs are close to "true" values and that end ICNC are similar
+        if IN_mode == "ABDINM"
+            TT.@test EKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test EKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(1.5)
+            TT.@test EKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        elseif IN_mode == "ABIFM"
+            TT.@test EKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test EKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
+            TT.@test EKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        elseif IN_mode == "ABHOM"
+            TT.@test EKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test EKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
+            TT.@test EKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        end
+    end
+
+    TT.@testset "UKI Perfect Model Calibrations on AIDA" begin
+        # test that coeffs are close to "true" values and that end ICNC are similar
+        if IN_mode == "ABDINM"
+            TT.@test UKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test UKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(1.5)
+            TT.@test UKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        elseif IN_mode == "ABIFM"
+            TT.@test UKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test UKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
+            TT.@test UKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        elseif IN_mode == "ABHOM"
+            TT.@test UKI_calibrated_parameters[1] ≈ coeff_true[1] rtol = FT(0.3)
+            TT.@test UKI_calibrated_parameters[2] ≈ coeff_true[2] rtol = FT(0.3)
+            TT.@test UKI_calibrated_soln[9, end] ≈ true_soln[9, end] rtol =
+                FT(0.3)
+        end
     end
 
 end