updated examples with MB

CliMA · Jun 21, 2024 · baff6a2 · baff6a2
1 parent 5b4786b
commit baff6a2
Show file tree

Hide file tree

Showing 7 changed files with 68 additions and 40 deletions.
diff --git a/examples/AerosolActivation/aerosol_activation.jl b/examples/AerosolActivation/aerosol_activation.jl
@@ -139,40 +139,35 @@ nothing # hide
 # calibrate is also used to generate observations.
 # We use the total number and mass of activated aerosol particles as
 # our observational data.
-observation_data_names = ["N_act", "M_act"];
+observation_data_names = ["N_act_and_M_act"];
 nothing # hide
 
 # We generate artificial truth samples based on the default values
 # of parameters we are calibrating.
-n_samples = 10
 G_t = run_activation_model(molar_mass_true, osmotic_coeff_true)
-y_t = zeros(length(G_t), n_samples)
 
 Γy = convert(Array, LinearAlgebra.Diagonal([0.01 * G_t[1], 0.01 * G_t[2]]))
 μ = zeros(length(G_t));
 nothing # hide
 
 # And add noise to the generated truth sample.
-for i in 1:n_samples
-    y_t[:, i] = G_t .+ rand(Distributions.MvNormal(μ, Γy))
-end
-
-truth_array = EKP.Observation(y_t, Γy, observation_data_names)
-nothing # hide
-
-# One could try for the truth to be a mean of the generated array.
-# Or do the calibration for all individual truth samples and then
-# compute the mean of calibrated parameters.
-# For now we are just taking one truth array member.
-truth_sample = truth_array.samples[1];
+y_t = G_t .+ rand(Distributions.MvNormal(μ, Γy))
+
+observation = EKP.Observation(
+    Dict(
+        "samples" => y_t,
+        "covariances" => Γy,
+        "names" => observation_data_names,
+    ),
+)
 nothing # hide
 
 # We use 50 ensemble members and do 10 iterations.
 N_ens = 50
 N_iter = 10
 
 initial_par = EKP.construct_initial_ensemble(rng, priors, N_ens)
-ekiobj = EKP.EnsembleKalmanProcess(initial_par, truth_sample, truth_array.obs_noise_cov, EKP.Inversion())
+ekiobj = EKP.EnsembleKalmanProcess(initial_par, observation, EKP.Inversion())
 nothing # hide
 
 # Finally, we can run the Ensemble Kalman Process calibration.

diff --git a/examples/Cloudy/Cloudy_example_eki.jl b/examples/Cloudy/Cloudy_example_eki.jl
@@ -70,7 +70,7 @@ priors = combine_distributions([prior_N0, prior_θ, prior_k])
 ###  Define the data from which we want to learn the parameters
 ###
 
-data_names = ["M0", "M1", "M2"]
+data_names = ["M0_M1_M2"]
 moments = [0.0, 1.0, 2.0]
 n_moments = length(moments)
 
@@ -110,8 +110,13 @@ for i in 1:n_samples
     y_t[:, i] = G_t .+ rand(MvNormal(μ, Γy))
 end
 
-truth = Observation(y_t, Γy, data_names)
-truth_sample = truth.mean
+truth = Observation(
+    Dict(
+        "samples" => vec(mean(y_t,dims=2))
+        "covariances" => Γy,
+        "names" => data_names,
+    ),
+)
 
 
 ###
@@ -125,8 +130,7 @@ N_iter = 8 # number of EKI iterations
 initial_par = construct_initial_ensemble(rng, priors, N_ens)
 ekiobj = EnsembleKalmanProcess(
     initial_par,
-    truth_sample,
-    truth.obs_noise_cov,
+    truth,
     Inversion(),
     timestepper = DefaultTimestepper(0.1),
 )

diff --git a/examples/Cloudy/Cloudy_example_uki.jl b/examples/Cloudy/Cloudy_example_uki.jl
@@ -71,7 +71,7 @@ priors = combine_distributions([prior_N0, prior_θ, prior_k])
 ###  Define the data from which we want to learn the parameters
 ###
 
-data_names = ["M0", "M1", "M2"]
+data_names = ["M0_M1_M2"]
 moments = [0.0, 1.0, 2.0]
 n_moments = length(moments)
 
@@ -110,8 +110,13 @@ for i in 1:n_samples
     y_t[:, i] = G_t .+ rand(MvNormal(μ, Γy))
 end
 
-truth = Observation(y_t, Γy, data_names)
-truth_sample = truth.mean
+truth = Observation(
+    Dict(
+        "samples" => vec(mean(y_t,dims=2))
+        "covariances" => Γy,
+        "names" => data_names,
+    ),
+)
 
 
 ###
@@ -129,7 +134,7 @@ N_iter = 50 # number of iterations
 update_freq = 1
 
 process = Unscented(mean(priors), cov(priors); α_reg = α_reg, update_freq = update_freq)
