Plot error from AIDA calibs

Author: amylu00

papers/ice_nucleation_2024/AIDA_calibrations.jl

@@ -129,15 +129,29 @@ 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 
+    # EKI_final_iter_spread = EKI_output[4]   # parameters for EACH ensemble in FINAL iteration
+    UKI_final_iter_spread = UKI_output[4]     # parameters for EACH ensemble in FINAL iteration
+
     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))
 
+    EKI_loss_batch = []
+    UKI_loss_batch = []
+
     ### Plot for individual experiments.
     for (exp_index, data_file) in enumerate(data_file_name_list)
         if batch_name == "HOM"
@@ -160,20 +174,40 @@ for (batch_index, batch_name) in enumerate(batch_names)
             push!(overview_data.UKI_calibrated_parcel, UKI_parcel)
         end
 
+        # Parcel runs using ensemble means of parameters at each iteration for each exp in batch
+        EKI_loss_exp = []
+        UKI_loss_exp = []
+        for j in 1:EKI_n_iterations
+            EKI_parcel_iteration = run_model([params_list[exp_index]], nuc_mode, EKI_mean_each_iter[j], FT, [IC_list[exp_index]], end_sim)
+            EKI_loss_iteration = (EKI_parcel_iteration[9, end] / Nₜ[exp_index]) - frozen_frac_moving_mean[exp_index][end]
+            append!(EKI_loss_exp, abs.(EKI_loss_iteration))
+        end
+        push!(EKI_loss_batch, EKI_loss_exp)
+        for k in 1:UKI_n_iterations
+            UKI_parcel_iteration = run_model([params_list[exp_index]], nuc_mode, UKI_mean_each_iter[k], FT, [IC_list[exp_index]], end_sim)
+            UKI_loss_iteration = (UKI_parcel_iteration[9, end] / Nₜ[exp_index]) - frozen_frac_moving_mean[exp_index][end]
+            append!(UKI_loss_exp, abs.(UKI_loss_iteration))
+        end
+        push!(UKI_loss_batch, UKI_loss_exp) 
+
         ## Plots.
         ## Plotting AIDA data.
         AIDA_data = data_file in edf_data_names ? unpack_data(data_file) : unpack_data(data_file, total_t = t_max[exp_index])
         (; 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_calibrated_ensemble_means,
+            UKI_n_iterations, UKI_calibrated_ensemble_means,
+        )
 
         ## Calibrated parcel simulations.
         #  Does the calibrated parcel give reasonable outputs?
@@ -192,21 +226,28 @@ 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
+    
+    ## Plotting loss function over calibration iterations
+    loss_fig = plot_loss_func(
+        batch_name,
+        EKI_n_iterations, UKI_n_iterations,
+        EKI_loss_batch, UKI_loss_batch,
+    )
+
+end  # iterating over batches
 
 # Plotting overview
 plot_ICNC_overview(overview_data)

papers/ice_nucleation_2024/calibration.jl

@@ -350,9 +350,14 @@ 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]
 
-    return [calibrated_coeffs, ϕ_n_values]
+    return [calibrated_coeffs, ϕ_n_values, mean_each_iter]
 end
 
 function calibrate_J_parameters_UKI(FT, IN_mode, params, IC, y_truth, end_sim, Γ,; perfect_model = false)
@@ -394,14 +399,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 +409,17 @@ 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
 
-    return [UKI_mean, ϕ_n, final_iter]
+    return [UKI_mean, all_ϕ_n, mean_each_iter, ϕ_n]
 end
 
 function ensemble_means(ϕ_n_values, N_iterations, N_ensemble)

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,29 +50,49 @@ 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_calibrated_ensemble_means,
+    UKI_n_iterations, UKI_calibrated_ensemble_means,
+)
 
-    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!(
         m_coeff_ax,
         collect(1:EKI_n_iterations),
-        calibrated_ensemble_means[1],
-        label = "ensemble mean",
+        EKI_calibrated_ensemble_means[1],
+        label = "EKI",
         color = :orange,
         linewidth = 2.5,
     )
     MK.lines!(
         c_coeff_ax,
         collect(1:EKI_n_iterations),
-        calibrated_ensemble_means[2],
-        label = "ensemble mean",
+        EKI_calibrated_ensemble_means[2],
+        label = "EKI",
         color = :orange,
         linewidth = 2.5,
     )
+    MK.lines!(
+        m_coeff_ax,
+        collect(1:UKI_n_iterations),
+        UKI_calibrated_ensemble_means[1],
+        label = "UKI",
+        color = :fuchsia,
+        linewidth = 2.5,
+    )
+    MK.lines!(
+        c_coeff_ax,
+        collect(1:UKI_n_iterations),
+        UKI_calibrated_ensemble_means[2],
+        label = "UKI",
+        color = :fuchsia,
+        linewidth = 2.5,
+    )
 
     MK.save("$batch_name" * "_calibrated_coeffs_fig.svg", calibrated_coeffs_fig)
 
@@ -293,7 +315,7 @@ 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
 
@@ -390,3 +412,49 @@ function plot_ICNC_overview(overview_data)
 
     return overview_fig
 end
+
+"""
+    plot_loss_func(
+        batch_name,
+        EKI_n_iterations, UKI_n_iterations,
+        EKI_loss_batch, UKI_loss_batch,
+    )
+
+Plots the difference in loss functions over each iteration of
+    calibration for a batch. The absolute value of the difference in
+    loss functions for all experiments are summed at each iteration.
+    Plot is saved and returned.
+"""
+function plot_loss_func(
+    batch_name,
+    EKI_n_iterations, UKI_n_iterations,
+    EKI_loss_batch, UKI_loss_batch,
+)
+    EKI_loss = vec(sum(reduce(hcat, EKI_loss_batch), dims=2))
+    UKI_loss = vec(sum(reduce(hcat, UKI_loss_batch), dims=2))
+
+    loss_fig = MK.Figure(size = (700, 600), fontsize = 24)
+    ax_loss = MK.Axis(loss_fig[1, 1], ylabel = "|Δ(Frozen Fraction)| [-]", xlabel = "Iteration [-]", title = "$batch_name: |Simulated Final FF - AIDA Final FF|")
+    MK.lines!(
+        ax_loss,
+        collect(1:EKI_n_iterations),
+        EKI_loss,
+        label = "EKI",
+        linewidth = 2.5,
+        color = :orange,
+    )
+    MK.lines!(
+        ax_loss,
+        collect(1:UKI_n_iterations),
+        UKI_loss,
+        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