Replacement of Emulator processing with new encoding framework

docs/make.jl

@@ -42,6 +42,7 @@ pages = [
     "Calibrate" => "calibrate.md",
     "Emulate" => emulate,
     "Sample" => "sample.md",
+    "Data Processing and Dimension reduction" => "data_processing.md",
     "Glossary" => "glossary.md",
     "API" => api,
 ]

docs/src/API/Emulators.md

@@ -9,9 +9,8 @@ Emulator
 optimize_hyperparameters!(::Emulator)
 Emulator(::MachineLearningTool, ::PairedDataContainer{FT}) where {FT <: AbstractFloat}
 predict
-normalize
-standardize
-reverse_standardize
-svd_transform
-svd_reverse_transform_mean_cov
+encode_data
+decode_data
+encode_structure_matrix
+decode_structure_matrix
 ```

docs/src/API/GaussianProcess.md

@@ -15,7 +15,7 @@ GaussianProcess(
     ::FT,
     ::PredictionType,
 ) where {GPPkg <: GaussianProcessesPackage, K <: GaussianProcesses.Kernel, KPy <: PyObject, AK <:AbstractGPs.Kernel, FT <: AbstractFloat}
-build_models!(::GaussianProcess{GPJL}, ::PairedDataContainer{FT}) where {FT <: AbstractFloat}
+build_models!(::GaussianProcess{GPJL}, ::PairedDataContainer{FT}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat}
 optimize_hyperparameters!(::GaussianProcess{GPJL})
 predict(::GaussianProcess{GPJL},  ::AbstractMatrix{FT}) where {FT <: AbstractFloat}
 ```

docs/src/API/MarkovChainMonteCarlo.md

@@ -50,9 +50,3 @@ EmulatorPosteriorModel
 MCMCState
 accept_ratio
 ```
-
-## Internals - Other
-
-```@docs
-to_decorrelated
-```

docs/src/API/RandomFeatures.md

@@ -22,7 +22,7 @@ build_default_prior
 ```@docs
 ScalarRandomFeatureInterface
 ScalarRandomFeatureInterface(::Int,::Int)
-build_models!(::ScalarRandomFeatureInterface, ::PairedDataContainer{FT}) where {FT <: AbstractFloat}
+build_models!(::ScalarRandomFeatureInterface, ::PairedDataContainer{FT}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat}
 predict(::ScalarRandomFeatureInterface, ::M) where {M <: AbstractMatrix}
 ```
 
@@ -31,7 +31,7 @@ predict(::ScalarRandomFeatureInterface, ::M) where {M <: AbstractMatrix}
 ```@docs
 VectorRandomFeatureInterface
 VectorRandomFeatureInterface(::Int, ::Int, ::Int)
