update Lorenz spatial dep example

Author: odunbar

examples/Lorenz/calibrate_spatial_dep.jl

@@ -169,12 +169,12 @@ x0 = x_spun_up[:, end]  #last element of the run is the initial condition for cr
 
 
 #Creating sythetic data
-T = 50.0
+T = 24.0
 ny = nx * 2   #number of data points
 lorenz_config_settings = LorenzConfig(dt, T)
 
 # construct how we compute Observations
-T_start = T - 25.0
+T_start = T - 20.0
 T_end = T
 observation_config = ObservationConfig(T_start, T_end)
 
@@ -196,9 +196,10 @@ y_ens = hcat(
 obs_noise_cov = cov(y_ens, dims = 2) + 1e-2 * I
 y = y_ens[:, 1]
 
+#Prior covariance
+
 pl = 4.0
 psig = 5.0
-#Prior covariance
 B = zeros(nx, nx)
 for ii in 1:nx
     for jj in 1:nx
@@ -207,6 +208,12 @@ for ii in 1:nx
 end
 B_sqrt = sqrt(B)
 
+#=
+psig = 5.0
+B = psig^2*I
+B_sqrt = sqrt(B)
+=#
+
 #Prior mean
 mu = 4.0 * ones(nx)
 

examples/Lorenz/emulate_sample_spatial_dep.jl

@@ -28,13 +28,14 @@ function main()
     ]
 
     #### CHOOSE YOUR CASE: 
-    mask = [1]# 1:1 # e.g. 1:2 or [2]
+    mask = [2]# 1:1 # e.g. 1:2 or [2]
     for (case) in cases[mask]
 
 
         println("case: ", case)
         min_iter = 1
-        max_iter = 7 # number of EKP iterations to use data from is at most this
+        skip_iter = 1
+        max_iter = 8 # number of EKP iterations to use data from is at most this
 
         ####
 
@@ -71,7 +72,7 @@ function main()
 
         # Emulate-sample settings
         # choice of machine-learning tool in the emulation stage
-        nugget = 1e-3
+        nugget = 1e-8
         if case == "GP"
             gppackage = Emulators.GPJL()
             pred_type = Emulators.YType()
@@ -106,7 +107,7 @@ function main()
         # Get training points from the EKP iteration number in the second input term  
         N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP
         min_iter = min(max_iter, max(1, min_iter))
-        input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:(N_iter - 1))
+        input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1))
         input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) #  "next" iterations
         # Save data
         @save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs

examples/Lorenz/plot_spatial_dep.jl

@@ -2,7 +2,7 @@
 
 # some context calues from emulate_sample_spatial_dep.jl
 cases = ["GP", "RF-scalar"]
-case = cases[1]
+case = cases[2]
 
 # load packages
 # CES 
@@ -19,6 +19,9 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses
 using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions
 const EKP = EnsembleKalmanProcesses
 
+cases = ["GP", "RF-scalar"]
+case = cases[2]
+
 # load
 homedir = pwd()
 println(homedir)
@@ -29,6 +32,8 @@ truth_params = load(data_file)["truth_params_constrained"]
 final_params = load(data_file)["final_params_constrained"]
 posterior = load(data_file)["posterior"]
 
+prior_file = joinpath(data_save_directory, "priors_spatial_dep.jld2")
+prior = load(prior_file)["prior"]
 
 ekp_file = joinpath(data_save_directory, "ekp_spatial_dep.jld2")
 ekpobj = load(ekp_file)["ekpobj"]
@@ -49,9 +54,9 @@ quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(co
 gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16)
 p1 = plot(
     range(0, nx - 1, step = 1),
-    [truth_params final_params],
-    label = ["solution" "EKI"],
-    color = [:black :lightgreen],
+    truth_params,
+    label = "solution",
+    color = :black,
     linewidth = 4,
     xlabel = "Spatial index",
     ylabel = "Forcing (input)",
@@ -72,14 +77,77 @@ p2 = plot(
     bottom_margin = 15mm,
     xticks = (Int.(0:10:ny), [0, 10, 20, 30, (40, 0), 10, 20, 30, 40]),
 )
-plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean-final-output", color = :lightgreen, linewidth = 4)
+
+l = @layout [a b]
+plt = plot(p1, p2, layout = l)
+
+savefig(plt, figure_save_directory * "data_spatial_dep.png")
+savefig(plt, figure_save_directory * "data_spatial_dep.pdf")
+
+p3 = deepcopy(p1)
+p4 = deepcopy(p2)
+
+plot!(p1,
+    range(0, nx - 1, step = 1),
+    final_params,
+      label = "mean ensemble input",
+    color = :lightgreen,
+    linewidth = 4,
+)
+plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4)
 
 l = @layout [a b]
 plt = plot(p1, p2, layout = l)
 
 savefig(plt, figure_save_directory * "solution_spatial_dep_ens$(N_ens).png")
 savefig(plt, figure_save_directory * "solution_spatial_dep_ens$(N_ens).pdf")
 
