[Updated] Likelihood-informed processors (#410)

* first port from PR 376

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

* rm linlinexp

* truncate based on relative log^2(1+li) not li itself

* add docs manifest

* remove data-as-samples (make default behaviour), first tests

* format

* add encoder_kwargs_from test

* typo

* typo

* logic typo

* increase shaky tol

* refactor private/public encode/decode functions

* format

* new docs manifest

* more coverage

* more coverage

* add cases

* rm trap rule switch for now

* format

* test pipeline gets better kwargs

* recreate manifest for docs and julia 1.11.9

* add coverage

* add apply_to for paired_data_container

* typo

* format

* nicer table, and bugfix ref with deepcopy

* format

* docs update

* refine example setup for LI 1:5 in  1:1 out

* add ref in as comparison to true parameters in example

* remove extra message

* docs improvement

* codecov

* updates from review

* codecov, and a bugfix for get_training_points

* format

Author: odunbar

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, 0.11"
 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"

docs/Manifest.toml

@@ -12,37 +12,72 @@ The framework works as follows
 An external API is also available using the `encode_data`, and `encode_structure_matrix` methods if needed to encode new data, or new covariance-type matrices.
 
 ## Defaults and recommendations
-### Default schedule (no explicit dimension reduction)
-When no schedule is provided, i.e. 
+
+### Recommended schedule: ("latest and greatest")
+
+Given `ekp` and `prior` from an ensemble kalman process run (e.g., the "Calibrate" stage). Assume we have chosen our machine learning tool and obtained training data. Our recommendeded family to balance efficiency and reduce dimension is to use the `retain_var` and `retain_info`. 
+
 ```julia
-emulator = Emulator(machine_learning_tool, input_output_pairs)
+retain_var = 0.999 # reduce dimension retaining 99.9% of variance (for PCA)
+retain_info = 0.999 # reduce dimension retaining 99.9% of information (for likelihood-informed)
+ekp_it = 5 # number of eki iteration (for likelihood-informed)
+
+encoder_schedule = [
+    (decorrelate_sample_cov(retain_var = retain_var), "in"),
+    (decorrelate_structure_mat(retain_var = retain_var), "out"),
+    (likelihood_informed(retain_info=retain_info, iters = 1:ekp_it), "in"),
+    (likelihood_informed(retain_info=retain_info, iters = 1:1), "out"),
+]
 ```
-The default schedule under-the-hood, is given by
+The motivation is that PCA is an efficient, but blunt, tool for reducing dimension from ((>100s) -> (100s)) dimension, then use likelihood-informed tools to hone down ((100s) -> (10s)) dimensions. `iters=1:X` is more performant in input than output dimension for `X>1` so we only do this here. This schedule is added to the emulator as follows
+```
+emulator = Emulator(
+    machine_learning_tool,       
+    input_output_pairs;
+    encoder_schedule = encoder_schedule,
+    encoder_kwargs = encoder_kwargs_from(ekp, prior), # kwargs from "Calibrate" objects
+)
+```
+
+## Recommended `encoder_kwargs`
+
+The `EnsembleKalmanProcesses` infrastructure naturally provides information to build the data processors. We provide some utilities to extract these into `kwargs` of the correct format, the recommended (but not only) option is:
+
 ```julia
-schedule = (decorrelate_sample_cov(), "in_and_out")
+# `prior::ParameterDistribution`
+# `ekp::EnsembleKalmanProcess`
+encoder_kwargs = encoder_kwargs_from(ekp, prior) 
 ```
-If the user only provides the observational noise covariance, 
+
+### Default schedule (to whiten with no dimension reduction)
+When the user provides no `encoder_schedule`, the following is created
 ```julia
 emulator = Emulator(
     machine_learning_tool,       
     input_output_pairs;
-    encoder_kwargs = (; obs_noise_cov = obs_noise_cov),
+    encoder_kwargs = encoder_kwargs,
 )
 ```
 The default schedule under-the-hood, is given by
 ```julia
 schedule = [
     (decorrelate_sample_cov(), "in"),
-    (decorrelate_structure_mat(), "out"), # uses obs_noise_cov
+    (decorrelate_structure_mat(), "out"), # uses obs_noise_cov kwarg
+]
+```
+With no kwargs also provided, the following schedule under-the-hood, is given by
+```julia
+schedule = [
+    (decorrelate_sample_cov(), "in_and_out"),
 ]
 ```
 
 !!! note "Switch encoding off"
     To ensure that no encoding is happening, the user must pass in an empty schedule `encoder_schedule = []`
 
-### Recommended: schedule for PCA dimension reduction
+## Other typical schedule: PCA dimension reduction
 
-Our recommendeded family to balance efficiency and reduce dimension is to use the `retain_var` kwargs.
+Robust and blunt dimension reduction can be done with only PCA
 ```julia
 retain_var_in = 0.99 # reduce dimension retaining 99% of input variance
 retain_var_out = 0.95 # reduce dimension retaining 95% of output variance
@@ -59,22 +94,21 @@ emulator = Emulator(
     encoder_kwargs = encoder_kwargs,
 )
 ```
-## Recommended: Get `encoder_kwargs` from `ObservationSeries` etc.
+## kwargs from other objects:
 
-To transition more smoothly from the `EnsembleKalmanProcesses` infrastructure, one can also get the kwargs in the desired format from any [`Observation` and `ObservationSeries`](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/observations/) objects, as well as [`ParameterDistribution`](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/parameter_distributions/) objects used in the `EnsembleKalmanProcesses.jl` package.
+One can also obtain individual kwarg sets from other objects, and can merge these as follows
 ```julia
 # `prior::ParameterDistribution`
-input_kwargs = get_kwargs_from(prior)
-
 # `observation_series::ObservationSeries`
+input_kwargs = get_kwargs_from(prior)
 output_kwargs = get_kwargs_from(observation_series)
-
 encoder_kwargs = merge(input_kwargs, output_kwargs)
 ```
+
+Though not recommended, you can also build your own keyword arguments.
 !!! note "Observational noise notes"
     1. The `obs_noise_cov` object does not need to be a constructed `AbstractMatrix`, rather it can be any type of compactly stored matrix compatible with `EnsembleKalmanProcesses` (e.g., [here](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/observations/)).
-    2. When `ObservationSeries` contains multiple observations, it is assumed that the covariance of the noise of all observations are the same. (If they are not, only the first will be taken by default in `get_kwargs_from(...)`
-    
+    2. When `ObservationSeries` contains multiple observations, it is assumed that the covariance of the noise of all observations are the same. (If they are not, only the first will be taken by default in `get_kwargs_from(...)`    
 
 ## Building and interpreting encoder schedules
 
@@ -145,7 +179,7 @@ The `DataContainerProcessor` encodes "input" data agnostic of the "output" data,
 
 The `PairedDataContainerProcessor` encodes inputs (or outputs) using information of the both inputs and outputs in pairs
 - `CanonicalCorrelation` - constructed with `canonical_correlation()`, which performs [canonical correlation analysis](https://en.wikipedia.org/wiki/Canonical_correlation) to process the pairs. In effect this performs PCA on the cross-correlation from input and output samples.
-- [Coming soon] `LikelihoodInformed` - this will use the data or likelihood from the inverse problem at hand to build diagnostic matrices that are used to find informative directions for dimension reduction. In particular we build generalizations of current frameworks (e.g., [Cui, Zahm 2021](http://doi.org/10.1088/1361-6420/abeafb), [Baptista, Marzouk, Zahm 2022](https://arxiv.org/abs/2207.08670)) to use the latest EKP iterations.
+- `LikelihoodInformed`- constructed with `likelihood_informed()`- we use the data or likelihood from the inverse problem at hand to build diagnostic matrices that are used to find informative directions for dimension reduction. In particular we build generalizations of current frameworks (e.g., [Cui, Zahm 2021](http://doi.org/10.1088/1361-6420/abeafb), [Baptista, Marzouk, Zahm 2022](https://arxiv.org/abs/2207.08670)) to use the latest EKP iterations. We in fact generalize this framework here [Preprint Coming soon] for use in partially-informed distributions (e.g. EKI iterations)
 
 This is an extensible framework, and so new data processors can be added to this library.
 

docs/src/data_processing.md

@@ -0,0 +1,18 @@
+[deps]
+AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
+CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
+ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+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/Project.toml

@@ -0,0 +1,65 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+
+
+include("./models.jl") # defines dims, files, mod_type
+# select the forward map
+@info "Executing model type $(mod_type)"
+
+function main()
+    if mod_type == "linear"
+        prior_cov, y, obs_noise_cov, model, true_parameter = lin(input_dim, output_dim, rng)
+    else
+        bad_model(mod_type, mod_types)
+    end
+
+    prior_obj = ParameterDistribution(
+        Parameterized(MvNormal(zeros(size(prior_cov, 1)), prior_cov)),
+        fill(no_constraint(), size(prior_cov, 1)),
+        "$(mod_type)_prior",
+    )
+
+    n_ensemble = 80
+    initial_ensemble = construct_initial_ensemble(rng, prior_obj, n_ensemble)
+    ekp = EnsembleKalmanProcess(initial_ensemble, y, obs_noise_cov, TransformInversion(); rng)
+
+    max_iter = 20
+    n_iters = 0
+    for i in 1:max_iter
+        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_$(mod_type).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
+main()

examples/DimensionReduction/calibrate.jl

@@ -0,0 +1,165 @@
+using Distributions
+using EnsembleKalmanProcesses
+using JLD2
+using LinearAlgebra
+using Random
+using MCMCChains
+using AdvancedMH
+using DataFrames # table print at end 
+# CES
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+using CalibrateEmulateSample.Utilities
+
+
+include("./models.jl")
+# select the forward map
+mod_types = ["linear"]
+mod_type = mod_types[1]
+@info "Executing model type $(mod_type)"
+
+function main()
+    loaded = load("datafiles/ekp_$(mod_type).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"]
+    mask = [1, length(αs)] # don't do all cases
+    αs = αs[mask]
+
+    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(), "in_and_out")]
+
+    # chosen some truncations with similar dimension
+    rvs_rinfos = [(0.995, 0.99995), (0.99, 0.9995), (0.98, 0.995), (0.95, 0.97), (0.9, 0.9)]
+
+    # 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", "PCA-in-out", "PCA(mat)-in-out", ["LI-in(1:$(i))-out(1:1)" for i in mask]...]
+
+
+    if mod_type == "linear"
+        ni_scaling = 0.0 # noise injection when eval forward map (def. 1)
+    else
+        bad_model(mod_type, mod_types)
+    end
+
+    all_errs = zeros(length(rvs_rinfos), length(names))
+    reduced_dims = zeros(Int, length(rvs_rinfos), length(names)) # store reduced dims for final table
+
+    for (idx, rv_rinfo) in enumerate(rvs_rinfos)
+        rv, rinfo = rv_rinfo
+        encoder_schedule_decorrelate_samp = [(decorrelate_sample_cov(retain_var = rv), "in_and_out")]
+        encoder_schedule_decorrelate = [(decorrelate_structure_mat(retain_var = rv), "in_and_out")]
+        encoder_schedules_li = [
+            [
+                (decorrelate_structure_mat(), "in_and_out"),
+                (likelihood_informed(retain_info = rinfo, iters = 1:i), "in"),
+                (likelihood_informed(retain_info = rinfo, iters = 1:1), "out"),
+            ] for i in mask
+        ]
+
+        em_ref = forward_map_wrapper(
+            param -> forward_map(param, model),
+            prior,
+            flat_io_pairs;
+            encoder_schedule = deepcopy(encoder_schedule_ref),
+            encoder_kwargs = deepcopy(encoder_kwargs),
+            noise_injector_scaling = ni_scaling, # shouldnt be needed
+        )
+
+        em_decorrelate_samp = forward_map_wrapper(
+            param -> forward_map(param, model),
+            prior,
+            flat_io_pairs;
+            encoder_schedule = deepcopy(encoder_schedule_decorrelate_samp),
+            encoder_kwargs = deepcopy(encoder_kwargs),
+            noise_injector_scaling = ni_scaling,
+        )
+
+        em_decorrelate = forward_map_wrapper(
+            param -> forward_map(param, model),
+            prior,
+            flat_io_pairs;
+            encoder_schedule = deepcopy(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 = deepcopy(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 (iidx, nn, em) in zip(1:length(names), names, vcat(em_ref, em_decorrelate_samp, em_decorrelate, ems_li...))
+
+            reduced_dims[idx, iidx] = get_encoded_dim(get_encoder_schedule(em), "in")
+            println(" ")
+            @info "Encoding name: $(nn)"
+            E, _ = get_encoder_from_schedule(get_encoder_schedule(em), "in")
+
+            if isnothing(E)
+                @info "No truncation"
+            else
+                ed = get_encoded_dim(get_encoder_schedule(em), "in")
+                @info "Truncation criteria (var>$(rv), or information > $(rinfo))"
+                @info "input dim reduced to: $(ed)"
+
+            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[:, 1] - v) / norm(post_means[:, 1]) for v in eachcol(post_means[:, 2:end])]'
+    end
+
+    df = DataFrame(all_errs, [Symbol(n) for n in names[1:end]])
+
+    # build dataframe
+    pairs = [(reduced_dims[i, j], all_errs[i, j]) for i in axes(all_errs, 1), j in axes(all_errs, 2)]
+
+    df = DataFrame(pairs, Symbol.(names[1:end]))  # columns = names
+    df.truncation = rvs_rinfos            # add row labels
+    return df, all_errs, reduced_dims[:, 1:end]
+end
+
+df, all_errs, red_dims = main()
+
+@info "(reduced_dim, error) of posterior mean to whitened \"reference\" solution,\n when using the truncation ({for PCA}, {for LI})"
+select!(df, :reference, :)