-build_models!(::VectorRandomFeatureInterface, ::PairedDataContainer{FT}) where {FT <: AbstractFloat}
+build_models!(::VectorRandomFeatureInterface, ::PairedDataContainer{FT}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat}
 predict(::VectorRandomFeatureInterface, ::M) where {M <: AbstractMatrix}
 ```
 

docs/src/data_processing.md

@@ -0,0 +1,69 @@
+# [Data Processing and Dimension Reduction](@id data-proc)
+
+## Overview 
+When working with high-dimensional problems with modest training data pairs, the bottleneck of CES procedure is the training of a competent emulator. It is often necessary to process and dimensionally-reduce the data to ease the learning task of the emulator. We provide a flexible (and extensible!) framework to create encoders and decoders for this purpose. The framework works as follows
+- An `encoder_schedule` defines the type of processing to be applied to input and/or output spaces
+- The `encoder_schedule` is passed into the `Emulator` where it is initialized and stored. 
+- The encoder will be used automatically to encode training data and predictions within the Emulate and Sample routines.
+
+An external API is also available using the `encode_data`, and `encode_structure_matrix` methods if needed.
+
+## Define an encoder schedule
+
+The user may provide an encoder schedule to transform data in a useful way. For example, mapping all output data each dimension to be bounded in [0,1]
+```julia
+simple_schedule = (minmax_scale(), "out")
+```
+The unit of the encoder contains a `DataProcessor`, the type of processing, and an string, whether to apply to input data,`"in"`, or output data `"out"`, or both `"in_and_out"`.
+
+The encoder schedule can be a vector of several units that apply multiple `DataProcessors` in order:
+```julia
+complex_schedule = [
+  (decorrelate_sample_cov(), "in"),
+  (quartile_scale(), "in"), 
+  (decorrelate_structure_mat(retain_var=0.95), "out"),
+  (canonical_correlation(), "in_and_out"),
+]
+```
+In this (rather unrealistic) chain;
+1. The inputs are decorrelated with their sample mean and covariance (and projected to low dimensional subspace if necessary) i.e PCA
+2. The scaled inputs are then subject to a "Robust" univariate scaling, mapping 1st-3rd quartiles to [0,1]
+3. The outputs are decorrelated using an "output structure matrix" (provided to the emulator `output_structure_matrix=`). Furthermore, apply a dimension-reduction to a space that retains 95% of the total variance. 
+4. In the reduced input-output space, a canonical correlation analysis is performed. Data is oriented and reduced (if necessary) maximize the joint correlation between inputs and outputs.
+
+!!! note "Default Encoder schedule"
+    The current default encoder schedule applies `decorrelate_structure_mat()` if a structure matrix (input or output) is provided, else it applies `decorrelate_sample_cov()`.
+
+!!! note "Switch encoding off"
+   To ensure that no encoding is happening, the user must pass in an empty schedule `encoder_schedule = []`
+
+
+## Creating an emulator with a schedule
+
+The schedule is then passed into the Emulator, along with the data and desired structure matrices
+```julia
+emulator = Emulator(
+    machine_learning_tool,       
+    input_output_pairs; 
+    output_structure_matrix = obs_noise_cov, 
+    encoder_schedule = complex_schedule,
+)
+```
+Note that due to the item `(decorrelate_structure_mat(retain_var=0.95), "out")` in the schedule, we must provide the `output_structure_matrix`.
+
+# Types of data processors
+
+We currently provide two main types of data processing: the `DataContainerProcessor` and `PairedDataContainerProcessor`.
+
+The `DataContainerProcessor` encodes "input" data agnostic of the "output" data, and vice versa, examples of current implementations are:
+- `UnivariateAffineScaling`: such as `quartile_scale()`, `minmax_scale()`, and `zscore_scale()`, which apply some basic univariate scaling to the data in each dimension
+- `Decorrelator`: such as `decorrelate_structure_mat()` and `decorrelate_sample_cov()`, or `decorrelate()` which perform [(truncated-)PCA](https://en.wikipedia.org/wiki/Singular_value_decomposition) using either the sample-estimated or user-provided covariance matrices (or their sum).
+
+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 (e.g., [Cui, Zahm 2021](http://doi.org/10.1088/1361-6420/abeafb), [Baptista, Marzouk, Zahm 2022](https://arxiv.org/abs/2207.08670))
+
+This is an extensible framework, and so new data processors can be added to this library.
+
+!!! note "Some experiments"
+    Some effects of the data processing (with older API) are outlined in a practical setting in the results and appendices of [Howland, Dunbar, Schneider, (2022)](https://doi.org/10.1029/2021MS002735).

docs/src/emulate.md

@@ -17,14 +17,11 @@ Wrapping a predefined machine learning tool, e.g. a Gaussian process `gauss_proc
 emulator = Emulator(
     gauss_proc, 
     input_output_pairs; # optional arguments after this
-    obs_noise_cov = Γy,
-    normalize_inputs = true,
-    standardize_outputs = true,
-    standardize_outputs_factors = factor_vector,
-    retained_svd_frac = 0.95,
+    output_structure_matrix = Γy,
+    encoder_schedule = encoder_schedule,
 )
 ```
-The optional arguments above relate to the data processing.
+The optional arguments above relate to the data processing, which is described [here](@ref data-proc)
 
 ### Emulator Training
 
@@ -40,34 +37,6 @@ y, cov = Emulator.predict(emulator, new_inputs)
 ```
 This returns both a mean value and a covariance.
 
