Plot error from AIDA calibs

amylu00 · amylu00 · commit 16b21915709d · 2025-04-23T11:02:14.000-07:00
diff --git a/papers/ice_nucleation_2024/AIDA_calibrations.jl b/papers/ice_nucleation_2024/AIDA_calibrations.jl
@@ -16,6 +16,32 @@ include(joinpath(pkgdir(CM), "papers", "ice_nucleation_2024", "calibration.jl"))
 include(joinpath(pkgdir(CM), "papers", "ice_nucleation_2024", "unpack_AIDA.jl"))
 include(joinpath(pkgdir(CM), "papers", "ice_nucleation_2024", "plots.jl"))
 
+function MSE(
+    n_iterations, n_ensembles,
+    all_params, data_file_name_list,
+    params_list, nuc_mode, FT, IC_list, end_sim,
+    y_truth, Nₜ,
+)
+
+    loss_vector = [] # should have length of n_iterations with each element being the MSE from the whole batch of ensembles
+    for iter in 1:n_iterations
+        loss_iteration = []
+        for ens in 1:n_ensembles
+            loss_ensemble = []
+            for exp in 1:length(data_file_name_list)
+                parcel_ensemble =
+                    run_model([params_list[exp]], nuc_mode, all_params[iter][:, ens], FT, [IC_list[exp]], end_sim)
+                loss_ensemble_exp = y_truth[exp][end] - (parcel_ensemble[9, end] / Nₜ[exp])
+                append!(loss_ensemble, loss_ensemble_exp)
+            end
+            append!(loss_iteration, loss_ensemble)
+        end
+        append!(loss_vector, (sum(loss_iteration)^2)/(length(loss_iteration))) # MSE
+    end
+
+    return loss_vector
+end
+
 # Defining data names, start/end times, etc.
 global edf_data_names = [
     "in05_17_aida.edf", "in05_18_aida.edf",
@@ -129,12 +155,24 @@ for (batch_index, batch_name) in enumerate(batch_names)
     EKI_output = calibrate_J_parameters_EKI(FT, nuc_mode, params_list, IC_list, y_truth, end_sim, Γ)
     UKI_output = calibrate_J_parameters_UKI(FT, nuc_mode, params_list, IC_list, y_truth, end_sim, Γ)
 
+    EKI_calibrated_parameters = EKI_output[1] # MEAN of parameters from ensembles in FINAL iteration
+    UKI_calibrated_parameters = UKI_output[1] # MEAN of parameters from ensembles in FINAL iteration 
+    EKI_all_params = EKI_output[2]            # parameters from EACH ensemble in EACH iteration
+    UKI_all_params = UKI_output[2]            # parameters from EACH ensemble in EACH iteration
+    EKI_mean_each_iter = EKI_output[3]        # MEAN of parameters from ensembles in EACH iteration 
+    UKI_mean_each_iter = UKI_output[3]        # MEAN of parameters from ensembles in EACH iteration 
+    UKI_final_iter_spread = UKI_output[4]     # parameters for EACH ensemble in FINAL iteration
+    EKI_error = EKI_output[4]
+    UKI_error = UKI_output[5]
+
     EKI_n_iterations = size(EKI_output[2])[1]
     EKI_n_ensembles = size(EKI_output[2][1])[2]
+    UKI_n_iterations = size(UKI_output[2])[1]
+    UKI_n_ensembles = size(UKI_output[2][1])[2]
+
+    EKI_calibrated_ensemble_means = ensemble_means(EKI_all_params, EKI_n_iterations, EKI_n_ensembles)
+    UKI_calibrated_ensemble_means = ensemble_means(UKI_all_params, UKI_n_iterations, UKI_n_ensembles)
 
-    EKI_calibrated_parameters = EKI_output[1]
-    UKI_calibrated_parameters = UKI_output[1]
-    calibrated_ensemble_means = ensemble_means(EKI_output[2], EKI_n_iterations, EKI_n_ensembles)
     merge!(EKI_calibrated_coeff_dict, Dict(batch_name => EKI_calibrated_parameters))
     merge!(UKI_calibrated_coeff_dict, Dict(batch_name => UKI_calibrated_parameters))
 
@@ -166,14 +204,18 @@ for (batch_index, batch_name) in enumerate(batch_names)
         (; AIDA_t_profile, AIDA_T_profile, AIDA_P_profile, AIDA_ICNC_profile, AIDA_e_profile) = AIDA_data
 
         AIDA_ICNC_data_fig = plot_AIDA_ICNC_data(
-            exp_names[exp_index],
+            exp_names[exp_index], start_time[exp_index],
             AIDA_t_profile, AIDA_ICNC_profile,
             t_profile[exp_index], ICNC_moving_avg[exp_index], frozen_frac_moving_mean[exp_index],
         )
 
         ## Calibrated coefficients.
         #  Did they converge?
-        calibrated_coeffs_fig = plot_calibrated_coeffs(batch_name, EKI_n_iterations, calibrated_ensemble_means)
+        calibrated_coeffs_fig = plot_calibrated_coeffs(
+            batch_name,
+            EKI_n_iterations, EKI_n_ensembles, EKI_all_params,
+            UKI_n_iterations, UKI_n_ensembles, UKI_all_params,
+        )
 
         ## Calibrated parcel simulations.
         #  Does the calibrated parcel give reasonable outputs?
@@ -192,21 +234,32 @@ for (batch_index, batch_name) in enumerate(batch_names)
         )
 
         ## Looking at spread in UKI calibrated parameters
