diff --git a/docs/make.jl b/docs/make.jl index fef114079..36e591bf4 100644 --- a/docs/make.jl +++ b/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, ] diff --git a/docs/src/API/Emulators.md b/docs/src/API/Emulators.md index 6ea3d19c1..e37aede0c 100644 --- a/docs/src/API/Emulators.md +++ b/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 ``` \ No newline at end of file diff --git a/docs/src/API/GaussianProcess.md b/docs/src/API/GaussianProcess.md index 6290887c0..7f4ad96ee 100644 --- a/docs/src/API/GaussianProcess.md +++ b/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} ``` \ No newline at end of file diff --git a/docs/src/API/MarkovChainMonteCarlo.md b/docs/src/API/MarkovChainMonteCarlo.md index 64d5e5a41..f60baf453 100644 --- a/docs/src/API/MarkovChainMonteCarlo.md +++ b/docs/src/API/MarkovChainMonteCarlo.md @@ -50,9 +50,3 @@ EmulatorPosteriorModel MCMCState accept_ratio ``` - -## Internals - Other - -```@docs -to_decorrelated -``` \ No newline at end of file diff --git a/docs/src/API/RandomFeatures.md b/docs/src/API/RandomFeatures.md index 7e112f479..623798a3e 100644 --- a/docs/src/API/RandomFeatures.md +++ b/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} ``` diff --git a/docs/src/data_processing.md b/docs/src/data_processing.md new file mode 100644 index 000000000..6742f52d5 --- /dev/null +++ b/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). diff --git a/docs/src/emulate.md b/docs/src/emulate.md index a36705369..a2239eafa 100644 --- a/docs/src/emulate.md +++ b/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 diff --git a/examples/Cloudy/Cloudy_calibrate.jl b/examples/Cloudy/Cloudy_calibrate.jl index 037de2d65..1b4fd34b5 100644 --- a/examples/Cloudy/Cloudy_calibrate.jl +++ b/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, :])) diff --git a/examples/Cloudy/Cloudy_emulate_sample.jl b/examples/Cloudy/Cloudy_emulate_sample.jl index 0f3deac48..030f503be 100644 --- a/examples/Cloudy/Cloudy_emulate_sample.jl +++ b/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)) ### diff --git a/examples/Cloudy/Project.toml b/examples/Cloudy/Project.toml index 9b6148406..7d8614819 100644 --- a/examples/Cloudy/Project.toml +++ b/examples/Cloudy/Project.toml @@ -17,4 +17,4 @@ StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" [compat] Cloudy = "0.2" -julia = "~1.6" +julia = "1.6" diff --git a/examples/Darcy/calibrate.jl b/examples/Darcy/calibrate.jl index f79d5dda7..b807a6766 100644 --- a/examples/Darcy/calibrate.jl +++ b/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,7 +103,7 @@ 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) @@ -111,11 +111,11 @@ function main() 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) diff --git a/examples/Darcy/emulate_sample.jl b/examples/Darcy/emulate_sample.jl index 5ba96b8fe..ca046dbb8 100644 --- a/examples/Darcy/emulate_sample.jl +++ b/examples/Darcy/emulate_sample.jl @@ -18,6 +18,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers + include("GModel.jl") function main() @@ -90,19 +91,15 @@ function main() # Save data @save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs - retained_svd_frac = 1.0 - - normalized = true - # do we want to use SVD to decorrelate outputs - decorrelate = true + # data processing + encoding_schedule = (decorrelate_structure_mat(), "in_and_out") emulator = Emulator( mlt, input_output_pairs; - obs_noise_cov = Γy, - normalize_inputs = normalized, - retained_svd_frac = retained_svd_frac, - decorrelate = decorrelate, + input_structure_matrix = cov(prior), + output_structure_matrix = Γy, + encoding_schedule = encoding_schedule, ) optimize_hyperparameters!(emulator, kernbounds = [fill(-1e2, n_params + 1), fill(1e2, n_params + 1)]) diff --git a/examples/EDMF_data/emulator-rank-test.jl b/examples/EDMF_data/emulator-rank-test.jl index f26a07ef8..4bf51fd4d 100644 --- a/examples/EDMF_data/emulator-rank-test.jl +++ b/examples/EDMF_data/emulator-rank-test.jl @@ -20,6 +20,8 @@ using CalibrateEmulateSample.Utilities rng_seed = 42424242 Random.seed!(rng_seed) +include("rebuild_priors.jl") + function main() @@ -32,19 +34,17 @@ function main() cases = [ "GP", # diagonalize, train scalar GP, assume diag inputs "RF-prior", - "RF-vector-svd-nonsep", - "RF-vector-svd-sep", #Bad kernel for comparison - "RF-vector-nosvd-nonsep", + "RF-nonsep", + "RF-lr-lr", #Bad kernel for comparison ] rank_cases = [ [0], # not used [10], # some rank > 0 1:10, # 1:5, # input rank always 5, output rank 1:5 - 1:10, ] - case_id = 2 + case_id = 3 case = cases[case_id] rank_test = rank_cases[case_id] n_repeats = 1 @@ -157,6 +157,18 @@ function main() train_pairs = DataContainers.PairedDataContainer(train_inputs, train_outputs) end + prior_config = get_prior_config(exp_name) + means = prior_config["prior_mean"] + std = prior_config["unconstrained_σ"] + constraints = prior_config["constraints"] + + prior = combine_distributions([ + ParameterDistribution( + Dict("name" => name, "distribution" => Parameterized(Normal(mean, std)), "constraint" => constraints[name]), + ) for (name, mean) in means + ]) + + @info "Completed data loading stage" println(" ") ############################################## @@ -188,7 +200,7 @@ function main() "n_iteration" => n_iteration, "n_features_opt" => Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works "localization" => SEC(0.05), - "n_ensemble" => 400, + "n_ensemble" => 100, "n_cross_val_sets" => n_cross_val_sets, ) if case == "RF-prior" @@ -209,7 +221,7 @@ function main() prediction_type = pred_type, noise_learn = false, ) - elseif case ∈ ["RF-vector-svd-sep"] + elseif case ∈ ["RF-lr-lr"] kernel_structure = SeparableKernel(LowRankFactor(5, nugget), LowRankFactor(rank_val, nugget)) n_features = 500 @@ -222,7 +234,7 @@ function main() optimizer_options = overrides, ) - elseif case ∈ ["RF-vector-svd-nonsep", "RF-prior"] + elseif case ∈ ["RF-nonsep", "RF-prior"] kernel_structure = NonseparableKernel(LowRankFactor(rank_val, nugget)) n_features = 500 @@ -234,31 +246,20 @@ function main() kernel_structure = kernel_structure, optimizer_options = overrides, ) - elseif case ∈ ["RF-vector-nosvd-nonsep"] - kernel_structure = NonseparableKernel(LowRankFactor(rank_val, nugget)) - n_features = 500 - - mlt = VectorRandomFeatureInterface( - n_features, - input_dim, - output_dim, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - decorrelate = false end - # Fit an emulator to the data - normalized = true + # data processing: + retain_var = 0.9 + encoder_schedule = (decorrelate_structure_mat(retain_var = retain_var), "in_and_out") + # Fit an emulator to the data ttt[rank_id, rep_idx] = @elapsed begin emulator = Emulator( mlt, train_pairs; - obs_noise_cov = truth_cov, - normalize_inputs = normalized, - decorrelate = decorrelate, + input_structure_matrix = cov(prior), + output_structure_matrix = truth_cov, + encoder_schedule = encoder_schedule, ) # Optimize the GP hyperparameters for better fit diff --git a/examples/EDMF_data/plot_posterior.jl b/examples/EDMF_data/plot_posterior.jl index 0133c00c1..c7500e99e 100644 --- a/examples/EDMF_data/plot_posterior.jl +++ b/examples/EDMF_data/plot_posterior.jl @@ -19,7 +19,7 @@ function main() # 5-parameter calibration exp exp_name = "ent-det-tked-tkee-stab-calibration" - date_of_run = Date(2024, 11, 03) + date_of_run = Date(2025, 07, 03) @info "plotting results found in $(date_of_run)" # Output figure read/write directory figure_save_directory = joinpath(@__DIR__, "output", exp_name, string(date_of_run)) @@ -29,13 +29,12 @@ function main() cases = [ "GP", # diagonalize, train scalar GP, assume diag inputs "RF-prior", - "RF-vector-svd-nonsep", - "RF-vector-svd-sep", - "RF-vector-nosvd-nonsep", + "RF-nonsep", + "RF-lr-lr", ] - case_rf = cases[3] - kernel_rank = 3 - prior_kernel_rank = 3 + case_rf = cases[4] + kernel_rank = 5 + prior_kernel_rank = 5 # load posterior_filepath = joinpath(data_save_directory, "$(case_rf)_$(kernel_rank)_posterior.jld2") if !isfile(posterior_filepath) diff --git a/examples/EDMF_data/rebuild_priors.jl b/examples/EDMF_data/rebuild_priors.jl new file mode 100644 index 000000000..8d652c496 --- /dev/null +++ b/examples/EDMF_data/rebuild_priors.jl @@ -0,0 +1,39 @@ +# code deprecated due to JLD2 +#= +prior_filepath = joinpath(exp_dir, "prior.jld2") +if !isfile(prior_filepath) + LoadError("prior file \"prior.jld2\" not found in directory \"" * exp_dir * "/\"") +else + prior_dict_raw = load(prior_filepath) #using JLD2 + prior = prior_dict_raw["prior"] +end +=# +# build prior if jld2 does not work +function get_prior_config(s::SS) where {SS <: AbstractString} + config = Dict() + if s == "ent-det-calibration" + config["constraints"] = + Dict("entrainment_factor" => [bounded(0.0, 1.0)], "detrainment_factor" => [bounded(0.0, 1.0)]) + config["prior_mean"] = Dict("entrainment_factor" => 0.13, "detrainment_factor" => 0.51) + config["unconstrained_σ"] = 1.0 + elseif s == "ent-det-tked-tkee-stab-calibration" + config["constraints"] = Dict( + "entrainment_factor" => [bounded(0.0, 1.0)], + "detrainment_factor" => [bounded(0.0, 1.0)], + "tke_ed_coeff" => [bounded(0.01, 1.0)], + "tke_diss_coeff" => [bounded(0.01, 1.0)], + "static_stab_coeff" => [bounded(0.01, 1.0)], + ) + config["prior_mean"] = Dict( + "entrainment_factor" => 0.13, + "detrainment_factor" => 0.51, + "tke_ed_coeff" => 0.14, + "tke_diss_coeff" => 0.22, + "static_stab_coeff" => 0.4, + ) + config["unconstrained_σ"] = 1.0 + else + throw(ArgumentError("prior for experiment $s not found, please implement in uq_for_edmf.jl")) + end + return config +end diff --git a/examples/EDMF_data/uq_for_edmf.jl b/examples/EDMF_data/uq_for_edmf.jl index 5deb78114..a708b00f5 100644 --- a/examples/EDMF_data/uq_for_edmf.jl +++ b/examples/EDMF_data/uq_for_edmf.jl @@ -25,6 +25,7 @@ using CalibrateEmulateSample.Utilities rng_seed = 42424242 Random.seed!(rng_seed) +include("rebuild_priors.jl") function main() @@ -37,11 +38,10 @@ function main() cases = [ "GP", # diagonalize, train scalar GP, assume diag inputs "RF-prior", - "RF-vector-svd-nonsep", - "RF-vector-svd-sep", - "RF-vector-nosvd-nonsep", + "RF-nonsep", + "RF-lr-lr", ] - case = cases[3] + case = cases[4] # Output figure save directory figure_save_directory = joinpath(@__DIR__, "output", exp_name, string(Dates.today())) @@ -130,45 +130,6 @@ function main() end # load and create prior distributions - # code deprecated due to JLD2 - #= - prior_filepath = joinpath(exp_dir, "prior.jld2") - if !isfile(prior_filepath) - LoadError("prior file \"prior.jld2\" not found in directory \"" * exp_dir * "/\"") - else - prior_dict_raw = load(prior_filepath) #using JLD2 - prior = prior_dict_raw["prior"] - end - =# - # build prior if jld2 does not work - function get_prior_config(s::SS) where {SS <: AbstractString} - config = Dict() - if s == "ent-det-calibration" - config["constraints"] = - Dict("entrainment_factor" => [bounded(0.0, 1.0)], "detrainment_factor" => [bounded(0.0, 1.0)]) - config["prior_mean"] = Dict("entrainment_factor" => 0.13, "detrainment_factor" => 0.51) - config["unconstrained_σ"] = 1.0 - elseif s == "ent-det-tked-tkee-stab-calibration" - config["constraints"] = Dict( - "entrainment_factor" => [bounded(0.0, 1.0)], - "detrainment_factor" => [bounded(0.0, 1.0)], - "tke_ed_coeff" => [bounded(0.01, 1.0)], - "tke_diss_coeff" => [bounded(0.01, 1.0)], - "static_stab_coeff" => [bounded(0.01, 1.0)], - ) - config["prior_mean"] = Dict( - "entrainment_factor" => 0.13, - "detrainment_factor" => 0.51, - "tke_ed_coeff" => 0.14, - "tke_diss_coeff" => 0.22, - "static_stab_coeff" => 0.4, - ) - config["unconstrained_σ"] = 1.0 - else - throw(ArgumentError("prior for experiment $s not found, please implement in uq_for_edmf.jl")) - end - return config - end prior_config = get_prior_config(exp_name) means = prior_config["prior_mean"] std = prior_config["unconstrained_σ"] @@ -199,7 +160,7 @@ function main() @info "Begin Emulation stage" # Create GP object train_frac = 0.9 - kernel_rank = 3 # svd-sep should be kr - 5 (as input rank is set to 5) + kernel_rank = 5 # svd-sep should be kr - 5 (as input rank is set to 5) n_cross_val_sets = 2 @info "Kernel rank: $(kernel_rank)" max_feature_size = @@ -210,9 +171,9 @@ function main() "scheduler" => DataMisfitController(terminate_at = 1000), "cov_sample_multiplier" => 0.2, "n_iteration" => 15, - "n_features_opt" => Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works + "n_features_opt" => 100, #Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works "localization" => SEC(0.05), - "n_ensemble" => 400, + "n_ensemble" => 100,#400, "n_cross_val_sets" => n_cross_val_sets, ) if case == "RF-prior" @@ -235,8 +196,8 @@ function main() prediction_type = pred_type, noise_learn = false, ) - elseif case ∈ ["RF-vector-svd-sep"] - kernel_structure = SeparableKernel(LowRankFactor(5, nugget), LowRankFactor(kernel_rank, nugget)) + elseif case ∈ ["RF-lr-lr"] + kernel_structure = SeparableKernel(LowRankFactor(1, nugget), LowRankFactor(kernel_rank, nugget)) n_features = 500 mlt = VectorRandomFeatureInterface( @@ -248,19 +209,7 @@ function main() optimizer_options = overrides, ) - elseif case ∈ ["RF-vector-svd-nonsep", "RF-prior"] - kernel_structure = NonseparableKernel(LowRankFactor(kernel_rank, nugget)) - n_features = 500 - - mlt = VectorRandomFeatureInterface( - n_features, - input_dim, - output_dim, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - elseif case ∈ ["RF-vector-nosvd-nonsep"] + elseif case ∈ ["RF-nonsep", "RF-prior"] kernel_structure = NonseparableKernel(LowRankFactor(kernel_rank, nugget)) n_features = 500 @@ -272,52 +221,27 @@ function main() kernel_structure = kernel_structure, optimizer_options = overrides, ) - decorrelate = false end - # Fit an emulator to the data - normalized = true + # data processing: + retain_var = 0.9 + encoder_schedule = (decorrelate_structure_mat(retain_var = retain_var), "in_and_out") + # Fit an emulator to the data emulator = Emulator( mlt, input_output_pairs; - obs_noise_cov = truth_cov, - normalize_inputs = normalized, - decorrelate = decorrelate, + input_structure_matrix = cov(prior), + output_structure_matrix = truth_cov, + encoder_schedule = encoder_schedule, ) # Optimize the GP hyperparameters for better fit optimize_hyperparameters!(emulator) - if case ∈ ["RF-vector-nosvd-nonsep", "RF-vector-svd-nonsep"] + if case ∈ ["RF-nonsep"] push!(opt_diagnostics, get_optimizer(mlt)[1]) #length-1 vec of vec -> vec end - # plot eki convergence plot - #= - if length(opt_diagnostics) > 0 - err_cols = reduce(hcat, opt_diagnostics) #error for each repeat as columns? - - #save data - error_filepath = joinpath(data_save_directory, "eki_conv_error.jld2") - save(error_filepath, "error", err_cols) - - # print all repeats - f5 = Figure(resolution = (1.618 * 300, 300), markersize = 4) - ax_conv = Axis(f5[1, 1], xlabel = "Iteration", ylabel = "max-normalized error") - if n_repeats == 1 - lines!(ax_conv, collect(1:size(err_cols, 1))[:], err_cols[:], solid_color = :blue) # If just one repeat - else - for idx in 1:size(err_cols, 1) - err_normalized = (err_cols' ./ err_cols[1, :])' # divide each series by the max, so all errors start at 1 - series!(ax_conv, err_normalized', solid_color = :blue) - end - end - save(joinpath(figure_save_directory, "eki-conv_$(case).png"), f5, px_per_unit = 3) - save(joinpath(figure_save_directory, "eki-conv_$(case).pdf"), f5, px_per_unit = 3) - - end - =# - emulator_filepath = joinpath(data_save_directory, "$(case)_$(kernel_rank)_emulator.jld2") save(emulator_filepath, "emulator", emulator) diff --git a/examples/Emulator/G-function/emulate-test-n-features.jl b/examples/Emulator/G-function/emulate-test-n-features.jl index 042694fbe..2dc1c5776 100644 --- a/examples/Emulator/G-function/emulate-test-n-features.jl +++ b/examples/Emulator/G-function/emulate-test-n-features.jl @@ -12,6 +12,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.Utilities using CairoMakie, ColorSchemes #for plots seed = 2589436 @@ -92,10 +93,9 @@ function main() TV = [(1 / (3 * (1 + ai)^2)) * prod_tmp2[i] / prod_tmp for (i, ai) in enumerate(a)] - cases = ["Prior", "GP", "RF-scalar"] case = cases[2] - decorrelate = true + encoder_schedule = (decorrelate_structure_mat(), "out") nugget = Float64(1e-12) n_features_vec = [25, 50, 100, 200, 400, 800] # < data points ideally as output dim = 1 @@ -116,7 +116,6 @@ function main() "n_iteration" => 20, "cov_sample_multiplier" => 1.0, # "localization" => SEC(0.1),#,Doesnt help much tbh - # "accelerator" => NesterovAccelerator(), "n_ensemble" => 100, "n_cross_val_sets" => n_cross_val_sets, ) @@ -152,7 +151,12 @@ function main() # Emulate ttt[f_idx, rep_idx] = @elapsed begin - emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate) + emulator = Emulator( + mlt, + iopairs; + output_structure_matrix = Γ * I, + encoder_schedule = deepcopy(encoder_schedule), + ) optimize_hyperparameters!(emulator) end diff --git a/examples/Emulator/G-function/emulate.jl b/examples/Emulator/G-function/emulate.jl index 23007f5d6..3cc3ab123 100644 --- a/examples/Emulator/G-function/emulate.jl +++ b/examples/Emulator/G-function/emulate.jl @@ -12,6 +12,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.Utilities using CairoMakie, ColorSchemes #for plots seed = 2589436 @@ -39,8 +40,8 @@ function main() rng = MersenneTwister(seed) - n_repeats = 20# repeat exp with same data. - n_dimensions = 20 + n_repeats = 5# repeat exp with same data. + n_dimensions = 10 # To create the sampling n_data_gen = 800 @@ -80,6 +81,7 @@ function main() input[:, i] = samples[ind[i], :] output[i] = y[ind[i]] + noise[i] end + encoder_schedule = (decorrelate_structure_mat(), "out") iopairs = PairedDataContainer(input, output) # analytic sobol indices taken from @@ -90,21 +92,17 @@ function main() prod_tmp2 = [prod(1 .+ 1 ./ (3 .* (1 .+ a[1:end .!== j]) .^ 2)) for j in 1:n_dimensions] TV = [(1 / (3 * (1 + ai)^2)) * prod_tmp2[i] / prod_tmp for (i, ai) in enumerate(a)] - - cases = ["Prior", "GP", "RF-scalar"] case = cases[3] - decorrelate = true nugget = Float64(1e-12) overrides = Dict( "verbose" => true, "scheduler" => DataMisfitController(terminate_at = 1e2), - "n_features_opt" => 150, + "n_features_opt" => 120, "n_iteration" => 10, "cov_sample_multiplier" => 3.0, #"localization" => SEC(0.1),#,Doesnt help much tbh - #"accelerator" => NesterovAccelerator(), "n_ensemble" => 100, #40*n_dimensions, "n_cross_val_sets" => 4, ) @@ -144,7 +142,12 @@ function main() # Emulate times[rep_idx] = @elapsed begin - emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate) + emulator = Emulator( + mlt, + iopairs; + output_structure_matrix = Γ * I, + encoder_schedule = deepcopy(encoder_schedule), + ) optimize_hyperparameters!(emulator) end @@ -307,7 +310,7 @@ function main() else for idx in 1:size(err_cols, 1) err_normalized = (err_cols' ./ err_cols[1, :])' # divide each series by the max, so all errors start at 1 - series!(ax_conv, err_normalized', color = :blue) + series!(ax_conv, err_normalized', solid_color = :blue) end end diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl index 45e04acd1..74421011d 100644 --- a/examples/Emulator/Ishigami/emulate.jl +++ b/examples/Emulator/Ishigami/emulate.jl @@ -11,6 +11,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.Utilities using JLD2 @@ -80,18 +81,17 @@ function main() output[i] = y[ind[i]] + noise[i] end iopairs = PairedDataContainer(input, output) - cases = ["Prior", "GP", "RF-scalar"] - case = cases[3] - decorrelate = true - nugget = Float64(1e-12) + case = cases[2] + encoder_schedule = (decorrelate_structure_mat(), "out") + nugget = Float64(1e-4) overrides = Dict( "scheduler" => DataMisfitController(terminate_at = 1e4), "n_features_opt" => 150, - "n_ensemble" => 30, + "n_ensemble" => 50, "n_iteration" => 20, - "accelerator" => NesterovAccelerator(), ) + if case == "Prior" # don't do anything overrides["n_iteration"] = 0 @@ -123,7 +123,8 @@ function main() end # Emulate - emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate) + emulator = + Emulator(mlt, iopairs; output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule)) optimize_hyperparameters!(emulator) # get EKP errors - just stored in "optimizer" box for now diff --git a/examples/Emulator/L63/emulate.jl b/examples/Emulator/L63/emulate.jl index 64ee4b6c3..b623a89b5 100644 --- a/examples/Emulator/L63/emulate.jl +++ b/examples/Emulator/L63/emulate.jl @@ -8,7 +8,7 @@ using JLD2 using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.DataContainers -using EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.Utilities function lorenz(du, u, p, t) du[1] = 10.0 * (u[2] - u[1]) @@ -26,7 +26,7 @@ function main() # rng rng = MersenneTwister(1232435) - n_repeats = 20 # repeat exp with same data. + n_repeats = 1 # repeat exp with same data. println("run experiment $n_repeats times") #for later plots @@ -107,7 +107,7 @@ function main() "RF-svd-sep", ] - case = cases[7] + case = cases[3] # 7 nugget = Float64(1e-8) u_test = [] @@ -121,8 +121,6 @@ function main() "cov_sample_multiplier" => 1.0, #5.0, "n_features_opt" => 150, "n_iteration" => 10, - #"accelerator" => DefaultAccelerator(), - #"localization" => EnsembleKalmanProcesses.Localizers.SECNice(0.01,1.0), # localization / s "n_ensemble" => 200, "verbose" => true, "n_cross_val_sets" => 2, @@ -213,11 +211,11 @@ function main() # Emulate if case ∈ ["RF-nosvd-nonsep", "RF-nosvd-sep"] - decorrelate = false + encoder_schedule = [] # no encoding else - decorrelate = true + encoder_schedule = (decorrelate_structure_mat(), "out") end - emulator = Emulator(mlt, iopairs; obs_noise_cov = Γy, decorrelate = decorrelate) + emulator = Emulator(mlt, iopairs; output_structure_matrix = Γy, encoder_schedule = deepcopy(encoder_schedule)) optimize_hyperparameters!(emulator) # diagnostics diff --git a/examples/Emulator/L63/emulate_diff-rank-test.jl b/examples/Emulator/L63/emulate_diff-rank-test.jl index 222147bf2..e60b7ebcc 100644 --- a/examples/Emulator/L63/emulate_diff-rank-test.jl +++ b/examples/Emulator/L63/emulate_diff-rank-test.jl @@ -8,6 +8,7 @@ using JLD2 using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Utilities using EnsembleKalmanProcesses.Localizers function lorenz(du, u, p, t) @@ -185,12 +186,17 @@ function main() # Emulate if case ∈ ["RF-nosvd-nonsep", "RF-nosvd-sep"] - decorrelate = false + encoder_schedule = [] else - decorrelate = true + encoder_schedule = nothing end ttt[rank_id, rep_idx] = @elapsed begin - emulator = Emulator(mlt, iopairs; obs_noise_cov = Γy, decorrelate = decorrelate) + emulator = Emulator( + mlt, + iopairs; + output_structure_matrix = Γy, + encoder_schedule = deepcopy(encoder_schedule), + ) optimize_hyperparameters!(emulator) end diff --git a/examples/Emulator/L63/plot_results.jl b/examples/Emulator/L63/plot_results.jl index 923eccf6a..bcf347b4f 100644 --- a/examples/Emulator/L63/plot_results.jl +++ b/examples/Emulator/L63/plot_results.jl @@ -7,7 +7,7 @@ function main() # filepaths output_directory = joinpath(@__DIR__, "output") cases = ["GP", "RF-prior", "RF-scalar", "RF-scalar-diagin", "RF-svd-nonsep", "RF-nosvd-nonsep", "RF-nosvd-sep"] #for paper, 1 2 & 5. - case = cases[1] + case = cases[3] @info "plotting case $case" #for later plots fontsize = 20 diff --git a/examples/Emulator/Regression_2d_2d/Project.toml b/examples/Emulator/Regression_2d_2d/Project.toml index 9d41c3993..d1ef6639c 100644 --- a/examples/Emulator/Regression_2d_2d/Project.toml +++ b/examples/Emulator/Regression_2d_2d/Project.toml @@ -11,5 +11,5 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [compat] -FiniteDiff = "~2.10" -julia = "~1.6" +FiniteDiff = "2.10" +julia = "1.6" diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl index df9282d32..fd55ca02d 100644 --- a/examples/Emulator/Regression_2d_2d/compare_regression.jl +++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl @@ -9,6 +9,7 @@ using Statistics using LinearAlgebra using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Utilities using CalibrateEmulateSample.EnsembleKalmanProcesses plot_flag = true if plot_flag @@ -41,24 +42,37 @@ function main() cases = [ "gp-skljl", "gp-gpjl", # Very slow prediction... - "rf-scalar", - "rf-svd-diag", - "rf-svd-nondiag", - "rf-nosvd-diag", - "rf-nosvd-nondiag", - "rf-svd-nonsep", - "rf-nosvd-nonsep", + "rf-lr-scalar", + "rf-lr-lr", + "rf-nonsep", ] - case_mask = [1, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite) + + encoder_names = ["no-proc", "sample-struct-proc", "sample-proc", "struct-mat-proc", "combined-proc", "cca-proc"] + encoders = [ + [], # no proc + nothing, # default proc + (decorrelate_sample_cov(), "in_and_out"), + (decorrelate_structure_mat(), "in_and_out"), + (decorrelate(), "in_and_out"), + (canonical_correlation(), "in_and_out"), + ] + + # USER CHOICES + encoder_mask = [3, 4, 6] # best performing + case_mask = [1, 3, 4, 5] # (KEEP set to 1:length(cases) when pushing for buildkite) + #problem n = 200 # number of training points p = 2 # input dim d = 2 # output dim - X = 2.0 * π * rand(p, n) + prior_cov = (pi^2) * I(p) + X = rand(MvNormal(zeros(p), prior_cov), n) # G(x1, x2) - g1(x) = sin.(x[1, :]) .+ 2 * cos.(2 * x[2, :]) - g2(x) = 3 * sin.(3 * x[1, :]) .- 4 * cos.(4 * x[2, :]) + g1(x) = sin.(x[1, :]) .+ cos.(2 * x[2, :]) + g2(x) = 2 * sin.(2 * x[1, :]) .- 3 * cos.(x[2, :]) + g1(x, y) = sin(x) + 2 * cos(2 * y) + g2(x, y) = 2 * sin(2 * x) - 3 * cos(y) g1x = g1(X) g2x = g2(X) gx = zeros(2, n) @@ -67,7 +81,7 @@ function main() # Add noise η μ = zeros(d) - Σ = 0.05 * [[0.8, 0.1] [0.1, 0.5]] # d x d + Σ = 0.1 * [[0.4, -0.3] [-0.3, 0.5]] # d x d noise_samples = rand(MvNormal(μ, Σ), n) # y = G(x) + η Y = gx .+ noise_samples @@ -78,258 +92,153 @@ function main() #plot training data with and without noise if plot_flag - p1 = plot( - X[1, :], - X[2, :], - g1x, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zguidefontrotation = 90, - ) - - figpath = joinpath(output_directory, "observed_y1nonoise.png") - savefig(figpath) + n_pts = 200 + x1 = range(-2 * pi, stop = 2 * π, length = n_pts) + x2 = range(-2 * pi, stop = 2 * π, length = n_pts) - p2 = plot( - X[1, :], - X[2, :], - g2x, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zguidefontrotation = 90, - ) - figpath = joinpath(output_directory, "observed_y2nonoise.png") - savefig(figpath) + p1 = contourf(x1, x2, g1.(x1', x2), c = :cividis, xlabel = "x1", ylabel = "x2", aspect_ratio = :equal) + scatter!(p1, X[1, :], X[2, :], c = :black, ms = 3, label = "train point") - p3 = plot( - X[1, :], - X[2, :], - Y[1, :], - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zguidefontrotation = 90, - ) - figpath = joinpath(output_directory, "observed_y1.png") + figpath = joinpath(output_directory, "g1_true.png") savefig(figpath) + p2 = contourf(x1, x2, g2.(x1', x2), c = :cividis, xlabel = "x1", ylabel = "x2", aspect_ratio = :equal) + scatter!(p2, X[1, :], X[2, :], c = :black, ms = 3, label = "train point") - p4 = plot( - X[1, :], - X[2, :], - Y[2, :], - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zguidefontrotation = 90, - ) - figpath = joinpath(output_directory, "observed_y2.png") + figpath = joinpath(output_directory, "g2_true.png") savefig(figpath) end for case in cases[case_mask] - - println(" ") - @info "running case $case" - - # common Gaussian feature setup - pred_type = YType() - - # common random feature setup - n_features = 300 - optimizer_options = - Dict("n_iteration" => 20, "n_features_opt" => 100, "n_ensemble" => 80, "cov_sample_multiplier" => 1.0) - nugget = 1e-12 - - - - if case == "gp-skljl" - gppackage = SKLJL() - gaussian_process = GaussianProcess(gppackage, noise_learn = true) - emulator = Emulator(gaussian_process, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - elseif case == "gp-gpjl" - @warn "gp-gpjl case is very slow at prediction" - gppackage = GPJL() - gaussian_process = GaussianProcess(gppackage, noise_learn = true) - emulator = Emulator(gaussian_process, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - elseif case == "rf-scalar" - srfi = ScalarRandomFeatureInterface( - n_features, - p, - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()), - optimizer_options = optimizer_options, - ) - emulator = Emulator(srfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - elseif case == "rf-svd-diag" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), DiagonalFactor()), - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - elseif case == "rf-svd-nondiag" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), LowRankFactor(2, nugget)), - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - elseif case == "rf-svd-nonsep" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = NonseparableKernel(LowRankFactor(4, nugget)), - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true) - - elseif case == "rf-nosvd-diag" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), DiagonalFactor()), # roughly normalize output by noise - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true, decorrelate = false) - elseif case == "rf-nosvd-nondiag" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), LowRankFactor(2, nugget)), # roughly normalize output by noise - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true, decorrelate = false) - elseif case == "rf-nosvd-nonsep" - vrfi = VectorRandomFeatureInterface( - n_features, - p, - d, - kernel_structure = NonseparableKernel(LowRankFactor(4, nugget)), - optimizer_options = deepcopy(optimizer_options), - ) - emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true, decorrelate = false) - - end - println("build RF with $n training points and $(n_features) random features.") - - optimize_hyperparameters!(emulator) # although RF already optimized - - # Plot mean and variance of the predicted observables y1 and y2 - # For this, we generate test points on a x1-x2 grid. - n_pts = 200 - x1 = range(0.0, stop = 2 * π, length = n_pts) - x2 = range(0.0, stop = 2 * π, length = n_pts) - X1, X2 = meshgrid(x1, x2) - # Input for predict has to be of size input_dim x N_samples - inputs = permutedims(hcat(X1[:], X2[:]), (2, 1)) - - em_mean, em_cov = predict(emulator, inputs, transform_to_real = true) - println("end predictions at ", n_pts * n_pts, " points") - - - #plot predictions - for y_i in 1:d - - em_var_temp = [diag(em_cov[j]) for j in 1:length(em_cov)] # (40000,) - em_var = permutedims(reduce(vcat, [x' for x in em_var_temp]), (2, 1)) # 2 x 40000 - - mean_grid = reshape(em_mean[y_i, :], n_pts, n_pts) # 2 x 40000 - if plot_flag - p5 = plot( - x1, - x2, - mean_grid, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zlabel = "mean of y" * string(y_i), - zguidefontrotation = 90, + for (enc_name, enc) in zip(encoder_names[encoder_mask], encoders[encoder_mask]) + + println(" ") + @info """ + + ------------- + running case $case with encoder $enc_name + ------------- + """ + + # common Gaussian feature setup + pred_type = YType() + + # common random feature setup + n_features = 200 + optimizer_options = + Dict("n_iteration" => 5, "n_features_opt" => 200, "n_ensemble" => 80, "cov_sample_multiplier" => 5.0) + nugget = 1e-8 + + # data processing schedule + + if case == "gp-skljl" + gppackage = SKLJL() + mlt = GaussianProcess(gppackage, noise_learn = true) + + elseif case == "gp-gpjl" + @warn "gp-gpjl case is very slow at prediction" + gppackage = GPJL() + mlt = GaussianProcess(gppackage, noise_learn = true) + + elseif case == "rf-lr-scalar" + mlt = ScalarRandomFeatureInterface( + n_features, + p, + kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()), + optimizer_options = optimizer_options, ) - end - var_grid = reshape(em_var[y_i, :], n_pts, n_pts) - if plot_flag - p6 = plot( - x1, - x2, - var_grid, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zlabel = "var of y" * string(y_i), - zguidefontrotation = 90, + elseif case == "rf-lr-lr" + mlt = VectorRandomFeatureInterface( + n_features, + p, + d, + kernel_structure = SeparableKernel(LowRankFactor(2, nugget), LowRankFactor(2, nugget)), + optimizer_options = deepcopy(optimizer_options), + ) + elseif case == "rf-nonsep" + mlt = VectorRandomFeatureInterface( + n_features, + p, + d, + kernel_structure = NonseparableKernel(LowRankFactor(4, nugget)), + optimizer_options = deepcopy(optimizer_options), ) - plot(p5, p6, layout = (1, 2), legend = false) - - savefig(joinpath(output_directory, case * "_y" * string(y_i) * "_predictions.png")) end - end - # Plot the true components of G(x1, x2) - g1_true = g1(inputs) - g1_true_grid = reshape(g1_true, n_pts, n_pts) - if plot_flag - p7 = plot( - x1, - x2, - g1_true_grid, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zlabel = "sin(x1) + cos(x2)", - zguidefontrotation = 90, + emulator = Emulator( + mlt, + iopairs, + encoder_schedule = enc, + input_structure_matrix = prior_cov, + output_structure_matrix = Σ, ) - savefig(joinpath(output_directory, case * "_true_g1.png")) - end - g2_true = g2(inputs) - g2_true_grid = reshape(g2_true, n_pts, n_pts) - if plot_flag - p8 = plot( - x1, - x2, - g2_true_grid, - st = :surface, - camera = (30, 60), - c = :cividis, - xlabel = "x1", - ylabel = "x2", - zlabel = "sin(x1) - cos(x2)", - zguidefontrotation = 90, - ) - g_true_grids = [g1_true_grid, g2_true_grid] - savefig(joinpath(output_directory, case * "_true_g2.png")) + optimize_hyperparameters!(emulator) # although RF already optimized + + # Plot mean and variance of the predicted observables y1 and y2 + # For this, we generate test points on a x1-x2 grid. + n_pts = 200 + x1 = range(-2 * π, stop = 2 * π, length = n_pts) + x2 = range(-2 * π, stop = 2 * π, length = n_pts) + X1, X2 = meshgrid(x1, x2) + # Input for predict has to be of size input_dim x N_samples + inputs = permutedims(hcat(X1[:], X2[:]), (2, 1)) + + em_mean, em_cov = predict(emulator, inputs, transform_to_real = true) + println("end predictions at ", n_pts * n_pts, " points") + + + #plot predictions + for y_i in 1:d + + em_var_temp = [diag(em_cov[j]) for j in 1:length(em_cov)] # (40000,) + em_var = permutedims(reduce(vcat, [x' for x in em_var_temp]), (2, 1)) # 2 x 40000 + + mean_grid = reshape(em_mean[y_i, :], n_pts, n_pts) # 2 x 40000 + if plot_flag + p5 = contourf( + x1, + x2, + mean_grid, + c = :cividis, + xlabel = "x1", + ylabel = "x2", + title = "mean output $y_i", + ) + end + var_grid = reshape(em_var[y_i, :], n_pts, n_pts) + if plot_flag + p6 = contourf( + x1, + x2, + var_grid, + c = :cividis, + xlabel = "x1", + ylabel = "x2", + title = "var output $y_i", + ) + + plot(p5, p6, layout = (1, 2), legend = false) + + savefig(joinpath(output_directory, case * "_" * enc_name * "_y" * string(y_i) * "_predictions.png")) + end + end - end - MSE = 1 / size(em_mean, 2) * sqrt(sum((em_mean[1, :] - g1_true) .^ 2 + (em_mean[2, :] - g2_true) .^ 2)) - println("L^2 error of mean and latent truth:", MSE) + # Plot the true components of G(x1, x2) + g1_true = g1(inputs) + g1_true_grid = reshape(g1_true, n_pts, n_pts) - # Plot the difference between the truth and the mean of the predictions - for y_i in 1:d + + g2_true = g2(inputs) + g2_true_grid = reshape(g2_true, n_pts, n_pts) + g_true_grids = [g1_true_grid, g2_true_grid] + MSE = 1 / size(em_mean, 2) * sqrt(sum((em_mean[1, :] - g1_true) .^ 2 + (em_mean[2, :] - g2_true) .^ 2)) + println("L^2 error of mean and latent truth:", MSE) + #= + # Plot the difference between the truth and the mean of the predictions + for y_i in 1:d # Reshape em_cov to size N_samples x output_dim em_var_temp = [diag(em_cov[j]) for j in 1:length(em_cov)] # (40000,) @@ -342,21 +251,21 @@ function main() if plot_flag zlabel = "1/var * (true_y" * string(y_i) * " - predicted_y" * string(y_i) * ")^2" - p9 = plot( + p9 = contourf( x1, x2, sqrt.(1.0 ./ var_grid .* (g_true_grids[y_i] .- mean_grid) .^ 2), - st = :surface, - camera = (30, 60), c = :magma, zlabel = zlabel, xlabel = "x1", ylabel = "x2", - zguidefontrotation = 90, + title = "weighted difference $y_i" ) - savefig(joinpath(output_directory, case * "_y" * string(y_i) * "_difference_truth_prediction.png")) + savefig(joinpath(output_directory, case*"_"* enc_name * "_y" * string(y_i) * "_difference_truth_prediction.png")) + end end + =# end end end diff --git a/examples/GCM/emulate_sample_script.jl b/examples/GCM/emulate_sample_script.jl index 4c0f9ca9c..07098b641 100644 --- a/examples/GCM/emulate_sample_script.jl +++ b/examples/GCM/emulate_sample_script.jl @@ -52,6 +52,10 @@ function main() truth = load(datafile)["truth"] # 96 obs_noise_cov = load(datafile)["obs_noise_cov"] # 96 x 96 + priorfile = "priors.jld2" + prior = load(priorfile)["prior"] + + #take only first 400 points iter_mask = [1, 2, 4] data_mask = 1:96 @@ -73,7 +77,6 @@ function main() stacked_inputs = reshape(permutedims(inputs, (1, 3, 2)), (N_ens * N_iter, input_dim)) stacked_outputs = reshape(permutedims(outputs, (1, 3, 2)), (N_ens * N_iter, output_dim)) input_output_pairs = PairedDataContainer(stacked_inputs, stacked_outputs, data_are_columns = false) #data are rows - normalized = true # setup random features eki_options_override = Dict( @@ -155,26 +158,24 @@ function main() end if emulator_type == "VectorRFR-nondiag" - #standardizing with data median for each data object seems reasonable in this setting. - standards = vec( - hcat( - median(stacked_outputs[:, 1:32], dims = 1), - median(stacked_outputs[:, 33:64], dims = 1), - median(stacked_outputs[:, 65:96], dims = 1), - ), - ) - println(standards) + + encoder_schedule = [(quartile_scale(), "in"), (decorrelate_structure_mat(), "out")] emulator = Emulator( mlt, input_output_pairs; - obs_noise_cov = obs_noise_cov, - normalize_inputs = normalized, - standardize_outputs = true, - standardize_outputs_factors = standards, - decorrelate = false, + output_structure_matrix = obs_noise_cov, + encoder_schedule = encoder_schedule, ) else - emulator = Emulator(mlt, input_output_pairs; obs_noise_cov = obs_noise_cov, normalize_inputs = normalized) + encoder_schedule = [(decorrelate_structure_mat(), "in_and_out")] + + emulator = Emulator( + mlt, + input_output_pairs; + input_structure_matrix = cov(prior), + output_structure_matrix = obs_noise_cov, + encoder_schedule = encoder_schedule, + ) end optimize_hyperparameters!(emulator) @@ -267,8 +268,6 @@ function main() ### MCMC - priorfile = "priors.jld2" - prior = load(priorfile)["prior"] ## ### Sample: Markov Chain Monte Carlo diff --git a/examples/Lorenz/calibrate.jl b/examples/Lorenz/calibrate.jl index c546d0491..bff342e4c 100644 --- a/examples/Lorenz/calibrate.jl +++ b/examples/Lorenz/calibrate.jl @@ -279,19 +279,19 @@ function main() # EKI iterations println("EKP inversion error:") err = zeros(N_iter) - final_iter = [N_iter] + n_iter = N_iter for i in 1:N_iter params_i = EKP.get_ϕ_final(priors, ekiobj) # the `ϕ` indicates that the `params_i` are in the constrained space g_ens = GModel.run_G_ensemble(params_i, lorenz_settings_G) terminated = EKP.update_ensemble!(ekiobj, g_ens) if !isnothing(terminated) - final_iter = i - 1 # final update was previous iteration + n_iter = i - 1 # final update was previous iteration break end err[i] = EKP.get_error(ekiobj)[end] #mean((params_true - mean(params_i,dims=2)).^2) println("Iteration: " * string(i) * ", Error: " * string(err[i])) end - N_iter = final_iter[1] #in case it terminated early + n_iter #in case it terminated early # EKI results: Has the ensemble collapsed toward the truth? println("True parameters: ") @@ -331,7 +331,7 @@ function main() yaxis = "A", ylims = extrema(ϕ_init[2, :]), ) - for i in 1:N_iter + for i in 1:n_iter ϕ_i = EKP.get_ϕ(priors, ekiobj, i) scatter!(p, ϕ_i[1, :], ϕ_i[2, :], color = :grey, label = false) end diff --git a/examples/Lorenz/calibrate_spatial_dep.jl b/examples/Lorenz/calibrate_spatial_dep.jl index d39ab9f6a..25152046c 100644 --- a/examples/Lorenz/calibrate_spatial_dep.jl +++ b/examples/Lorenz/calibrate_spatial_dep.jl @@ -169,12 +169,12 @@ x0 = x_spun_up[:, end] #last element of the run is the initial condition for cr #Creating sythetic data -T = 14.0 +T = 50.0 ny = nx * 2 #number of data points lorenz_config_settings = LorenzConfig(dt, T) # construct how we compute Observations -T_start = 4.0 +T_start = T - 25.0 T_end = T observation_config = ObservationConfig(T_start, T_end) @@ -193,11 +193,11 @@ y_ens = hcat( ) # estimate noise as the variability of these pushed-forward samples on the attractor -obs_noise_cov = cov(y_ens, dims = 2) +obs_noise_cov = cov(y_ens, dims = 2) + 1e-2 * I y = y_ens[:, 1] -pl = 2.0 -psig = 3.0 +pl = 4.0 +psig = 5.0 #Prior covariance B = zeros(nx, nx) for ii in 1:nx @@ -208,7 +208,7 @@ end B_sqrt = sqrt(B) #Prior mean -mu = 8.0 * ones(nx) +mu = 4.0 * ones(nx) #Creating prior distribution distribution = Parameterized(MvNormal(mu, B)) @@ -232,13 +232,15 @@ rng = MersenneTwister(rng_seed) # initial parameters: N_params x N_ens initial_params = construct_initial_ensemble(rng, prior, N_ens) -method = Inversion() +method = Inversion(prior) @info "Ensemble size: $(N_ens)" ekpobj = EKP.EnsembleKalmanProcess(initial_params, y, obs_noise_cov, method; rng = copy(rng), verbose = true) count = 0 -n_iter = 0 +n_iter = N_iter +n_samples_exp = N_iter * N_ens +x_on_attractor = x_spun_up[:, shuffled_ids[1:n_samples_exp]] # randomly select points from second half of spin up for i in 1:N_iter params_i = get_ϕ_final(prior, ekpobj) @@ -247,7 +249,7 @@ for i in 1:N_iter [ lorenz_forward( EnsembleMemberConfig(params_i[:, j]), - (x0 .+ ic_cov_sqrt * randn(rng, nx)), + x_on_attractor[:, (i - 1) * N_ens + j], lorenz_config_settings, observation_config, ) for j in 1:N_ens diff --git a/examples/Lorenz/emulate_sample.jl b/examples/Lorenz/emulate_sample.jl index c9fb04feb..02ed90d5a 100644 --- a/examples/Lorenz/emulate_sample.jl +++ b/examples/Lorenz/emulate_sample.jl @@ -5,8 +5,8 @@ include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) using Distributions # probability distributions and associated functions using LinearAlgebra ENV["GKSwstype"] = "100" -using StatsPlots using Plots +using StatsPlots using Random using JLD2 @@ -18,36 +18,17 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.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) - return norm_factor -end - -function get_standardizing_factors(data::Array{FT, 1}) where {FT} - # Input: data size: N_data*N_ensembles (splatted) - # Ensemble median of the data - norm_factor = median(data) - return norm_factor -end - - function main() cases = [ - "GP", # diagonalize, train scalar GP, assume diag inputs - "RF-scalar-diagin", # diagonalize, train scalar RF, assume diag inputs (most comparable to GP) - "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs - "RF-vector-svd-diag", - "RF-vector-svd-nondiag", - "RF-vector-nosvd-diag", - "RF-vector-nosvd-nondiag", - "RF-vector-svd-nonsep", + "gp-gpjl", # diagonalize, train scalar GP, assume diag inputs + "rf-lr-scalar", + "rf-lr-lr", + "rf-nonsep", ] #### CHOOSE YOUR CASE: - mask = 1:8 # 1:8 # e.g. 1:8 or [7] + mask = 1:4 # for (case) in cases[mask] @@ -55,7 +36,6 @@ function main() min_iter = 1 max_iter = 5 # number of EKP iterations to use data from is at most this overrides = Dict( - "verbose" => true, "scheduler" => DataMisfitController(terminate_at = 100.0), "cov_sample_multiplier" => 1.0, "n_iteration" => 20, @@ -106,7 +86,7 @@ function main() # Emulate-sample settings # choice of machine-learning tool in the emulation stage nugget = 0.001 - if case == "GP" + if case == "gp-gpjl" gppackage = Emulators.GPJL() pred_type = Emulators.YType() mlt = GaussianProcess( @@ -115,11 +95,9 @@ function main() prediction_type = pred_type, noise_learn = false, ) - elseif case ∈ ["RF-scalar", "RF-scalar-diagin"] + elseif case == "rf-lr-scalar" n_features = 100 - kernel_structure = - case == "RF-scalar-diagin" ? SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) : - SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()) + kernel_structure = SeparableKernel(LowRankFactor(1, nugget), OneDimFactor()) mlt = ScalarRandomFeatureInterface( n_features, n_params, @@ -127,12 +105,9 @@ function main() kernel_structure = kernel_structure, optimizer_options = overrides, ) - elseif case ∈ ["RF-vector-svd-diag", "RF-vector-nosvd-diag", "RF-vector-svd-nondiag", "RF-vector-nosvd-nondiag"] + elseif case == "rf-lr-lr" # do we want to assume that the outputs are decorrelated in the machine-learning problem? - kernel_structure = - case ∈ ["RF-vector-svd-diag", "RF-vector-nosvd-diag"] ? - SeparableKernel(LowRankFactor(2, nugget), DiagonalFactor(nugget)) : - SeparableKernel(LowRankFactor(2, nugget), LowRankFactor(3, nugget)) + kernel_structure = SeparableKernel(LowRankFactor(2, nugget), LowRankFactor(2, nugget)) n_features = 500 mlt = VectorRandomFeatureInterface( @@ -143,7 +118,7 @@ function main() kernel_structure = kernel_structure, optimizer_options = overrides, ) - elseif case ∈ ["RF-vector-svd-nonsep"] + elseif case == "rf-nonsep" kernel_structure = NonseparableKernel(LowRankFactor(3, nugget)) n_features = 500 @@ -158,11 +133,6 @@ function main() end # Standardize the output data - # Use median over all data since all data are the same type - truth_sample_norm = vcat(truth_sample...) - norm_factor = get_standardizing_factors(truth_sample_norm) - #norm_factor = vcat(norm_factor...) - norm_factor = fill(norm_factor, size(truth_sample)) # Get training points from the EKP iteration number in the second input term N_iter = min(max_iter, length(get_u(ekiobj)) - 1) # number of paired iterations taken from EKP @@ -196,23 +166,11 @@ function main() savefig(p, joinpath(figure_save_directory, "training_points.png")) end - standardize = false - retained_svd_frac = 1.0 - normalized = true - # do we want to use SVD to decorrelate outputs - decorrelate = case ∈ ["RF-vector-nosvd-diag", "RF-vector-nosvd-nondiag"] ? false : true - - - emulator = Emulator( - mlt, - input_output_pairs; - obs_noise_cov = Γy, - normalize_inputs = normalized, - standardize_outputs = standardize, - standardize_outputs_factors = norm_factor, - retained_svd_frac = retained_svd_frac, - decorrelate = decorrelate, - ) + # Process data + retain_var = 0.95 + encoder_schedule = [(quartile_scale(), "in"), (decorrelate_structure_mat(retain_var = retain_var), "out")] + + emulator = Emulator(mlt, input_output_pairs; output_structure_matrix = Γy, encoder_schedule = encoder_schedule) optimize_hyperparameters!(emulator) # Check how well the Gaussian Process regression predicts on the @@ -266,13 +224,25 @@ function main() posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns constrained_posterior_samples = mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + xlims = extrema(constrained_posterior_samples[1, :]) + ylims = extrema(constrained_posterior_samples[2, :]) + + # plot with PairPlots + gr(dpi = 300, size = (300, 300)) - p = cornerplot(permutedims(constrained_posterior_samples, (2, 1)), label = param_names, compact = true) + p = cornerplot( + permutedims(constrained_posterior_samples, (2, 1)), + label = param_names, + compact = true, + xlims = xlims, + ylims = ylims, + ) plot!(p.subplots[1], [truth_params_constrained[1]], seriestype = "vline", w = 1.5, c = :steelblue, ls = :dash) # vline on top histogram plot!(p.subplots[3], [truth_params_constrained[2]], seriestype = "hline", w = 1.5, c = :steelblue, ls = :dash) # hline on right histogram plot!(p.subplots[2], [truth_params_constrained[1]], seriestype = "vline", w = 1.5, c = :steelblue, ls = :dash) # v & h line on scatter. plot!(p.subplots[2], [truth_params_constrained[2]], seriestype = "hline", w = 1.5, c = :steelblue, ls = :dash) + figpath = joinpath(figure_save_directory, "posterior_2d-" * case * ".pdf") savefig(figpath) figpath = joinpath(figure_save_directory, "posterior_2d-" * case * ".png") diff --git a/examples/Lorenz/emulate_sample_spatial_dep.jl b/examples/Lorenz/emulate_sample_spatial_dep.jl index cd19f7c57..50b634ddc 100644 --- a/examples/Lorenz/emulate_sample_spatial_dep.jl +++ b/examples/Lorenz/emulate_sample_spatial_dep.jl @@ -19,38 +19,22 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.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) - return norm_factor -end - -function get_standardizing_factors(data::Array{FT, 1}) where {FT} - # Input: data size: N_data*N_ensembles (splatted) - # Ensemble median of the data - norm_factor = median(data) - return norm_factor -end - function main() cases = [ - "GP", # SLOW + "GP", "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs ] #### CHOOSE YOUR CASE: - mask = [2]# 1:1 # e.g. 1:2 or [2] + mask = [1]# 1:1 # e.g. 1:2 or [2] for (case) in cases[mask] println("case: ", case) min_iter = 1 - max_iter = 6 # number of EKP iterations to use data from is at most this - - # we do not want termination, as our priors have relatively little interpretation + max_iter = 7 # number of EKP iterations to use data from is at most this #### @@ -79,7 +63,6 @@ function main() truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) Γy = get_obs_noise_cov(ekpobj) - n_params = length(truth_params) # "input dim" output_dim = size(Γy, 1) ### @@ -88,7 +71,7 @@ function main() # Emulate-sample settings # choice of machine-learning tool in the emulation stage - nugget = 0.001 + nugget = 1e-3 if case == "GP" gppackage = Emulators.GPJL() pred_type = Emulators.YType() @@ -100,14 +83,14 @@ function main() ) elseif case == "RF-scalar" overrides = Dict( - "verbose" => true, - "scheduler" => DataMisfitController(terminate_at = 100.0), + # "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 1000.0), "cov_sample_multiplier" => 1.0, "n_iteration" => 8, "n_features_opt" => 40, ) - n_features = 200 - kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()) + n_features = 80 + kernel_structure = SeparableKernel(LowRankFactor(1, nugget), OneDimFactor()) mlt = ScalarRandomFeatureInterface( n_features, n_params, @@ -119,12 +102,6 @@ function main() throw(ArgumentError("case $(case) not recognised, please choose from the list $(cases)")) end - # Standardize the output data - # Use median over all data since all data are the same type - truth_sample_norm = vcat(truth_sample...) - norm_factor = get_standardizing_factors(truth_sample_norm) - #norm_factor = vcat(norm_factor...) - norm_factor = fill(norm_factor, size(truth_sample)) # Get training points from the EKP iteration number in the second input term N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP @@ -136,50 +113,45 @@ function main() # plot training points in constrained space if case == cases[mask[1]] + i1 = 1 + i2 = 10 gr(dpi = 300, size = (400, 400)) inputs_unconstrained = get_inputs(input_output_pairs) inputs_constrained = transform_unconstrained_to_constrained(priors, inputs_unconstrained) p = plot( title = "training points", - xlims = extrema(inputs_constrained[1, :]), - xaxis = "F", - yaxis = "A", - ylims = extrema(inputs_constrained[2, :]), + xlims = extrema(inputs_constrained[i1, :]), + xaxis = "x$(i1)", + yaxis = "x$(i2)", + ylims = extrema(inputs_constrained[i2, :]), ) - scatter!(p, inputs_constrained[1, :], inputs_constrained[2, :], color = :magenta, label = false) + scatter!(p, inputs_constrained[i1, :], inputs_constrained[i2, :], color = :magenta, label = "train") inputs_test_unconstrained = get_inputs(input_output_pairs_test) inputs_test_constrained = transform_unconstrained_to_constrained(priors, inputs_test_unconstrained) - scatter!(p, inputs_test_constrained[1, :], inputs_test_constrained[2, :], color = :black, label = false) + scatter!(p, inputs_test_constrained[i1, :], inputs_test_constrained[i2, :], color = :black, label = "test") - vline!(p, [truth_params_constrained[1]], linestyle = :dash, linecolor = :red, label = false) - hline!(p, [truth_params_constrained[2]], linestyle = :dash, linecolor = :red, label = false) + vline!(p, [truth_params_constrained[i1]], linestyle = :dash, linecolor = :red, label = false) + hline!(p, [truth_params_constrained[i2]], linestyle = :dash, linecolor = :red, label = false) savefig(p, joinpath(figure_save_directory, "training_points.pdf")) savefig(p, joinpath(figure_save_directory, "training_points.png")) end - standardize = false - retained_svd_frac = 0.99 # keep 99% of the singular values - normalized = true - # do we want to use SVD to decorrelate outputs - decorrelate = true - + # data processing configuration + retain_var = 0.95 + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] emulator = Emulator( mlt, input_output_pairs; - obs_noise_cov = Γy, - normalize_inputs = normalized, - standardize_outputs = standardize, - standardize_outputs_factors = norm_factor, - retained_svd_frac = retained_svd_frac, - decorrelate = decorrelate, + input_structure_matrix = cov(priors), + output_structure_matrix = Γy, + encoder_schedule = encoder_schedule, ) optimize_hyperparameters!(emulator) # Check how well the Gaussian Process regression predicts on the # true parameters - #if retained_svd_frac==1.0 y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1), transform_to_real = true) y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test), transform_to_real = true) @@ -256,6 +228,8 @@ function main() input_output_pairs, "truth_params", truth_params, + "encoder_schedule", + get_encoder_schedule(emulator), ) end end diff --git a/examples/Lorenz/plot_spatial_dep.jl b/examples/Lorenz/plot_spatial_dep.jl index 7ea4d456c..4265469b7 100644 --- a/examples/Lorenz/plot_spatial_dep.jl +++ b/examples/Lorenz/plot_spatial_dep.jl @@ -1,14 +1,8 @@ # some context values (from calibrate_spatial_dep.jl) # some context calues from emulate_sample_spatial_dep.jl -cases = [ - "GP", # SLOW - "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs - "RF-vector-svd-diag", # inaccurate - "RF-vector-svd-nondiag", - "RF-vector-svd-nonsep", -] -case = cases[2] +cases = ["GP", "RF-scalar"] +case = cases[1] # load packages # CES @@ -35,10 +29,13 @@ truth_params = load(data_file)["truth_params_constrained"] final_params = load(data_file)["final_params_constrained"] posterior = load(data_file)["posterior"] + ekp_file = joinpath(data_save_directory, "ekp_spatial_dep.jld2") ekpobj = load(ekp_file)["ekpobj"] N_ens = get_N_ens(ekpobj) - +nx = length(truth_params) +y = get_obs(ekpobj) +ny = length(y) # get samples and quantiles from posterior param_names = get_name(posterior) posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns @@ -63,7 +60,7 @@ p1 = plot( ) p2 = plot( - 1:length(y), + 1:ny, y, ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), label = "data", diff --git a/examples/Sinusoid/calibrate.jl b/examples/Sinusoid/calibrate.jl index ac610e533..330ef84be 100644 --- a/examples/Sinusoid/calibrate.jl +++ b/examples/Sinusoid/calibrate.jl @@ -87,12 +87,17 @@ ensemble_kalman_process = EKP.EnsembleKalmanProcess(initial_ensemble, y_obs, Γ, # We are now ready to carry out the inversion. At each iteration, we get the # ensemble from the last iteration, apply ``G(\theta)`` to each ensemble member, # and apply the Kalman update to the ensemble. +n_iter = N_iterations for i in 1:N_iterations params_i = EKP.get_ϕ_final(prior, ensemble_kalman_process) G_ens = hcat([G(params_i[:, i]; rng = rng) for i in 1:N_ensemble]...) - EKP.update_ensemble!(ensemble_kalman_process, G_ens) + terminate = EKP.update_ensemble!(ensemble_kalman_process, G_ens) + if !isnothing(terminate) + n_iter = i - 1 + break + end end @@ -123,8 +128,8 @@ vline!([theta_true[1]], color = :red, style = :dash, label = :false) hline!([theta_true[2]], color = :red, style = :dash, label = :false) -iteration_colormap = colormap("Blues", N_iterations) -for j in 1:N_iterations +iteration_colormap = colormap("Blues", n_iter) +for j in 1:n_iter ensemble_j = EKP.get_ϕ(prior, ensemble_kalman_process, j) plot!( ensemble_j[1, :], @@ -165,7 +170,7 @@ save( "N_ensemble", N_ensemble, "N_iterations", - N_iterations, + n_iter, "rng", rng, ) diff --git a/examples/Sinusoid/emulate.jl b/examples/Sinusoid/emulate.jl index dcf0027c0..0b02a6259 100644 --- a/examples/Sinusoid/emulate.jl +++ b/examples/Sinusoid/emulate.jl @@ -20,6 +20,7 @@ using JLD2 using CalibrateEmulateSample using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.Utilities const CES = CalibrateEmulateSample const EKP = CalibrateEmulateSample.EnsembleKalmanProcesses @@ -70,8 +71,11 @@ outputs = CES.Utilities.get_outputs(input_output_pairs) gppackage = Emulators.GPJL() gauss_proc = Emulators.GaussianProcess(gppackage, noise_learn = false) +# Create a data encoding to be performed on the input-output pairs +encoder_schedule = [(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")] + # Build emulator with data -emulator_gp = Emulator(gauss_proc, input_output_pairs, normalize_inputs = true, obs_noise_cov = Γ) +emulator_gp = Emulator(gauss_proc, input_output_pairs, output_structure_matrix = Γ, encoder_schedule = encoder_schedule) optimize_hyperparameters!(emulator_gp) # We have built the Gaussian process emulator and we can now use it for prediction. We will validate the emulator @@ -95,10 +99,9 @@ kernel_structure = NonseparableKernel(LowRankFactor(2, nugget)) optimizer_options = Dict( "n_ensemble" => 50, "cov_sample_multiplier" => 10, - "scheduler" => EKP.DataMisfitController(on_terminate = "continue"), - "n_iteration" => 50, + "scheduler" => EKP.DataMisfitController(terminate_at = 1e3), + "n_iteration" => 10, "rng" => rng, - "verbose" => true, ) random_features = VectorRandomFeatureInterface( n_features, @@ -108,8 +111,9 @@ random_features = VectorRandomFeatureInterface( kernel_structure = kernel_structure, optimizer_options = optimizer_options, ) +encoder_schedule_rf = [(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")] emulator_random_features = - Emulator(random_features, input_output_pairs, normalize_inputs = true, obs_noise_cov = Γ, decorrelate = false) + Emulator(random_features, input_output_pairs, encoder_schedule = encoder_schedule_rf, output_structure_matrix = Γ) optimize_hyperparameters!(emulator_random_features) diff --git a/src/Emulator.jl b/src/Emulator.jl index 1904fd9be..1c836a295 100644 --- a/src/Emulator.jl +++ b/src/Emulator.jl @@ -1,6 +1,12 @@ module Emulators using ..DataContainers +import ..Utilities.encode_with_schedule +import ..Utilities.decode_with_schedule +import ..Utilities.encode_data +import ..Utilities.decode_data +import ..Utilities.encode_structure_matrix +import ..Utilities.decode_structure_matrix using DocStringExtensions using Statistics @@ -14,10 +20,10 @@ export Emulator export calculate_normalization export build_models! export optimize_hyperparameters! -export predict - +export predict, encode_data, decode_data, encode_structure_matrix, decode_structure_matrix +export get_machine_learning_tool, get_io_pairs, get_encoded_io_pairs, get_encoder_schedule """ -$(DocStringExtensions.TYPEDEF) +$(TYPEDEF) Type to dispatch different emulators: @@ -45,503 +51,287 @@ function predict(mlt, new_inputs; mlt_kwargs...) throw_define_mlt() end - - # We will define the different emulator types after the general statements """ -$(DocStringExtensions.TYPEDEF) +$(TYPEDEF) Structure used to represent a general emulator, independently of the algorithm used. # Fields -$(DocStringExtensions.TYPEDFIELDS) +$(TYPEDFIELDS) """ -struct Emulator{FT <: AbstractFloat} +struct Emulator{FT <: AbstractFloat, VV <: AbstractVector} "Machine learning tool, defined as a struct of type MachineLearningTool." machine_learning_tool::MachineLearningTool - "Normalized, standardized, transformed pairs given the Booleans normalize\\_inputs, standardize\\_outputs, retained\\_svd\\_frac." - training_pairs::PairedDataContainer{FT} - "Mean of input; length *input\\_dim*." - input_mean::AbstractVector{FT} - "If normalizing: whether to fit models on normalized inputs (`(inputs - input_mean) * sqrt_inv_input_cov`)." - normalize_inputs::Bool - "(Linear) normalization transformation; size *input\\_dim* × *input\\_dim*." - normalization::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} - "Whether to fit models on normalized outputs: `outputs / standardize_outputs_factor`." - standardize_outputs::Bool - "If standardizing: Standardization factors (characteristic values of the problem)." - standardize_outputs_factors::Union{AbstractVector{FT}, Nothing} - "The singular value decomposition of *obs\\_noise\\_cov*, such that *obs\\_noise\\_cov* = decomposition.U * Diagonal(decomposition.S) * decomposition.Vt. NB: the SVD may be reduced in dimensions." - decomposition::Union{SVD, Nothing} - "Fraction of singular values kept in decomposition. A value of 1 implies full SVD spectrum information." - retained_svd_frac::FT + "original training data" + io_pairs::PairedDataContainer{FT} + "encoded training data" + encoded_io_pairs::PairedDataContainer{FT} + "Store of the pipeline to encode (/decode) the data" + encoder_schedule::VV end -get_machine_learning_tool(emulator::Emulator) = emulator.machine_learning_tool - -# Constructor for the Emulator Object """ -$(DocStringExtensions.TYPEDSIGNATURES) +$(TYPEDSIGNATURES) -Positional Arguments - - `machine_learning_tool` ::MachineLearningTool, - - `input_output_pairs` ::PairedDataContainer -Keyword Arguments - - `obs_noise_cov`: A matrix/uniform scaling to provide the observational noise covariance of the data - used for data processing (default `nothing`), - - `normalize_inputs`: Normalize the inputs to be unit Gaussian, in the smallest full-rank space of the data (default `true`), - - `standardize_outputs`: Standardize outputs with by dividing by a vector of provided factors (default `false`), - - `standardize_outputs_factors`: If standardizing, the provided dim_output-length vector of factors, - - `decorrelate`: Apply (truncated) SVD to the outputs. Predictions are returned in the decorrelated space, (default `true`) - - `retained_svd_frac`: The cumulative sum of singular values retained after output SVD truncation (default 1.0 - no truncation) +Gets the `machine_learning_tool` field of the `Emulator` """ -function Emulator( - machine_learning_tool::MachineLearningTool, - input_output_pairs::PairedDataContainer{FT}; - obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing, - normalize_inputs::Bool = true, - standardize_outputs::Bool = false, - standardize_outputs_factors::Union{AbstractVector{FT}, Nothing} = nothing, - decorrelate::Bool = true, - retained_svd_frac::FT = 1.0, - mlt_kwargs..., -) where {FT <: AbstractFloat} - - # For Consistency checks - input_dim, output_dim = size(input_output_pairs, 1) - if isa(obs_noise_cov, UniformScaling) - # Cast UniformScaling to Diagonal, since UniformScaling incompatible with - # SVD and standardize() - obs_noise_cov = Diagonal(obs_noise_cov, output_dim) - end - if obs_noise_cov !== nothing - err2 = "obs_noise_cov must be of size ($output_dim, $output_dim), got $(size(obs_noise_cov))" - size(obs_noise_cov) == (output_dim, output_dim) || throw(ArgumentError(err2)) - else - @warn "The covariance of the observational noise (a.k.a obs_noise_cov) is useful for data processing. Large approximation errors can occur without it. If possible, please provide it using the keyword obs_noise_cov." - end - - - # [1.] Normalize the inputs? - input_mean = vec(mean(get_inputs(input_output_pairs), dims = 2)) #column vector - normalization = nothing - if normalize_inputs - normalization = calculate_normalization(get_inputs(input_output_pairs)) - training_inputs = normalize(get_inputs(input_output_pairs), input_mean, normalization) - # new input_dim < input_dim when inputs lie in a proper linear subspace. - else - training_inputs = get_inputs(input_output_pairs) - end +get_machine_learning_tool(emulator::Emulator) = emulator.machine_learning_tool - # [2.] Standardize the outputs? - if standardize_outputs - training_outputs, obs_noise_cov = - standardize(get_outputs(input_output_pairs), obs_noise_cov, standardize_outputs_factors) - else - training_outputs = get_outputs(input_output_pairs) - end +""" +$(TYPEDSIGNATURES) - # [3.] Decorrelating the outputs, not performed for vector RF - if decorrelate - #Transform data if obs_noise_cov available - # (if obs_noise_cov==nothing, transformed_data is equal to data) - decorrelated_training_outputs, decomposition = - svd_transform(training_outputs, obs_noise_cov, retained_svd_frac = retained_svd_frac) +Gets the `io_pairs` field of the `Emulator` +""" +get_io_pairs(emulator::Emulator) = emulator.io_pairs - training_pairs = PairedDataContainer(training_inputs, decorrelated_training_outputs) - # [4.] build an emulator +""" +$(TYPEDSIGNATURES) - build_models!(machine_learning_tool, training_pairs; mlt_kwargs...) - else - if decorrelate || !isa(machine_learning_tool, VectorRandomFeatureInterface) - throw(ArgumentError("$machine_learning_tool is incompatible with option Emulator(...,decorrelate = false)")) - end - decomposition = nothing - training_pairs = PairedDataContainer(training_inputs, training_outputs) - # [4.] build an emulator - providing the noise covariance as a Tikhonov regularizer - build_models!(machine_learning_tool, training_pairs; regularization_matrix = obs_noise_cov, mlt_kwargs...) - end +Gets the `encoded_io_pairs` field of the `Emulator` +""" +get_encoded_io_pairs(emulator::Emulator) = emulator.encoded_io_pairs - return Emulator{FT}( - machine_learning_tool, - training_pairs, - input_mean, - normalize_inputs, - normalization, - standardize_outputs, - standardize_outputs_factors, - decomposition, - retained_svd_frac, - ) -end +""" +$(TYPEDSIGNATURES) +Gets the `encoder_schedul` field of the `Emulator` """ -$(DocStringExtensions.TYPEDSIGNATURES) +get_encoder_schedule(emulator::Emulator) = emulator.encoder_schedule -Optimizes the hyperparameters in the machine learning tool. +# Constructor for the Emulator Object """ -function optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) where {FT <: AbstractFloat} - optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...) -end +$(TYPEDSIGNATURES) +Constructor of the Emulator object, -""" -$(DocStringExtensions.TYPEDSIGNATURES) +Positional Arguments + - `machine_learning_tool`: the selected machine learning tool object (e.g. Gaussian process / Random feature interface) + - `input_output_pairs`: the paired input-output data points stored in a `PairedDataContainer` -Makes a prediction using the emulator on new inputs (each new inputs given as data columns). -Default is to predict in the decorrelated space. +Keyword Arguments + - `encoder_schedule`[=`nothing`]: the schedule of data encoding/decoding. This will be passed into the method `create_encoder_schedule` internally. `nothing` sets sets a default schedule `(decorrelate_samples_cov(), "in_and_out")`. Pass `[]` for no encoding. + - `input_structure_matrix`[=`nothing`]: Some encoders make use of an input structure (e.g., the prior covariance matrix). Particularly useful for few samples. + - `output_structure_matrix` [=`nothing`] Some encoders make use of an input structure (e.g., the prior covariance matrix). Particularly useful for few samples. +Other keywords are passed to the machine learning tool initialization """ -function predict( - emulator::Emulator{FT}, - new_inputs::AM; - transform_to_real = false, - vector_rf_unstandardize = true, +function Emulator( + machine_learning_tool::MachineLearningTool, + input_output_pairs::PairedDataContainer{FT}; + encoder_schedule = nothing, + input_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing, + output_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing, + obs_noise_cov = nothing, # temporary mlt_kwargs..., -) where {FT <: AbstractFloat, AM <: AbstractMatrix} - # Check if the size of new_inputs is consistent with the GP model's input - # dimension. - input_dim, output_dim = size(emulator.training_pairs, 1) +) where {FT <: AbstractFloat} - N_samples = size(new_inputs, 2) + # For Consistency checks + input_dim, output_dim = size(input_output_pairs, 1) - # check sizing against normalization - if emulator.normalize_inputs - size(new_inputs, 1) == size(emulator.normalization, 2) || throw( - ArgumentError( - "Emulator object and input observations do not have consistent dimensions, expected $(size(emulator.normalization,2)), received $(size(new_inputs,1))", - ), + if !isnothing(obs_noise_cov) && isnothing(output_structure_matrix) + @warn( + "Keyword `obs_noise_cov=` is now deprecated, and replaced with `output_structure_matrix`. \n Continuing by setting `output_structure_matrix=obs_noise_cov`." ) - else - size(new_inputs, 1) == input_dim || throw( - ArgumentError( - "Emulator object and input observations do not have consistent dimensions, expected $(input_dim), received $(size(new_inputs,1))", - ), + output_structure_matrix = obs_noise_cov + elseif !isnothing(obs_noise_cov) && !isnothing(output_structure_matrix) + @warn( + "Keyword `obs_noise_cov=` is now deprecated and will be ignored. \n Continuing with value of `output_structure_matrix=`" ) end - # [1.] normalize - normalized_new_inputs = normalize(emulator, new_inputs) - # [2.] predict. Note: ds = decorrelated, standard - ds_outputs, ds_output_var = predict(emulator.machine_learning_tool, normalized_new_inputs, mlt_kwargs...) - - # [3.] transform back to real coordinates or remain in decorrelated coordinates - if !isa(get_machine_learning_tool(emulator), VectorRandomFeatureInterface) - if transform_to_real && emulator.decomposition === nothing - throw(ArgumentError("""Need SVD decomposition to transform back to original space, - but GaussianProcess.decomposition == nothing. - Try setting transform_to_real=false""")) - elseif transform_to_real && emulator.decomposition !== nothing - - #transform back to real coords - cov becomes dense - s_outputs, s_output_cov = svd_reverse_transform_mean_cov(ds_outputs, ds_output_var, emulator.decomposition) - if output_dim == 1 - s_output_cov = reshape([s_output_cov[i][1] for i in 1:N_samples], 1, :) - end - - # [4.] unstandardize - return reverse_standardize(emulator, s_outputs, s_output_cov) + # [1.] Initializes and performs data encoding schedule + # Default processing: decorrelate_sample_cov() where no structure matrix provided, and decorrelate_structure_mat() where provided. + if isnothing(encoder_schedule) + encoder_schedule = [] + if isnothing(input_structure_matrix) + push!(encoder_schedule, (decorrelate_sample_cov(), "in")) else - # remain in decorrelated, standardized coordinates (cov remains diagonal) - # Convert to vector of matrices to match the format - # when transform_to_real=true - ds_output_diagvar = vec([Diagonal(ds_output_var[:, j]) for j in 1:N_samples]) - if output_dim == 1 - ds_output_diagvar = reshape([ds_output_diagvar[i][1] for i in 1:N_samples], 1, :) - end - return ds_outputs, ds_output_diagvar + push!(encoder_schedule, (decorrelate_structure_mat(), "in")) end - else - # create a vector of covariance matrices - ds_output_covvec = vec([ds_output_var[:, :, j] for j in 1:N_samples]) - - if emulator.decomposition === nothing # without applying SVD - if vector_rf_unstandardize - # [4.] unstandardize - return reverse_standardize(emulator, ds_outputs, ds_output_covvec) - else - return ds_outputs, ds_output_covvec - end - - else #if we applied SVD - if transform_to_real # ...and want to return coordinates - s_outputs, s_output_cov = - svd_reverse_transform_mean_cov(ds_outputs, ds_output_covvec, emulator.decomposition) - if output_dim == 1 - s_output_cov = reshape([s_output_cov[i][1] for i in 1:N_samples], 1, :) - end - # [4.] unstandardize - return reverse_standardize(emulator, s_outputs, s_output_cov) - else #... and want to stay in decorrelated standardized coordinates - - # special cases where you only return variances and not covariances, could be better acheived with dispatch... - kernel_structure = get_kernel_structure(get_machine_learning_tool(emulator)) - if nameof(typeof(kernel_structure)) == :SeparableKernel - output_cov_structure = get_output_cov_structure(kernel_structure) - if nameof(typeof(output_cov_structure)) == :DiagonalFactor - ds_output_diagvar = vec([Diagonal(ds_output_covvec[j]) for j in 1:N_samples]) # extracts diagonal - return ds_outputs, ds_output_diagvar - elseif nameof(typeof(output_cov_structure)) == :OneDimFactor - ds_output_diagvar = ([ds_output_covvec[i][1, 1] for i in 1:N_samples], 1, :) # extracts value - return ds_outputs, ds_output_diagvar - else - return ds_outputs, ds_output_covvec - end - else - return ds_outputs, ds_output_covvec - end - end + if isnothing(output_structure_matrix) + push!(encoder_schedule, (decorrelate_sample_cov(), "out")) + else + push!(encoder_schedule, (decorrelate_structure_mat(), "out")) end - end + enc_schedule = create_encoder_schedule(encoder_schedule) + (encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) = + initialize_and_encode_with_schedule!( + enc_schedule, + input_output_pairs, + input_structure_matrix, + output_structure_matrix, + ) + + # build the machine learning tool in the encoded space + build_models!( + machine_learning_tool, + encoded_io_pairs, + encoded_input_structure_matrix, + encoded_output_structure_matrix; + mlt_kwargs..., + ) + return Emulator{FT, typeof(enc_schedule)}(machine_learning_tool, input_output_pairs, encoded_io_pairs, enc_schedule) end -# Normalization, Standardization, and Decorrelation """ -$(DocStringExtensions.TYPEDSIGNATURES) +$(TYPEDSIGNATURES) -Calculate the normalization of inputs. +Optimizes the hyperparameters in the machine learning tool. Note that some machine learning packages train hyperparameters on construction so this call is not necessary """ -function calculate_normalization(inputs::VOrM) where {VOrM <: AbstractVecOrMat} - input_mean = vec(mean(inputs, dims = 2)) - input_cov = cov(inputs, dims = 2) - - if rank(input_cov) == size(input_cov, 1) - normalization = sqrt(inv(input_cov)) - else # if not full rank, normalize non-zero singular values - svd_in = svd(input_cov) - sqrt_inv_sv = 1 ./ sqrt.(svd_in.S[1:rank(input_cov)]) - normalization = Diagonal(sqrt_inv_sv) * svd_in.Vt[1:rank(input_cov), :] #non-square - @info "reducing input dimension from $(size(input_cov,1)) to $rank(input_cov) during low rank in normalization" - end - return normalization +function optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) where {FT <: AbstractFloat} + optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...) end """ -$(DocStringExtensions.TYPEDSIGNATURES) +$(TYPEDSIGNATURES) -Normalize the input data, with a normalizing function. +Encode the new data (a `DataContainer`, or matrix where data are columns) representing inputs (`"in"`) or outputs (`"out"`). with the stored and initialized encoder schedule. """ -function normalize(emulator::Emulator, inputs::VOrM) where {VOrM <: AbstractVecOrMat} - if emulator.normalize_inputs - return normalize(inputs, emulator.input_mean, emulator.normalization) +function encode_data( + emulator::Emulator, + data::MorDC, + in_or_out::AS, +) where {AS <: AbstractString, MorDC <: Union{AbstractMatrix, DataContainer}} + if isa(data, AbstractMatrix) + return get_data(encode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out)) else - return inputs + return encode_with_schedule(get_encoder_schedule(emulator), data, in_or_out) end end -function normalize( - inputs::VOrM, - input_mean::V, - normalization::M, -) where {VOrM <: AbstractVecOrMat, V <: AbstractVector, M <: AbstractMatrix} - return normalization * (inputs .- input_mean) -end """ -$(DocStringExtensions.TYPEDSIGNATURES) +$(TYPEDSIGNATURES) -Standardize with a vector of factors (size equal to output dimension). +Encode a new structure matrix in the input space (`"in"`) or output space (`"out"`). with the stored and initialized encoder schedule. """ -function standardize( - outputs::VorM, - output_covs::VofMUS, - factors::V, -) where { - VorM <: AbstractVecOrMat, - VofMUS <: Union{AbstractVector{<:AbstractMatrix}, AbstractVector{<:UniformScaling}}, - V <: AbstractVector, -} - # Case where `output_covs` is a Vector of covariance matrices - n = length(factors) - outputs = outputs ./ factors - cov_factors = factors .* factors' - for i in 1:length(output_covs) - if isa(output_covs[i], Matrix) - output_covs[i] = output_covs[i] ./ cov_factors - elseif isa(output_covs[i], UniformScaling) - # assert all elements of factors are same, otherwise can't cast back to - # UniformScaling - @assert all(y -> y == first(factors), factors) - output_covs[i] = output_covs[i] / cov_factors[1, 1] - else - # Diagonal, TriDiagonal etc. case - # need to do ./ as dense matrix and cast back to original type - # https://discourse.julialang.org/t/get-generic-constructor-of-parametric-type/57189/8 - T = typeof(output_covs[i]) - cast_T = getfield(parentmodule(T), nameof(T)) - output_covs[i] = cast_T(Matrix(output_covs[i]) ./ cov_factors) - end - end - return outputs, output_covs +function encode_structure_matrix( + emulator::Emulator, + structure_mat::USorMorN, + in_or_out::AS, +) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} + return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out) end -function standardize( - outputs::VorM, - output_cov::MorVorUS, - factors::V, -) where { - VorM <: AbstractVecOrMat, - MorVorUS <: Union{AbstractMatrix, AbstractVector{<:AbstractFloat}, UniformScaling}, # must be vector of floats or could cause inf. loop with other method definition. - V <: AbstractVector, -} - - # Case where `output_cov` is a single covariance matrix - stdized_out, stdized_covs = standardize(outputs, [output_cov], factors) - return stdized_out, stdized_covs[1] -end """ -$(DocStringExtensions.TYPEDSIGNATURES) +$(TYPEDSIGNATURES) -Reverse a previous standardization with the stored vector of factors (size equal to output -dimension). `output_cov` is a Vector of covariance matrices, such as is returned by -[`svd_reverse_transform_mean_cov`](@ref). +Decode the new data (a `DataContainer`, or matrix where data are columns) representing inputs (`"in"`) or outputs (`"out"`). with the stored and initialized encoder schedule. """ -function reverse_standardize( - emulator::Emulator{FT}, - outputs::VorM, - output_covs::VorMtwo, -) where { - VorM <: AbstractVecOrMat, - VorMtwo <: AbstractVecOrMat, # can be different type - FT <: AbstractFloat, -} - if emulator.standardize_outputs - return standardize(outputs, output_covs, 1.0 ./ emulator.standardize_outputs_factors) +function decode_data( + emulator::Emulator, + data::MorDC, + in_or_out::AS, +) where {AS <: AbstractString, MorDC <: Union{AbstractMatrix, DataContainer}} + if isa(data, AbstractMatrix) + return get_data(decode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out)) else - return outputs, output_covs + return decode_with_schedule(get_encoder_schedule(emulator), data, in_or_out) end end - - """ -$(DocStringExtensions.TYPEDSIGNATURES) - -Apply a singular value decomposition (SVD) to the data - - `data` - GP training data/targets; size *output\\_dim* × *N\\_samples* - - `obs_noise_cov` - covariance of observational noise - - `truncate_svd` - Project onto this fraction of the largest principal components. Defaults - to 1.0 (no truncation). - -Returns the transformed data and the decomposition, which is a matrix -factorization of type LinearAlgebra.SVD. - -Note: If `F::SVD` is the factorization object, U, S, V and Vt can be obtained via -F.U, F.S, F.V and F.Vt, such that A = U * Diagonal(S) * Vt. The singular values -in S are sorted in descending order. +$(TYPEDSIGNATURES) + +Decode a new structure matrix in the input space (`"in"`) or output space (`"out"`). with the stored and initialized encoder schedule. """ -function svd_transform( - data::AbstractMatrix{FT}, - obs_noise_cov::Union{AbstractMatrix{FT}, Nothing}; - retained_svd_frac::FT = 1.0, -) where {FT <: AbstractFloat} - if obs_noise_cov === nothing - # no-op case - return data, nothing - end - # actually have a matrix to take the SVD of - decomp = svd(obs_noise_cov) - sqrt_singular_values_inv = Diagonal(1.0 ./ sqrt.(decomp.S)) - if retained_svd_frac < 1.0 - # Truncate the SVD as a form of regularization - # Find cutoff - S_cumsum = cumsum(decomp.S) / sum(decomp.S) - ind = findall(x -> (x > retained_svd_frac), S_cumsum) - k = ind[1] - n = size(data)[1] - println("SVD truncated at k: ", k, "/", n) - # Apply truncated SVD - transformed_data = sqrt_singular_values_inv[1:k, 1:k] * decomp.Vt[1:k, :] * data - decomposition = SVD(decomp.U[:, 1:k], decomp.S[1:k], decomp.Vt[1:k, :]) - else - transformed_data = sqrt_singular_values_inv * decomp.Vt * data - decomposition = decomp - end - return transformed_data, decomposition +function decode_structure_matrix( + emulator::Emulator, + structure_mat::USorMorN, + in_or_out::AS, +) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} + return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out) end -function svd_transform( - data::AbstractVector{FT}, - obs_noise_cov::Union{AbstractMatrix{FT}, Nothing}; - retained_svd_frac::FT = 1.0, -) where {FT <: AbstractFloat} - # method for 1D data - transformed_data, decomposition = - svd_transform(reshape(data, :, 1), obs_noise_cov; retained_svd_frac = retained_svd_frac) - return vec(transformed_data), decomposition -end +""" +$(TYPEDSIGNATURES) -function svd_transform( - data::AbstractVecOrMat{FT}, - obs_noise_cov::UniformScaling{FT}; - retained_svd_frac::FT = 1.0, -) where {FT <: AbstractFloat} - # method for UniformScaling - return svd_transform(data, Diagonal(obs_noise_cov, size(data, 1)); retained_svd_frac = retained_svd_frac) -end +Makes a prediction using the emulator on new inputs (each new inputs given as data columns). +Default is to predict in the decorrelated space. +Return type of N inputs: (in the output space) + - 1-D: mean [1 x N], cov [1 x N] + - p-D: mean [p x N], cov N x [p x p] """ -$(DocStringExtensions.TYPEDSIGNATURES) - -Transform the mean and covariance back to the original (correlated) coordinate system - - `μ` - predicted mean; size *output\\_dim* × *N\\_predicted\\_points*. - - `σ2` - predicted variance; size *output\\_dim* × *N\\_predicted\\_points*. - - predicted covariance; size *N\\_predicted\\_points* array of size *output\\_dim* × *output\\_dim*. - - `decomposition` - SVD decomposition of *obs\\_noise\\_cov*. - -Returns the transformed mean (size *output\\_dim* × *N\\_predicted\\_points*) and variance. -Note that transforming the variance back to the original coordinate system -results in non-zero off-diagonal elements, so instead of just returning the -elements on the main diagonal (i.e., the variances), we return the full -covariance at each point, as a vector of length *N\\_predicted\\_points*, where -each element is a matrix of size *output\\_dim* × *output\\_dim*. -""" -function svd_reverse_transform_mean_cov( - μ::AbstractMatrix{FT}, - σ2::AbstractVector,# vector of output_dim x output_dim matrices - decomposition::SVD, -) where {FT <: AbstractFloat} +function predict( + emulator::Emulator{FT}, + new_inputs::AM; + transform_to_real = false, + mlt_kwargs..., +) where {FT <: AbstractFloat, AM <: AbstractMatrix} + # Check if the size of new_inputs is consistent with the GP model's input + # dimension. + # un-encoded data to get dimensions + input_dim, output_dim = size(get_io_pairs(emulator), 1) + encoded_input_dim, encoded_output_dim = size(get_encoded_io_pairs(emulator), 1) + + N_samples = size(new_inputs, 2) - N_predicted_points = length(σ2) - if !(eltype(σ2) <: AbstractMatrix) + if size(new_inputs, 1) != input_dim throw( ArgumentError( - "input σ2 must be a vector of eltype <: AbstractMatrix, instead eltype(σ2) = ", - string(eltype(σ2)), + "Emulator object and input observations do not have consistent dimensions, expected $(input_dim), received $(size(new_inputs,1))", ), ) end - # We created meanvGP = D_inv * Vt * mean_v so meanv = V * D * meanvGP - sqrt_singular_values = Diagonal(sqrt.(decomposition.S)) - transformed_μ = decomposition.V * sqrt_singular_values * μ + # encode the new input data + encoded_inputs = encode_data(emulator, new_inputs, "in") + # predict in encoding space + # returns outputs: [enc_out_dim x n_samples] + # Scalar-methods uncertainties=variances: [enc_out_dim x n_samples] + # Vector-methods uncertainties=covariances: [enc_out_dim x enc_out_dim x n_samples) + encoded_outputs, encoded_uncertainties = predict(get_machine_learning_tool(emulator), encoded_inputs, mlt_kwargs...) - transformed_σ2 = Vector{Symmetric{FT, Matrix{FT}}}(undef, N_predicted_points) - # Back transformation - - for j in 1:N_predicted_points - σ2_j = decomposition.V * sqrt_singular_values * σ2[j] * sqrt_singular_values * decomposition.Vt - transformed_σ2[j] = Symmetric(σ2_j) - end + var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov" - return transformed_μ, transformed_σ2 + # return decoded or encoded? + if transform_to_real + decoded_outputs = decode_data(emulator, encoded_outputs, "out") -end + decoded_covariances = zeros(eltype(encoded_outputs), output_dim, output_dim, size(encoded_uncertainties)[end]) + if var_or_cov == "var" + for (i, col) in enumerate(eachcol(encoded_uncertainties)) + decoded_covariances[:, :, i] .= decode_structure_matrix(emulator, Diagonal(col), "out") + end + else # == "cov" + for (i, mat) in enumerate(eachslice(encoded_uncertainties, dims = 3)) + decoded_covariances[:, :, i] .= decode_structure_matrix(emulator, mat, "out") + end + end -function svd_reverse_transform_mean_cov( - μ::AbstractMatrix{FT}, - σ2::AbstractMatrix{FT}, - decomposition::SVD, -) where {FT <: AbstractFloat} - @assert size(μ) == size(σ2) - output_dim, N_predicted_points = size(σ2) - σ2_as_cov = vec([Diagonal(σ2[:, j]) for j in 1:N_predicted_points]) + if output_dim > 1 + return decoded_outputs, eachslice(decoded_covariances, dims = 3) + else + # here the covs are [1 x 1 x samples] just return [1 x samples] + return decoded_outputs, decoded_covariances[1, :, :] + end - return svd_reverse_transform_mean_cov(μ, σ2_as_cov, decomposition) -end + else + encoded_covariances_mat = + zeros(eltype(encoded_outputs), encoded_output_dim, encoded_output_dim, size(encoded_uncertainties)[end]) + if var_or_cov == "var" + for (i, col) in enumerate(eachcol(encoded_uncertainties)) + encoded_covariances_mat[:, :, i] = Diagonal(col) + end + else # =="cov" + for (i, mat) in enumerate(eachslice(encoded_uncertainties, dims = 3)) + encoded_covariances_mat[:, :, i] = mat + end + end + if encoded_output_dim > 1 + return encoded_outputs, eachslice(encoded_covariances_mat, dims = 3) + else + # here the covs are [1 x 1 x samples] just return [1 x samples] + return encoded_outputs, encoded_covariances_mat[1, :, :] + end + end +end end diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl index 6a59669ba..3fd1ae3b8 100644 --- a/src/GaussianProcess.jl +++ b/src/GaussianProcess.jl @@ -65,7 +65,7 @@ models. $(DocStringExtensions.TYPEDFIELDS) """ -struct GaussianProcess{GPPackage, FT} <: MachineLearningTool +struct GaussianProcess{GPPackage, FT, VV <: AbstractVector} <: MachineLearningTool "The Gaussian Process (GP) Regression model(s) that are fitted to the given input-data pairs." models::Vector{Union{<:GaussianProcesses.GPE, <:PyObject, <:AbstractGPs.PosteriorGP, Nothing}} "Kernel object." @@ -76,7 +76,8 @@ struct GaussianProcess{GPPackage, FT} <: MachineLearningTool alg_reg_noise::FT "Prediction type (`y` to predict the data, `f` to predict the latent function)." prediction_type::PredictionType - + "Regularization vector for each output dimension (based on alg_reg_noise" + regularization::VV end @@ -112,7 +113,16 @@ function GaussianProcess( alg_reg_noise = 1.0 end - return GaussianProcess{typeof(package), FT}(models, kernel, noise_learn, alg_reg_noise, prediction_type) + vv = typeof(alg_reg_noise)[] + + return GaussianProcess{typeof(package), FT, typeof(vv)}( + models, + kernel, + noise_learn, + alg_reg_noise, + prediction_type, + vv, + ) end # First we create the GPJL implementation @@ -138,7 +148,9 @@ Method to build Gaussian process models based on the package. """ function build_models!( gp::GaussianProcess{GPJL}, - input_output_pairs::PairedDataContainer{FT}; + input_output_pairs::PairedDataContainer{FT}, + input_structure_matrix, + output_structure_matrix; kwargs..., ) where {FT <: AbstractFloat} # get inputs and outputs @@ -170,6 +182,7 @@ function build_models!( println("Using user-defined kernel", kern) end + if gp.noise_learn # Add white noise to kernel white_logstd = log(1.0) @@ -177,9 +190,19 @@ function build_models!( kern = kern + white println("Learning additive white noise") end - logstd_regularization_noise = log(sqrt(gp.alg_reg_noise)) + # use the output_structure_matrix to scale regularization scale + regularization = if isnothing(output_structure_matrix) + 1.0 * ones(N_models) + elseif isa(output_structure_matrix, UniformScaling) + output_structure_matrix.λ * ones(N_models) + else + diag(output_structure_matrix) + end + + logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise)) for i in 1:N_models + logstd_regularization_i = logstd_regularization_noise[i] # Make a copy of the kernel (because it gets altered in every # iteration) kernel_i = deepcopy(kern) @@ -194,12 +217,13 @@ function build_models!( kmean = MeanZero() # Instantiate GP model - m = GaussianProcesses.GPE(input_values, output_values[i, :], kmean, kernel_i, logstd_regularization_noise) + m = GaussianProcesses.GPE(input_values, output_values[i, :], kmean, kernel_i, logstd_regularization_i) println("created GP: ", i) push!(models, m) end + append!(gp.regularization, logstd_regularization_noise) end @@ -258,7 +282,9 @@ predict(gp::GaussianProcess{GPJL}, new_inputs::AbstractMatrix{FT}) where {FT <: #now we build the SKLJL implementation function build_models!( gp::GaussianProcess{SKLJL}, - input_output_pairs::PairedDataContainer{FT}; + input_output_pairs::PairedDataContainer{FT}, + input_structure_matrix, + output_structure_matrix; kwargs..., ) where {FT <: AbstractFloat} # get inputs and outputs @@ -295,12 +321,20 @@ function build_models!( kern = kern + white println("Learning additive white noise") end - regularization_noise = gp.alg_reg_noise - + # use the output_structure_matrix to scale regularization scale + regularization = if isnothing(output_structure_matrix) + 1.0 * ones(N_models) + elseif isa(output_structure_matrix, UniformScaling) + output_structure_matrix.λ * ones(N_models) + else + diag(output_structure_matrix) + end + regularization_noise_vec = gp.alg_reg_noise .* regularization for i in 1:N_models + regularization_noise_i = regularization_noise_vec[i] kernel_i = deepcopy(kern) data_i = output_values[i, :] - m = pyGP.GaussianProcessRegressor(kernel = kernel_i, n_restarts_optimizer = 10, alpha = regularization_noise) + m = pyGP.GaussianProcessRegressor(kernel = kernel_i, n_restarts_optimizer = 10, alpha = regularization_noise_i) # ScikitLearn.fit! arguments: # input_values: (N_samples × input_dim) # data_i: (N_samples,) @@ -310,6 +344,8 @@ function build_models!( push!(models, m) @info(m.kernel) end + append!(gp.regularization, regularization_noise_vec) + end @@ -327,7 +363,9 @@ function predict(gp::GaussianProcess{SKLJL}, new_inputs::AbstractMatrix{FT}) whe # for SKLJL does not return the observational noise (even if return_std = true) # we must add contribution depending on whether we learnt the noise or not. - σ2[:, :] = σ2[:, :] .+ gp.alg_reg_noise + for i in 1:size(σ2, 2) + σ2[:, i] = σ2[:, i] + gp.regularization + end return μ, σ2 end @@ -336,7 +374,9 @@ end #We build the AGPJL implementation function build_models!( gp::GaussianProcess{AGPJL}, - input_output_pairs::PairedDataContainer{FT}; + input_output_pairs::PairedDataContainer{FT}, + input_structure_matrix, + output_structure_matrix; kernel_params = nothing, kwargs..., ) where {FT <: AbstractFloat} @@ -386,7 +426,15 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow N_models = size(output_values, 1) #size(transformed_data)[1] - regularization_noise = gp.alg_reg_noise + # use the output_structure_matrix to scale regularization scale + regularization = if isnothing(output_structure_matrix) + 1.0 * ones(N_models) + elseif isa(output_structure_matrix, UniformScaling) + output_structure_matrix.λ * ones(N_models) + else + diag(output_structure_matrix) + end + regularization_noise = gp.alg_reg_noise .* regularization # now obtain the values of the hyperparameters if N_models == 1 && !(isa(kernel_params, AbstractVector)) # i.e. just a Dict @@ -399,12 +447,13 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow var_sqexp = exp.(2 .* kernel_params_vec[i]["log_std_sqexp"]) # Float var_noise = exp.(2 .* kernel_params_vec[i]["log_std_noise"]) # Float rbf_invlen = 1 ./ exp.(kernel_params_vec[i]["log_rbf_len"])# Vec + regularization_noise_i = regularization_noise[i] opt_kern = var_sqexp * (KernelFunctions.SqExponentialKernel() ∘ ARDTransform(rbf_invlen[:])) + var_noise * KernelFunctions.WhiteKernel() opt_f = AbstractGPs.GP(opt_kern) - opt_fx = opt_f(input_values', regularization_noise) + opt_fx = opt_f(input_values', regularization_noise_i; obsdim = 2) data_i = output_values[i, :] opt_post_fx = posterior(opt_fx, data_i) @@ -412,6 +461,8 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow push!(models, opt_post_fx) println(opt_post_fx.prior.kernel) end + append!(gp.regularization, regularization_noise) + end @@ -428,11 +479,12 @@ function predict(gp::GaussianProcess{AGPJL}, new_inputs::AM) where {AM <: Abstra σ2 = zeros(FTorD, N_models, N_samples) for i in 1:N_models pred_gp = gp.models[i] - pred = pred_gp(new_inputs) + pred = pred_gp(new_inputs; obsdim = 2) μ[i, :] = mean(pred) σ2[i, :] = var(pred) end - - σ2[:, :] .= σ2[:, :] .+ gp.alg_reg_noise + for i in 1:size(σ2, 2) + σ2[:, i] .= σ2[:, i] + gp.regularization + end return μ, σ2 end diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 46c30051f..31b10ab43 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -36,52 +36,6 @@ export EmulatorPosteriorModel, sample, esjd -# ------------------------------------------------------------------------------------------ -# Output space transformations between original and SVD-decorrelated coordinates. -# Redundant with what's in Emulators.jl, but need to reimplement since we don't have -# access to obs_noise_cov - -""" -$(DocStringExtensions.TYPEDSIGNATURES) - -Transform samples from the original (correlated) coordinate system to the SVD-decorrelated -coordinate system used by [`Emulator`](@ref). Used in the constructor for [`MCMCWrapper`](@ref). - -The keyword `single_vec` wraps the output in a vector if `true` (default). -""" -function to_decorrelated(data::AbstractVector{FT}, em::Emulator{FT}; single_vec = true) where {FT <: AbstractFloat} - # method for a single sample - if em.standardize_outputs && em.standardize_outputs_factors !== nothing - # standardize() data by scale factors, if they were given - data = data ./ em.standardize_outputs_factors - end - decomp = em.decomposition - if decomp !== nothing - # Use SVD decomposition of obs noise cov, if given, to transform data to - # decorrelated coordinates. - inv_sqrt_singvals = Diagonal(1.0 ./ sqrt.(decomp.S)) - return single_vec ? [vec(inv_sqrt_singvals * decomp.Vt * data)] : inv_sqrt_singvals * decomp.Vt * data - else - return single_vec ? [vec(data)] : data - end -end - -function to_decorrelated(data::AbstractMatrix{FT}, em::Emulator{FT}) where {FT <: AbstractFloat} - # method for Matrix with columns that are samples - return [vec(to_decorrelated(cd, em, single_vec = false)) for cd in eachcol(data)] - -end - - -function to_decorrelated(data::AVV, em::Emulator{FT}) where {AVV <: AbstractVector, FT <: AbstractFloat} - # method for vector of samples - if isa(data[1], AbstractVector) - return [vec(to_decorrelated(d, em, single_vec = false)) for d in data] - else # turns out it is just one vector of a non-float type - return to_decorrelated(convert.(FT, data), em) - end -end - # ------------------------------------------------------------------------------------------ # Sampler extensions to differentiate vanilla RW and pCN algorithms @@ -296,8 +250,7 @@ function emulator_log_density_model( # Vector of N_samples covariance matrices. For MH, N_samples is always 1, so we # have to reshape()/re-cast input/output; simpler to do here than add a # predict() method. - g, g_cov = Emulators.predict(em, reshape(θ, :, 1), transform_to_real = false, vector_rf_unstandardize = false) - #TODO vector_rf will always unstandardize, but other methods will not, so we require this additional flag. + g, g_cov = Emulators.predict(em, reshape(θ, :, 1), transform_to_real = false) if isa(g_cov[1], Real) @@ -542,13 +495,13 @@ AbstractMCMC's terminology). # Fields $(DocStringExtensions.TYPEDFIELDS) """ -struct MCMCWrapper{AMorAV <: Union{AbstractVector, AbstractMatrix}, AV <: AbstractVector} +struct MCMCWrapper{AVV <: AbstractVector, AV <: AbstractVector} "[`ParameterDistribution`](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/parameter_distributions/) object describing the prior distribution on parameter values." prior::ParameterDistribution - "A vector or [Nx1] matrix, describing a single observation data (or NxM column-matrix / vector or vectors for multiple observations) provided by the user." - observations::AMorAV + "[output_dim x N_samples] matrix, of given observation data." + observations::AVV "Vector of observations describing the data samples to actually used during MCMC sampling (that have been transformed into a space consistent with emulator outputs)." - decorrelated_observations::AV + encoded_observations::AV "`AdvancedMH.DensityModel` object, used to evaluate the posterior density being sampled from." log_posterior_map::AbstractMCMC.AbstractModel "Object describing a MCMC sampling algorithm and its settings." @@ -575,8 +528,7 @@ decorrelation) that was applied in the Emulator. It creates and wraps an instanc - [`BarkerSampling`](@ref): Metropolis-Hastings sampling using the Barker proposal, which has a robustness to choosing step-size parameters. -- `obs_sample`: A single sample from the observations. Can, e.g., be picked from an - Observation struct using `get_obs_sample`. +- `obs_sample`: Vector (for one sample) or matrix with columns as samples from the observation. Can, e.g., be picked from an Observation struct using `get_obs_sample`. - `prior`: [`ParameterDistribution`](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/parameter_distributions/) object containing the parameters' prior distributions. - `emulator`: [`Emulator`](@ref) to sample from. @@ -594,10 +546,18 @@ function MCMCWrapper( kwargs..., ) where {AV <: AbstractVector, AMorAV <: Union{AbstractVector, AbstractMatrix}} - # decorrelate observations into a vector - decorrelated_obs = to_decorrelated(observation, em) + # make into iterable over vectors + obs_slice = if observation isa AbstractVector{<:AbstractVector} + observation + else # NB a vector is treated as a column here: + eachcol(observation) + end + + # encoding works on columns but mcmc wants vec-of-vec + encoded_obs = [vec(encode_data(em, reshape(obs, :, 1), "out")) for obs in obs_slice] - log_posterior_map = EmulatorPosteriorModel(prior, em, decorrelated_obs) + + log_posterior_map = EmulatorPosteriorModel(prior, em, encoded_obs) mh_proposal_sampler = MetropolisHastingsSampler(mcmc_alg, prior) # parameter names are needed in every dimension in a MCMCChains object needed for diagnostics @@ -617,7 +577,7 @@ function MCMCWrapper( :chain_type => MCMCChains.Chains, ) sample_kwargs = merge(sample_kwargs, kwargs) # override defaults with any explicit values - return MCMCWrapper(prior, observation, decorrelated_obs, log_posterior_map, mh_proposal_sampler, sample_kwargs) + return MCMCWrapper(prior, obs_slice, encoded_obs, log_posterior_map, mh_proposal_sampler, sample_kwargs) end """ @@ -777,6 +737,7 @@ function get_posterior(mcmc::MCMCWrapper, chain::MCMCChains.Chains) p_chain = Array(Chains(chain, :parameters)) # discard internal/diagnostic data p_samples = [Samples(p_chain[:, slice, 1], params_are_columns = false) for slice in p_slices] + # live in same space as prior # checks if a function distribution, by looking at if the distribution is nested p_constraints = [ diff --git a/src/RandomFeature.jl b/src/RandomFeature.jl index 2b379ef4e..7a07eae8a 100644 --- a/src/RandomFeature.jl +++ b/src/RandomFeature.jl @@ -742,7 +742,7 @@ function calculate_ensemble_mean_and_coeffnorm( moc_tmp = similar(mean_of_covs) mtmp = zeros(output_dim, n_test) - println("calculating " * string(N_ens) * " ensemble members...") + # println("calculating " * string(N_ens) * " ensemble members...") for i in ProgressBar(1:N_ens) for j in collect(1:repeats) @@ -947,7 +947,7 @@ function calculate_ensemble_mean_and_coeffnorm( n_test = length(test_idx) - println("calculating " * string(N_ens) * " ensemble members...") + # println("calculating " * string(N_ens) * " ensemble members...") nthreads = Threads.nthreads() diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl index 2333a0226..b23983f17 100644 --- a/src/ScalarRandomFeature.jl +++ b/src/ScalarRandomFeature.jl @@ -24,6 +24,8 @@ struct ScalarRandomFeatureInterface{S <: AbstractString, RNG <: AbstractRNG, KST input_dim::Int "choice of random number generator" rng::RNG + "regularization" + regularization::Vector{Union{Matrix, UniformScaling, Diagonal}} "Kernel structure type (e.g. Separable or Nonseparable)" kernel_structure::KST "Random Feature decomposition, choose from \"svd\" or \"cholesky\" (default)" @@ -76,6 +78,14 @@ gets the rng field """ EKP.get_rng(srfi::ScalarRandomFeatureInterface) = srfi.rng +""" +$(DocStringExtensions.TYPEDSIGNATURES) + +Gets the regularization field +""" +get_regularization(srfi::ScalarRandomFeatureInterface) = srfi.regularization + + """ $(DocStringExtensions.TYPEDSIGNATURES) @@ -144,6 +154,7 @@ function ScalarRandomFeatureInterface( # Initialize vector for GP models rfms = Vector{RF.Methods.RandomFeatureMethod}(undef, 0) fitted_features = Vector{RF.Methods.Fit}(undef, 0) + regularization = Vector{Union{Matrix, UniformScaling, Nothing}}(undef, 0) if isnothing(kernel_structure) kernel_structure = SeparableKernel(cov_structure_from_string("lowrank", input_dim), OneDimFactor()) @@ -191,6 +202,7 @@ function ScalarRandomFeatureInterface( n_features, input_dim, rng, + regularization, kernel_structure, feature_decomposition, optimizer_opts, @@ -294,7 +306,7 @@ RFM_from_hyperparameters( ) where { RNG <: AbstractRNG, ForVM <: Union{Real, AbstractVecOrMat}, - MorUS <: Union{Matrix, UniformScaling}, + MorUS <: Union{AbstractMatrix, UniformScaling}, S <: AbstractString, MT <: MultithreadType, } = RFM_from_hyperparameters(srfi, rng, l, regularization, n_features, batch_sizes, input_dim, multithread_type) @@ -306,14 +318,16 @@ Builds the random feature method from hyperparameters. We use cosine activation """ function build_models!( srfi::ScalarRandomFeatureInterface, - input_output_pairs::PairedDataContainer{FT}; + input_output_pairs::PairedDataContainer{FT}, + input_structure_matrix, + output_structure_matrix; kwargs..., ) where {FT <: AbstractFloat} + # get inputs and outputs input_values = get_inputs(input_output_pairs) output_values = get_outputs(input_output_pairs) n_rfms, n_data = size(output_values) - noise_sd = 1.0 input_dim = size(input_values, 1) @@ -371,10 +385,13 @@ function build_models!( end end - - - #regularization = I = 1.0 in scalar case - regularization = I + regularization = if isnothing(output_structure_matrix) + 1.0 * I(n_rfms) + elseif isa(output_structure_matrix, UniformScaling) + output_structure_matrix + else + Diagonal(output_structure_matrix) + end @info ( "hyperparameter learning for $n_rfms models using $n_train training points, $n_test validation points and $n_features_opt features" @@ -382,6 +399,7 @@ function build_models!( n_iteration = optimizer_options["n_iteration"] diagnostics = zeros(n_iteration, n_rfms) for i in 1:n_rfms + regularization_i = regularization[i, i] * I io_pairs_opt = PairedDataContainer(input_values, reshape(output_values[i, :], 1, size(output_values, 2))) @@ -422,7 +440,7 @@ function build_models!( srfi, rng, μ_hp, - regularization, + regularization_i, n_features_opt, train_idx[cv_idx], test_idx[cv_idx], @@ -434,7 +452,7 @@ function build_models!( cov_correction = cov_correction, ) Γ = internal_Γ - Γ[1:n_test, 1:n_test] += regularization # + approx_σ2 + Γ[1:n_test, 1:n_test] += regularization_i # + approx_σ2 Γ[(n_test + 1):end, (n_test + 1):end] += I if !isposdef(Γ) Γ = posdef_correct(Γ) @@ -480,7 +498,7 @@ function build_models!( srfi, rng, lvec, - regularization, + regularization_i, n_features_opt, train_idx[cv_idx], test_idx[cv_idx], @@ -536,7 +554,7 @@ function build_models!( srfi, rng, hp_optimal, - regularization, + regularization_i, n_features, batch_sizes, input_dim, @@ -548,6 +566,7 @@ function build_models!( push!(fitted_features, fitted_features_i) end + push!(get_regularization(srfi), regularization) push!(optimizer, diagnostics) end @@ -564,7 +583,7 @@ end """ $(DocStringExtensions.TYPEDSIGNATURES) -Prediction of data observation (not latent function) at new inputs (passed in as columns in a matrix). That is, we add the observational noise into predictions. +Prediction of emulator mean at new inputs (passed in as columns in a matrix), and a prediction of the total covariance at new inputs equal to (emulator covariance + noise covariance). """ function predict( srfi::ScalarRandomFeatureInterface, @@ -591,8 +610,13 @@ function predict( ) end - # add the noise contribution from the regularization - σ2[:, :] = σ2[:, :] .+ 1.0 + # add the noise contribution stored within the regularization + reg = get_regularization(srfi)[1] + reg_diag = isa(reg, UniformScaling) ? reg.λ * ones(M) : diag(reg) + + for i in 1:M + σ2[i, :] .+= reg_diag[i] + end return μ, σ2 end diff --git a/src/Utilities.jl b/src/Utilities.jl index 7a85a811d..fa36c5988 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -12,22 +12,16 @@ EnsembleKalmanProcess = EnsembleKalmanProcesses.EnsembleKalmanProcess using ..DataContainers export get_training_points -export orig2zscore -export zscore2orig export PairedDataContainerProcessor, DataContainerProcessor -export UnivariateAffineScaling, ElementwiseScaler, QuartileScaling, MinMaxScaling, ZScoreScaling -export quartile_scale, minmax_scale, zscore_scale -export Decorrelater, decorrelate_sample_cov, decorrelate_structure_mat, decorrelate -export get_type, - get_shift, get_scale, get_data_mean, get_encoder_mat, get_decoder_mat, get_retain_var, get_decorrelate_with -export create_encoder_schedule, encode_with_schedule, decode_with_schedule -export initialize_processor!, - initialize_and_encode_data!, encode_data, decode_data, encode_structure_matrix, decode_structure_matrix - -export CanonicalCorrelation -export canonical_correlation -export get_apply_to +export create_encoder_schedule, + initialize_and_encode_with_schedule!, + encode_with_schedule, + decode_with_schedule, + encode_data, + encode_structure_matrix, + decode_data, + decode_structure_matrix @@ -67,743 +61,74 @@ function get_training_points( return training_points end -function orig2zscore(X::AbstractVector{FT}, mean::AbstractVector{FT}, std::AbstractVector{FT}) where {FT} - # Compute the z scores of a vector X using the given mean - # and std - Z = zeros(size(X)) - for i in 1:length(X) - Z[i] = (X[i] - mean[i]) / std[i] - end - return Z -end - -function orig2zscore(X::AbstractMatrix{FT}, mean::AbstractVector{FT}, std::AbstractVector{FT}) where {FT} - # Compute the z scores of matrix X using the given mean and - # std. Transformation is applied column-wise. - Z = zeros(size(X)) - n_cols = size(X)[2] - for i in 1:n_cols - Z[:, i] = (X[:, i] .- mean[i]) ./ std[i] - end - return Z -end - -function zscore2orig(Z::AbstractVector{FT}, mean::AbstractVector{FT}, std::AbstractVector{FT}) where {FT} - # Transform X (a vector of z scores) back to the original - # values - X = zeros(size(Z)) - for i in 1:length(X) - X[i] = Z[i] .* std[i] .+ mean[i] - end - return X -end - -function zscore2orig(Z::AbstractMatrix{FT}, mean::AbstractVector{FT}, std::AbstractVector{FT}) where {FT} - X = zeros(size(Z)) - # Transform X (a matrix of z scores) back to the original - # values. Transformation is applied column-wise. - n_cols = size(Z)[2] - for i in 1:n_cols - X[:, i] = Z[:, i] .* std[i] .+ mean[i] - end - return X -end # Data processing tooling: -abstract type PairedDataContainerProcessor end # tools that operate on inputs and outputs -abstract type DataContainerProcessor end # tools that operate on only inputs or outputs - -abstract type UnivariateAffineScaling end -abstract type QuartileScaling <: UnivariateAffineScaling end -abstract type MinMaxScaling <: UnivariateAffineScaling end -abstract type ZScoreScaling <: UnivariateAffineScaling end - -# Processors -""" -$(TYPEDEF) - -The ElementwiseScaler{T} will create an encoding of the data_container via elementwise affine transformations. - -Different methods `T` will build different transformations: -- [`quartile_scale`](@ref) : creates `QuartileScaling`, -- [`minmax_scale`](@ref) : creates `MinMaxScaling` -- [`zscore_scale`](@ref) : creates `ZScoreScaling` - -and are accessed with [`get_type`](@ref) -""" -struct ElementwiseScaler{T, VV <: AbstractVector} <: DataContainerProcessor - "storage for the shift applied to data" - shift::VV - "storage for the scaling" - scale::VV -end - -""" -$(TYPEDSIGNATURES) - -Constructs `ElementwiseScaler{QuartileScaling}` processor. -As part of an encoder schedule, it will apply the transform ``\\frac{x - Q2(x)}{Q3(x) - Q1(x)}`` to each data dimension. -Also known as "robust scaling" -""" -quartile_scale() = ElementwiseScaler(QuartileScaling) - -""" -$(TYPEDSIGNATURES) - -Constructs `ElementwiseScaler{MinMaxScaling}` processor. -As part of an encoder schedule, this will apply the transform ``\\frac{x - \\min(x)}{\\max(x) - \\min(x)}`` to each data dimension. -""" -minmax_scale() = ElementwiseScaler(MinMaxScaling) - -""" -$(TYPEDSIGNATURES) - -Constructs `ElementwiseScaler{ZScoreScaling}` processor. -As part of an encoder schedule, this will apply the transform ``\\frac{x-\\mu}{\\sigma}``, (where ``x\\sim N(\\mu,\\sigma)``), to each data dimension. -For multivariate standardization, see [`Decorrelater`](@ref) -""" -zscore_scale() = ElementwiseScaler(ZScoreScaling) - -ElementwiseScaler(::Type{UAS}) where {UAS <: UnivariateAffineScaling} = - ElementwiseScaler{UAS, Vector{Float64}}(Float64[], Float64[]) - -""" -$(TYPEDSIGNATURES) - -Gets the UnivariateAffineScaling type `T` -""" -get_type(es::ElementwiseScaler{T}) where {T} = T - -""" -$(TYPEDSIGNATURES) - -Gets the `shift` field of the `ElementwiseScaler` -""" -get_shift(es::ElementwiseScaler) = es.shift - -""" -$(TYPEDSIGNATURES) - -Gets the `scale` field of the `ElementwiseScaler` -""" -get_scale(es::ElementwiseScaler) = es.scale - -function Base.show(io::IO, es::ElementwiseScaler) - out = "ElementwiseScaler: $(get_type(es))" - print(io, out) -end - -function initialize_processor!( - es::ElementwiseScaler, - data::MM, - T::Type{QS}, -) where {MM <: AbstractMatrix, QS <: QuartileScaling} - quartiles_vec = [quantile(dd, [0.25, 0.5, 0.75]) for dd in eachrow(data)] - quartiles_mat = reduce(hcat, quartiles_vec) # 3 rows: Q1, Q2, and Q3 - append!(get_shift(es), quartiles_mat[2, :]) - append!(get_scale(es), (quartiles_mat[3, :] - quartiles_mat[1, :])) -end - -function initialize_processor!( - es::ElementwiseScaler, - data::MM, - T::Type{MMS}, -) where {MM <: AbstractMatrix, MMS <: MinMaxScaling} - minmax_vec = [[minimum(dd), maximum(dd)] for dd in eachrow(data)] - minmax_mat = reduce(hcat, minmax_vec) # 2 rows: min max - append!(get_shift(es), minmax_mat[1, :]) - append!(get_scale(es), (minmax_mat[2, :] - minmax_mat[1, :])) -end - -function initialize_processor!( - es::ElementwiseScaler, - data::MM, - T::Type{ZSS}, -) where {MM <: AbstractMatrix, ZSS <: ZScoreScaling} - stat_vec = [[mean(dd), std(dd)] for dd in eachrow(data)] - stat_mat = reduce(hcat, stat_vec) # 2 rows: mean, std - append!(get_shift(es), stat_mat[1, :]) - append!(get_scale(es), stat_mat[2, :]) -end - -function initialize_processor!(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} - if length(get_shift(es)) == 0 - T = get_type(es) - initialize_processor!(es, data, T) - end -end - -""" -$(TYPEDSIGNATURES) - -Apply the `ElementwiseScaler` encoder, on a columns-are-data matrix -""" -function encode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} - out = deepcopy(data) - for i in 1:size(out, 1) - out[i, :] .-= get_shift(es)[i] - out[i, :] /= get_scale(es)[i] - end - return out -end - -""" -$(TYPEDSIGNATURES) - -Apply the `ElementwiseScaler` decoder, on a columns-are-data matrix -""" -function decode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} - out = deepcopy(data) - for i in 1:size(out, 1) - out[i, :] *= get_scale(es)[i] - out[i, :] .+= get_shift(es)[i] - end - return out -end - -""" -$(TYPEDSIGNATURES) - -Computes and populates the `shift` and `scale` fields for the `ElementwiseScaler` -""" -initialize_processor!(es::ElementwiseScaler, data::MM, structure_matrix) where {MM <: AbstractMatrix} = - initialize_processor!(es, data) - - -""" -$(TYPEDSIGNATURES) - -Apply the `ElementwiseScaler` encoder to a provided structure matrix -""" -function encode_structure_matrix(es::ElementwiseScaler, structure_matrix::MM) where {MM <: AbstractMatrix} - return Diagonal(1 ./ get_scale(es)) * structure_matrix * Diagonal(1 ./ get_scale(es)) -end - -""" -$(TYPEDSIGNATURES) - -Apply the `ElementwiseScaler` decoder to a provided structure matrix -""" -function decode_structure_matrix(es::ElementwiseScaler, enc_structure_matrix::MM) where {MM <: AbstractMatrix} - return Diagonal(get_scale(es)) * enc_structure_matrix * Diagonal(get_scale(es)) -end - - -""" -$(TYPEDEF) - -Decorrelate the data via taking an SVD decomposition and projecting onto the singular-vectors. - -Preferred construction is with the methods -- [`decorrelate_structure_mat`](@ref) -- [`decorrelate_sample_cov`](@ref) -- [`decorrelate`](@ref) - -For `decorrelate_structure_mat`: -The SVD is taken over a structure matrix (e.g., `prior_cov` for inputs, `obs_noise_cov` for outputs). The structure matrix will become exactly `I` after processing. - -For `decorrelate_sample_cov`: -The SVD is taken over the estimated covariance of the data. The data samples will have a `Normal(0,I)` distribution after processing, - -For `decorrelate(;decorrelate_with="combined")` (default): -The SVD is taken to be the sum of structure matrix and estimated covariance. This may be more robust to ill-specification of structure matrix, or poor estimation of the sample covariance. - -# Fields -$(TYPEDFIELDS) -""" -struct Decorrelater{VV1, VV2, VV3, FT, AS <: AbstractString} <: DataContainerProcessor - "storage for the data mean" - data_mean::VV1 - "the matrix used to perform encoding" - encoder_mat::VV2 - "the inverse of the the matrix used to perform encoding" - decoder_mat::VV3 - "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" - retain_var::FT - "Switch to choose what form of matrix to use to decorrelate the data" - decorrelate_with::AS -end - -""" -$(TYPEDSIGNATURES) - -Constructs the `Decorrelater` struct. Users can add optional keyword arguments: -- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained -- `decorrelate_with` [=`"combined"`]: from which matrix do we provide subspace directions, options are - - `"structure_mat"`, see [`decorrelate_structure_mat`](@ref) - - `"sample_cov"`, see [`decorrelate_sample_cov`](@ref) - - `"combined"`, sums the `"sample_cov"` and `"structure_mat"` matrices -""" -decorrelate(; retain_var::FT = Float64(1.0), decorrelate_with = "combined") where {FT} = - Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), decorrelate_with) - -""" -$(TYPEDSIGNATURES) - -Constructs the `Decorrelater` struct, setting decorrelate_with = "sample_cov". Encoding data with this will ensure that the distribution of data samples after encoding will be `Normal(0,I)`. One can additionally add keywords: -- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained -""" -decorrelate_sample_cov(; retain_var::FT = Float64(1.0)) where {FT} = - Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), "sample_cov") - -""" -$(TYPEDSIGNATURES) - -Constructs the `Decorrelater` struct, setting decorrelate_with = "structure_mat". This encoding will transform a provided structure matrix into `I`. One can additionally add keywords: -- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained -""" -decorrelate_structure_mat(; retain_var::FT = Float64(1.0)) where {FT} = - Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), "structure_mat") - -""" -$(TYPEDSIGNATURES) - -returns the `data_mean` field of the `Decorrelater`. -""" -get_data_mean(dd::Decorrelater) = dd.data_mean - -""" -$(TYPEDSIGNATURES) - -returns the `encoder_mat` field of the `Decorrelater`. -""" -get_encoder_mat(dd::Decorrelater) = dd.encoder_mat - -""" -$(TYPEDSIGNATURES) - -returns the `decoder_mat` field of the `Decorrelater`. -""" -get_decoder_mat(dd::Decorrelater) = dd.decoder_mat - -""" -$(TYPEDSIGNATURES) - -returns the `retain_var` field of the `Decorrelater`. -""" -get_retain_var(dd::Decorrelater) = dd.retain_var - -""" -$(TYPEDSIGNATURES) - -returns the `decorrelate_with` field of the `Decorrelater`. -""" -get_decorrelate_with(dd::Decorrelater) = dd.decorrelate_with - -function Base.show(io::IO, dd::Decorrelater) - out = "Decorrelater" - out *= ": decorrelate_with=$(get_decorrelate_with(dd))" - if get_retain_var(dd) < 1.0 - out *= ", retain_var=$(get_retain_var(dd))" - end - print(io, out) -end - -""" -$(TYPEDSIGNATURES) - -Computes and populates the `data_mean` and `encoder_mat` and `decoder_mat` fields for the `Decorrelater` -""" -function initialize_processor!( - dd::Decorrelater, - data::MM, - structure_matrix::USorM, -) where {MM <: AbstractMatrix, USorM <: Union{UniformScaling, AbstractMatrix}} - if length(get_data_mean(dd)) == 0 - push!(get_data_mean(dd), vec(mean(data, dims = 2))) - end - - if length(get_encoder_mat(dd)) == 0 - - # Can do tsvd here for large matrices - decorrelate_with = get_decorrelate_with(dd) - if decorrelate_with == "structure_mat" - svdA = svd(structure_matrix) - rk = rank(structure_matrix) - elseif decorrelate_with == "sample_cov" - cd = cov(data, dims = 2) - svdA = svd(cd) - rk = rank(cd) - elseif decorrelate_with == "combined" - spluscd = structure_matrix + cov(data, dims = 2) - svdA = svd(spluscd) - rk = rank(spluscd) - else - throw( - ArgumentError( - "Keyword `decorrelate_with` must be taken from [\"sample_cov\", \"structure_mat\", \"combined\"]. Received $(decorrelate_with)", - ), - ) - end - ret_var = get_retain_var(dd) - if ret_var < 1.0 - sv_cumsum = cumsum(svdA.S .^ 2) / sum(svdA.S .^ 2) # variance contributions are (sing_val)^2 - trunc = minimum(findall(x -> (x > ret_var), sv_cumsum)) - @info " truncating at $(trunc)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc])% of the variance of the structure matrix" - else - trunc = rk - end - sqrt_inv_sv = Diagonal(1.0 ./ sqrt.(svdA.S[1:trunc])) - sqrt_sv = Diagonal(sqrt.(svdA.S[1:trunc])) - - # as we have svd of cov-matrix we can use U or Vt - encoder_mat = sqrt_inv_sv * svdA.Vt[1:trunc, :] - decoder_mat = svdA.Vt[1:trunc, :]' * sqrt_sv - - push!(get_encoder_mat(dd), encoder_mat) - push!(get_decoder_mat(dd), decoder_mat) - end -end - - -""" -$(TYPEDSIGNATURES) - -Apply the `Decorrelater` encoder, on a columns-are-data matrix -""" -function encode_data(dd::Decorrelater, data::MM) where {MM <: AbstractMatrix} - data_mean = get_data_mean(dd)[1] - encoder_mat = get_encoder_mat(dd)[1] - return encoder_mat * (data .- data_mean) -end - -""" -$(TYPEDSIGNATURES) - -Apply the `Decorrelater` decoder, on a columns-are-data matrix -""" -function decode_data(dd::Decorrelater, data::MM) where {MM <: AbstractMatrix} - data_mean = get_data_mean(dd)[1] - decoder_mat = get_decoder_mat(dd)[1] - return decoder_mat * data .+ data_mean -end - -""" -$(TYPEDSIGNATURES) - -Apply the `Decorrelater` encoder to a provided structure matrix -""" -function encode_structure_matrix(dd::Decorrelater, structure_matrix::MM) where {MM <: AbstractMatrix} - encoder_mat = get_encoder_mat(dd)[1] - return encoder_mat * structure_matrix * encoder_mat' -end - -""" -$(TYPEDSIGNATURES) - -Apply the `Decorrelater` decoder to a provided structure matrix -""" -function decode_structure_matrix(dd::Decorrelater, enc_structure_matrix::MM) where {MM <: AbstractMatrix} - decoder_mat = get_decoder_mat(dd)[1] - return decoder_mat * enc_structure_matrix * decoder_mat' -end - -# ...struct VariationalAutoEncoder <: DataContainerProcessor end - -# PDCProcessors -# struct InverseProblemInformed <: PairedDataContainerProcessor end -# struct LikelihoodInformed <: PairedDataContainerProcessor end - - -""" -$(TYPEDEF) - -Uses both input and output data to learn a subspace of maximal correlation between inputs and outputs. The subspace for a pair (X,Y) will be of size minimum(rank(X),rank(Y)), computed using SVD-based method -e.g. See e.g., https://numerical.recipes/whp/notes/CanonCorrBySVD.pdf - -Preferred construction is with the [`canonical_correlation`](@ref) method - -# Fields -$(TYPEDFIELDS) -""" -struct CanonicalCorrelation{VV1, VV2, VV3, FT, VV4} <: PairedDataContainerProcessor - "storage for the input or output data mean" - data_mean::VV1 - "the encoding matrix of input or output canonical correlations" - encoder_mat::VV2 - "the decoding matrix of input or output canonical correlations" - decoder_mat::VV3 - "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" - retain_var::FT - "Stores whether this is an input or output encoder (vector with string \"in\" or \"out\")" - apply_to::VV4 -end - -""" -$(TYPEDSIGNATURES) - -Constructs the `CanonicalCorrelation` struct. Can optionally provide the keyword -- `retain_var`[=1.0]: to project onto the leading singular vectors (of the input-output product) such that `retain_var` variance is retained. -""" -canonical_correlation(; retain_var::FT = Float64(1.0)) where {FT} = - CanonicalCorrelation(Any[], Any[], Any[], clamp(retain_var, FT(0), FT(1)), AbstractString[]) - -""" -$(TYPEDSIGNATURES) - -returns the `data_mean` field of the `CanonicalCorrelation`. -""" -get_data_mean(cc::CanonicalCorrelation) = cc.data_mean - -""" -$(TYPEDSIGNATURES) - -returns the `encoder_mat` field of the `CanonicalCorrelation`. -""" -get_encoder_mat(cc::CanonicalCorrelation) = cc.encoder_mat +abstract type DataProcessor end +abstract type PairedDataContainerProcessor <: DataProcessor end # tools that operate on inputs and outputs +abstract type DataContainerProcessor <: DataProcessor end # tools that operate on only inputs or outputs -""" -$(TYPEDSIGNATURES) +# define how to have equality +Base.:(==)(a::DCP, b::DCP) where {DCP <: DataContainerProcessor} = + all(getfield(a, f) == getfield(b, f) for f in fieldnames(DCP)) +Base.:(==)(a::PDCP, b::PDCP) where {PDCP <: PairedDataContainerProcessor} = + all(getfield(a, f) == getfield(b, f) for f in fieldnames(PDCP)) -returns the `decoder_mat` field of the `CanonicalCorrelation`. -""" -get_decoder_mat(cc::CanonicalCorrelation) = cc.decoder_mat -""" -$(TYPEDSIGNATURES) - -returns the `retain_var` field of the `CanonicalCorrelation`. -""" -get_retain_var(cc::CanonicalCorrelation) = cc.retain_var - -""" -$(TYPEDSIGNATURES) - -returns the `apply_to` field of the `CanonicalCorrelation`. -""" -get_apply_to(cc::CanonicalCorrelation) = cc.apply_to +#### -function Base.show(io::IO, cc::CanonicalCorrelation) - out = "CanonicalCorrelation:" - if length(get_apply_to(cc)) > 0 - out *= " apply_to=$(get_apply_to(cc)[1])" - end - if get_retain_var(cc) < 1.0 - out *= " retain_var=$(get_retain_var(cc))" +function _encode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString} + input_data, output_data = get_data(data) + if apply_to == "in" + return encode_data(proc, input_data) + elseif apply_to == "out" + return encode_data(proc, output_data) + else + bad_apply_to(apply_to) end - print(io, out) end +function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString} + input_data, output_data = get_data(data) -function initialize_processor!( - cc::CanonicalCorrelation, - in_data::MM, - out_data::MM, - apply_to::AS, -) where {MM <: AbstractMatrix, AS <: AbstractString} - - if apply_to ∉ ["in", "out"] + if apply_to == "in" + return decode_data(proc, input_data) + elseif apply_to == "out" + return decode_data(proc, output_data) + else bad_apply_to(apply_to) end - - if length(get_apply_to(cc)) == 0 - push!(get_apply_to(cc), apply_to) - end - - if length(get_data_mean(cc)) == 0 - if apply_to == "in" - push!(get_data_mean(cc), vec(mean(in_data, dims = 2))) - elseif apply_to == "out" - push!(get_data_mean(cc), vec(mean(out_data, dims = 2))) - end - end - - if length(get_encoder_mat(cc)) == 0 - - if size(in_data, 2) < size(in_data, 1) || size(out_data, 2) < size(out_data, 1) - throw( - ArgumentError( - "CanonicalCorrelation implementation not defined for # data samples < dimensions, please obtain more samples, or perform prior dimension reduction approaches until this is satisfied", - ), - ) - end - - # Individually decompose in and out - svdi = svd(in_data .- mean(in_data, dims = 2)) - svdo = svd(out_data .- mean(out_data, dims = 2)) - - # ensure correct shaping (in_mat = (in_dim x n_samples), out_mat = (out_dim x n_samples)) - in_mat_sq, in_mat_nonsq = (size(svdi.U, 1) == size(svdi.U, 2)) ? (svdi.U, svdi.Vt) : (svdi.Vt, svdi.U) - out_mat_sq, out_mat_nonsq = (size(svdo.U, 1) == size(svdo.U, 2)) ? (svdo.U, svdo.Vt) : (svdo.Vt, svdo.U) - - svdio = svd(out_mat_nonsq * in_mat_nonsq') - - # retain variance - ret_var = get_retain_var(cc) - if ret_var < 1.0 - sv_cumsum = cumsum(svdio.S .^ 2) / sum(svdio.S .^ 2) # variance contributions are (sing_val)^2 - trunc = minimum(findall(x -> (x > ret_var), sv_cumsum)) - @info " truncating at $(trunc)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc])% of the variance in the joint space" - else - trunc = min(rank(in_data), rank(out_data)) - end - - if apply_to == "in" - in_dim = size(in_data, 1) - svdio_mat = (size(svdio.U, 1) == in_dim) ? svdio.U : svdio.V - # mat' * Sx⁻¹ * Uxt - encoder_mat = svdio_mat[:, 1:trunc]' * Diagonal(1 ./ svdi.S) * in_mat_sq' - decoder_mat = in_mat_sq * Diagonal(svdi.S) * svdio_mat[:, 1:trunc] - elseif apply_to == "out" - out_dim = size(out_data, 1) - svdio_mat = (size(svdio.U, 1) == out_dim) ? svdio.U : svdio.V - # Vt * Sy⁻¹ * Uyt - encoder_mat = svdio_mat[:, 1:trunc]' * Diagonal(1 ./ svdo.S) * out_mat_sq' - decoder_mat = out_mat_sq * Diagonal(svdo.S) * svdio_mat[:, 1:trunc] - end - - push!(get_encoder_mat(cc), encoder_mat) - push!(get_decoder_mat(cc), decoder_mat) - - # Note: To check CCA: - # u = in_encoder * (in - mean(in)) - # v = out_encoder * (out - mean(out)) - # u * u' = v * v' = I, - # v * u' = u * v' = Diagonal(svdio.S[1:trunc]) - end end -""" -$(TYPEDSIGNATURES) - -Computes and populates the `data_mean`, `encoder_mat`, `decoder_mat` and `apply_to` fields for the `CanonicalCorrelation` -""" -initialize_processor!( - cc::CanonicalCorrelation, - in_data::MM, - out_data::MM, - structure_matrix, +function _initialize_and_encode_data!( + proc::PairedDataContainerProcessor, + data, + structure_mats, apply_to::AS, -) where {MM <: AbstractMatrix, AS <: AbstractString} = initialize_processor!(cc, in_data, out_data, apply_to) - - -""" -$(TYPEDSIGNATURES) - -Apply the `CanonicalCorrelation` encoder, on a columns-are-data matrix -""" -function encode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} - data_mean = get_data_mean(cc)[1] - encoder_mat = get_encoder_mat(cc)[1] - return encoder_mat * (data .- data_mean) +) where {AS <: AbstractString} + initialize_processor!(proc, get_data(data)..., structure_mats..., apply_to) + return _encode_data(proc, data, apply_to) end -""" -$(TYPEDSIGNATURES) - -Apply the `CanonicalCorrelation` decoder, on a columns-are-data matrix -""" -function decode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} - data_mean = get_data_mean(cc)[1] - decoder_mat = get_decoder_mat(cc)[1] - return decoder_mat * data .+ data_mean -end - -""" -$(TYPEDSIGNATURES) - -Apply the `CanonicalCorrelation` encoder to a provided structure matrix -""" -function encode_structure_matrix(cc::CanonicalCorrelation, structure_matrix::MM) where {MM <: AbstractMatrix} - encoder_mat = get_encoder_mat(cc)[1] - return encoder_mat * structure_matrix * encoder_mat' -end - -""" -$(TYPEDSIGNATURES) - -Apply the `CanonicalCorrelation` decoder to a provided structure matrix -""" -function decode_structure_matrix(cc::CanonicalCorrelation, enc_structure_matrix::MM) where {MM <: AbstractMatrix} - decoder_mat = get_decoder_mat(cc)[1] - return decoder_mat * enc_structure_matrix * decoder_mat' -end - - - - - -#### - - -# generic functions to initialize, encode, and decode data. With Data processors or PairedData processors -function initialize_and_encode_data!( - dcp::DCP, - data::MM, - structure_mat::USorM, -) where {DCP <: DataContainerProcessor, USorM <: Union{UniformScaling, AbstractMatrix}, MM <: AbstractMatrix} - initialize_processor!(dcp, data, structure_mat) - return encode_data(dcp, data) -end - -""" -$(TYPEDSIGNATURES) - -Initializes the `DataContainerProcessor` encoder (often requires data, and structure matrices), then encodes the provided columns-are-data matrix -""" -initialize_and_encode_data!( - dcp::DCP, - data::MM, - structure_mat::USorM, - apply_to::AS, -) where { - DCP <: DataContainerProcessor, - MM <: AbstractMatrix, - USorM <: Union{UniformScaling, AbstractMatrix}, - AS <: AbstractString, -} = initialize_and_encode_data!(dcp, data, structure_mat) - - -""" -$(TYPEDSIGNATURES) - -decodes the columns-are-data matrix with the processor. -""" -decode_data( - dcp::DCP, - data::MM, - apply_to::AS, -) where {DCP <: DataContainerProcessor, MM <: AbstractMatrix, AS <: AbstractString} = decode_data(dcp, data) - -""" -$(TYPEDSIGNATURES) - -Initializes the `PairedDataContainerProcesser` encoder (often requires input & output data, and structure matrices), then encodes either the input or output data (pair of columns-are-data matrices) based on `apply_to`. -""" -function initialize_and_encode_data!( - dcp::PDCP, +function _initialize_and_encode_data!( + proc::DataContainerProcessor, data, - structure_mat::USorM, + structure_mats, apply_to::AS, -) where {PDCP <: PairedDataContainerProcessor, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} - input_data, output_data = data - initialize_processor!(dcp, input_data, output_data, structure_mat, apply_to) +) where {AS <: AbstractString} + input_data, output_data = get_data(data) + input_structure_mat, output_structure_mat = structure_mats if apply_to == "in" - return encode_data(dcp, input_data) + initialize_processor!(proc, input_data, input_structure_mat) elseif apply_to == "out" - return encode_data(dcp, output_data) + initialize_processor!(proc, output_data, output_structure_mat) else bad_apply_to(apply_to) end -end -""" -$(TYPEDSIGNATURES) - -decodes the input or output dat (pair of columns-are-data matrices) a with the processor, based on `apply_to`. -""" -function decode_data(dcp::PDCP, data, apply_to::AS) where {PDCP <: PairedDataContainerProcessor, AS <: AbstractString} - input_data, output_data = data - if apply_to == "in" - return decode_data(dcp, input_data) - elseif apply_to == "out" - return decode_data(dcp, output_data) - else - bad_apply_to(apply_to) - end + return _encode_data(proc, data, apply_to) end """ @@ -819,58 +144,32 @@ enc_schedule = [ (DataProcessor4(...), "in_and_out"), ] ``` -This function creates the encoder scheduler that is also machine readable +This function creates the encoder scheduler that is also machine readable. E.g., ```julia enc_schedule = [ - (DataProcessor1(...), x -> get_inputs(x), "in"), - (DataProcessor2(...), x -> get_outputs(x), "out"), - (DataProcessor2(...), x -> get_outputs(x), "out"), - (PairedDataProcessor3(...), x -> (get_outputs(x), get_outputs(x)), "in"), - (DataProcessor4(...), x -> get_inputs(x), "in"), - (DataProcessor4(...), x -> get_outputs(x), "out"), + (DataProcessor1(...), "in"), + (DataProcessor2(...), "out"), + (DataProcessor2(...), "out"), + (PairedDataProcessor3(...),"in"), + (DataProcessor4(...), "in"), + (DataProcessor4(...), "out"), ] ``` and the decoder schedule is a copy of the encoder schedule reversed (and processors copied) """ function create_encoder_schedule(schedule_in::VV) where {VV <: AbstractVector} - encoder_schedule = [] for (processor, apply_to) in schedule_in # converts the string into the extraction of data - if isa(processor, DataContainerProcessor) - if apply_to == "in" - func = x -> get_inputs(x) - push!(encoder_schedule, (processor, func, apply_to)) - - elseif apply_to == "out" - func = x -> get_outputs(x) - push!(encoder_schedule, (processor, func, apply_to)) - - elseif apply_to == "in_and_out" - func1 = x -> get_inputs(x) - func2 = x -> get_outputs(x) - push!(encoder_schedule, (processor, func1, "in")) - push!(encoder_schedule, (deepcopy(processor), func2, "out")) - else - @warn( - "Expected schedule keywords ∈ {\"in\",\"out\",\"in_and_out\"}. Received $(apply_to), ignoring processor $(processor)..." - ) - end - - # extract all the data (needed for the paired processor, but still pass through what you apply to) - elseif isa(processor, PairedDataContainerProcessor) - if apply_to ∈ ["in", "out"] - func = x -> (get_inputs(x), get_outputs(x)) - push!(encoder_schedule, (processor, func, apply_to)) - elseif apply_to == "in_and_out" - func = x -> (get_inputs(x), get_outputs(x)) - push!(encoder_schedule, (processor, func, "in")) - push!(encoder_schedule, (deepcopy(processor), func, "out")) - else - @warn( - "Expected schedule keywords ∈ {\"in\",\"out\",\"in_and_out\"}. Received $(apply_to), ignoring processor $(processor)..." - ) - end + if apply_to ∈ ["in", "out"] + push!(encoder_schedule, (processor, apply_to)) + elseif apply_to == "in_and_out" + push!(encoder_schedule, (processor, "in")) + push!(encoder_schedule, (deepcopy(processor), "out")) + else + @warn( + "Expected schedule keywords ∈ {\"in\",\"out\",\"in_and_out\"}. Received $(apply_to), ignoring processor $(processor)..." + ) end end @@ -882,52 +181,43 @@ Create size-1 encoder schedule with a tuple of `(DataProcessor1(...), apply_to)` """ create_encoder_schedule(schedule_in::TT) where {TT <: Tuple} = create_encoder_schedule([schedule_in]) - -function bad_apply_to(apply_to::AS) where {AS <: AbstractString} - throw( - ArgumentError( - "processer can only be applied to inputs (\"in\") or outputs (\"out\"). received $(apply_to). \n Please use `create_encoder_schedule` prior to encoding/decoding to ensure correct schedule format", - ), - ) -end - # Functions to encode/decode with uninitialized schedule (require structure matrices as input) """ $TYPEDSIGNATURES Takes in the created encoder schedule (See [`create_encoder_schedule`](@ref)), and initializes it, and encodes the paired data container, and structure matrices with it. """ -function encode_with_schedule( +function initialize_and_encode_with_schedule!( encoder_schedule::VV, io_pairs::PDC, - input_structure_mat::USorM1, - output_structure_mat::USorM2, + input_structure_mat::USorMorN1, + output_structure_mat::USorMorN2, ) where { VV <: AbstractVector, PDC <: PairedDataContainer, - USorM1 <: Union{UniformScaling, AbstractMatrix}, - USorM2 <: Union{UniformScaling, AbstractMatrix}, + USorMorN1 <: Union{UniformScaling, AbstractMatrix, Nothing}, + USorMorN2 <: Union{UniformScaling, AbstractMatrix, Nothing}, } processed_io_pairs = deepcopy(io_pairs) processed_input_structure_mat = deepcopy(input_structure_mat) processed_output_structure_mat = deepcopy(output_structure_mat) # apply_to is the string "in", "out" etc. - for (processor, extract_data, apply_to) in encoder_schedule + for (processor, apply_to) in encoder_schedule @info "Initialize encoding of data: \"$(apply_to)\" with $(processor)" + + processed = _initialize_and_encode_data!( + processor, + processed_io_pairs, + (processed_input_structure_mat, processed_output_structure_mat), + apply_to, + ) + if apply_to == "in" - structure_matrix = processed_input_structure_mat - elseif apply_to == "out" - structure_matrix = processed_output_structure_mat - else - bad_apply_to(apply_to) - end - processed = initialize_and_encode_data!(processor, extract_data(processed_io_pairs), structure_matrix, apply_to) - if apply_to == "in" - processed_input_structure_mat = encode_structure_matrix(processor, structure_matrix) + processed_input_structure_mat = encode_structure_matrix(processor, processed_input_structure_mat) processed_io_pairs = PairedDataContainer(processed, get_outputs(processed_io_pairs)) elseif apply_to == "out" - processed_output_structure_mat = encode_structure_matrix(processor, structure_matrix) + processed_output_structure_mat = encode_structure_matrix(processor, processed_output_structure_mat) processed_io_pairs = PairedDataContainer(get_inputs(processed_io_pairs), processed) end end @@ -936,6 +226,21 @@ function encode_with_schedule( end # Functions to encode/decode with initialized schedule + +# cases when structure_matrix is Nothing: +encode_structure_matrix(dp::DP, n::Nothing) where {DP <: DataProcessor} = nothing +decode_structure_matrix(dp::DP, n::Nothing) where {DP <: DataProcessor} = nothing +encode_with_schedule( + encoder_schedule::VV, + n::Nothing, + in_or_out::AS, +) where {VV <: AbstractVector, AS <: AbstractString} = nothing +decode_with_schedule( + encoder_schedule::VV, + n::Nothing, + in_or_out::AS, +) where {VV <: AbstractVector, AS <: AbstractString} = nothing + """ $TYPEDSIGNATURES @@ -946,33 +251,22 @@ function encode_with_schedule( data_container::DC, in_or_out::AS, ) where {VV <: AbstractVector, DC <: DataContainer, AS <: AbstractString} - if in_or_out ∉ ["in", "out"] bad_in_or_out(in_or_out) end processed_container = deepcopy(data_container) # apply_to is the string "in", "out" etc. - for (processor, extract_data, apply_to) in encoder_schedule + for (processor, apply_to) in encoder_schedule if apply_to == in_or_out processed = encode_data(processor, get_data(processed_container)) processed_container = DataContainer(processed) - end end return processed_container end -function bad_in_or_out(in_or_out::AS) where {AS <: AbstractString} - throw( - ArgumentError( - "`in_or_out` must be either \"in\" (data is an input) or \"out\" (data is an output). Received $(in_or_out)", - ), - ) -end - - """ $TYPEDSIGNATURES @@ -983,14 +277,13 @@ function encode_with_schedule( structure_matrix::USorM, in_or_out::AS, ) where {VV <: AbstractVector, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} - - if !(in_or_out ∈ ["in", "out"]) + if in_or_out ∉ ["in", "out"] bad_in_or_out(in_or_out) end processed_structure_matrix = deepcopy(structure_matrix) # apply_to is the string "in", "out" etc. - for (processor, extract_data, apply_to) in encoder_schedule + for (processor, apply_to) in encoder_schedule if apply_to == in_or_out processed_structure_matrix = encode_structure_matrix(processor, processed_structure_matrix) end @@ -1015,15 +308,14 @@ function decode_with_schedule( USorM2 <: Union{UniformScaling, AbstractMatrix}, PDC <: PairedDataContainer, } - processed_io_pairs = deepcopy(io_pairs) processed_input_structure_mat = deepcopy(input_structure_mat) processed_output_structure_mat = deepcopy(output_structure_mat) # apply_to is the string "in", "out" etc. for idx in reverse(eachindex(encoder_schedule)) - (processor, extract_data, apply_to) = encoder_schedule[idx] - processed = decode_data(processor, extract_data(processed_io_pairs), apply_to) + (processor, apply_to) = encoder_schedule[idx] + processed = _decode_data(processor, processed_io_pairs, apply_to) if apply_to == "in" processed_input_structure_mat = decode_structure_matrix(processor, processed_input_structure_mat) @@ -1049,7 +341,6 @@ function decode_with_schedule( data_container::DC, in_or_out::AS, ) where {VV <: AbstractVector, DC <: DataContainer, AS <: AbstractString} - if in_or_out ∉ ["in", "out"] bad_in_or_out(in_or_out) end @@ -1057,11 +348,10 @@ function decode_with_schedule( # apply_to is the string "in", "out" etc. for idx in reverse(eachindex(encoder_schedule)) - (processor, extract_data, apply_to) = encoder_schedule[idx] + (processor, apply_to) = encoder_schedule[idx] if apply_to == in_or_out processed = decode_data(processor, get_data(processed_container)) processed_container = DataContainer(processed) - end end @@ -1078,7 +368,6 @@ function decode_with_schedule( structure_matrix::USorM, in_or_out::AS, ) where {VV <: AbstractVector, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} - if in_or_out ∉ ["in", "out"] bad_in_or_out(in_or_out) end @@ -1086,7 +375,7 @@ function decode_with_schedule( # apply_to is the string "in", "out" etc. for idx in reverse(eachindex(encoder_schedule)) - (processor, extract_data, apply_to) = encoder_schedule[idx] + (processor, apply_to) = encoder_schedule[idx] if apply_to == in_or_out processed_structure_matrix = decode_structure_matrix(processor, processed_structure_matrix) end @@ -1096,5 +385,29 @@ function decode_with_schedule( end +# Errors + +function bad_apply_to(apply_to::AS) where {AS <: AbstractString} + throw( + ArgumentError( + "processer can only be applied to inputs (\"in\") or outputs (\"out\"). received $(apply_to). \n Please use `create_encoder_schedule` prior to encoding/decoding to ensure correct schedule format", + ), + ) +end + +function bad_in_or_out(in_or_out::AS) where {AS <: AbstractString} + throw( + ArgumentError( + "`in_or_out` must be either \"in\" (data is an input) or \"out\" (data is an output). Received $(in_or_out)", + ), + ) +end + + +# Processors + +include("Utilities/canonical_correlation.jl") +include("Utilities/decorrelator.jl") +include("Utilities/elementwise_scaler.jl") end # module diff --git a/src/Utilities/canonical_correlation.jl b/src/Utilities/canonical_correlation.jl new file mode 100644 index 000000000..e8be4a611 --- /dev/null +++ b/src/Utilities/canonical_correlation.jl @@ -0,0 +1,225 @@ +# included in Utilities.jl + +export CanonicalCorrelation, canonical_correlation, get_apply_to +export get_shift, get_scale, get_data_mean, get_encoder_mat, get_decoder_mat, get_retain_var, get_decorrelate_with + +""" +$(TYPEDEF) + +Uses both input and output data to learn a subspace of maximal correlation between inputs and outputs. The subspace for a pair (X,Y) will be of size minimum(rank(X),rank(Y)), computed using SVD-based method +e.g. See e.g., https://numerical.recipes/whp/notes/CanonCorrBySVD.pdf + +Preferred construction is with the [`canonical_correlation`](@ref) method + +# Fields +$(TYPEDFIELDS) +""" +struct CanonicalCorrelation{VV1, VV2, VV3, FT, VV4} <: PairedDataContainerProcessor + "storage for the input or output data mean" + data_mean::VV1 + "the encoding matrix of input or output canonical correlations" + encoder_mat::VV2 + "the decoding matrix of input or output canonical correlations" + decoder_mat::VV3 + "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" + retain_var::FT + "Stores whether this is an input or output encoder (vector with string \"in\" or \"out\")" + apply_to::VV4 +end + +""" +$(TYPEDSIGNATURES) + +Constructs the `CanonicalCorrelation` struct. Can optionally provide the keyword +- `retain_var`[=1.0]: to project onto the leading singular vectors (of the input-output product) such that `retain_var` variance is retained. +""" +canonical_correlation(; retain_var::FT = Float64(1.0)) where {FT} = + CanonicalCorrelation(Any[], Any[], Any[], clamp(retain_var, FT(0), FT(1)), AbstractString[]) + +""" +$(TYPEDSIGNATURES) + +returns the `data_mean` field of the `CanonicalCorrelation`. +""" +get_data_mean(cc::CanonicalCorrelation) = cc.data_mean + +""" +$(TYPEDSIGNATURES) + +returns the `encoder_mat` field of the `CanonicalCorrelation`. +""" +get_encoder_mat(cc::CanonicalCorrelation) = cc.encoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `decoder_mat` field of the `CanonicalCorrelation`. +""" +get_decoder_mat(cc::CanonicalCorrelation) = cc.decoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `retain_var` field of the `CanonicalCorrelation`. +""" +get_retain_var(cc::CanonicalCorrelation) = cc.retain_var + +""" +$(TYPEDSIGNATURES) + +returns the `apply_to` field of the `CanonicalCorrelation`. +""" +get_apply_to(cc::CanonicalCorrelation) = cc.apply_to + +function Base.show(io::IO, cc::CanonicalCorrelation) + + out = "CanonicalCorrelation:" + if length(get_apply_to(cc)) > 0 + out *= " apply_to=$(get_apply_to(cc)[1])" + end + if get_retain_var(cc) < 1.0 + out *= " retain_var=$(get_retain_var(cc))" + end + print(io, out) +end + + +function initialize_processor!( + cc::CanonicalCorrelation, + in_data::MM, + out_data::MM, + apply_to::AS, +) where {MM <: AbstractMatrix, AS <: AbstractString} + + if apply_to ∉ ["in", "out"] + bad_apply_to(apply_to) + end + + if length(get_apply_to(cc)) == 0 + push!(get_apply_to(cc), apply_to) + end + + if length(get_data_mean(cc)) == 0 + if apply_to == "in" + push!(get_data_mean(cc), vec(mean(in_data, dims = 2))) + elseif apply_to == "out" + push!(get_data_mean(cc), vec(mean(out_data, dims = 2))) + end + end + + if length(get_encoder_mat(cc)) == 0 + + if size(in_data, 2) < size(in_data, 1) || size(out_data, 2) < size(out_data, 1) + throw( + ArgumentError( + "CanonicalCorrelation implementation not defined for # data samples < dimensions, please obtain more samples, or perform prior dimension reduction approaches until this is satisfied", + ), + ) + end + + # Individually decompose in and out + svdi = svd(in_data .- mean(in_data, dims = 2)) + svdo = svd(out_data .- mean(out_data, dims = 2)) + + # ensure correct shaping (in_mat = (in_dim x n_samples), out_mat = (out_dim x n_samples)) + in_mat_sq, in_mat_nonsq = (size(svdi.U, 1) == size(svdi.U, 2)) ? (svdi.U, svdi.Vt) : (svdi.Vt, svdi.U) + out_mat_sq, out_mat_nonsq = (size(svdo.U, 1) == size(svdo.U, 2)) ? (svdo.U, svdo.Vt) : (svdo.Vt, svdo.U) + + svdio = svd(in_mat_nonsq * out_mat_nonsq') + + # retain variance + ret_var = get_retain_var(cc) + if ret_var < 1.0 + sv_cumsum = cumsum(svdio.S .^ 2) / sum(svdio.S .^ 2) # variance contributions are (sing_val)^2for these matrices + trunc_val = minimum(findall(x -> (x > ret_var), sv_cumsum)) + @info " truncating at $(trunc_val)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc_val])% of the variance in the joint space" + else + trunc_val = min(rank(in_data), rank(out_data)) + end + in_dim = size(in_data, 1) + in_svdio_mat, out_svdio_mat = (size(svdio.U, 1) == in_dim) ? (svdio.U, svdio.V) : (svdio.V, svdio.U') + if apply_to == "in" + # mat' * Sx⁻¹ * Uxt + encoder_mat = in_svdio_mat[:, 1:trunc_val]' * Diagonal(1 ./ svdi.S) * in_mat_sq' + decoder_mat = in_mat_sq * Diagonal(svdi.S) * in_svdio_mat[:, 1:trunc_val] + + elseif apply_to == "out" + out_dim = size(out_data, 1) + # Vt * Sy⁻¹ * Uyt + encoder_mat = out_svdio_mat[:, 1:trunc_val]' * Diagonal(1 ./ svdo.S) * out_mat_sq' + decoder_mat = out_mat_sq * Diagonal(svdo.S) * out_svdio_mat[:, 1:trunc_val] + end + + push!(get_encoder_mat(cc), encoder_mat) + push!(get_decoder_mat(cc), decoder_mat) + + # Note: To check CCA: + # u = in_encoder * (in_data .- mean(in_data, dims=2)) + # v = out_encoder * (out_data .- mean(out_data, dims=2)) + # u * u' = v * v' = I, + # v * u' = u * v' = Diagonal(svdio.S[1:trunc_val]) + end +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `data_mean`, `encoder_mat`, `decoder_mat` and `apply_to` fields for the `CanonicalCorrelation` +""" +initialize_processor!( + cc::CanonicalCorrelation, + in_data::MM, + out_data::MM, + input_structure_matrix, + output_structure_matrix, + apply_to::AS, +) where {MM <: AbstractMatrix, AS <: AbstractString} = initialize_processor!(cc, in_data, out_data, apply_to) + + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` encoder, on a columns-are-data matrix +""" +function encode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(cc)[1] + encoder_mat = get_encoder_mat(cc)[1] + return encoder_mat * (data .- data_mean) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` decoder, on a columns-are-data matrix +""" +function decode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(cc)[1] + decoder_mat = get_decoder_mat(cc)[1] + return decoder_mat * data .+ data_mean +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` encoder to a provided structure matrix +""" +function encode_structure_matrix( + cc::CanonicalCorrelation, + structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + encoder_mat = get_encoder_mat(cc)[1] + return encoder_mat * structure_matrix * encoder_mat' +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` decoder to a provided structure matrix +""" +function decode_structure_matrix( + cc::CanonicalCorrelation, + enc_structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + decoder_mat = get_decoder_mat(cc)[1] + return decoder_mat * enc_structure_matrix * decoder_mat' +end diff --git a/src/Utilities/decorrelator.jl b/src/Utilities/decorrelator.jl new file mode 100644 index 000000000..62b51da1a --- /dev/null +++ b/src/Utilities/decorrelator.jl @@ -0,0 +1,234 @@ +# included in Utilities.jl + +export Decorrelator, decorrelate_sample_cov, decorrelate_structure_mat, decorrelate +export get_data_mean, get_encoder_mat, get_decoder_mat, get_retain_var, get_decorrelate_with + +""" +$(TYPEDEF) + +Decorrelate the data via taking an SVD decomposition and projecting onto the singular-vectors. + +Preferred construction is with the methods +- [`decorrelate_structure_mat`](@ref) +- [`decorrelate_sample_cov`](@ref) +- [`decorrelate`](@ref) + +For `decorrelate_structure_mat`: +The SVD is taken over a structure matrix (e.g., `prior_cov` for inputs, `obs_noise_cov` for outputs). The structure matrix will become exactly `I` after processing. + +For `decorrelate_sample_cov`: +The SVD is taken over the estimated covariance of the data. The data samples will have a `Normal(0,I)` distribution after processing, + +For `decorrelate(;decorrelate_with="combined")` (default): +The SVD is taken to be the sum of structure matrix and estimated covariance. This may be more robust to ill-specification of structure matrix, or poor estimation of the sample covariance. + +# Fields +$(TYPEDFIELDS) +""" +struct Decorrelator{VV1, VV2, VV3, FT, AS <: AbstractString} <: DataContainerProcessor + "storage for the data mean" + data_mean::VV1 + "the matrix used to perform encoding" + encoder_mat::VV2 + "the inverse of the the matrix used to perform encoding" + decoder_mat::VV3 + "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" + retain_var::FT + "Switch to choose what form of matrix to use to decorrelate the data" + decorrelate_with::AS +end + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelator` struct. Users can add optional keyword arguments: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +- `decorrelate_with` [=`"combined"`]: from which matrix do we provide subspace directions, options are + - `"structure_mat"`, see [`decorrelate_structure_mat`](@ref) + - `"sample_cov"`, see [`decorrelate_sample_cov`](@ref) + - `"combined"`, sums the `"sample_cov"` and `"structure_mat"` matrices +""" +decorrelate(; retain_var::FT = Float64(1.0), decorrelate_with = "combined") where {FT} = + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), decorrelate_with) + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelator` struct, setting decorrelate_with = "sample_cov". Encoding data with this will ensure that the distribution of data samples after encoding will be `Normal(0,I)`. One can additionally add keywords: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +""" +decorrelate_sample_cov(; retain_var::FT = Float64(1.0)) where {FT} = + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), "sample_cov") + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelator` struct, setting decorrelate_with = "structure_mat". This encoding will transform a provided structure matrix into `I`. One can additionally add keywords: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +""" +decorrelate_structure_mat(; retain_var::FT = Float64(1.0)) where {FT} = + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), "structure_mat") + +""" +$(TYPEDSIGNATURES) + +returns the `data_mean` field of the `Decorrelator`. +""" +get_data_mean(dd::Decorrelator) = dd.data_mean + +""" +$(TYPEDSIGNATURES) + +returns the `encoder_mat` field of the `Decorrelator`. +""" +get_encoder_mat(dd::Decorrelator) = dd.encoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `decoder_mat` field of the `Decorrelator`. +""" +get_decoder_mat(dd::Decorrelator) = dd.decoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `retain_var` field of the `Decorrelator`. +""" +get_retain_var(dd::Decorrelator) = dd.retain_var + +""" +$(TYPEDSIGNATURES) + +returns the `decorrelate_with` field of the `Decorrelator`. +""" +get_decorrelate_with(dd::Decorrelator) = dd.decorrelate_with + +function Base.show(io::IO, dd::Decorrelator) + out = "Decorrelator" + out *= ": decorrelate_with=$(get_decorrelate_with(dd))" + if get_retain_var(dd) < 1.0 + out *= ", retain_var=$(get_retain_var(dd))" + end + print(io, out) +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `data_mean` and `encoder_mat` and `decoder_mat` fields for the `Decorrelator` +""" +function initialize_processor!( + dd::Decorrelator, + data::MM, + structure_matrix::USorMorN, +) where {MM <: AbstractMatrix, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} + if length(get_data_mean(dd)) == 0 + push!(get_data_mean(dd), vec(mean(data, dims = 2))) + end + + if length(get_encoder_mat(dd)) == 0 + + # Can do tsvd here for large matrices + decorrelate_with = get_decorrelate_with(dd) + if decorrelate_with == "structure_mat" + if isnothing(structure_matrix) + throw( + ArgumentError( + "DataProcessor `decorrelate_structure_mat` requires a user-provided structure matrix: received `nothing`. \n please provide (for input or output as needed) as a keyword argument `Emulator(...; input_structure_matrix=..., output_structure_matrix=...)` ", + ), + ) + elseif isa(structure_matrix, UniformScaling) + data_dim = size(data, 1) + svdA = svd(structure_matrix(data_dim)) + rk = data_dim + else + svdA = svd(structure_matrix) + rk = rank(structure_matrix) + end + elseif decorrelate_with == "sample_cov" + cd = cov(data, dims = 2) + svdA = svd(cd) + rk = rank(cd) + elseif decorrelate_with == "combined" + spluscd = structure_matrix + cov(data, dims = 2) + svdA = svd(spluscd) + rk = rank(spluscd) + else + throw( + ArgumentError( + "Keyword `decorrelate_with` must be taken from [\"sample_cov\", \"structure_mat\", \"combined\"]. Received $(decorrelate_with)", + ), + ) + end + ret_var = get_retain_var(dd) + if ret_var < 1.0 + sv_cumsum = cumsum(svdA.S) / sum(svdA.S) # variance contributions are (sing_val) for these matrices + trunc_val = minimum(findall(x -> (x > ret_var), sv_cumsum)) + @info " truncating at $(trunc_val)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc_val])% of the variance of the structure matrix" + else + trunc_val = rk + if rk < size(data, 1) + @info " truncating at $(trunc_val)/$(size(data,1)), as low-rank data detected" + end + end + + sqrt_inv_sv = Diagonal(1.0 ./ sqrt.(svdA.S[1:trunc_val])) + sqrt_sv = Diagonal(sqrt.(svdA.S[1:trunc_val])) + # as we have svd of cov-matrix we can use U or Vt + encoder_mat = sqrt_inv_sv * svdA.Vt[1:trunc_val, :] + decoder_mat = svdA.Vt[1:trunc_val, :]' * sqrt_sv + + push!(get_encoder_mat(dd), encoder_mat) + push!(get_decoder_mat(dd), decoder_mat) + end +end + + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelator` encoder, on a columns-are-data matrix +""" +function encode_data(dd::Decorrelator, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(dd)[1] + encoder_mat = get_encoder_mat(dd)[1] + return encoder_mat * (data .- data_mean) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelator` decoder, on a columns-are-data matrix +""" +function decode_data(dd::Decorrelator, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(dd)[1] + decoder_mat = get_decoder_mat(dd)[1] + return decoder_mat * data .+ data_mean +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelator` encoder to a provided structure matrix +""" +function encode_structure_matrix( + dd::Decorrelator, + structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + encoder_mat = get_encoder_mat(dd)[1] + return encoder_mat * structure_matrix * encoder_mat' +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelator` decoder to a provided structure matrix +""" +function decode_structure_matrix( + dd::Decorrelator, + enc_structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + decoder_mat = get_decoder_mat(dd)[1] + return decoder_mat * enc_structure_matrix * decoder_mat' +end diff --git a/src/Utilities/elementwise_scaler.jl b/src/Utilities/elementwise_scaler.jl new file mode 100644 index 000000000..e139be4b3 --- /dev/null +++ b/src/Utilities/elementwise_scaler.jl @@ -0,0 +1,185 @@ +# included in Utilities.jl + +export UnivariateAffineScaling, ElementwiseScaler, QuartileScaling, MinMaxScaling, ZScoreScaling +export quartile_scale, minmax_scale, zscore_scale +export get_type, get_shift, get_scale + +""" +$(TYPEDEF) + +The ElementwiseScaler{T} will create an encoding of the data_container via elementwise affine transformations. + +Different methods `T` will build different transformations: +- [`quartile_scale`](@ref) : creates `QuartileScaling`, +- [`minmax_scale`](@ref) : creates `MinMaxScaling` +- [`zscore_scale`](@ref) : creates `ZScoreScaling` + +and are accessed with [`get_type`](@ref) +""" +struct ElementwiseScaler{T, VV <: AbstractVector} <: DataContainerProcessor + "storage for the shift applied to data" + shift::VV + "storage for the scaling" + scale::VV +end + +abstract type UnivariateAffineScaling end +abstract type QuartileScaling <: UnivariateAffineScaling end +abstract type MinMaxScaling <: UnivariateAffineScaling end +abstract type ZScoreScaling <: UnivariateAffineScaling end + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{QuartileScaling}` processor. +As part of an encoder schedule, it will apply the transform ``\\frac{x - Q2(x)}{Q3(x) - Q1(x)}`` to each data dimension. +Also known as "robust scaling" +""" +quartile_scale() = ElementwiseScaler(QuartileScaling) + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{MinMaxScaling}` processor. +As part of an encoder schedule, this will apply the transform ``\\frac{x - \\min(x)}{\\max(x) - \\min(x)}`` to each data dimension. +""" +minmax_scale() = ElementwiseScaler(MinMaxScaling) + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{ZScoreScaling}` processor. +As part of an encoder schedule, this will apply the transform ``\\frac{x-\\mu}{\\sigma}``, (where ``x\\sim N(\\mu,\\sigma)``), to each data dimension. +For multivariate standardization, see [`Decorrelator`](@ref) +""" +zscore_scale() = ElementwiseScaler(ZScoreScaling) + +ElementwiseScaler(::Type{UAS}) where {UAS <: UnivariateAffineScaling} = + ElementwiseScaler{UAS, Vector{Float64}}(Float64[], Float64[]) + +""" +$(TYPEDSIGNATURES) + +Gets the UnivariateAffineScaling type `T` +""" +get_type(es::ElementwiseScaler{T}) where {T} = T + +""" +$(TYPEDSIGNATURES) + +Gets the `shift` field of the `ElementwiseScaler` +""" +get_shift(es::ElementwiseScaler) = es.shift + +""" +$(TYPEDSIGNATURES) + +Gets the `scale` field of the `ElementwiseScaler` +""" +get_scale(es::ElementwiseScaler) = es.scale + +function Base.show(io::IO, es::ElementwiseScaler) + out = "ElementwiseScaler: $(get_type(es))" + print(io, out) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{QS}, +) where {MM <: AbstractMatrix, QS <: QuartileScaling} + quartiles_vec = [quantile(dd, [0.25, 0.5, 0.75]) for dd in eachrow(data)] + quartiles_mat = reduce(hcat, quartiles_vec) # 3 rows: Q1, Q2, and Q3 + append!(get_shift(es), quartiles_mat[2, :]) + append!(get_scale(es), (quartiles_mat[3, :] - quartiles_mat[1, :])) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{MMS}, +) where {MM <: AbstractMatrix, MMS <: MinMaxScaling} + minmax_vec = [[minimum(dd), maximum(dd)] for dd in eachrow(data)] + minmax_mat = reduce(hcat, minmax_vec) # 2 rows: min max + append!(get_shift(es), minmax_mat[1, :]) + append!(get_scale(es), (minmax_mat[2, :] - minmax_mat[1, :])) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{ZSS}, +) where {MM <: AbstractMatrix, ZSS <: ZScoreScaling} + stat_vec = [[mean(dd), std(dd)] for dd in eachrow(data)] + stat_mat = reduce(hcat, stat_vec) # 2 rows: mean, std + append!(get_shift(es), stat_mat[1, :]) + append!(get_scale(es), stat_mat[2, :]) +end + +function initialize_processor!(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + if length(get_shift(es)) == 0 + T = get_type(es) + initialize_processor!(es, data, T) + end +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` encoder, on a columns-are-data matrix +""" +function encode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + out = deepcopy(data) + for i in 1:size(out, 1) + out[i, :] .-= get_shift(es)[i] + out[i, :] /= get_scale(es)[i] + end + return out +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` decoder, on a columns-are-data matrix +""" +function decode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + out = deepcopy(data) + for i in 1:size(out, 1) + out[i, :] *= get_scale(es)[i] + out[i, :] .+= get_shift(es)[i] + end + return out +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `shift` and `scale` fields for the `ElementwiseScaler` +""" +initialize_processor!(es::ElementwiseScaler, data::MM, structure_matrix) where {MM <: AbstractMatrix} = + initialize_processor!(es, data) + + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` encoder to a provided structure matrix +""" +function encode_structure_matrix( + es::ElementwiseScaler, + structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + return Diagonal(1 ./ get_scale(es)) * structure_matrix * Diagonal(1 ./ get_scale(es)) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` decoder to a provided structure matrix +""" +function decode_structure_matrix( + es::ElementwiseScaler, + enc_structure_matrix::USorM, +) where {USorM <: Union{UniformScaling, AbstractMatrix}} + return Diagonal(get_scale(es)) * enc_structure_matrix * Diagonal(get_scale(es)) +end diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl index cad00bd6d..50eedc316 100644 --- a/src/VectorRandomFeature.jl +++ b/src/VectorRandomFeature.jl @@ -352,10 +352,11 @@ Build Vector Random Feature model for the input-output pairs subject to regulari """ function build_models!( vrfi::VectorRandomFeatureInterface, - input_output_pairs::PairedDataContainer{FT}; - regularization_matrix::Union{M, Nothing} = nothing, + input_output_pairs::PairedDataContainer{FT}, + input_structure_matrix, + output_structure_matrix; kwargs..., -) where {FT <: AbstractFloat, M <: AbstractMatrix} +) where {FT <: AbstractFloat} # get inputs and outputs input_values = get_inputs(input_output_pairs) @@ -443,22 +444,16 @@ function build_models!( "hyperparameter learning using $n_train training points, $n_test validation points and $n_features_opt features" ) - # regularization_matrix = nothing when we use scaled SVD to decorrelate the space, - # in this setting, noise = I - if regularization_matrix === nothing - regularization = I - + # think of the output_structure_matrix as the observational noise covariance, or a related quantity + if isnothing(output_structure_matrix) + regularization = 1.0 * I + elseif !isposdef(output_structure_matrix) + regularization = posdef_correct(output_structure_matrix) + println("RF output structure matrix is not positive definite, correcting for use as a regularizer") else - # think of the regularization_matrix as the observational noise covariance, or a related quantity - if !isposdef(regularization_matrix) - regularization = posdef_correct(regularization_matrix) - println("RF regularization matrix is not positive definite, correcting") - - else - regularization = regularization_matrix - end - + regularization = output_structure_matrix end + # [2.] Estimate covariance at mean value μ_hp = transform_unconstrained_to_constrained(prior, mean(prior)) cov_sample_multiplier = optimizer_options["cov_sample_multiplier"] diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl index 370e419bd..503aedd06 100644 --- a/test/Emulator/runtests.jl +++ b/test/Emulator/runtests.jl @@ -9,83 +9,11 @@ using LinearAlgebra using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Utilities #build an unknown type struct MLTester <: Emulators.MachineLearningTool end -function constructor_tests( - iopairs::PairedDataContainer, - Σ::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing}, - norm_factors, - decomposition, -) where {FT <: AbstractFloat} - - input_mean = vec(mean(get_inputs(iopairs), dims = 2)) #column vector - normalization = sqrt(inv(Symmetric(cov(get_inputs(iopairs), dims = 2)))) - normalization = Emulators.calculate_normalization(get_inputs(iopairs)) - norm_inputs = Emulators.normalize(get_inputs(iopairs), input_mean, normalization) - # [4.] test emulator preserves the structures - mlt = MLTester() - @test_throws ErrorException emulator = Emulator( - mlt, - iopairs, - obs_noise_cov = Σ, - normalize_inputs = true, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) - - #build a known type, with defaults - gp = GaussianProcess(GPJL()) - emulator = Emulator( - gp, - iopairs, - obs_noise_cov = Σ, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) - - # compare SVD/norm/stand with stored emulator version - test_decomp = emulator.decomposition - @test test_decomp.V == decomposition.V #(use [:,:] to make it an array) - @test test_decomp.Vt == decomposition.Vt - @test test_decomp.S == decomposition.S - - gp = GaussianProcess(GPJL()) - emulator2 = Emulator( - gp, - iopairs, - obs_noise_cov = Σ, - normalize_inputs = true, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) - train_inputs = get_inputs(emulator2.training_pairs) - @test train_inputs == norm_inputs - train_inputs2 = Emulators.normalize(emulator2, get_inputs(iopairs)) - @test train_inputs2 == norm_inputs - - # reverse standardise - gp = GaussianProcess(GPJL()) - emulator3 = Emulator( - gp, - iopairs, - obs_noise_cov = Σ, - normalize_inputs = false, - standardize_outputs = true, - standardize_outputs_factors = norm_factors, - retained_svd_frac = 1.0, - ) - - #standardized and decorrelated (sd) data - s_y, s_y_cov = Emulators.standardize(get_outputs(iopairs), Σ, norm_factors) - sd_train_outputs = get_outputs(emulator3.training_pairs) - sqrt_singular_values_inv = Diagonal(1.0 ./ sqrt.(emulator3.decomposition.S)) - decorrelated_s_y = sqrt_singular_values_inv * emulator3.decomposition.Vt * s_y - @test decorrelated_s_y == sd_train_outputs -end - @testset "Emulators" begin #build some quick data + noise m = 50 @@ -101,117 +29,91 @@ end noise_samples = rand(MvNormal(μ, Σ), m) y += noise_samples - iopairs = PairedDataContainer(x, y, data_are_columns = true) - @test get_inputs(iopairs) == x - @test get_outputs(iopairs) == y - - # [1.] test SVD (2D y version) - test_SVD = svd(Σ) - transformed_y, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0) - @test_throws MethodError Emulators.svd_transform(y, Σ[:, 1], retained_svd_frac = 1.0) - @test test_SVD.V[:, :] == decomposition.V #(use [:,:] to make it an array) - @test test_SVD.Vt == decomposition.Vt - @test test_SVD.S == decomposition.S - @test size(transformed_y) == size(y) + io_pairs = PairedDataContainer(x, y, data_are_columns = true) + @test get_inputs(io_pairs) == x + @test get_outputs(io_pairs) == y - # 1D y version - transformed_y, decomposition2 = Emulators.svd_transform(y[:, 1], Σ, retained_svd_frac = 1.0) - @test_throws MethodError Emulators.svd_transform(y[:, 1], Σ[:, 1], retained_svd_frac = 1.0) - @test size(transformed_y) == size(y[:, 1]) + # Unknown ML type + mlt = MLTester() + @test_throws ErrorException emulator = Emulator(mlt, io_pairs) - # Reverse SVD - y_new, y_cov_new = - Emulators.svd_reverse_transform_mean_cov(reshape(transformed_y, (d, 1)), ones(d, 1), decomposition2) - @test_throws MethodError Emulators.svd_reverse_transform_mean_cov(transformed_y, ones(d, 1), decomposition2) - @test_throws MethodError Emulators.svd_reverse_transform_mean_cov( - reshape(transformed_y, (d, 1)), - ones(d), - decomposition2, + # test getters & defaults + gp = GaussianProcess(GPJL()) + em = Emulator(gp, io_pairs) + @test get_machine_learning_tool(em) == gp + @test get_io_pairs(em) == io_pairs + default_encoder = (decorrelate_sample_cov(), "in_and_out") # for these inputs this is the default + enc_sch = create_encoder_schedule(default_encoder) + enc_io_pairs, enc_I_in, enc_I_out = initialize_and_encode_with_schedule!(enc_sch, io_pairs, nothing, 1.0 * I) + @test get_encoder_schedule(em)[1][1] == enc_sch[1][1] # inputs: proc + @test get_encoder_schedule(em)[1][2] == enc_sch[1][2] # inputs: apply_to + @test get_encoder_schedule(em)[2][1] == enc_sch[2][1] # outputs... + @test get_encoder_schedule(em)[2][2] == enc_sch[2][2] + @test get_data(get_encoded_io_pairs(em)) == get_data(enc_io_pairs) + @test get_data(get_encoded_io_pairs(em)) == get_data(enc_io_pairs) + @test isnothing(enc_I_in) + + #NB - encoders all tested in Utilities here just testing some API + encoded_mat = encode_data(em, x, "in") + encoded_dc = encode_data(em, DataContainer(x), "in") + encoded_nothing = encode_structure_matrix(em, nothing, "in") + encoded_I = encode_structure_matrix(em, 1.0 * I, "out") + tol = 1e-14 + @test isapprox(norm(encoded_mat - get_data(encoded_dc)), 0, atol = tol * p * m) + @test isapprox(norm(encoded_mat - get_inputs(enc_io_pairs)), 0, atol = tol * p * m) + @test isnothing(encoded_nothing) + @test isapprox(norm(enc_I_out - encoded_I), 0, atol = tol * d * d) + + decoded_dc = decode_data(em, encoded_dc, "in") + decoded_mat = decode_data(em, encoded_mat, "in") + decoded_nothing = decode_structure_matrix(em, nothing, "in") + decoded_I = decode_structure_matrix(em, encoded_I, "out") + @test isapprox(norm(get_data(decoded_dc) - decoded_mat), 0, atol = tol * p * m) + @test isapprox(norm(decoded_mat - x), 0, atol = tol * p * m) + @test isnothing(decoded_nothing) + @test isapprox(norm(decoded_I - 1.0 * I), 0, atol = tol * d * d) + + # test obs_noise_cov (check the warning at the start) + @test_logs (:warn,) (:info,) (:info,) (:warn,) Emulator(gp, io_pairs, obs_noise_cov = Σ) + @test_logs (:warn,) (:info,) (:info,) (:warn,) Emulator( + gp, + io_pairs, + input_structure_matrix = 4.0 * I, + obs_noise_cov = 3.0 * I, + output_structure_matrix = 2.0 * I, ) - @test size(y_new)[1] == size(y[:, 1])[1] - @test y_new ≈ y[:, 1] - @test y_cov_new[1] ≈ Σ - - - # Truncation - transformed_y, trunc_decomposition = Emulators.svd_transform(y[:, 1], Σ, retained_svd_frac = 0.95) - trunc_size = size(trunc_decomposition.S)[1] - @test test_SVD.S[1:trunc_size] == trunc_decomposition.S - @test size(transformed_y)[1] == trunc_size - - # [2.] test Normalization - # full rank - input_mean = vec(mean(get_inputs(iopairs), dims = 2)) #column vector - normalization = sqrt(inv(Symmetric(cov(get_inputs(iopairs), dims = 2)))) - @test all(isapprox.(Emulators.calculate_normalization(get_inputs(iopairs)), normalization, atol = 1e-12)) - - - norm_inputs = Emulators.normalize(get_inputs(iopairs), input_mean, normalization) - @test all(isapprox.(norm_inputs, normalization * (get_inputs(iopairs) .- input_mean), atol = 1e-12)) - @test isapprox.(norm(cov(norm_inputs, dims = 2) - I), 0.0, atol = 1e-8) - - # reduced rank - reduced_inputs = get_inputs(iopairs)[:, 1:(p - 1)] - input_mean = vec(mean(reduced_inputs, dims = 2)) #column vector - input_cov = cov(reduced_inputs, dims = 2) - r = rank(input_cov) # = p-2 - svd_in = svd(input_cov) - sqrt_inv_sv = 1 ./ sqrt.(svd_in.S[1:r]) - normalization = Diagonal(sqrt_inv_sv) * svd_in.Vt[1:r, :] # size r x p - @test all(isapprox.(Emulators.calculate_normalization(reduced_inputs), normalization, atol = 1e-12)) + em1 = Emulator(gp, io_pairs, obs_noise_cov = Σ) - norm_inputs = Emulators.normalize(reduced_inputs, input_mean, normalization) - @test size(norm_inputs) == (r, p - 1) - @test isapprox.(norm(cov(norm_inputs, dims = 2)[1:r, 1:r] - I(r)), 0.0, atol = 1e-12) - - # [3.] test Standardization - norm_factors = 10.0 - norm_factors = fill(norm_factors, size(y[:, 1])) # must be size of output dim - - s_y, s_y_cov = Emulators.standardize(get_outputs(iopairs), Σ, norm_factors) - @test s_y == get_outputs(iopairs) ./ norm_factors - @test s_y_cov == Σ ./ (norm_factors .* norm_factors') - - @testset "Emulators - dense Array Σ" begin - constructor_tests(iopairs, Σ, norm_factors, decomposition) - - # truncation - only test here, where singular value of Σ differ - gp = GaussianProcess(GPJL()) - emulator4 = Emulator( - gp, - iopairs, - obs_noise_cov = Σ, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 0.9, - ) - trunc_size = size(emulator4.decomposition.S)[1] - @test test_SVD.S[1:trunc_size] == emulator4.decomposition.S - end - - @testset "Emulators - Diagonal Σ" begin - x = rand(3, m) #R^3 - y = rand(d, m) #R^5 - # "noise" - Σ = Diagonal(rand(d) .+ 1.0) - noise_samples = rand(MvNormal(zeros(d), Σ), m) - y += noise_samples - iopairs = PairedDataContainer(x, y, data_are_columns = true) - _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0) + em2 = Emulator( + gp, + io_pairs, + input_structure_matrix = 4.0 * I(p), + obs_noise_cov = Σ, + output_structure_matrix = 2.0 * I, + ) + Γ = randn(p, p) + Γ = Γ * Γ' + em3 = Emulator(gp, io_pairs, input_structure_matrix = Γ) + enc_sch1 = create_encoder_schedule([(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")]) + enc_sch2 = create_encoder_schedule([(decorrelate_structure_mat(), "in"), (decorrelate_structure_mat(), "out")]) + enc_sch3 = create_encoder_schedule([(decorrelate_structure_mat(), "in"), (decorrelate_sample_cov(), "out")]) + (_, _, _) = initialize_and_encode_with_schedule!( + enc_sch1, + io_pairs, + 1.0 * I(p), + Σ, # obs noise cov becomes the output structure matrix + ) + @test get_encoder_schedule(em1) == enc_sch1 - constructor_tests(iopairs, Σ, norm_factors, decomposition) - end + (_, _, _) = initialize_and_encode_with_schedule!( + enc_sch2, + io_pairs, + 4.0 * I(p), + 2.0 * I(d), # out_struct_mat overrides obs noise cov + ) + @test get_encoder_schedule(em2) == enc_sch2 - @testset "Emulators - UniformScaling Σ" begin - x = rand(3, m) #R^3 - y = rand(d, m) #R^5 - # "noise" - Σ = UniformScaling(1.3) - noise_samples = rand(MvNormal(zeros(d), Σ), m) - y += noise_samples - iopairs = PairedDataContainer(x, y, data_are_columns = true) - _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0) + (_, _, _) = initialize_and_encode_with_schedule!(enc_sch3, io_pairs, Γ, I(d)) + @test get_encoder_schedule(em3) == enc_sch3 - constructor_tests(iopairs, Σ, norm_factors, decomposition) - end end diff --git a/test/GaussianProcess/runtests.jl b/test/GaussianProcess/runtests.jl index f92f24393..13b43d9fc 100644 --- a/test/GaussianProcess/runtests.jl +++ b/test/GaussianProcess/runtests.jl @@ -13,6 +13,7 @@ copy!(pykernels, pyimport_conda("sklearn.gaussian_process.kernels", "scikit-lear using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Utilities @testset "GaussianProcess" begin @@ -62,33 +63,11 @@ using CalibrateEmulateSample.DataContainers @test gp1.prediction_type == pred_type @test gp1.alg_reg_noise == 1e-4 + em1 = Emulator(gp1, iopairs, encoder_schedule = []) - em1 = Emulator( - gp1, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + @test_logs (:warn,) Emulator(gp1, iopairs, encoder_schedule = []) # check that gp1 does not get more models added under second call - @test_logs (:warn,) (:warn,) Emulator( - gp1, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) # check that gp1 does not get more models added under second call - - Emulator( - gp1, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + Emulator(gp1, iopairs, encoder_schedule = []) @test length(gp1.models) == 1 Emulators.optimize_hyperparameters!(em1) @@ -101,14 +80,7 @@ using CalibrateEmulateSample.DataContainers # GaussianProcess 1b: use GPJL to create an abstractGP dist. agp = GaussianProcess(AGPJL(); noise_learn = true, alg_reg_noise = 1e-4, prediction_type = pred_type) - @test_throws ArgumentError Emulator( - agp, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + @test_throws ArgumentError Emulator(agp, iopairs, encoder_schedule = []) gp1_opt_params = Emulators.get_params(gp1)[1] # one model only gp1_opt_param_names = get_param_names(gp1)[1] # one model only @@ -119,26 +91,10 @@ using CalibrateEmulateSample.DataContainers "log_std_noise" => gp1_opt_params[end], ) - em_agp_from_gp1 = Emulator( - agp, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - kernel_params = kernel_params, - ) + em_agp_from_gp1 = Emulator(agp, iopairs, encoder_schedule = [], kernel_params = kernel_params) optimize_hyperparameters!(em_agp_from_gp1) # skip rebuild: - @test_logs (:warn,) (:warn,) Emulator( - agp, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - kernel_params = kernel_params, - ) + @test_logs (:warn,) Emulator(agp, iopairs, encoder_schedule = [], kernel_params = kernel_params) μ1b, σ1b² = Emulators.predict(em_agp_from_gp1, new_inputs) @@ -155,14 +111,7 @@ using CalibrateEmulateSample.DataContainers gp2 = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) - em2 = Emulator( - gp2, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + em2 = Emulator(gp2, iopairs, encoder_schedule = []) Emulators.optimize_hyperparameters!(em2) @@ -179,39 +128,13 @@ using CalibrateEmulateSample.DataContainers GPkernel = var * se gp3 = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) - em3 = Emulator( - gp3, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) - @test_logs (:warn,) (:warn,) Emulator( - gp3, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) - Emulator( - gp3, - iopairs, - obs_noise_cov = nothing, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + em3 = Emulator(gp3, iopairs, encoder_schedule = []) + @test_logs (:warn,) Emulator(gp3, iopairs, encoder_schedule = []) + Emulator(gp3, iopairs, encoder_schedule = []) @test length(gp3.models) == 1 # check that gp3 does not get more models added under repeated calls Emulators.optimize_hyperparameters!(em3) - #gp3 = GaussianProcess(iopairs, gppackage; GPkernel=GPkernel, obs_noise_cov=nothing, - # normalized=false, noise_learn=true, - # retained_svd_frac=1.0, standardize=false, - # prediction_type=pred_type, norm_factor=nothing) - μ3, σ3² = Emulators.predict(em3, new_inputs) @test vec(μ3) ≈ [0.0, 1.0, 0.0, -1.0, 0.0] atol = 0.3 @test vec(σ3²) ≈ [0.016, 0.002, 0.003, 0.004, 0.003] atol = 1e-2 @@ -250,28 +173,15 @@ using CalibrateEmulateSample.DataContainers @test get_inputs(iopairs2) == X @test get_outputs(iopairs2) == Y + # with noise learning - e.g. we add a kernel to learn the noise (even though we provide the Sigma to the emulator) gp4_noise_learnt = GaussianProcess(gppackage; kernel = nothing, noise_learn = true, prediction_type = pred_type) # without noise learning, just use the SVD transform to deal with observational noise - em4_noise_learnt = Emulator( - gp4_noise_learnt, - iopairs2, - obs_noise_cov = Σ, - normalize_inputs = true, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + em4_noise_learnt = Emulator(gp4_noise_learnt, iopairs2, output_structure_matrix = Σ) gp4 = GaussianProcess(gppackage; kernel = nothing, noise_learn = false, prediction_type = pred_type) - em4_noise_from_Σ = Emulator( - gp4, - iopairs2, - obs_noise_cov = Σ, - normalize_inputs = true, - standardize_outputs = false, - retained_svd_frac = 1.0, - ) + em4_noise_from_Σ = Emulator(gp4, iopairs2, output_structure_matrix = Σ) Emulators.optimize_hyperparameters!(em4_noise_learnt) Emulators.optimize_hyperparameters!(em4_noise_from_Σ) @@ -316,14 +226,7 @@ using CalibrateEmulateSample.DataContainers ) for model_params in gp4_opt_params ] - em_agp_from_gp4 = Emulator( - agp4, - iopairs2, - obs_noise_cov = Σ, - normalize_inputs = true, - retained_svd_frac = 1.0, - kernel_params = kernel_params, - ) + em_agp_from_gp4 = Emulator(agp4, iopairs2, output_structure_matrix = Σ, kernel_params = kernel_params) μ4b, σ4b² = Emulators.predict(em_agp_from_gp4, new_inputs, transform_to_real = true) diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index e0e4c3793..27621f09d 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -9,6 +9,7 @@ const MCMC = MarkovChainMonteCarlo using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Utilities function test_data(; rng_seed = 41, n = 20, var_y = 0.05, rest...) # Seed for pseudo-random number generator @@ -49,19 +50,19 @@ function test_data_mv(; rng_seed = 41, n = 20, var_y = 0.05, input_dim = 10, res return y, σ2_y, PairedDataContainer(x, y, data_are_columns = true), rng end -function test_gp_mv(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing) +function test_gp_mv(y, σ2_y, iopairs::PairedDataContainer) gppackage = GPJL() pred_type = YType() # Construct kernel: # Squared exponential kernel (note that hyperparameters are on log scale) # with observational noise gp = GaussianProcess(gppackage; noise_learn = true, prediction_type = pred_type) - em = Emulator(gp, iopairs; obs_noise_cov = σ2_y) + em = Emulator(gp, iopairs; output_structure_matrix = σ2_y) Emulators.optimize_hyperparameters!(em) return em end -function test_gp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing) +function test_gp_1(y, σ2_y, iopairs::PairedDataContainer) gppackage = GPJL() pred_type = YType() # Construct kernel: @@ -69,20 +70,12 @@ function test_gp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing # with observational noise GPkernel = SE(log(1.0), log(1.0)) gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) - em = Emulator( - gp, - iopairs; - obs_noise_cov = σ2_y, - normalize_inputs = false, - standardize_outputs = false, - standardize_outputs_factors = norm_factor, - retained_svd_frac = 1.0, - ) + em = Emulator(gp, iopairs; output_structure_matrix = σ2_y, encoder_schedule = []) Emulators.optimize_hyperparameters!(em) return em end -function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing) +function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer) gppackage = GPJL() pred_type = YType() # Construct kernel: @@ -90,15 +83,7 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = # with observational noise GPkernel = SE(log(1.0), log(1.0)) gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) - em = Emulator( - gp, - iopairs; - obs_noise_cov = σ2_y, - normalize_inputs = false, - standardize_outputs = false, - standardize_outputs_factors = norm_factor, - retained_svd_frac = 1.0, - ) + em = Emulator(gp, iopairs; output_structure_matrix = σ2_y, encoder_schedule = []) Emulators.optimize_hyperparameters!(em) # now make agp from gp @@ -115,22 +100,14 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = ) for model_params in gp_opt_params ] - em_agp = Emulator( - agp, - iopairs, - obs_noise_cov = σ2_y, - normalize_inputs = false, - standardize_outputs = false, - retained_svd_frac = 1.0, - standardize_outputs_factors = norm_factor, - kernel_params = kernel_params, - ) + em_agp = + Emulator(agp, iopairs, output_structure_matrix = σ2_y, encoder_schedule = [], kernel_params = kernel_params) return em, em_agp end -function test_gp_2(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing) +function test_gp_2(y, σ2_y, iopairs::PairedDataContainer) gppackage = GPJL() pred_type = YType() # Construct kernel: @@ -138,15 +115,9 @@ function test_gp_2(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing # with observational noise GPkernel = SE(log(1.0), log(1.0)) gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) - em = Emulator( - gp, - iopairs; - obs_noise_cov = σ2_y, - normalize_inputs = false, - standardize_outputs = true, - standardize_outputs_factors = norm_factor, - retained_svd_frac = 0.9, - ) + retain_var = 0.95 + encoder_schedule = [(quartile_scale(), "out"), (decorrelate_structure_mat(retain_var = retain_var), "out")] + em = Emulator(gp, iopairs; output_structure_matrix = σ2_y, encoder_schedule = encoder_schedule) Emulators.optimize_hyperparameters!(em) return em @@ -157,7 +128,7 @@ function mcmc_test_template( σ2_y, em::Emulator; mcmc_alg = RWMHSampling(), - obs_sample = 1.0, + obs_sample = [1.0], init_params = 3.0, step = 0.25, rng = Random.GLOBAL_RNG, @@ -189,14 +160,6 @@ end Dict(:mcmc_alg => RWMHSampling(), :obs_sample => obs_sample, :init_params => [3.0], :step => 0.5, :rng => rng) - @testset "Constructor: standardize" begin - em = test_gp_1(y, σ2_y, iopairs) - test_obs = MarkovChainMonteCarlo.to_decorrelated(obs_sample, em) - # The MCMC stored a SVD-transformed sample, in a vector - # 1.0/sqrt(0.05) * obs_sample ≈ 4.472 - @test isapprox(test_obs[1], (obs_sample ./ sqrt(σ2_y[1, 1])); atol = 1e-2) - end - @testset "MV priors" begin # 10D dist with 1 name, just build the wrapper for test prior_mv = test_prior_mv() @@ -228,9 +191,7 @@ end @test all(isapprox.(esjd1, esjd1b, rtol = 0.2)) # now test SVD normalization - norm_factor = 10.0 - norm_factor = fill(norm_factor, size(y[:, 1])) # must be size of output dim - em_2 = test_gp_2(y, σ2_y, iopairs; norm_factor = norm_factor) + em_2 = test_gp_2(y, σ2_y, iopairs) _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...) # difference between mean_1 and mean_2 only from MCMC convergence @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1) @@ -301,9 +262,7 @@ end @test all(isapprox.(esjd1, esjd1b, rtol = 0.2)) # now test SVD normalization - norm_factor = 10.0 - norm_factor = fill(norm_factor, size(y[:, 1])) # must be size of output dim - em_2 = test_gp_2(y, σ2_y, iopairs; norm_factor = norm_factor) + em_2 = test_gp_2(y, σ2_y, iopairs) _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...) # difference between mean_1 and mean_2 only from MCMC convergence @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1) diff --git a/test/RandomFeature/runtests.jl b/test/RandomFeature/runtests.jl index baa23f674..1201faa8e 100644 --- a/test/RandomFeature/runtests.jl +++ b/test/RandomFeature/runtests.jl @@ -306,9 +306,9 @@ rng = Random.MersenneTwister(seed) ) # build emulators - em_srfi = Emulator(srfi, iopairs, obs_noise_cov = obs_noise_cov) + em_srfi = Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov) n_srfi = length(get_rfms(srfi)) - em_vrfi = Emulator(vrfi, iopairs, obs_noise_cov = obs_noise_cov) + em_vrfi = Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) n_vrfi = length(get_rfms(vrfi)) # test bad case @@ -324,11 +324,11 @@ rng = Random.MersenneTwister(seed) @test_throws ArgumentError Emulator(srfi_bad, iopairs) # test under repeats - @test_logs (:warn,) Emulator(srfi, iopairs, obs_noise_cov = obs_noise_cov) - Emulator(srfi, iopairs, obs_noise_cov = obs_noise_cov) + @test_logs (:info,) (:info,) (:warn,) Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov) + Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov) @test length(get_rfms(srfi)) == n_srfi - @test_logs (:warn,) Emulator(vrfi, iopairs, obs_noise_cov = obs_noise_cov) - Emulator(vrfi, iopairs, obs_noise_cov = obs_noise_cov) + @test_logs (:info,) (:info,) (:warn,) Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) + Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) @test length(get_rfms(vrfi)) == n_vrfi @@ -422,14 +422,14 @@ rng = Random.MersenneTwister(seed) # build emulators # svd: scalar, scalar + diag in, vector + diag out, vector # no-svd: vector - em_srfi_svd_diagin = Emulator(srfi_diagin, iopairs, obs_noise_cov = Σ) - em_srfi_svd = Emulator(srfi, iopairs, obs_noise_cov = Σ) + em_srfi_svd_diagin = Emulator(srfi_diagin, iopairs, output_structure_matrix = Σ) + em_srfi_svd = Emulator(srfi, iopairs, output_structure_matrix = Σ) - em_vrfi_svd_diagout = Emulator(vrfi_diagout, iopairs, obs_noise_cov = Σ) - em_vrfi_svd = Emulator(vrfi, iopairs, obs_noise_cov = Σ) + em_vrfi_svd_diagout = Emulator(vrfi_diagout, iopairs, output_structure_matrix = Σ) + em_vrfi_svd = Emulator(vrfi, iopairs, output_structure_matrix = Σ) - em_vrfi = Emulator(vrfi, iopairs, obs_noise_cov = Σ, decorrelate = false) - em_vrfi_nonsep = Emulator(vrfi_nonsep, iopairs, obs_noise_cov = Σ, decorrelate = false) + em_vrfi = Emulator(vrfi, iopairs, output_structure_matrix = Σ, decorrelate = false) + em_vrfi_nonsep = Emulator(vrfi_nonsep, iopairs, output_structure_matrix = Σ, decorrelate = false) #TODO truncated SVD option for vector (involves resizing prior) @@ -452,8 +452,8 @@ rng = Random.MersenneTwister(seed) @test isapprox.(norm(μ_ss - outx), 0, atol = tol_μ) @test isapprox.(norm(μ_vsd - outx), 0, atol = tol_μ) @test isapprox.(norm(μ_vs - outx), 0, atol = tol_μ) - @test isapprox.(norm(μ_v - outx), 0, atol = 5 * tol_μ) # approximate option so likely less good approx - @test isapprox.(norm(μ_vns - outx), 0, atol = 5 * tol_μ) # approximate option so likely less good approx + @test isapprox.(norm(μ_v - outx), 0, atol = 10 * tol_μ) # approximate option so likely less good approx + @test isapprox.(norm(μ_vns - outx), 0, atol = 10 * tol_μ) # approximate option so likely less good approx # An example with the other threading option vrfi_tul = VectorRandomFeatureInterface( @@ -464,7 +464,7 @@ rng = Random.MersenneTwister(seed) optimizer_options = Dict("train_fraction" => 0.7, "multithread" => "tullio"), rng = rng, ) - em_vrfi_svd_tul = Emulator(vrfi_tul, iopairs, obs_noise_cov = Σ) + em_vrfi_svd_tul = Emulator(vrfi_tul, iopairs, output_structure_matrix = Σ) μ_vs_tul, σ²_vs_tul = Emulators.predict(em_vrfi_svd_tul, new_inputs, transform_to_real = true) @test isapprox.(norm(μ_vs_tul - outx), 0, atol = tol_μ) diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl index 687532b22..6871b9608 100644 --- a/test/Utilities/runtests.jl +++ b/test/Utilities/runtests.jl @@ -13,34 +13,6 @@ using CalibrateEmulateSample.DataContainers # Seed for pseudo-random number generator rng = Random.MersenneTwister(41) - arr = vcat([i * ones(3)' for i in 1:5]...) - arr_t = permutedims(arr, (2, 1)) - - mean_arr = dropdims(mean(arr, dims = 1), dims = 1) - std_arr = dropdims(std(arr, dims = 1), dims = 1) - @test mean_arr == [3.0, 3.0, 3.0] - @test std_arr ≈ [1.58, 1.58, 1.58] atol = 1e-2 - z_arr = orig2zscore(arr, mean_arr, std_arr) - z_arr_test = [ - -1.265 -1.265 -1.265 - -0.632 -0.632 -0.632 - 0.0 0.0 0.0 - 0.632 0.632 0.632 - 1.265 1.265 1.265 - ] - @test z_arr ≈ z_arr_test atol = 1e-2 - orig_arr = zscore2orig(z_arr, mean_arr, std_arr) - @test orig_arr ≈ arr atol = 1e-5 - - v = vec([1.0 2.0 3.0]) - mean_v = vec([0.5 1.5 2.5]) - std_v = vec([0.5 0.5 0.5]) - z_v = orig2zscore(v, mean_v, std_v) - println(z_v) - @test z_v ≈ [1.0, 1.0, 1.0] atol = 1e-5 - orig_v = zscore2orig(z_v, mean_v, std_v) - @test orig_v ≈ v atol = 1e-5 - # test get_training_points # first create the EnsembleKalmanProcess n_ens = 10 @@ -100,10 +72,31 @@ end dd3 = decorrelate_structure_mat(retain_var = 0.7) @test get_retain_var(dd3) == 0.7 @test get_decorrelate_with(dd3) == "structure_mat" - DD = Decorrelater([1], [2], [3], 1.0, "test") + DD = Decorrelator([1], [2], [3], 1.0, "test") @test get_data_mean(DD) == [1] @test get_encoder_mat(DD) == [2] @test get_decoder_mat(DD) == [3] + + + cc = canonical_correlation() + @test get_retain_var(cc) == 1.0 + cc2 = canonical_correlation(retain_var = 0.7) + @test get_retain_var(cc2) == 0.7 + cc3 = CanonicalCorrelation([1], [2], [3], 1.0, "test") + @test get_data_mean(cc3) == [1] + @test get_encoder_mat(cc3) == [2] + @test get_decoder_mat(cc3) == [3] + @test get_apply_to(cc3) == "test" + + + # test equalities + cc = canonical_correlation() + cc_copy = canonical_correlation() + dd = decorrelate() + dd_copy = decorrelate() + @test cc == cc_copy + @test dd == dd_copy + # get some data as IO pairs for functional tests in_dim = 10 @@ -150,7 +143,7 @@ end for (name, sch, ll_flag) in zip(test_names, schedules, lossless) encoder_schedule = create_encoder_schedule(sch) (encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) = - encode_with_schedule(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + initialize_and_encode_with_schedule!(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) (decoded_io_pairs, decoded_prior_cov, decoded_obs_noise_cov) = decode_with_schedule(encoder_schedule, encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) @@ -237,7 +230,7 @@ end # Multivariate lossy dim-reduction tests if name == "decorrelate-structure-mat-retain-0.95-var" svdc = svd(test_covv) - var_cumsum = cumsum(svdc.S .^ 2) ./ sum(svdc.S .^ 2) + var_cumsum = cumsum(svdc.S) ./ sum(svdc.S) @test var_cumsum[dimm] > 0.95 @test var_cumsum[dimm - 1] < 0.95 @test all(isapprox.(pop_mean, zeros(dimm), atol = tol)) @@ -279,6 +272,14 @@ end end + # test throws on lack of sample_mat + sch2a = (decorrelate_structure_mat(), "in") + schedule2a = create_encoder_schedule(sch2a) + @test_throws ArgumentError initialize_and_encode_with_schedule!(schedule2a, io_pairs, nothing, obs_noise_cov) # nothing + sch2b = (decorrelate_structure_mat(), "out") + schedule2b = create_encoder_schedule(sch2b) + @test_throws ArgumentError initialize_and_encode_with_schedule!(schedule2b, io_pairs, prior_cov, nothing) + # combine a few lossless encoding schedules (lossless requires samples>dims) @@ -300,19 +301,21 @@ end @test_logs (:warn,) create_encoder_schedule((canonical_correlation(), "bad")) @test_logs (:warn,) create_encoder_schedule((zscore_scale(), "bad")) - func = x -> (get_inputs(x), get_outputs(x)) - bad_encoder_schedule = [(canonical_correlation(), func, "bad")] - @test_throws ArgumentError encode_with_schedule(bad_encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + bad_encoder_schedule = [(canonical_correlation(), "bad")] + @test_throws ArgumentError initialize_and_encode_with_schedule!( + bad_encoder_schedule, + io_pairs, + prior_cov, + obs_noise_cov, + ) @test_throws ArgumentError decode_with_schedule(bad_encoder_schedule, io_pairs, prior_cov, obs_noise_cov) - @test_throws ArgumentError initialize_and_encode_data!(canonical_correlation(), func(io_pairs), prior_cov, "bad") - @test_throws ArgumentError decode_data(canonical_correlation(), func(io_pairs), "bad") encoder_schedule = create_encoder_schedule(schedule_builder) # encode the data using the schedule (encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) = - encode_with_schedule(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + initialize_and_encode_with_schedule!(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) # decode the data using the schedule (decoded_io_pairs, decoded_prior_cov, decoded_obs_noise_cov) =