Start adding example

Author: ArneBouillon

examples/DimensionReduction/Project.toml

@@ -0,0 +1,17 @@
+[deps]
+AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
+CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
+ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
+Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"

examples/DimensionReduction/calibrate_linlinexp.jl

@@ -0,0 +1,89 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+
+include("./models.jl")
+
+## Define custom learning-rate scheduler for better control over timepoints reached
+mutable struct CheckpointScheduler <: EnsembleKalmanProcesses.LearningRateScheduler
+    αs::Vector{Float64}
+    scheduler
+
+    current_index
+
+    CheckpointScheduler(αs, scheduler) = new(αs, scheduler, 2)
+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]
+        scheduler.current_index += 1
+    end
+
+    nothing
+end
+function get_algorithm_time(ekp::EnsembleKalmanProcess) # This is not defined in older package versions; this is temporary here
+    return accumulate(+, ekp.Δt)
+end
+
+rng = Random.MersenneTwister(123)
+input_dim = 100
+output_dim = 50
+αs = [0.0, 0.25, 0.5, 0.75, 1.0]
+
+num_trials = 1
+for trial in 1:num_trials
+    @info "Trial $trial"
+
+    prior_cov, y, obs_noise_cov, model, true_parameter = linlinexp(input_dim, output_dim, rng)
+    prior_obj = ParameterDistribution(
+        Parameterized(MvNormal(zeros(size(prior_cov, 1)), prior_cov)), fill(no_constraint(), size(prior_cov, 1)), "linlinexp_prior",
+    )
+
+    n_ensemble = 200
+    initial_ensemble = construct_initial_ensemble(rng, prior_obj, n_ensemble)
+    ekp = EnsembleKalmanProcess(
+        initial_ensemble,
+        y,
+        obs_noise_cov,
+        TransformInversion();
+        rng,
+        scheduler = CheckpointScheduler(αs, EKSStableScheduler(2.0, 0.01)),
+    )
+
+    n_iters = 0
+    while vcat([0.0], get_algorithm_time(ekp))[end] < maximum(αs)
+        n_iters += 1
+        G_ens = hcat([forward_map(param, model) for param in eachcol(get_ϕ_final(prior_obj, ekp))]...)
+        terminate = update_ensemble!(ekp, G_ens)
+        if !isnothing(terminate)
+            throw("EKI terminated prematurely: $(terminate)! Shouldn't happen...")
+        end
+    end
+    @info "EKP iterations: $n_iters"
+    @info "Loss over iterations: $(get_error(ekp))"
+
+    ekp_samples = Dict()
+    for α in αs
+        closest_iter = argmin(0:n_iters) do i
+            abs(α - (i == 0 ? 0.0 : get_algorithm_time(ekp)[i]))
+        end + 1
+        ekp_samples[α] = (get_u(ekp, closest_iter), closest_iter == n_iters + 1 ? hcat([forward_map(u, model) for u in eachcol(get_u(ekp, closest_iter))]...) : get_g(ekp, closest_iter))
+    end
+
+
+    #! format: off
+    save(
+        "datafiles/ekp_linlinexp_$(trial).jld2",
+        "ekpobj", ekp,
+        "ekp_samples", ekp_samples,
+        "prior_obj", prior_obj,
+        "y", y,
+        "obs_noise_cov", obs_noise_cov,
+        "model", model,
+        "true_parameter", true_parameter,
+    )
+    #! format: on
+end

examples/DimensionReduction/datafiles/.gitkeep

@@ -0,0 +1,77 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+
+# CES
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+using CalibrateEmulateSample.Utilities
+
+include("./models.jl")
+
+rng = Random.MersenneTwister(123)
+input_dim = 100
+output_dim = 50
+αs = [0.0, 0.25, 0.5, 0.75, 1.0]
+
+num_trials = 1
+for trial in 1:num_trials
+    loaded = load("datafiles/ekp_linlinexp_$(trial).jld2")
+    ekpobj = loaded["ekpobj"]
+    ekp_samp = loaded["ekp_samples"]
+    prior_obj = loaded["prior_obj"]
+    obs_noise_cov = loaded["obs_noise_cov"]
+    y = loaded["y"]
+    model = loaded["model"]
+
+
+    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),
+            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 = 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)
+
+        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)
+
+        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))
+    end
+
+    println(post_means[1:10,:])
+end

