first port from PR 376

odunbar · odunbar · commit ae228128cee4 · 2026-04-02T15:20:22.000-07:00
typos

add pkgs

encoder_kwargs in nice form

remove dim criterion for now

docstring

API refinement

add example

builds encoder, but mcmc bug

examples run, bugs for dimensionality and product order resolved

noise injection reduced

add some get_encoded_dim functions

add consistency for both in-and-out dimensions for li, when using multiple distributions

runs through with output reduction, logic simplified for now
diff --git a/Project.toml b/Project.toml
@@ -19,6 +19,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 LowRankApprox = "898213cb-b102-5a47-900c-97e73b919f73"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
+Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
@@ -47,6 +49,8 @@ KernelFunctions = "0.10.64"
 LinearMaps = "3.11.4"
 LowRankApprox = "0.5.5"
 MCMCChains = "7"
+Manifolds = "0.11.19"
+Manopt = "0.5.34"
 Plots = "1.41.1"
 Printf = "1"
 ProgressBars = "1"
diff --git a/examples/DimensionReduction/Project.toml b/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"
diff --git a/examples/DimensionReduction/calibrate_linlinexp.jl b/examples/DimensionReduction/calibrate_linlinexp.jl
@@ -0,0 +1,72 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+
+include("./models.jl")
+
+datadir = joinpath(@__DIR__,"datafiles")
+if !isdir(datadir)
+    mkpath(datadir)
+end
+
+rng = Random.MersenneTwister(41)
+input_dim = 100
+output_dim = 100
+
+num_trials = 5
+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 = 50
+    initial_ensemble = construct_initial_ensemble(rng, prior_obj, n_ensemble)
+    ekp = EnsembleKalmanProcess(
+        initial_ensemble,
+        y,
+        obs_noise_cov,
+        TransformInversion();
+        rng,
+    )
+
+    n_iters = 0
+    for i in 1:50
+        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)
+            break
+        end
+    end
+    @info "EKP iterations: $n_iters"
+    @info "Loss over iterations: $(get_error(ekp))"
+    @info "Timesteps: $(ekp.Δt)"
+    @info "Checkpoints: $(get_algorithm_time(ekp))"
+    αs = vec([0 get_algorithm_time(ekp)[1:end]...])
+    ekp_samples = Dict(α => (get_u(ekp, i), get_g(ekp, i)) for (i,α) in enumerate(αs[1:end-1]))
+
+    # evaluate at final time
+    G_ens = hcat([forward_map(param, model) for param in eachcol(get_ϕ_final(prior_obj, ekp))]...)
+    ekp_samples[αs[end]] = (get_u(ekp,length(αs)), G_ens)
+    
+    #! format: off
+    save(
+        joinpath(datadir,"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,
+        "alphas", αs,
+    )
+    #! format: on
+end
diff --git a/examples/DimensionReduction/emulate_sample_linlinexp.jl b/examples/DimensionReduction/emulate_sample_linlinexp.jl
@@ -0,0 +1,140 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+using MCMCChains
+using AdvancedMH
+
+# CES
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+using CalibrateEmulateSample.Utilities
+
+
+include("./models.jl")
+
+rng = Random.MersenneTwister(41)
+input_dim = 100
+output_dim = 100
+
+num_trials = 1
+for trial in 1:num_trials
+    loaded = load("datafiles/ekp_linlinexp_$(trial).jld2")
+    ekpobj = loaded["ekpobj"]
+    ekp_samples = loaded["ekp_samples"]
+    prior = loaded["prior_obj"]
+    obs_noise_cov = loaded["obs_noise_cov"]
+    y = loaded["y"]
+    model = loaded["model"]
+    true_parameter = loaded["true_parameter"]
+    αs = loaded["alphas"]
+
+    min_iter = 1
+    max_iter = min(10, length(αs)) # number of EKP iterations to use data from is at most this
+    training_points = [ekp_samples[α] for α in αs[min_iter:max_iter]]
+
+    encoder_schedule_ref = [(decorrelate_structure_mat(; retain_var = 1.0), "in_and_out")]
+   
+    rvs = collect(0.975:0.0125:1.0)
+    rkls = collect(0.999:0.0005:1.0)
+    all_errs = zeros(length(rvs), 1 + length(αs))
+
+    # as we have more input samples than out stored in ekp. we provide the extended set of outputs to be used in the likelihood informed processor
+    final_samples_out = ekp_samples[αs[end]][2]
+    encoder_kwargs = encoder_kwargs_from(ekpobj, prior; final_samples_out=final_samples_out)
+
+    flat_io_pairs = get_training_points(ekpobj, min_iter:max_iter, g_final=final_samples_out)
+    
+    names = ["reference", "decorrelate, PCA-in", ["decorrelate, LI-in 1:$(i)" for i in 1:length(αs)]...]
+
+    ni_scaling = 0.01 # noise injection into null space scaling (def. 1)
+    
+    for (idx, (rv, rkl)) in enumerate(zip(rvs, rkls))
+        encoder_schedule_decorrelate = [
+            (decorrelate_structure_mat(retain_var = rv), "in"),
+            (decorrelate_structure_mat(), "out"),
+        ]
+        encoder_schedules_li = [
+            [
+                (decorrelate_structure_mat(), "in_and_out"),
+                (likelihood_informed(retain_kl=rkl, iters=1:i), "in"),
+            ] for i in 1:length(αs)
+        ]
+
+        em_ref = forward_map_wrapper(
+            param -> forward_map(param, model),
+            prior,
+            flat_io_pairs;
+            encoder_schedule = encoder_schedule_ref,
+            encoder_kwargs = deepcopy(encoder_kwargs),
+            noise_injector_scaling=ni_scaling, # shouldnt be needed
+        )
+
+        em_decorrelate = forward_map_wrapper(
+            param -> forward_map(param, model),
+            prior,
+            flat_io_pairs;
+            encoder_schedule = encoder_schedule_decorrelate,
+            encoder_kwargs = deepcopy(encoder_kwargs),
+            noise_injector_scaling=ni_scaling,
+        )
+
+        ems_li = [
+            forward_map_wrapper(
+                param -> forward_map(param, model),
+                prior,               
+                flat_io_pairs;
+                encoder_schedule = encoder_schedule,
+                encoder_kwargs = deepcopy(encoder_kwargs),
+                noise_injector_scaling=ni_scaling,
+            ) for encoder_schedule in encoder_schedules_li
+        ]
+
+        post_means = reshape(true_parameter, input_dim, 1)
+        post_covs = []
+        
+        for (nn, em) in zip(names,vcat(em_ref, em_decorrelate, ems_li...))
+            println(" ")
+            @info "Encoding name: $(nn)"            
+            E,_ = get_encoder_from_schedule(get_encoder_schedule(em), "in")
+            if isnothing(E)
+                @info "No truncation"                
+            else
+
+                @info "Truncation criteria (var>$(rv), or kl > $(rkl))"
+                @info "input dim reduced to: $(size(E,1))"
+        
+            end 
+            println(" ")
+            u0 = rand(MvNormal(mean(prior), cov(prior)))
+            mcmc = MCMCWrapper(RWMHSampling(), y, prior, 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)
+            mcmc = MCMCWrapper(RWMHSampling(), y, prior, em; init_params = u0)
+            chain = MarkovChainMonteCarlo.sample(
+                mcmc,
+                40_000;
+                stepsize = new_step,
+                discard_initial = 5_000,
+            )
+            posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain)
+
+            post_mean = mean(posterior)
+            post_cov = cov(posterior)
+
+            post_means = hcat(post_means, reshape(post_mean, input_dim, 1))
+            push!(post_covs, post_cov)
+        end
+        
+        # post_means[:,1] = true parameter
+        # post_means[:,2] = ref
+        # post_means[:,3] = decor, LI etc.
+        # so err cols are normalized diff to ref of (decor, LI etc.) (lower is better)
+        all_errs[idx, :] =
+            [norm(post_means[:, 2] - v) / norm(post_means[:, 2]) for v in eachcol(post_means[:, 3:end])]'
+    end
+    @info "error of posterior mean to whitened \"reference\" solution. for $(names[2:end])"
+    display(all_errs)
+end
diff --git a/examples/DimensionReduction/models.jl b/examples/DimensionReduction/models.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(rng, 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
diff --git a/src/Emulator.jl b/src/Emulator.jl
@@ -41,17 +41,17 @@ include(joinpath("MachineLearningTools", "RandomFeature.jl")) # Random Freatures
 # etc.
 
 # defaults in error, all MachineLearningTools require these functions.
-function throw_define_mlt()
-    throw(ErrorException("Unknown MachineLearningTool defined, please use a known implementation"))
+function throw_define_mlt(mlt)
+    throw(ErrorException("Unknown MachineLearningTool defined, please use a known implementation. Please check all methods are defined for the MLT received: \n $mlt"))
 end
 function build_models!(mlt, iopairs, input_structure_mats, output_structure_mats, mlt_kwargs...)
-    throw_define_mlt()
+    throw_define_mlt(mlt)
 end
 function optimize_hyperparameters!(mlt)
-    throw_define_mlt()
+    throw_define_mlt(mlt)
 end
 function predict(mlt, new_inputs; mlt_kwargs...)
-    throw_define_mlt()
+    throw_define_mlt(mlt)
 end
 
 # We will define the different emulator types after the general statements
@@ -502,8 +502,6 @@ function forward_map_wrapper(
         else
             push!(encoder_schedule, (decorrelate_sample_cov(), "out"))
         end
-    else
-        @warn "Please note that only the output encoder is used in this implementation. \nThe input encoder will be initialized if provided, but not used during sampling, which is completed in the full parameter space."
     end
 
     encoder_schedule = create_encoder_schedule(encoder_schedule)
diff --git a/src/Utilities.jl b/src/Utilities.jl
diff --git a/src/Utilities/likelihood_informed.jl b/src/Utilities/likelihood_informed.jl