-ukiobj = EnsembleKalmanProcess(truth_sample, truth.obs_noise_cov, process)
+ukiobj = EnsembleKalmanProcess(truth, process)
 
 # Initialize a ParticleDistribution with dummy parameters. The parameters 
 # will then be set within `run_dyn_model`

diff --git a/examples/Lorenz/Lorenz_example.jl b/examples/Lorenz/Lorenz_example.jl
@@ -84,7 +84,7 @@ end
 ###
 ###  Define the data from which we want to learn the parameters
 ###
-data_names = ["y0", "y1"]
+data_names = ["y0_y1"]
 
 
 ###
@@ -177,8 +177,13 @@ end
 
 
 # Construct observation object
-truth = Observation(yt, Γy, data_names)
-truth_sample = truth.mean
+truth = Observation(
+    Dict(
+        "samples" => vec(mean(yt,dims=2))
+        "covariances" => Γy,
+        "names" => data_names,
+    ),
+)
 ###
 ###  Calibrate: Ensemble Kalman Inversion
 ###
@@ -193,7 +198,7 @@ N_iter = 5 # number of EKI iterations
 # initial parameters: N_params x N_ens
 initial_params = construct_initial_ensemble(rng, priors, N_ens)
 
-ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, Inversion())
+ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth, Inversion())
 
 # EKI iterations
 println("EKP inversion error:")

diff --git a/examples/Lorenz/Lorenz_example_ukp.jl b/examples/Lorenz/Lorenz_example_ukp.jl
@@ -82,7 +82,7 @@ end
 ###
 ###  Define the data from which we want to learn the parameters
 ###
-data_names = ["y0", "y1"]
+data_names = ["y0_y1"]
 
 
 ###
@@ -174,8 +174,13 @@ else
 end
 
 # Construct observation object
-truth = Observation(yt, Γy, data_names)
-truth_sample = truth.mean
+truth = Observation(
+    Dict(
+        "samples" => vec(mean(yt,dims=2))
+        "covariances" => Γy,
+        "names" => data_names,
+    ),
+)
 ###
 ###  Calibrate: Ensemble Kalman Inversion
 ###
@@ -198,7 +203,7 @@ update_freq = 0
 #       n_param+2  : sample n_param+2 sigma points
 N_ens = 2n_param + 1
 process = Unscented(mean(priors), cov(priors); α_reg = α_reg, update_freq = update_freq, sigma_points = "symmetric")
-ukiobj = EKP.EnsembleKalmanProcess(truth_sample, truth.obs_noise_cov, process)
+ukiobj = EKP.EnsembleKalmanProcess(truth, process)
 
 
 # UKI iterations

diff --git a/examples/Lorenz/distributed_Lorenz_example.jl b/examples/Lorenz/distributed_Lorenz_example.jl
@@ -108,7 +108,7 @@ end
 ###
 ###  Define the data from which we want to learn the parameters
 ###
-data_names = ["y0", "y1"]
+data_names = ["y0_y1"]
 
 
 ###
@@ -201,8 +201,13 @@ end
 
 
 # Construct observation object
-truth = Observation(yt, Γy, data_names)
-truth_sample = truth.mean
+truth = Observation(
+    Dict(
+        "samples" => vec(mean(yt,dims=2))
+        "covariances" => Γy,
+        "names" => data_names,
+    ),
+)
 ###
 ###  Calibrate: Ensemble Kalman Inversion
 ###
@@ -217,7 +222,7 @@ N_iter = 5 # number of EKI iterations
 # initial parameters: N_params x N_ens
 initial_params = construct_initial_ensemble(rng, priors, N_ens)
 
-ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, Inversion())
+ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth, Inversion())
 
 # EKI iterations
 println("EKP inversion error:")

diff --git a/examples/example_template.jl b/examples/example_template.jl
@@ -29,9 +29,18 @@ noise_covariance = prescribed_observational_noise_covariance
 ##
 ## Store the observation and covariance
 ##
+truth_sample = truth_data[1] 
+
+truth = Observation(
+    Dict(
+        "samples" => truth_sample,
+        "covariances" => observational_noise_covariance,
+        "names" => data_names,
+    ),
+)
+# note that you can stack data(i.e. different observed variables) conveniently with combine_observations([truth1,truth2])
+# if you have a lot of data you can also minibatch it conveniently by using an ObservationSeries - see the docs
 
-truth = Observation(truth_data, observational_noise_covariance, data_names)
-truth_sample = truth.samples[1]
 
 ##
 ## Now define the priors:
@@ -48,7 +57,7 @@ params = construct_initial_ensemble(priors, N_ens)
 
 # constructor for the EKI object
 ekiobj =
-    EnsembleKalmanProcessModule.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, Inversion())
+    EnsembleKalmanProcessModule.EnsembleKalmanProcess(initial_params, truth, Inversion())
 
 #outer EKI loop
 for i in 1:N_iterations