-
-## Data processing
-
-Some effects of the following are outlined in a practical setting in the results and appendices of [Howland, Dunbar, Schneider, (2022)](https://doi.org/10.1029/2021MS002735).
-
-### Diagonalization and output dimension reduction
-
-This arises from the optional arguments
-- `obs_noise_cov = Γy` (default: `nothing`)
-We always use singular value decomposition to diagonalize the output space, requiring output covariance `Γy`. *Why?* If we need to train a $$\mathbb{R}^{10} \to \mathbb{R}^{100}$$ emulator, diagonalization allows us to instead train 100 $$\mathbb{R}^{10} \to \mathbb{R}^{1}$$ emulators (far cheaper).
-- `retained_svd_frac = 0.95` (default `1.0`)
-Performance is increased further by throwing away less informative output dimensions, if 95% of the information (i.e., variance) is in the first 40 diagonalized output dimensions then setting `retained_svd_frac=0.95` will train only 40 emulators.
-
-!!! note
-    Diagonalization is an approximation. It is however a good approximation when the observational covariance varies slowly in the parameter space.
-!!! warn
-    Severe approximation errors can occur if `obs_noise_cov` is not provided.
-
-
-### Normalization and standardization
-
-This arises from the optional arguments
-- `normalize_inputs = true` (default: `true`)
-We normalize the input data in a standard way by centering, and scaling with the empirical covariance
-- `standardize_outputs = true` (default: `false`)
-- `standardize_outputs_factors = factor_vector` (default: `nothing`)
-To help with poor conditioning of the covariance matrix, users can also standardize each output dimension with by a multiplicative factor given by the elements of `factor_vector`.
-
 ## [Modular interface](@id modular-interface)
 
 Developers may contribute new tools by performing the following

examples/Cloudy/Cloudy_calibrate.jl

@@ -160,13 +160,19 @@ dummy = ones(n_params)
 dist_type = ParticleDistributions.GammaPrimitiveParticleDistribution(dummy...)
 model_settings = DynamicalModel.ModelSettings(kernel, dist_type, moments, tspan)
 # EKI iterations
+n_iter = N_iter
 for n in 1:N_iter
     # Return transformed parameters in physical/constrained space
     ϕ_n = get_ϕ_final(priors, ekiobj)
     # Evaluate forward map
     G_n = [DynamicalModel.run_dyn_model(ϕ_n[:, i], model_settings) for i in 1:N_ens]
     G_ens = hcat(G_n...)  # reformat
-    EnsembleKalmanProcesses.update_ensemble!(ekiobj, G_ens)
+    terminate = EnsembleKalmanProcesses.update_ensemble!(ekiobj, G_ens)
+    if !isnothing(terminate)
+        n_iter = n - 1
+        break
+    end
+
 end
 
 
@@ -202,7 +208,7 @@ save(
 gr(size = (1200, 400))
 
 u_init = get_u_prior(ekiobj)
-anim_eki_unconst_cloudy = @animate for i in 1:(N_iter - 1)
+anim_eki_unconst_cloudy = @animate for i in 1:(n_iter - 1)
     u_i = get_u(ekiobj, i)
 
     p1 = plot(u_i[1, :], u_i[2, :], seriestype = :scatter, xlims = extrema(u_init[1, :]), ylims = extrema(u_init[2, :]))
@@ -259,7 +265,7 @@ gif(anim_eki_unconst_cloudy, joinpath(output_directory, "cloudy_eki_unconstr.gif
 
 # Plots in the constrained space
 ϕ_init = transform_unconstrained_to_constrained(priors, u_init)
-anim_eki_cloudy = @animate for i in 1:(N_iter - 1)
+anim_eki_cloudy = @animate for i in 1:(n_iter - 1)
     ϕ_i = get_ϕ(priors, ekiobj, i)
 
     p1 = plot(ϕ_i[1, :], ϕ_i[2, :], seriestype = :scatter, xlims = extrema(ϕ_init[1, :]), ylims = extrema(ϕ_init[2, :]))

examples/Cloudy/Cloudy_emulate_sample.jl

@@ -16,16 +16,7 @@ using CairoMakie, PairPlots
 using CalibrateEmulateSample.Emulators
 using CalibrateEmulateSample.MarkovChainMonteCarlo
 using CalibrateEmulateSample.Utilities
-using EnsembleKalmanProcesses
-using EnsembleKalmanProcesses.ParameterDistributions
-using EnsembleKalmanProcesses.DataContainers
-
-function get_standardizing_factors(data::Array{FT, 2}) where {FT}
-    # Input: data size: N_data x N_ensembles
-    # Ensemble median of the data
-    norm_factor = median(data, dims = 2) # N_data x 1 array
-    return norm_factor
-end
+using CalibrateEmulateSample.EnsembleKalmanProcesses
 
 ################################################################################
 #                                                                              #
@@ -100,7 +91,7 @@ function main()
     cases = [
         "rf-scalar",
         "gp-gpjl",  # Veeeery slow predictions
-        "rf-nosvd-nonsep",
+        "rf-nonsep",
     ]
 
     # Specify cases to run (e.g., case_mask = [2] only runs the second case)
@@ -116,13 +107,13 @@ function main()
         "verbose" => true,
         "scheduler" => DataMisfitController(terminate_at = 100.0),
         "cov_sample_multiplier" => 1.0,
-        "n_iteration" => 20,
+        "n_iteration" => 10,
     )
 
     # We use the same input-output-pairs and normalization factors for
     # Gaussian Process and Random Feature cases
     input_output_pairs = get_training_points(ekiobj, length(get_u(ekiobj)) - 2)
-    norm_factors = get_standardizing_factors(get_outputs(input_output_pairs))
+    test_pairs = get_training_points(ekiobj, (length(get_u(ekiobj)) - 1):(length(get_u(ekiobj)) - 1))
     for case in cases[case_mask]
 
         println(" ")
@@ -140,12 +131,9 @@ function main()
             # Define machine learning tool
             mlt = GaussianProcess(gppackage; kernel = gp_kernel, prediction_type = pred_type, noise_learn = false)
 
-            decorrelate = true
-            standardize_outputs = true
-
         elseif case == "rf-scalar"
 
-            kernel_rank = 3
+            kernel_rank = 1
             kernel_structure = SeparableKernel(LowRankFactor(kernel_rank, nugget), OneDimFactor())
 
             # Define machine learning tool
@@ -156,10 +144,7 @@ function main()
                 optimizer_options = optimizer_options,
             )
 
-            decorrelate = true
-            standardize_outputs = true
-
-        elseif case == "rf-nosvd-nonsep"
+        elseif case == "rf-nonsep"
 
             # Define machine learning tool
             kernel_rank = 4
@@ -171,34 +156,32 @@ function main()
                 optimizer_options = optimizer_options,
             )
 
-            # Vector RF does not require decorrelation of outputs
-            decorrelate = false
-            standardize_outputs = false
-
-
         else
             error("Case $case is not implemented yet.")
 
         end
 
-        # The data processing normalizes input data, and decorrelates
-        # output data with information from Γy, if required
-        # Note: The `standardize_outputs_factors` are only used under the
-        # condition that `standardize_outputs` is true.
+        # Data processing
+        if case == "rf-nonsep"
+            encoder_schedule = []
+        else
+            encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+        end
+
+        # build emulator
         emulator = Emulator(
             mlt,
-            input_output_pairs,
-            obs_noise_cov = Γy,
-            decorrelate = decorrelate,
-            standardize_outputs = standardize_outputs,
-            standardize_outputs_factors = vcat(norm_factors...),
+            input_output_pairs;
+            input_structure_matrix = cov(priors),
+            output_structure_matrix = Γy,
+            encoder_schedule = encoder_schedule,
         )
 
         optimize_hyperparameters!(emulator)
 
         # Check how well the emulator predicts on the true parameters
         y_mean, y_var = Emulators.predict(emulator, reshape(θ_true, :, 1); transform_to_real = true)
-
+        y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(test_pairs); transform_to_real = true)
         println("Emulator ($(case)) prediction on true parameters: ")
         println(vec(y_mean))
         println("true data: ")
