[Updated] Likelihood-informed processors

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/src/data_processing.md

@@ -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.
 

examples/DimensionReduction/Project.toml

@@ -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/calibrate.jl

@@ -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/emulate_sample.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, :)