-        ϕ_UKI = UKI_output[2]
-        UKI_parcel_1 = run_model([params_list[exp_index]], nuc_mode, [ϕ_UKI[1,1], ϕ_UKI[2,1]], FT, [IC_list[exp_index]], end_sim)
-        UKI_parcel_2 = run_model([params_list[exp_index]], nuc_mode, [ϕ_UKI[1,1], ϕ_UKI[2,1]], FT, [IC_list[exp_index]], end_sim)
-        UKI_parcel_3 = run_model([params_list[exp_index]], nuc_mode, [ϕ_UKI[1,1], ϕ_UKI[2,1]], FT, [IC_list[exp_index]], end_sim)
-        UKI_parcel_4 = run_model([params_list[exp_index]], nuc_mode, [ϕ_UKI[1,1], ϕ_UKI[2,1]], FT, [IC_list[exp_index]], end_sim)
-        UKI_parcel_5 = run_model([params_list[exp_index]], nuc_mode, [ϕ_UKI[1,1], ϕ_UKI[2,1]], FT, [IC_list[exp_index]], end_sim)
+        UKI_parcel_1 = run_model([params_list[exp_index]], nuc_mode, [UKI_final_iter_spread[1,1], UKI_final_iter_spread[2,1]], FT, [IC_list[exp_index]], end_sim)
+        UKI_parcel_2 = run_model([params_list[exp_index]], nuc_mode, [UKI_final_iter_spread[1,2], UKI_final_iter_spread[2,2]], FT, [IC_list[exp_index]], end_sim)
+        UKI_parcel_3 = run_model([params_list[exp_index]], nuc_mode, [UKI_final_iter_spread[1,3], UKI_final_iter_spread[2,3]], FT, [IC_list[exp_index]], end_sim)
+        UKI_parcel_4 = run_model([params_list[exp_index]], nuc_mode, [UKI_final_iter_spread[1,4], UKI_final_iter_spread[2,4]], FT, [IC_list[exp_index]], end_sim)
+        UKI_parcel_5 = run_model([params_list[exp_index]], nuc_mode, [UKI_final_iter_spread[1,5], UKI_final_iter_spread[2,5]], FT, [IC_list[exp_index]], end_sim)
 
         UKI_spread_fig = plot_UKI_spread(
             exp_names[exp_index],
             t_profile[exp_index], frozen_frac_moving_mean[exp_index], Nₜ[exp_index],
             UKI_parcel_1, UKI_parcel_2, UKI_parcel_3, UKI_parcel_4, UKI_parcel_5, UKI_parcel,
         )
 
-    end # plotting individual experiments
-end
+    end # iterating over experiments in batch
+
+    # Parcel runs using ensemble means of parameters at each iteration for each exp in batch
+    # EKI_loss_batch = MSE(EKI_n_iterations, EKI_n_ensembles, EKI_all_params, data_file_name_list, params_list, nuc_mode, FT, IC_list, end_sim, y_truth, Nₜ)
+    # UKI_loss_batch = MSE(UKI_n_iterations, UKI_n_ensembles, UKI_all_params, data_file_name_list, params_list, nuc_mode, FT, IC_list, end_sim, y_truth, Nₜ)
+
+    ## Plotting loss function over calibration iterations
+    loss_fig = plot_loss_func(
+        batch_name,
+        EKI_n_iterations, UKI_n_iterations,
+        EKI_error, UKI_error,
+    )
+
+end  # iterating over batches
 
 # Plotting overview
 plot_ICNC_overview(overview_data)
