Update tests

Author: ArneBouillon

examples/DimensionReduction/calibrate_linlinexp.jl

@@ -18,7 +18,7 @@ end
 function EnsembleKalmanProcesses.calculate_timestep!(ekp, g, Δt_new, scheduler::CheckpointScheduler)
     EnsembleKalmanProcesses.calculate_timestep!(ekp, g, Δt_new, scheduler.scheduler)
     if scheduler.current_index <= length(scheduler.αs) && get_algorithm_time(ekp)[end] > scheduler.αs[scheduler.current_index]
-        get_algorithm_time(ekp)[end] = scheduler.αs[scheduler.current_index]
+        ekp.Δt[end] -= get_algorithm_time(ekp)[end] - scheduler.αs[scheduler.current_index]
         scheduler.current_index += 1
     end
 
@@ -28,9 +28,9 @@ function get_algorithm_time(ekp::EnsembleKalmanProcess) # This is not defined in
     return accumulate(+, ekp.Δt)
 end
 
-rng = Random.MersenneTwister(123)
+rng = Random.MersenneTwister(41)
 input_dim = 100
-output_dim = 50
+output_dim = 100
 αs = [0.0, 0.25, 0.5, 0.75, 1.0]
 
 num_trials = 1
@@ -42,7 +42,7 @@ for trial in 1:num_trials
         Parameterized(MvNormal(zeros(size(prior_cov, 1)), prior_cov)), fill(no_constraint(), size(prior_cov, 1)), "linlinexp_prior",
     )
 