+#
+plot!(p3,
+      range(0, nx - 1, step = 1),
+      get_ϕ_final(prior, ekpobj)[:,1],
+      label = "ensemble inputs",
+      color = :lightgreen,
+      linewidth = 4,
+      linealpha=0.1,
+      )
+
+plot!(p3,
+      range(0, nx - 1, step = 1),
+      get_ϕ_final(prior, ekpobj)[:,2:end],
+      label = "",
+      color = :lightgreen,
+      linewidth = 4,
+      linealpha=0.1,
+      )
+plot!(p3,
+      range(0, nx - 1, step = 1),
+      get_ϕ(prior, ekpobj,1)[:,1],
+      label = "",
+      color = :lightgreen,
+      linewidth = 4,
+      linealpha=0.1,
+      )
+
+plot!(p3,
+      range(0, nx - 1, step = 1),
+      get_ϕ(prior, ekpobj,1)[:,2:end],
+      label = "",
+      color = :lightgreen,
+      linewidth = 4,
+      linealpha=0.1,
+      )
+
+plot!(p4, 1:length(y), get_g_final(ekpobj)[:,1], color = :lightgreen, label = "ensemble outputs", linewidth = 4, linealpha=0.1)
+plot!(p4, 1:length(y), get_g_final(ekpobj)[:,2:end], color = :lightgreen, label = "", linewidth = 4, linealpha=0.1)
+plot!(p4, 1:length(y), get_g(ekpobj,1)[:,1], color = :lightgreen, label = "", linewidth = 4, linealpha=0.1)
+plot!(p4, 1:length(y), get_g(ekpobj,1)[:,2:end], color = :lightgreen, label = "", linewidth = 4, linealpha=0.1)
+
+
+plt = plot(p3, p4, layout = l)
+
+savefig(plt, figure_save_directory * "solution_spatial_dep_full_ens$(N_ens).png")
+savefig(plt, figure_save_directory * "solution_spatial_dep_full_ens$(N_ens).pdf")
 
 # plot - UQ results
 
@@ -107,7 +175,90 @@ plot!(
     fillalpha = 0.1,
 )
 
-
 figpath = joinpath(figure_save_directory, "posterior_ribbons_" * case)
 savefig(figpath * ".png")
 savefig(figpath * ".pdf")
+
+##########################
+cases = ["GP", "RF-scalar"]
+
+# load
+homedir = pwd()
+println(homedir)
+data_file_GP = joinpath(data_save_directory, "posterior_$(cases[1]).jld2")
+data_file_RF = joinpath(data_save_directory, "posterior_$(cases[2]).jld2")
+truth_params_GP = load(data_file_GP)["truth_params_constrained"]
+final_params_GP = load(data_file_GP)["final_params_constrained"]
+posterior_GP = load(data_file_GP)["posterior"]
+truth_params_RF = load(data_file_RF)["truth_params_constrained"]
+final_params_RF = load(data_file_RF)["final_params_constrained"]
+posterior_RF = load(data_file_RF)["posterior"]
+
+ekp_file = joinpath(data_save_directory, "ekp_spatial_dep.jld2")
+ekpobj = load(ekp_file)["ekpobj"]
+N_ens = get_N_ens(ekpobj)
+nx = length(truth_params)
+y = get_obs(ekpobj)
+ny = length(y)
+# get samples and quantiles from posterior
+param_names_GP = get_name(posterior_GP)
+posterior_samples_GP = vcat([get_distribution(posterior_GP)[name] for name in get_name(posterior_GP)]...) #samples are columns
+constrained_posterior_samples_GP =
+    mapslices(x -> transform_unconstrained_to_constrained(posterior_GP, x), posterior_samples_GP, dims = 1)
+
+quantiles_GP = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_GP)])' # rows are quantiles for row of posterior samples
+
+param_names_RF = get_name(posterior_RF)
+posterior_samples_RF = vcat([get_distribution(posterior_RF)[name] for name in get_name(posterior_RF)]...) #samples are columns
+constrained_posterior_samples_RF =
+    mapslices(x -> transform_unconstrained_to_constrained(posterior_RF, x), posterior_samples_RF, dims = 1)
+
+quantiles_RF = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_RF)])' # rows are quantiles for row of posterior samples
+
+# plot - UQ results - both
+
+gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16)
+p1 = plot(
+    range(0, nx - 1, step = 1),
+    [truth_params final_params],
+    label = ["solution" "EKI-opt"],
+    color = [:black :lightgreen],
+    linewidth = 4,
+    xlabel = "Spatial index",
+    ylabel = "Forcing (input)",
+    left_margin = 15mm,
+    bottom_margin = 15mm,
+)
+
+plot!(
+    p1,
+    range(0, nx - 1, step = 1),
+    quantiles_GP[:, 2], # median of all vals
+    color = :blue,
+    label = "GP posterior",
+    linewidth = 4,
+    ribbon = [quantiles_GP[:, 2] - quantiles_GP[:, 1] quantiles_GP[:, 3] - quantiles_GP[:, 2]],
+    linealpha=0.5,
+    fillalpha = 0.1,
+)
+
+plot!(
+    p1,
+    range(0, nx - 1, step = 1),
+    quantiles_RF[:, 2], # median of all vals
+    color = :red,
+    label = "posterior_RF",
+    linewidth = 4,
+    ribbon = [quantiles_RF[:, 2] - quantiles_RF[:, 1] quantiles_RF[:, 3] - quantiles_RF[:, 2]],
+    linealpha=0.5,
+    fillalpha = 0.1,
+)
+
+
+figpath = joinpath(figure_save_directory, "posterior_ribbons_both")
+savefig(figpath * ".png")
+savefig(figpath * ".pdf")
+
+
+
+