diff --git a/papers/ice_nucleation_2024/calibration.jl b/papers/ice_nucleation_2024/calibration.jl
@@ -249,8 +249,8 @@ function create_prior(FT, IN_mode, ; perfect_model = false, aerosol_type = nothi
                 m_stats = [FT(100), FT(50), FT(0), Inf]
                 c_stats = [FT(0.7), FT(20), -Inf, Inf]
             elseif aerosol_type == CMP.Dust(FT)
-                m_stats = [FT(50), FT(50), FT(0), Inf]
-                c_stats = [FT(0.7), FT(20), -Inf, Inf]
+                m_stats = [FT(50), FT(70), FT(0), Inf]
+                c_stats = [FT(0.7), FT(40), -Inf, Inf]
             else
                 error("Aerosol type not supported for ABDINM. Check create_prior function.")
             end
@@ -350,9 +350,15 @@ function calibrate_J_parameters_EKI(FT, IN_mode, params, IC, y_truth, end_sim, 
         digits = 6,
     )
 
+    mean_each_iter = []
+    for i in 1:final_iter[1]
+        push!(mean_each_iter, EKP.get_u_mean(EKI_obj, i))
+    end
+
     calibrated_coeffs = [m_coeff_ekp, c_coeff_ekp]
+    error = EKP.get_error(EKI_obj)
 
-    return [calibrated_coeffs, ϕ_n_values]
+    return [calibrated_coeffs, ϕ_n_values, mean_each_iter, error]
 end
 
 function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, end_sim, Γ,; perfect_model = false)
@@ -394,14 +400,7 @@ function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, end_sim, 
         # 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
@@ -411,9 +410,19 @@ function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, end_sim, 
     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)
+    ϕ_n = EKP.get_ϕ_final(prior, UKI_obj) # final iteraiton for each ensemble
+    all_ϕ_n = EKP.get_ϕ(prior, UKI_obj; return_array=true) # set of params for each ensembles AND each iteration
+
+    mean_each_iter = []
+    for i in 1:(final_iter[1] + 1)
+        mean_u_space = EKP.get_u_mean(UKI_obj, i)
+        mean_u_space_transformed = EKP.transform_unconstrained_to_constrained(prior, mean_u_space)
+        push!(mean_each_iter, mean_u_space_transformed)
+    end
+
+    error = EKP.get_error(UKI_obj)
 
-    return [UKI_mean, ϕ_n, final_iter]
+    return [UKI_mean, all_ϕ_n, mean_each_iter, ϕ_n, error]
 end
 
 function ensemble_means(ϕ_n_values, N_iterations, N_ensemble)
diff --git a/papers/ice_nucleation_2024/plots.jl b/papers/ice_nucleation_2024/plots.jl
@@ -2,6 +2,8 @@ import CairoMakie as MK
 
 include(joinpath(pkgdir(CM), "parcel", "ParcelCommon.jl"))
 
+#! format: off
+
 """
     plot_AIDA_ICNC_data(
         exp_name,
@@ -12,7 +14,7 @@ include(joinpath(pkgdir(CM), "parcel", "ParcelCommon.jl"))
 Plots raw AIDA ICNC data. Plot is saved and returned.
 """
 function plot_AIDA_ICNC_data(
-    exp_name,
+    exp_name, start_time,
     AIDA_t_profile, AIDA_ICNC_profile,
     t_profile, ICNC_moving_avg, frozen_frac_moving_mean,
 )
@@ -34,8 +36,8 @@ function plot_AIDA_ICNC_data(
         linestyle = :dash,
         linewidth = 2,
     )
-    MK.lines!(data_ax1, t_profile, ICNC_moving_avg, label = "AIDA ICNC moving mean", linewidth = 2.5, color = :blue)
-    MK.lines!(data_ax2, t_profile, frozen_frac_moving_mean, linewidth = 2.5, color = :blue)
+    MK.lines!(data_ax1, t_profile .+ start_time, ICNC_moving_avg, label = "AIDA ICNC moving mean", linewidth = 2.5, color = :blue)
+    MK.lines!(data_ax2, t_profile .+ start_time, frozen_frac_moving_mean, linewidth = 2.5, color = :blue)
     MK.axislegend(data_ax1, framevisible = true, labelsize = 12, position = :rc)
     MK.save("$exp_name" * "_ICNC.svg", AIDA_data_fig)
 