-    n_ensemble = 200
+    n_ensemble = 400
     initial_ensemble = construct_initial_ensemble(rng, prior_obj, n_ensemble)
     ekp = EnsembleKalmanProcess(
         initial_ensemble,
@@ -64,6 +64,8 @@ for trial in 1:num_trials
     end
     @info "EKP iterations: $n_iters"
     @info "Loss over iterations: $(get_error(ekp))"
+    @info "Timesteps: $(ekp.Δt)"
+    @info "Checkpoints: $(cumsum(ekp.Δt))"
 
     ekp_samples = Dict()
     for α in αs

examples/DimensionReduction/emulate_sample_linlinexp.jl

@@ -3,6 +3,8 @@ using EnsembleKalmanProcesses
 using JLD2
 using LinearAlgebra
 using Random
+using MCMCChains
+using AdvancedMH
 
 # CES
 using CalibrateEmulateSample.Emulators
@@ -11,9 +13,29 @@ using CalibrateEmulateSample.Utilities
 
 include("./models.jl")
 
-rng = Random.MersenneTwister(123)
+
+mutable struct NoEmulation <: Emulators.MachineLearningTool
+    f
+    output_structure_mats
+
+    NoEmulation(f) = new(f, nothing)
+end
+function Emulators.build_models!(ne::NoEmulation, iopairs, input_structure_mats, output_structure_mats; mlt_kwargs...)
+    ne.output_structure_mats = output_structure_mats
+end
+function Emulators.predict(ne::NoEmulation, new_inputs; mlt_kwargs...)
+    encoder_schedule = mlt_kwargs[:encoder_schedule]
+
+    decoded_inputs = decode_with_schedule(encoder_schedule, EnsembleKalmanProcesses.DataContainer(new_inputs), "in")
+    decoded_outputs = hcat(map(ne.f, eachcol(decoded_inputs.data))...)
+    encoded_outputs = encode_with_schedule(encoder_schedule, EnsembleKalmanProcesses.DataContainer(decoded_outputs), "out").data
+    encoded_outputs, cat([Utilities.get_structure_mat(ne.output_structure_mats) for _ in eachcol(new_inputs)]...; dims = 3)
+end
+
+
+rng = Random.MersenneTwister(41)
 input_dim = 100
-output_dim = 50
+output_dim = 100
 αs = [0.0, 0.25, 0.5, 0.75, 1.0]
 
 num_trials = 1
@@ -25,53 +47,69 @@ for trial in 1:num_trials
     obs_noise_cov = loaded["obs_noise_cov"]
     y = loaded["y"]
     model = loaded["model"]
+    true_parameter = loaded["true_parameter"]
 
 
     min_iter = 1
-    max_iter = 7 # number of EKP iterations to use data from is at most this
-
-    encoder_schedule_decorrelate = [(decorrelate_structure_mat(; retain_var = 0.7), "in_and_out")]
-    encoder_schedules_li = [
-        [(likelihood_informed(; retain_KL = 0.9, alpha = α, use_data_as_samples = false), "in_and_out")] for α in αs
-    ]
-
-    em_decorrelate = Emulator(
-        GaussianProcess(Emulators.GPJL(); kernel = nothing, prediction_type = Emulators.YType(), noise_learn = false),
-        Utilities.get_training_points(ekpobj, min_iter:max_iter);
-        encoder_schedule = encoder_schedule_decorrelate,
-        encoder_kwargs = (; prior_cov = cov(prior_obj), obs_noise_cov = obs_noise_cov),
-    )
-
-    ems_li = [
-        Emulator(
-            GaussianProcess(Emulators.GPJL(); kernel = nothing, prediction_type = Emulators.YType(), noise_learn = false),
+    max_iter = 10 # number of EKP iterations to use data from is at most this
+
+    encoder_schedule_ref = [(decorrelate_structure_mat(; retain_var = 1.0), "in_and_out")]
+
+    dims = [2, 4, 6, 8, 10]
+    all_errs = zeros(length(dims), 1 + length(αs))
+    for (dim_i, dim) in enumerate(dims)
+        encoder_schedule_decorrelate = [(Decorrelator([], [], [], (:dimension, dim), "structure_mat", nothing), "in"), (decorrelate_structure_mat(; retain_var = 1.0), "out")]
+        encoder_schedules_li = [
+            [(decorrelate_structure_mat(; retain_var = 1.0), "in_and_out"), (LikelihoodInformed(nothing, nothing, nothing, (:dimension, dim), α, :linreg, false), "in")] for α in αs
+        ]
+
+        em_ref = Emulator(
+            NoEmulation(param -> forward_map(param, model)),
+            Utilities.get_training_points(ekpobj, min_iter:max_iter);
+            encoder_schedule = encoder_schedule_ref,
+            encoder_kwargs = (; prior_cov = cov(prior_obj), obs_noise_cov = obs_noise_cov),
+        )
+
+        em_decorrelate = Emulator(
+            NoEmulation(param -> forward_map(param, model)),
             Utilities.get_training_points(ekpobj, min_iter:max_iter);
-            encoder_schedule = encoder_schedule,
-            encoder_kwargs = (; prior_cov = cov(prior_obj), obs_noise_cov = obs_noise_cov, samples_in = ekp_samp[α][1], samples_out = ekp_samp[α][2], observation = y),
+            encoder_schedule = encoder_schedule_decorrelate,
+            encoder_kwargs = (; prior_cov = cov(prior_obj), obs_noise_cov = obs_noise_cov),
         )
-        for (encoder_schedule, α) in zip(encoder_schedules_li, αs)
-    ]
 
-    post_means = zeros(input_dim, 0)
-    for em in vcat(em_decorrelate, ems_li...)
-        u0 = rand(MvNormal(mean(prior_obj), cov(prior_obj)))
-        mcmc = MCMCWrapper(RWMHSampling(), y, prior_obj, em; init_params = u0)
-        new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0)
+        ems_li = [
+            Emulator(
+                NoEmulation(param -> forward_map(param, model)),
+                Utilities.get_training_points(ekpobj, min_iter:max_iter);
+                encoder_schedule = encoder_schedule,
+                encoder_kwargs = (; prior_cov = cov(prior_obj), obs_noise_cov = obs_noise_cov, samples_in = ekp_samp[α][1], samples_out = ekp_samp[α][2], observation = y),
+            )
+            for (encoder_schedule, α) in zip(encoder_schedules_li, αs)
+        ]
+
+        post_means = reshape(true_parameter, input_dim, 1)
+        post_covs = []
+        for em in vcat(em_ref, em_decorrelate, ems_li...)
+            u0 = rand(MvNormal(mean(prior_obj), cov(prior_obj)))
+            mcmc = MCMCWrapper(RWMHSampling(), y, prior_obj, em; init_params = u0)
+            new_step = optimize_stepsize(mcmc; init_stepsize = 0.05, N = 2000, discard_initial = 0)
+
+            println("Begin MCMC - with step size ", new_step)
+            num_chains = 16
+            mcmc = MCMCWrapper(RWMHSampling(), y, prior_obj, em; init_params = [u0 for _ in 1:num_chains])
+            chain = MarkovChainMonteCarlo.sample(mcmc, MCMCThreads(), 40_000, num_chains; chain_type = Chains, stepsize = new_step, discard_initial = 5_000)
 
-        println("Begin MCMC - with step size ", new_step)
-        chain = MarkovChainMonteCarlo.sample(mcmc, 10_000; stepsize = new_step, discard_initial = 2_000)
+            posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain)
 
-        posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain)
+            post_mean = mean(posterior)
+            post_cov = cov(posterior)
 
-        post_mean = mean(posterior)
-        post_cov = cov(posterior)
-        println("post_mean")
-        println(post_mean)
-        println("post_cov")
-        println(post_cov)
+            post_means = hcat(post_means, reshape(post_mean, input_dim, 1))
+            push!(post_covs, post_cov)
+        end
 
-        post_means = hcat(post_means, reshape(post_mean, input_dim, 1))
+        all_errs[dim_i, :] = [norm(post_means[:,2] - v)/norm(post_means[:,2]) for v in eachcol(post_means[:,3:end])]'
     end
 
-    println(post_means[1:10,:])
+    display(all_errs)
 end

examples/DimensionReduction/models.jl

@@ -26,7 +26,7 @@ function linlinexp(input_dim, output_dim, rng)
     obs_noise_cov = Diagonal([Float64(j)^(-1 / 2) for j in 1:output_dim])
     noise = rand(rng, MvNormal(zeros(output_dim), obs_noise_cov))
     # true_parameter = reshape(ones(input_dim), :, 1)
-    true_parameter = rand(MvNormal(zeros(input_dim), Γ))
+    true_parameter = rand(rng, MvNormal(zeros(input_dim), Γ))
     y = vec(forward_map(true_parameter, model) + noise)
     return Γ, y, obs_noise_cov, model, true_parameter
 end