examples/DimensionReduction/emulate_sample_linlinexp.jl

@@ -0,0 +1,50 @@
+using Distributions
+using LinearAlgebra
+using Statistics
+
+using EnsembleKalmanProcesses
+using EnsembleKalmanProcesses.ParameterDistributions
+
+abstract type ForwardMapType end
+
+## Linear-times-exponential model
+function linlinexp(input_dim, output_dim, rng)
+    # prior
+    γ0 = 4.0
+    β_γ = -2
+    Γ = Diagonal([γ0 * (1.0 * j)^β_γ for j in 1:input_dim])
+
+    U = qr(randn(rng, (output_dim, output_dim))).Q
+    V = qr(randn(rng, (input_dim, input_dim))).Q
+    λ0 = 100.0
+    β_λ = -1
+    Λ = Diagonal([λ0 * (1.0 * j)^β_λ for j in 1:output_dim])
+    A = U * Λ * V[1:output_dim, :] # output x input
+    model = LinLinExp(input_dim, output_dim, A)
+
+    # generate data sample
+    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), Γ))
+    y = vec(forward_map(true_parameter, model) + noise)
+    return Γ, y, obs_noise_cov, model, true_parameter
+end
+
+struct LinLinExp{AM <: AbstractMatrix} <: ForwardMapType
+    input_dim::Int
+    output_dim::Int
+    G::AM
+end
+
+function forward_map(X::AVorM, model::LinLinExp) where {AVorM <: AbstractVecOrMat}
+    return model.G * (X .* exp.(0.05X))
+end
+
+function jac_forward_map(X::AbstractVector, model::LinLinExp)
+    return model.G * Diagonal(exp.(0.05X) .* (1 .+ 0.05X))
+end
+
+function jac_forward_map(X::AbstractMatrix, model::LinLinExp)
+    return [jac_forward_map(x, model) for x in eachcol(X)]
+end

examples/DimensionReduction/models.jl

@@ -243,8 +243,8 @@ function initialize_and_encode_with_schedule!(
     prior_cov::Union{Nothing, StructureMatrix} = nothing,
     obs_noise_cov::Union{Nothing, StructureMatrix} = nothing,
     observation::Union{Nothing, StructureVector} = nothing,
-    prior_samples_in::Union{Nothing, StructureVector} = nothing,
-    prior_samples_out::Union{Nothing, StructureVector} = nothing,
+    samples_in::Union{Nothing, StructureVector} = nothing,
+    samples_out::Union{Nothing, StructureVector} = nothing,
 ) where {VV <: AbstractVector, PDC <: PairedDataContainer}
     processed_io_pairs = deepcopy(io_pairs)
 