@@ -48,30 +50,70 @@ end
 
 Plots evolution of calibrated coefficients over the EKI iterations. Plot is saved and returned.
 """
-function plot_calibrated_coeffs(batch_name, EKI_n_iterations, calibrated_ensemble_means)
+function plot_calibrated_coeffs(
+    batch_name,
+    EKI_n_iterations, EKI_n_ensembles, EKI_all_params,
+    UKI_n_iterations, UKI_n_ensembles, UKI_all_params,
+)
+    EKI_x_axis = []
+    UKI_x_axis = []
+    EKI_m = []
+    UKI_m = []
+    EKI_c = []
+    UKI_c = []
+
+    for iter in 1:EKI_n_iterations
+        for ensemble_n in 1:EKI_n_ensembles
+            append!(EKI_x_axis, iter)
+            append!(EKI_m, EKI_all_params[iter][1, ensemble_n])
+            append!(EKI_c, EKI_all_params[iter][2, ensemble_n])
+        end
+    end
+    for iter in 1:UKI_n_iterations
+        for ensemble_n in 1:UKI_n_ensembles
+            append!(UKI_x_axis, iter)
+            append!(UKI_m, UKI_all_params[iter][1, ensemble_n])
+            append!(UKI_c, UKI_all_params[iter][2, ensemble_n])
+        end
+    end
 
-    calibrated_coeffs_fig = MK.Figure(size = (1100, 900), fontsize = 24)
+    calibrated_coeffs_fig = MK.Figure(size = (900, 600), fontsize = 24)
     m_coeff_ax = MK.Axis(calibrated_coeffs_fig[1, 1], ylabel = "m coefficient [-]", title = "$batch_name")
     c_coeff_ax =
-        MK.Axis(calibrated_coeffs_fig[1, 2], ylabel = "c coefficient [-]", xlabel = "iteration #", title = "EKI")
+        MK.Axis(calibrated_coeffs_fig[1, 2], ylabel = "c coefficient [-]", xlabel = "iteration #", title = "Ensemble Mean")
 
-    MK.lines!(
+    MK.scatter!(
         m_coeff_ax,
-        collect(1:EKI_n_iterations),
-        calibrated_ensemble_means[1],
-        label = "ensemble mean",
-        color = :orange,
-        linewidth = 2.5,
+        EKI_x_axis,
+        EKI_m,
+        label = "EKI",
+        color = (:orange, 0.5),
+        # marker = :xcross,
     )
-    MK.lines!(
+    MK.scatter!(
         c_coeff_ax,
-        collect(1:EKI_n_iterations),
-        calibrated_ensemble_means[2],
-        label = "ensemble mean",
-        color = :orange,
-        linewidth = 2.5,
+        EKI_x_axis,
+        EKI_c,
+        label = "EKI",
+        color = (:orange, 0.5),
+        # marker = :xcross
+    )
+    MK.scatter!(
+        m_coeff_ax,
+        UKI_x_axis,
+        UKI_m,
+        label = "UKI",
+        color = (:fuchsia, 0.5),
+    )
+    MK.scatter!(
+        c_coeff_ax,
+        UKI_x_axis,
+        UKI_c,
+        label = "UKI",
+        color = (:fuchsia, 0.5),
     )
 
+    MK.axislegend(m_coeff_ax, framevisible = false, labelsize = 18, position = :rt)
     MK.save("$batch_name" * "_calibrated_coeffs_fig.svg", calibrated_coeffs_fig)
 
     return calibrated_coeffs_fig
@@ -293,14 +335,12 @@ function plot_UKI_spread(
     MK.errorbars!(ax_spread, t_profile, frozen_frac_moving_mean, error, color = (:blue, 0.3))
     MK.save("$exp_name" * "_UKI_spread_fig.svg", UKI_spread_fig)
 
-    return plot_UKI_spread
+    return UKI_spread_fig
 
 end
 
 """
-    plot_ICNC_overview(
-        EKI_calibrated_parcel, UKI_parcel,
-    )
+    plot_ICNC_overview(overview_data)
 
 Plots frozen fraction evolution for EKI and UKI calibrated parcel simulations
 and AIDA data for all batch experiments. Plot is saved and returned.
@@ -390,3 +430,45 @@ function plot_ICNC_overview(overview_data)
 
     return overview_fig
 end
+
+"""
+    plot_loss_func(
+        batch_name,
+        EKI_n_iterations, UKI_n_iterations,
+        EKI_error, UKI_error,
+    )
+
+Plots the difference in loss functions over each iteration of
+    calibration for a batch. Plot is saved and returned.
+"""
+function plot_loss_func(
+    batch_name,
+    EKI_n_iterations, UKI_n_iterations,
+    EKI_error, UKI_error,
+)
+
+    loss_fig = MK.Figure(size = (700, 600), fontsize = 24)
+    ax_loss = MK.Axis(loss_fig[1, 1], ylabel = "Error [-]", xlabel = "Iteration [-]", title = "$batch_name: Error")
+    MK.lines!(
+        ax_loss,
+        collect(1:EKI_n_iterations),
+        EKI_error,
+        label = "EKI",
+        linewidth = 2.5,
+        color = :orange,
+    )
+    MK.lines!(
+        ax_loss,
+        collect(1:(UKI_n_iterations - 1)),
+        UKI_error,
+        label = "UKI",
+        linewidth = 2.5,
+        color = :fuchsia,
+    )
+    MK.axislegend(ax_loss, framevisible = false, labelsize = 18, position = :rc)
+    MK.save("$batch_name" * "_loss_fig.svg", loss_fig)
+
+    return loss_fig
+
+end
+#! format: on