@@ -207,6 +190,8 @@ function main()
         println(sqrt.(diag(y_var[1], 0)))
         println("Emulator ($(case)) MSE (truth): ")
         println(mean((truth_sample - vec(y_mean)) .^ 2))
+        println("Emulator ($(case)) MSE (next ensemble): ")
+        println(mean((get_outputs(test_pairs) - y_mean_test) .^ 2))
 
 
         ###

examples/Cloudy/Project.toml

@@ -17,4 +17,4 @@ StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
 
 [compat]
 Cloudy = "0.2"
-julia = "~1.6"
+julia = "1.6"

examples/Darcy/calibrate.jl

@@ -55,7 +55,7 @@ function main()
 
     smoothness = 0.1
     corr_length = 1.0
-    dofs = 5
+    dofs = 8
 
     grf = GRF.GaussianRandomField(
         GRF.CovarianceFunction(dim, GRF.Matern(smoothness, corr_length)),
@@ -103,19 +103,19 @@ function main()
     # We perform the inversion loop. Remember that within calls to `get_ϕ_final` the EKP transformations are applied, thus the ensemble that is returned will be the positively-bounded permeability field evaluated at all the discretization points. 
     println("Begin inversion")
     err = []
-    final_it = [N_iter]
+    n_iter = N_iter
     for i in 1:N_iter
         params_i = get_ϕ_final(prior, ekiobj)
         g_ens = run_G_ensemble(darcy, params_i)
         terminate = EKP.update_ensemble!(ekiobj, g_ens)
         push!(err, get_error(ekiobj)[end]) #mean((params_true - mean(params_i,dims=2)).^2)
         println("Iteration: " * string(i) * ", Error: " * string(err[i]))
         if !isnothing(terminate)
-            final_it[1] = i - 1
+            n_iter = i - 1
             break
         end
     end
-    n_iter = final_it[1]
+
     # We plot first the prior ensemble mean and pointwise variance of the permeability field, and also the pressure field solved with the ensemble mean. Each ensemble member is stored as a column and therefore for uses such as plotting one needs to reshape to the desired dimension.
     if PLOT_FLAG
         gr(size = (1500, 400), legend = false)