@@ -259,16 +259,16 @@ function initialize_and_encode_with_schedule!(
     end
 
     input_structure_vecs = deepcopy(input_structure_vecs)
-    if !isnothing(prior_samples_in)
-        (input_structure_vecs[:prior_samples_in] = prior_samples_in)
+    if !isnothing(samples_in)
+        (input_structure_vecs[:samples_in] = samples_in)
     end
 
     output_structure_vecs = deepcopy(output_structure_vecs)
     if !isnothing(observation)
         (output_structure_vecs[:observation] = observation)
     end
-    if !isnothing(prior_samples_out)
-        (output_structure_vecs[:prior_samples_out] = prior_samples_out)
+    if !isnothing(samples_out)
+        (output_structure_vecs[:samples_out] = samples_out)
     end
 
     # apply_to is the string "in", "out" etc.

src/Utilities.jl

@@ -11,15 +11,15 @@ mutable struct LikelihoodInformed{FT <: Real} <: PairedDataContainerProcessor
     dim_criterion::Tuple{Symbol, <:Number}
     α::FT
     grad_type::Symbol
-    use_prior_samples::Bool
+    use_data_as_samples::Bool
 end
 
-function likelihood_informed(retain_KL; alpha = 0.0, grad_type = :localsl, use_prior_samples = true)
+function likelihood_informed(; retain_KL, alpha = 0.0, grad_type = :localsl, use_data_as_samples = false)
     if grad_type ∉ [:linreg, :localsl]
         @error "Unknown grad_type=$grad_type"
     end
 
-    LikelihoodInformed(nothing, nothing, nothing, (:retain_KL, retain_KL), alpha, grad_type, use_prior_samples)
+    LikelihoodInformed(nothing, nothing, nothing, (:retain_KL, retain_KL), alpha, grad_type, use_data_as_samples)
 end
 
 get_encoder_mat(li::LikelihoodInformed) = li.encoder_mat
@@ -45,14 +45,13 @@ function initialize_processor!(
         else
             get_structure_vec(output_structure_vectors, :observation)
         end
-        samples_in, samples_out = if li.use_prior_samples
-            @assert α ≈ 0.0
+        samples_in, samples_out = if li.use_data_as_samples
+            (in_data, out_data)
+        else
             (
-                get_structure_vec(input_structure_vectors, :prior_samples_in),
-                get_structure_vec(output_structure_vectors, :prior_samples_out),
+                get_structure_vec(input_structure_vectors, :samples_in),
+                get_structure_vec(output_structure_vectors, :samples_out),
             )
-        else
-            (in_data, out_data)
         end
         obs_noise_cov = get_structure_mat(output_structure_matrices, :obs_noise_cov)
         noise_cov_inv = inv(obs_noise_cov)
@@ -79,10 +78,10 @@ function initialize_processor!(
 
         li.encoder_mat = if apply_to == "in" || α ≈ 0
             decomp = if apply_to == "in"
-                eigen(mean(grad' * noise_cov_inv * ((1-α)obs_noise_cov + α^2 * (y - g) * (y - g)') * noise_cov_inv * grad for (g, grad) in zip(eachcol(samples_out), grads)), sortby = (-))
+                eigen(hermitianpart(mean(grad' * noise_cov_inv * ((1-α)obs_noise_cov + α^2 * (y - g) * (y - g)') * noise_cov_inv * grad for (g, grad) in zip(eachcol(samples_out), grads))), sortby = (-))
             else
                 @assert apply_to == "out"
-                eigen(mean(grad * grad' for grad in grads), obs_noise_cov, sortby = (-))
+                eigen(hermitianpart(mean(grad * grad' for grad in grads)), obs_noise_cov, sortby = (-))
             end
 
             if li.dim_criterion[1] == :retain_KL
@@ -93,6 +92,7 @@ function initialize_processor!(
                 @assert li.dim_criterion[1] == :dimension
                 trunc_val = li.dim_criterion[2]
             end
+            @info "    truncating at $trunc_val/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc_val])% of the KL divergence reduction"
             li.encoder_mat = decomp.vectors[:, 1:trunc_val]'
         else
             @assert apply_to == "out"
@@ -135,13 +135,18 @@ function initialize_processor!(
                 if li.dim_criterion[1] == :retain_KL
                     retain_KL = li.dim_criterion[2]
                     ref = f(M, zeros(output_dim, 0))
-                    if f(M, Vs) / ref ≤ 1 - retain_KL
+                    val = f(M, Vs)
+                    if val / ref ≤ 1 - retain_KL
+                        @info "    truncating at $k/$output_dim retaining $(100.0*(1-val/ref))% of the KL divergence reduction"
                         break # TODO: Start bisecting?
                     else
                         k *= 2
                     end
                 else
                     @assert li.dim_criterion[1] == :dimension
+                    ref = f(M, zeros(output_dim, 0))
+                    val = f(M, Vs)
+                    @info "    truncating at $(li.dim_criterion[2])/$output_dim retaining $(100.0*val/ref)% of the KL divergence reduction"
                     break
                 end
             end