diff --git a/docs/src/API/GaussianProcess.md b/docs/src/API/GaussianProcess.md index 7f4ad96ee..065edec0f 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}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat} +build_models!(::GaussianProcess{GPJL}, ::PairedDataContainer{FT}, input_structure_mats, output_structure_mats) 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/RandomFeatures.md b/docs/src/API/RandomFeatures.md index 623798a3e..7f41d8287 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}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat} +build_models!(::ScalarRandomFeatureInterface, ::PairedDataContainer{FT}, input_structure_mats, output_structure_mats) 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}, input_structure_matrix, output_structure_matrix) where {FT <: AbstractFloat} +build_models!(::VectorRandomFeatureInterface, ::PairedDataContainer{FT}, input_structure_mats, output_structure_mats) where {FT <: AbstractFloat} predict(::VectorRandomFeatureInterface, ::M) where {M <: AbstractMatrix} ``` diff --git a/docs/src/data_processing.md b/docs/src/data_processing.md index 6742f52d5..1e17464dc 100644 --- a/docs/src/data_processing.md +++ b/docs/src/data_processing.md @@ -28,7 +28,7 @@ complex_schedule = [ 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. +3. The outputs are decorrelated using an "output structure matrix" (provided to the emulator in the `encoder_kwargs` keyword parameter, e.g. as `(; obs_cov_noise =)`). 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" @@ -44,12 +44,12 @@ The schedule is then passed into the Emulator, along with the data and desired s ```julia emulator = Emulator( machine_learning_tool, - input_output_pairs; - output_structure_matrix = obs_noise_cov, + input_output_pairs; encoder_schedule = complex_schedule, + encoder_kwargs = (; obs_noise_cov = obs_noise_cov), ) ``` -Note that due to the item `(decorrelate_structure_mat(retain_var=0.95), "out")` in the schedule, we must provide the `output_structure_matrix`. +Note that due to the item `(decorrelate_structure_mat(retain_var=0.95), "out")` in the schedule, we must provide an output structure matrix. In this case, we provide `obs_noise_cov`. # Types of data processors diff --git a/docs/src/emulate.md b/docs/src/emulate.md index a2239eafa..f53cb4f42 100644 --- a/docs/src/emulate.md +++ b/docs/src/emulate.md @@ -17,8 +17,8 @@ Wrapping a predefined machine learning tool, e.g. a Gaussian process `gauss_proc emulator = Emulator( gauss_proc, input_output_pairs; # optional arguments after this - output_structure_matrix = Γy, encoder_schedule = encoder_schedule, + encoder_kwargs = (; obs_noise_cov = Γy), ) ``` The optional arguments above relate to the data processing, which is described [here](@ref data-proc) @@ -44,7 +44,7 @@ Developers may contribute new tools by performing the following 2. Create a struct `MyMLTool <: MachineLearningTool`, containing any arguments or optimizer options 3. Create the following three methods to build, train, and predict with your tool (use `GaussianProcess.jl` as a guide) ``` -build_models!(mlt::MyMLTool, iopairs::PairedDataContainer) -> Nothing +build_models!(mlt::MyMLTool, iopairs::PairedDataContainer, input_structure_mats::Dict{Symbol, <:StructureMatrix}, output_structure_mats::Dict{Symbol, <:StructureMatrix}) -> Nothing optimize_hyperparameters!(mlt::MyMLTool, args...; kwargs...) -> Nothing function predict(mlt::MyMLTool, new_inputs::Matrix; kwargs...) -> Matrix, Union{Matrix, Array{,3} ``` diff --git a/examples/Cloudy/Cloudy_emulate_sample.jl b/examples/Cloudy/Cloudy_emulate_sample.jl index 030f503be..c6b04232e 100644 --- a/examples/Cloudy/Cloudy_emulate_sample.jl +++ b/examples/Cloudy/Cloudy_emulate_sample.jl @@ -172,9 +172,8 @@ function main() emulator = Emulator( mlt, input_output_pairs; - input_structure_matrix = cov(priors), - output_structure_matrix = Γy, encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(priors), obs_noise_cov = Γy), ) optimize_hyperparameters!(emulator) diff --git a/examples/Darcy/emulate_sample.jl b/examples/Darcy/emulate_sample.jl index ca046dbb8..b7874a3c2 100644 --- a/examples/Darcy/emulate_sample.jl +++ b/examples/Darcy/emulate_sample.jl @@ -92,14 +92,13 @@ function main() @save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs # data processing - encoding_schedule = (decorrelate_structure_mat(), "in_and_out") + encoder_schedule = (decorrelate_structure_mat(), "in_and_out") emulator = Emulator( mlt, input_output_pairs; - input_structure_matrix = cov(prior), - output_structure_matrix = Γy, - encoding_schedule = encoding_schedule, + encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(prior), obs_noise_cov = Γy), ) 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 4bf51fd4d..9817de632 100644 --- a/examples/EDMF_data/emulator-rank-test.jl +++ b/examples/EDMF_data/emulator-rank-test.jl @@ -256,10 +256,9 @@ function main() ttt[rank_id, rep_idx] = @elapsed begin emulator = Emulator( mlt, - train_pairs; - input_structure_matrix = cov(prior), - output_structure_matrix = truth_cov, + train_pairs, encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(prior), obs_noise_cov = truth_cov), ) # Optimize the GP hyperparameters for better fit diff --git a/examples/EDMF_data/uq_for_edmf.jl b/examples/EDMF_data/uq_for_edmf.jl index a708b00f5..fb19910c2 100644 --- a/examples/EDMF_data/uq_for_edmf.jl +++ b/examples/EDMF_data/uq_for_edmf.jl @@ -230,10 +230,9 @@ function main() # Fit an emulator to the data emulator = Emulator( mlt, - input_output_pairs; - input_structure_matrix = cov(prior), - output_structure_matrix = truth_cov, + input_output_pairs, encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(prior), obs_noise_cov = truth_cov), ) # Optimize the GP hyperparameters for better fit diff --git a/examples/Emulator/G-function/emulate-test-n-features.jl b/examples/Emulator/G-function/emulate-test-n-features.jl index 2dc1c5776..ecb324a94 100644 --- a/examples/Emulator/G-function/emulate-test-n-features.jl +++ b/examples/Emulator/G-function/emulate-test-n-features.jl @@ -154,8 +154,8 @@ function main() emulator = Emulator( mlt, iopairs; - output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = (; obs_noise_cov = Γ * I), ) optimize_hyperparameters!(emulator) end diff --git a/examples/Emulator/G-function/emulate.jl b/examples/Emulator/G-function/emulate.jl index 3cc3ab123..91ad687fe 100644 --- a/examples/Emulator/G-function/emulate.jl +++ b/examples/Emulator/G-function/emulate.jl @@ -145,8 +145,8 @@ function main() emulator = Emulator( mlt, iopairs; - output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = (; obs_noise_cov = Γ * I), ) optimize_hyperparameters!(emulator) end diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl index 74421011d..105304a48 100644 --- a/examples/Emulator/Ishigami/emulate.jl +++ b/examples/Emulator/Ishigami/emulate.jl @@ -123,8 +123,12 @@ function main() end # Emulate - emulator = - Emulator(mlt, iopairs; output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule)) + emulator = Emulator( + mlt, + iopairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = (; obs_noise_cov = Γ * I), + ) 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 b623a89b5..3702c58ec 100644 --- a/examples/Emulator/L63/emulate.jl +++ b/examples/Emulator/L63/emulate.jl @@ -215,7 +215,12 @@ function main() else encoder_schedule = (decorrelate_structure_mat(), "out") end - emulator = Emulator(mlt, iopairs; output_structure_matrix = Γy, encoder_schedule = deepcopy(encoder_schedule)) + emulator = Emulator( + mlt, + iopairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = (; obs_noise_cov = Γy), + ) 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 e60b7ebcc..38c896a1f 100644 --- a/examples/Emulator/L63/emulate_diff-rank-test.jl +++ b/examples/Emulator/L63/emulate_diff-rank-test.jl @@ -194,8 +194,8 @@ function main() emulator = Emulator( mlt, iopairs; - output_structure_matrix = Γy, encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = (; obs_noise_cov = Γy), ) optimize_hyperparameters!(emulator) end diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl index fd55ca02d..116515d45 100644 --- a/examples/Emulator/Regression_2d_2d/compare_regression.jl +++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl @@ -168,10 +168,9 @@ function main() emulator = Emulator( mlt, - iopairs, + iopairs; encoder_schedule = enc, - input_structure_matrix = prior_cov, - output_structure_matrix = Σ, + encoder_kwargs = (; prior_cov = prior_cov, obs_noise_cov = Σ), ) diff --git a/examples/GCM/emulate_sample_script.jl b/examples/GCM/emulate_sample_script.jl index 07098b641..0c53cf4a5 100644 --- a/examples/GCM/emulate_sample_script.jl +++ b/examples/GCM/emulate_sample_script.jl @@ -163,18 +163,17 @@ function main() emulator = Emulator( mlt, input_output_pairs; - output_structure_matrix = obs_noise_cov, encoder_schedule = encoder_schedule, + encoder_kwargs = (; obs_noise_cov = obs_noise_cov), ) else 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, + input_output_pairs, encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(prior), obs_noise_cov = obs_noise_cov), ) end diff --git a/examples/Lorenz/emulate_sample.jl b/examples/Lorenz/emulate_sample.jl index 02ed90d5a..b204e1df9 100644 --- a/examples/Lorenz/emulate_sample.jl +++ b/examples/Lorenz/emulate_sample.jl @@ -170,7 +170,12 @@ function main() 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) + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = encoder_schedule, + encoder_kwargs = (; obs_noise_cov = Γy), + ) optimize_hyperparameters!(emulator) # Check how well the Gaussian Process regression predicts on the diff --git a/examples/Lorenz/emulate_sample_spatial_dep.jl b/examples/Lorenz/emulate_sample_spatial_dep.jl index 50b634ddc..04178f72a 100644 --- a/examples/Lorenz/emulate_sample_spatial_dep.jl +++ b/examples/Lorenz/emulate_sample_spatial_dep.jl @@ -144,9 +144,8 @@ function main() emulator = Emulator( mlt, input_output_pairs; - input_structure_matrix = cov(priors), - output_structure_matrix = Γy, encoder_schedule = encoder_schedule, + encoder_kwargs = (; prior_cov = cov(priors), obs_noise_cov = Γy), ) optimize_hyperparameters!(emulator) diff --git a/examples/Sinusoid/emulate.jl b/examples/Sinusoid/emulate.jl index 0b02a6259..06fdfca00 100644 --- a/examples/Sinusoid/emulate.jl +++ b/examples/Sinusoid/emulate.jl @@ -75,7 +75,7 @@ gauss_proc = Emulators.GaussianProcess(gppackage, noise_learn = false) encoder_schedule = [(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")] # Build emulator with data -emulator_gp = Emulator(gauss_proc, input_output_pairs, output_structure_matrix = Γ, encoder_schedule = encoder_schedule) +emulator_gp = Emulator(gauss_proc, input_output_pairs, (obs_cov_noise = Γ,); 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 @@ -113,7 +113,7 @@ random_features = VectorRandomFeatureInterface( ) encoder_schedule_rf = [(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")] emulator_random_features = - Emulator(random_features, input_output_pairs, encoder_schedule = encoder_schedule_rf, output_structure_matrix = Γ) + Emulator(random_features, input_output_pairs, (obs_cov_noise = Γ,); encoder_schedule = encoder_schedule_rf) optimize_hyperparameters!(emulator_random_features) diff --git a/src/Emulator.jl b/src/Emulator.jl index 1c836a295..56bfe83a6 100644 --- a/src/Emulator.jl +++ b/src/Emulator.jl @@ -17,7 +17,6 @@ using Random export Emulator -export calculate_normalization export build_models! export optimize_hyperparameters! export predict, encode_data, decode_data, encode_structure_matrix, decode_structure_matrix @@ -41,7 +40,7 @@ include("RandomFeature.jl") function throw_define_mlt() throw(ErrorException("Unknown MachineLearningTool defined, please use a known implementation")) end -function build_models!(mlt, iopairs, mlt_kwargs...) +function build_models!(mlt, iopairs, input_structure_mats, output_structure_mats, mlt_kwargs...) throw_define_mlt() end function optimize_hyperparameters!(mlt) @@ -96,7 +95,7 @@ get_encoded_io_pairs(emulator::Emulator) = emulator.encoded_io_pairs """ $(TYPEDSIGNATURES) -Gets the `encoder_schedul` field of the `Emulator` +Gets the `encoder_schedule` field of the `Emulator` """ get_encoder_schedule(emulator::Emulator) = emulator.encoder_schedule @@ -111,69 +110,56 @@ Positional Arguments - `input_output_pairs`: the paired input-output data points stored in a `PairedDataContainer` 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. + - `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_sample_cov(), "in_and_out")]`, or `[(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")]` if an `encoder_kwargs` has a key `:obs_noise_cov`. Pass `[]` for no encoding. + - `encoder_kwargs`[=`NamedTuple()`]: a Dict or NamedTuple with keyword arguments to be passed to `initialize_and_encode_with_schedule!` Other keywords are passed to the machine learning tool initialization """ 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 + encoder_kwargs = NamedTuple(), + obs_noise_cov = nothing, # temporary mlt_kwargs..., ) where {FT <: AbstractFloat} - # For Consistency checks - input_dim, output_dim = size(input_output_pairs, 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`." - ) - 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=`" - ) + if !isnothing(obs_noise_cov) + if haskey(encoder_kwargs, :obs_noise_cov) + @warn "Keyword argument `obs_noise_cov=` is deprecated and will be ignored in favor of `encoder_kwargs[:obs_noise_cov]`." + else + @warn "Keyword argument `obs_noise_cov=` is deprecated. Please use `encoder_kwargs[:obs_noise_cov]` instead." + end end # [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")) + + push!(encoder_schedule, (decorrelate_sample_cov(), "in")) + if haskey(encoder_kwargs, :obs_noise_cov) || !isnothing(obs_noise_cov) + push!(encoder_schedule, (decorrelate_structure_mat(), "out")) else - push!(encoder_schedule, (decorrelate_structure_mat(), "in")) - 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, - ) + encoder_schedule = create_encoder_schedule(encoder_schedule) + (encoded_io_pairs, input_structure_mats, output_structure_mats, _, _) = initialize_and_encode_with_schedule!( + encoder_schedule, + input_output_pairs; + obs_noise_cov = obs_noise_cov, + encoder_kwargs..., + ) # build the machine learning tool in the encoded space - build_models!( + build_models!(machine_learning_tool, encoded_io_pairs, input_structure_mats, output_structure_mats; mlt_kwargs...) + return Emulator{FT, typeof(encoder_schedule)}( machine_learning_tool, + input_output_pairs, encoded_io_pairs, - encoded_input_structure_matrix, - encoded_output_structure_matrix; - mlt_kwargs..., + encoder_schedule, ) - return Emulator{FT, typeof(enc_schedule)}(machine_learning_tool, input_output_pairs, encoded_io_pairs, enc_schedule) end """ @@ -209,9 +195,9 @@ Encode a new structure matrix in the input space (`"in"`) or output space (`"out """ function encode_structure_matrix( emulator::Emulator, - structure_mat::USorMorN, + structure_mat::USorM, in_or_out::AS, -) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} +) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}} return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out) end @@ -240,9 +226,9 @@ Decode a new structure matrix in the input space (`"in"`) or output space (`"out """ function decode_structure_matrix( emulator::Emulator, - structure_mat::USorMorN, + structure_mat::USorM, in_or_out::AS, -) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} +) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}} return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out) end diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl index 3fd1ae3b8..9d6d403df 100644 --- a/src/GaussianProcess.jl +++ b/src/GaussianProcess.jl @@ -2,6 +2,8 @@ using EnsembleKalmanProcesses.DataContainers using DocStringExtensions +using ..Utilities: get_structure_mat + # [1] For GaussianProcesses import GaussianProcesses: predict, get_params, get_param_names using GaussianProcesses @@ -149,8 +151,8 @@ Method to build Gaussian process models based on the package. function build_models!( gp::GaussianProcess{GPJL}, input_output_pairs::PairedDataContainer{FT}, - input_structure_matrix, - output_structure_matrix; + input_structure_mats, + output_structure_mats; kwargs..., ) where {FT <: AbstractFloat} # get inputs and outputs @@ -191,12 +193,15 @@ function build_models!( println("Learning additive white noise") end # use the output_structure_matrix to scale regularization scale - regularization = if isnothing(output_structure_matrix) + regularization = if isempty(output_structure_mats) 1.0 * ones(N_models) - elseif isa(output_structure_matrix, UniformScaling) - output_structure_matrix.λ * ones(N_models) else - diag(output_structure_matrix) + output_structure_mat = get_structure_mat(output_structure_mats) + if isa(output_structure_mat, UniformScaling) + output_structure_mat.λ * ones(N_models) + else + diag(output_structure_mat) + end end logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise)) @@ -283,9 +288,8 @@ predict(gp::GaussianProcess{GPJL}, new_inputs::AbstractMatrix{FT}) where {FT <: function build_models!( gp::GaussianProcess{SKLJL}, input_output_pairs::PairedDataContainer{FT}, - input_structure_matrix, - output_structure_matrix; - kwargs..., + input_structure_mats, + output_structure_mats, ) where {FT <: AbstractFloat} # get inputs and outputs input_values = permutedims(get_inputs(input_output_pairs), (2, 1)) @@ -322,12 +326,15 @@ function build_models!( println("Learning additive white noise") end # use the output_structure_matrix to scale regularization scale - regularization = if isnothing(output_structure_matrix) + regularization = if isempty(output_structure_mats) 1.0 * ones(N_models) - elseif isa(output_structure_matrix, UniformScaling) - output_structure_matrix.λ * ones(N_models) else - diag(output_structure_matrix) + output_structure_mat = get_structure_mat(output_structure_mats) + if isa(output_structure_mat, UniformScaling) + output_structure_mat.λ * ones(N_models) + else + diag(output_structure_mat) + end end regularization_noise_vec = gp.alg_reg_noise .* regularization for i in 1:N_models @@ -375,10 +382,9 @@ end function build_models!( gp::GaussianProcess{AGPJL}, input_output_pairs::PairedDataContainer{FT}, - input_structure_matrix, - output_structure_matrix; + input_structure_mats, + output_structure_mats; kernel_params = nothing, - kwargs..., ) where {FT <: AbstractFloat} # get inputs and outputs input_values = permutedims(get_inputs(input_output_pairs), (2, 1)) @@ -427,12 +433,15 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow N_models = size(output_values, 1) #size(transformed_data)[1] # use the output_structure_matrix to scale regularization scale - regularization = if isnothing(output_structure_matrix) + regularization = if isempty(output_structure_mats) 1.0 * ones(N_models) - elseif isa(output_structure_matrix, UniformScaling) - output_structure_matrix.λ * ones(N_models) else - diag(output_structure_matrix) + output_structure_mat = get_structure_mat(output_structure_mats) + if isa(output_structure_mat, UniformScaling) + output_structure_mat.λ * ones(N_models) + else + diag(output_structure_mat) + end end regularization_noise = gp.alg_reg_noise .* regularization diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl index b23983f17..fa59c6741 100644 --- a/src/ScalarRandomFeature.jl +++ b/src/ScalarRandomFeature.jl @@ -319,9 +319,8 @@ Builds the random feature method from hyperparameters. We use cosine activation function build_models!( srfi::ScalarRandomFeatureInterface, input_output_pairs::PairedDataContainer{FT}, - input_structure_matrix, - output_structure_matrix; - kwargs..., + input_structure_mats, + output_structure_mats, ) where {FT <: AbstractFloat} # get inputs and outputs @@ -385,12 +384,15 @@ function build_models!( end end - regularization = if isnothing(output_structure_matrix) + regularization = if isempty(output_structure_mats) 1.0 * I(n_rfms) - elseif isa(output_structure_matrix, UniformScaling) - output_structure_matrix else - Diagonal(output_structure_matrix) + output_structure_mat = get_structure_mat(output_structure_mats) + if isa(output_structure_mat, UniformScaling) + output_structure_mat + else + Diagonal(output_structure_mat) + end end @info ( diff --git a/src/Utilities.jl b/src/Utilities.jl index fa36c5988..298926c93 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -24,6 +24,8 @@ export create_encoder_schedule, decode_structure_matrix +const StructureMatrix = Union{UniformScaling, AbstractMatrix} +const StructureVector = Union{AbstractVector, AbstractMatrix} # In case of a matrix, the columns should be seen as vectors """ $(DocStringExtensions.TYPEDSIGNATURES) @@ -77,6 +79,49 @@ Base.:(==)(a::PDCP, b::PDCP) where {PDCP <: PairedDataContainerProcessor} = #### +function get_structure_vec(structure_vecs, name = nothing) + if isnothing(name) + if length(structure_vecs) == 1 + return only(values(structure_vecs)) + elseif isempty(structure_vecs) + throw(ArgumentError("Please provide a structure vector.")) + else + throw( + ArgumentError( + "Structure vectors $(collect(keys(structure_vecs))) are present. I received argument `name = nothing`, so I don't know which to use.", + ), + ) + end + else + if haskey(structure_vecs, name) + return structure_vecs[name] + else + throw(ArgumentError("Structure vector $name not found. Options: $(collect(keys(structure_vecs))).")) + end + end +end + +function get_structure_mat(structure_mats, name = nothing) + if isnothing(name) + if length(structure_mats) == 1 + return only(values(structure_mats)) + elseif isempty(structure_mats) + throw(ArgumentError("Please provide a structure matrix.")) + else + throw( + ArgumentError( + "Structure matrices $(collect(keys(structure_mats))) are present. I received argument `name = nothing`, so I don't know which to use.", + ), + ) + end + else + if haskey(structure_mats, name) + return structure_mats[name] + else + throw(ArgumentError("Structure matrix $name not found. Options: $(collect(keys(structure_mats))).")) + end + end +end function _encode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString} input_data, output_data = get_data(data) @@ -105,9 +150,10 @@ function _initialize_and_encode_data!( proc::PairedDataContainerProcessor, data, structure_mats, + structure_vecs, apply_to::AS, ) where {AS <: AbstractString} - initialize_processor!(proc, get_data(data)..., structure_mats..., apply_to) + initialize_processor!(proc, get_data(data)..., structure_mats..., structure_vecs..., apply_to) return _encode_data(proc, data, apply_to) end @@ -115,15 +161,17 @@ function _initialize_and_encode_data!( proc::DataContainerProcessor, data, structure_mats, + structure_vecs, apply_to::AS, ) where {AS <: AbstractString} input_data, output_data = get_data(data) - input_structure_mat, output_structure_mat = structure_mats + input_structure_mats, output_structure_mats = structure_mats + input_structure_vecs, output_structure_vecs = structure_vecs if apply_to == "in" - initialize_processor!(proc, input_data, input_structure_mat) + initialize_processor!(proc, input_data, input_structure_mats, input_structure_vecs) elseif apply_to == "out" - initialize_processor!(proc, output_data, output_structure_mat) + initialize_processor!(proc, output_data, output_structure_mats, output_structure_vecs) else bad_apply_to(apply_to) end @@ -189,18 +237,41 @@ Takes in the created encoder schedule (See [`create_encoder_schedule`](@ref)), a """ function initialize_and_encode_with_schedule!( encoder_schedule::VV, - io_pairs::PDC, - input_structure_mat::USorMorN1, - output_structure_mat::USorMorN2, -) where { - VV <: AbstractVector, - PDC <: PairedDataContainer, - USorMorN1 <: Union{UniformScaling, AbstractMatrix, Nothing}, - USorMorN2 <: Union{UniformScaling, AbstractMatrix, Nothing}, -} + io_pairs::PDC; + input_structure_mats = Dict{Symbol, StructureMatrix}(), + output_structure_mats = Dict{Symbol, StructureMatrix}(), + input_structure_vecs = Dict{Symbol, StructureVector}(), + output_structure_vecs = Dict{Symbol, StructureVector}(), + prior_cov::Union{Nothing, StructureMatrix} = nothing, + obs_noise_cov::Union{Nothing, StructureMatrix} = nothing, + observation::Union{Nothing, StructureVector} = nothing, + prior_samples_in::Union{Nothing, StructureVector} = nothing, + prior_samples_out::Union{Nothing, StructureVector} = nothing, +) where {VV <: AbstractVector, PDC <: PairedDataContainer} processed_io_pairs = deepcopy(io_pairs) - processed_input_structure_mat = deepcopy(input_structure_mat) - processed_output_structure_mat = deepcopy(output_structure_mat) + + input_structure_mats = deepcopy(input_structure_mats) + if !isnothing(prior_cov) + (input_structure_mats[:prior_cov] = prior_cov) + end + + output_structure_mats = deepcopy(output_structure_mats) + if !isnothing(obs_noise_cov) + (output_structure_mats[:obs_noise_cov] = obs_noise_cov) + end + + input_structure_vecs = deepcopy(input_structure_vecs) + if !isnothing(prior_samples_in) + (input_structure_vecs[:prior_samples_in] = prior_samples_in) + end + + output_structure_vecs = deepcopy(output_structure_vecs) + if !isnothing(observation) + (output_structure_vecs[:observation] = observation) + end + if !isnothing(prior_samples_out) + (output_structure_vecs[:prior_samples_out] = prior_samples_out) + end # apply_to is the string "in", "out" etc. for (processor, apply_to) in encoder_schedule @@ -209,38 +280,35 @@ function initialize_and_encode_with_schedule!( processed = _initialize_and_encode_data!( processor, processed_io_pairs, - (processed_input_structure_mat, processed_output_structure_mat), + (input_structure_mats, output_structure_mats), + (input_structure_vecs, output_structure_vecs), apply_to, ) if apply_to == "in" - processed_input_structure_mat = encode_structure_matrix(processor, processed_input_structure_mat) + input_structure_mats = Dict{Symbol, StructureMatrix}( + name => encode_structure_matrix(processor, mat) for (name, mat) in input_structure_mats + ) + input_structure_vecs = Dict{Symbol, StructureVector}( + name => encode_data(processor, vec) for (name, vec) in input_structure_vecs + ) processed_io_pairs = PairedDataContainer(processed, get_outputs(processed_io_pairs)) elseif apply_to == "out" - processed_output_structure_mat = encode_structure_matrix(processor, processed_output_structure_mat) + output_structure_mats = Dict{Symbol, StructureMatrix}( + name => encode_structure_matrix(processor, mat) for (name, mat) in output_structure_mats + ) + output_structure_vecs = Dict{Symbol, StructureVector}( + name => encode_data(processor, vec) for (name, vec) in output_structure_vecs + ) processed_io_pairs = PairedDataContainer(get_inputs(processed_io_pairs), processed) end end - return processed_io_pairs, processed_input_structure_mat, processed_output_structure_mat + return processed_io_pairs, input_structure_mats, output_structure_mats, input_structure_vecs, output_structure_vecs 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 diff --git a/src/Utilities/canonical_correlation.jl b/src/Utilities/canonical_correlation.jl index e8be4a611..c7ab6a47c 100644 --- a/src/Utilities/canonical_correlation.jl +++ b/src/Utilities/canonical_correlation.jl @@ -170,8 +170,10 @@ initialize_processor!( cc::CanonicalCorrelation, in_data::MM, out_data::MM, - input_structure_matrix, - output_structure_matrix, + input_structure_matrices, + output_structure_matrices, + input_structure_vectors, + output_structure_vectors, apply_to::AS, ) where {MM <: AbstractMatrix, AS <: AbstractString} = initialize_processor!(cc, in_data, out_data, apply_to) @@ -203,10 +205,7 @@ $(TYPEDSIGNATURES) Apply the `CanonicalCorrelation` encoder to a provided structure matrix """ -function encode_structure_matrix( - cc::CanonicalCorrelation, - structure_matrix::USorM, -) where {USorM <: Union{UniformScaling, AbstractMatrix}} +function encode_structure_matrix(cc::CanonicalCorrelation, structure_matrix::SM) where {SM <: StructureMatrix} encoder_mat = get_encoder_mat(cc)[1] return encoder_mat * structure_matrix * encoder_mat' end @@ -216,10 +215,7 @@ $(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}} +function decode_structure_matrix(cc::CanonicalCorrelation, enc_structure_matrix::SM) where {SM <: StructureMatrix} 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 index 62b51da1a..9f237a827 100644 --- a/src/Utilities/decorrelator.jl +++ b/src/Utilities/decorrelator.jl @@ -36,6 +36,8 @@ struct Decorrelator{VV1, VV2, VV3, FT, AS <: AbstractString} <: DataContainerPro retain_var::FT "Switch to choose what form of matrix to use to decorrelate the data" decorrelate_with::AS + "When given, use the structure matrix by this name if `decorrelate_with` uses structure matrices. When `nothing`, try to use the only present structure matrix instead." + structure_mat_name::Union{Nothing, Symbol} end """ @@ -48,8 +50,8 @@ Constructs the `Decorrelator` struct. Users can add optional keyword arguments: - `"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) +decorrelate(; retain_var::FT = Float64(1.0), decorrelate_with = "combined", structure_mat_name = nothing) where {FT} = + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), decorrelate_with, structure_mat_name) """ $(TYPEDSIGNATURES) @@ -58,7 +60,7 @@ Constructs the `Decorrelator` struct, setting decorrelate_with = "sample_cov". E - `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") + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), "sample_cov", nothing) """ $(TYPEDSIGNATURES) @@ -66,8 +68,8 @@ $(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") +decorrelate_structure_mat(; retain_var::FT = Float64(1.0), structure_mat_name = nothing) where {FT} = + Decorrelator([], [], [], clamp(retain_var, FT(0), FT(1)), "structure_mat", structure_mat_name) """ $(TYPEDSIGNATURES) @@ -121,8 +123,9 @@ Computes and populates the `data_mean` and `encoder_mat` and `decoder_mat` field function initialize_processor!( dd::Decorrelator, data::MM, - structure_matrix::USorMorN, -) where {MM <: AbstractMatrix, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}} + structure_matrices::Dict{Symbol, <:StructureMatrix}, + ::Dict{Symbol, <:StructureVector}, +) where {MM <: AbstractMatrix} if length(get_data_mean(dd)) == 0 push!(get_data_mean(dd), vec(mean(data, dims = 2))) end @@ -132,13 +135,8 @@ function initialize_processor!( # 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) + structure_matrix = get_structure_mat(structure_matrices, dd.structure_mat_name) + if isa(structure_matrix, UniformScaling) data_dim = size(data, 1) svdA = svd(structure_matrix(data_dim)) rk = data_dim @@ -151,6 +149,7 @@ function initialize_processor!( svdA = svd(cd) rk = rank(cd) elseif decorrelate_with == "combined" + structure_matrix = get_structure_mat(structure_matrices, dd.structure_mat_name) spluscd = structure_matrix + cov(data, dims = 2) svdA = svd(spluscd) rk = rank(spluscd) @@ -212,10 +211,7 @@ $(TYPEDSIGNATURES) Apply the `Decorrelator` encoder to a provided structure matrix """ -function encode_structure_matrix( - dd::Decorrelator, - structure_matrix::USorM, -) where {USorM <: Union{UniformScaling, AbstractMatrix}} +function encode_structure_matrix(dd::Decorrelator, structure_matrix::SM) where {SM <: StructureMatrix} encoder_mat = get_encoder_mat(dd)[1] return encoder_mat * structure_matrix * encoder_mat' end @@ -225,10 +221,7 @@ $(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}} +function decode_structure_matrix(dd::Decorrelator, enc_structure_matrix::SM) where {SM <: StructureMatrix} 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 index e139be4b3..6cca8a9d2 100644 --- a/src/Utilities/elementwise_scaler.jl +++ b/src/Utilities/elementwise_scaler.jl @@ -86,7 +86,7 @@ end function initialize_processor!( es::ElementwiseScaler, data::MM, - T::Type{QS}, + ::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 @@ -97,7 +97,7 @@ end function initialize_processor!( es::ElementwiseScaler, data::MM, - T::Type{MMS}, + ::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 @@ -108,7 +108,7 @@ end function initialize_processor!( es::ElementwiseScaler, data::MM, - T::Type{ZSS}, + ::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 @@ -156,8 +156,12 @@ $(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) +initialize_processor!( + es::ElementwiseScaler, + data::MM, + structure_matrices, + structure_vectors, +) where {MM <: AbstractMatrix} = initialize_processor!(es, data) """ @@ -165,10 +169,7 @@ $(TYPEDSIGNATURES) Apply the `ElementwiseScaler` encoder to a provided structure matrix """ -function encode_structure_matrix( - es::ElementwiseScaler, - structure_matrix::USorM, -) where {USorM <: Union{UniformScaling, AbstractMatrix}} +function encode_structure_matrix(es::ElementwiseScaler, structure_matrix::SM) where {SM <: StructureMatrix} return Diagonal(1 ./ get_scale(es)) * structure_matrix * Diagonal(1 ./ get_scale(es)) end @@ -177,9 +178,6 @@ $(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}} +function decode_structure_matrix(es::ElementwiseScaler, enc_structure_matrix::SM) where {SM <: StructureMatrix} return Diagonal(get_scale(es)) * enc_structure_matrix * Diagonal(get_scale(es)) end diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl index 50eedc316..6abcc83dd 100644 --- a/src/VectorRandomFeature.jl +++ b/src/VectorRandomFeature.jl @@ -353,9 +353,8 @@ Build Vector Random Feature model for the input-output pairs subject to regulari function build_models!( vrfi::VectorRandomFeatureInterface, input_output_pairs::PairedDataContainer{FT}, - input_structure_matrix, - output_structure_matrix; - kwargs..., + input_structure_mats, + output_structure_mats, ) where {FT <: AbstractFloat} # get inputs and outputs @@ -445,13 +444,16 @@ function build_models!( ) # 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") + regularization = if isempty(output_structure_mats) + 1.0 * I else - regularization = output_structure_matrix + output_structure_mat = get_structure_mat(output_structure_mats) + if !isposdef(output_structure_mat) + println("RF output structure matrix is not positive definite, correcting for use as a regularizer") + posdef_correct(output_structure_mat) + else + output_structure_mat + end end # [2.] Estimate covariance at mean value diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl index 503aedd06..5eec9a6c0 100644 --- a/test/Emulator/runtests.jl +++ b/test/Emulator/runtests.jl @@ -44,76 +44,48 @@ struct MLTester <: Emulators.MachineLearningTool end @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) + enc_io_pairs, enc_I_in, enc_I_out = initialize_and_encode_with_schedule!(enc_sch, io_pairs; obs_noise_cov = 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) + @test isempty(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) + @test isapprox(norm(enc_I_out[:obs_noise_cov] - 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, + io_pairs; + encoder_kwargs = (; prior_cov = 4.0 * I, obs_noise_cov = 2.0 * I), obs_noise_cov = 3.0 * I, - output_structure_matrix = 2.0 * I, ) - em1 = Emulator(gp, io_pairs, obs_noise_cov = Σ) + em1 = Emulator(gp, io_pairs; encoder_kwargs = (; obs_noise_cov = Σ)) - 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!( + initialize_and_encode_with_schedule!( enc_sch1, - io_pairs, - 1.0 * I(p), - Σ, # obs noise cov becomes the output structure matrix + io_pairs; + prior_cov = 1.0 * I(p), + obs_noise_cov = Σ, # obs noise cov becomes the output structure matrix ) @test get_encoder_schedule(em1) == enc_sch1 - (_, _, _) = 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 - - (_, _, _) = initialize_and_encode_with_schedule!(enc_sch3, io_pairs, Γ, I(d)) - @test get_encoder_schedule(em3) == enc_sch3 - end diff --git a/test/GaussianProcess/runtests.jl b/test/GaussianProcess/runtests.jl index 13b43d9fc..371e9df1e 100644 --- a/test/GaussianProcess/runtests.jl +++ b/test/GaussianProcess/runtests.jl @@ -139,6 +139,10 @@ using CalibrateEmulateSample.Utilities @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 + gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) + Γ = 0.05I + em = Emulator(gp, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = Γ)) + @test gp.regularization[end] == gp.alg_reg_noise * Γ.λ # ------------------------------------------------------------------------- # Test case 2: 2D input, 2D output @@ -177,11 +181,11 @@ using CalibrateEmulateSample.Utilities # 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, output_structure_matrix = Σ) + em4_noise_learnt = Emulator(gp4_noise_learnt, iopairs2; encoder_kwargs = (; obs_noise_cov = Σ)) gp4 = GaussianProcess(gppackage; kernel = nothing, noise_learn = false, prediction_type = pred_type) - em4_noise_from_Σ = Emulator(gp4, iopairs2, output_structure_matrix = Σ) + em4_noise_from_Σ = Emulator(gp4, iopairs2; encoder_kwargs = (; obs_noise_cov = Σ)) Emulators.optimize_hyperparameters!(em4_noise_learnt) Emulators.optimize_hyperparameters!(em4_noise_from_Σ) @@ -226,7 +230,7 @@ using CalibrateEmulateSample.Utilities ) for model_params in gp4_opt_params ] - em_agp_from_gp4 = Emulator(agp4, iopairs2, output_structure_matrix = Σ, kernel_params = kernel_params) + em_agp_from_gp4 = Emulator(agp4, iopairs2; encoder_kwargs = (; obs_noise_cov = Σ), 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 27621f09d..975dbec36 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -57,7 +57,7 @@ function test_gp_mv(y, σ2_y, iopairs::PairedDataContainer) # 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; output_structure_matrix = σ2_y) + em = Emulator(gp, iopairs; encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) return em end @@ -70,7 +70,7 @@ function test_gp_1(y, σ2_y, iopairs::PairedDataContainer) # 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; output_structure_matrix = σ2_y, encoder_schedule = []) + em = Emulator(gp, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) return em end @@ -83,7 +83,7 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer) # 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; output_structure_matrix = σ2_y, encoder_schedule = []) + em = Emulator(gp, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) # now make agp from gp @@ -100,8 +100,13 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer) ) for model_params in gp_opt_params ] - em_agp = - Emulator(agp, iopairs, output_structure_matrix = σ2_y, encoder_schedule = [], kernel_params = kernel_params) + em_agp = Emulator( + agp, + iopairs; + encoder_schedule = [], + encoder_kwargs = (; obs_noise_cov = σ2_y), + kernel_params = kernel_params, + ) return em, em_agp end @@ -117,7 +122,7 @@ function test_gp_2(y, σ2_y, iopairs::PairedDataContainer) gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) 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) + em = Emulator(gp, iopairs; encoder_schedule = encoder_schedule, encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) return em diff --git a/test/RandomFeature/runtests.jl b/test/RandomFeature/runtests.jl index 1201faa8e..90feb3a38 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, output_structure_matrix = obs_noise_cov) + em_srfi = Emulator(srfi, iopairs; encoder_kwargs = (; obs_noise_cov = obs_noise_cov)) n_srfi = length(get_rfms(srfi)) - em_vrfi = Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) + em_vrfi = Emulator(vrfi, iopairs; encoder_kwargs = (; obs_noise_cov = obs_noise_cov)) n_vrfi = length(get_rfms(vrfi)) # test bad case @@ -324,11 +324,19 @@ rng = Random.MersenneTwister(seed) @test_throws ArgumentError Emulator(srfi_bad, iopairs) # test under repeats - @test_logs (:info,) (:info,) (:warn,) Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov) - Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov) + @test_logs (:info,) (:info,) (:warn,) Emulator( + srfi, + iopairs; + encoder_kwargs = (; obs_noise_cov = obs_noise_cov), + ) + Emulator(srfi, iopairs; encoder_kwargs = (; obs_noise_cov = obs_noise_cov)) @test length(get_rfms(srfi)) == n_srfi - @test_logs (:info,) (:info,) (:warn,) Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) - Emulator(vrfi, iopairs, output_structure_matrix = obs_noise_cov) + @test_logs (:info,) (:info,) (:warn,) Emulator( + vrfi, + iopairs; + encoder_kwargs = (; obs_noise_cov = obs_noise_cov), + ) + Emulator(vrfi, iopairs; encoder_kwargs = (; obs_noise_cov = obs_noise_cov)) @test length(get_rfms(vrfi)) == n_vrfi @@ -422,14 +430,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, output_structure_matrix = Σ) - em_srfi_svd = Emulator(srfi, iopairs, output_structure_matrix = Σ) + em_srfi_svd_diagin = Emulator(srfi_diagin, iopairs; encoder_kwargs = (; obs_noise_cov = Σ)) + em_srfi_svd = Emulator(srfi, iopairs; encoder_kwargs = (; 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_svd_diagout = Emulator(vrfi_diagout, iopairs; encoder_kwargs = (; obs_noise_cov = Σ)) + em_vrfi_svd = Emulator(vrfi, iopairs; encoder_kwargs = (; obs_noise_cov = Σ)) - em_vrfi = Emulator(vrfi, iopairs, output_structure_matrix = Σ, decorrelate = false) - em_vrfi_nonsep = Emulator(vrfi_nonsep, iopairs, output_structure_matrix = Σ, decorrelate = false) + em_vrfi = Emulator(vrfi, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = Σ)) + em_vrfi_nonsep = Emulator(vrfi_nonsep, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = Σ)) #TODO truncated SVD option for vector (involves resizing prior) @@ -464,7 +472,7 @@ rng = Random.MersenneTwister(seed) optimizer_options = Dict("train_fraction" => 0.7, "multithread" => "tullio"), rng = rng, ) - em_vrfi_svd_tul = Emulator(vrfi_tul, iopairs, output_structure_matrix = Σ) + em_vrfi_svd_tul = Emulator(vrfi_tul, iopairs; encoder_kwargs = (; obs_noise_cov = Σ)) μ_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 6871b9608..f235ed0f3 100644 --- a/test/Utilities/runtests.jl +++ b/test/Utilities/runtests.jl @@ -49,6 +49,25 @@ end # Seed for pseudo-random number generator rng = Random.MersenneTwister(4154) + # Tests for get_structure_vec and get_structure_mat + structure_vecs = Dict("a" => [1, 2, 3], "b" => [4, 5, 6]) + structure_mats = Dict("c" => [1 2; 3 4], "d" => [5 6; 7 8]) + @test Utilities.get_structure_vec(structure_vecs, "a") == [1, 2, 3] + @test Utilities.get_structure_vec(structure_vecs, "b") == [4, 5, 6] + @test Utilities.get_structure_mat(structure_mats, "c") == [1 2; 3 4] + @test Utilities.get_structure_mat(structure_mats, "d") == [5 6; 7 8] + @test_throws ArgumentError Utilities.get_structure_vec(structure_vecs, "c") + @test_throws ArgumentError Utilities.get_structure_mat(structure_mats, "a") + @test_throws ArgumentError Utilities.get_structure_vec(structure_vecs) + @test_throws ArgumentError Utilities.get_structure_mat(structure_mats) + + structure_vecs = Dict("a" => [1, 2, 3]) + @test Utilities.get_structure_vec(structure_vecs) == [1, 2, 3] + @test Utilities.get_structure_vec(structure_vecs, "a") == [1, 2, 3] + structure_mats = Dict("c" => [1 2; 3 4]) + @test Utilities.get_structure_mat(structure_mats) == [1 2; 3 4] + @test Utilities.get_structure_mat(structure_mats, "c") == [1 2; 3 4] + # quick build tests and test getters zs = zscore_scale() mm = minmax_scale() @@ -72,7 +91,7 @@ end dd3 = decorrelate_structure_mat(retain_var = 0.7) @test get_retain_var(dd3) == 0.7 @test get_decorrelate_with(dd3) == "structure_mat" - DD = Decorrelator([1], [2], [3], 1.0, "test") + DD = Decorrelator([1], [2], [3], 1.0, "test", nothing) @test get_data_mean(DD) == [1] @test get_encoder_mat(DD) == [2] @test get_decoder_mat(DD) == [3] @@ -142,17 +161,21 @@ 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) = - 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) + (encoded_io_pairs, encoded_input_structure_mats, encoded_output_structure_mats, _, _) = + initialize_and_encode_with_schedule!(encoder_schedule, io_pairs; prior_cov, obs_noise_cov) + + (decoded_io_pairs, decoded_input_structure_mat, decoded_output_structure_mat) = decode_with_schedule( + encoder_schedule, + encoded_io_pairs, + encoded_input_structure_mats[:prior_cov], + encoded_output_structure_mats[:obs_noise_cov], + ) for (enc_dat, dec_dat, test_dat, enc_covv, dec_covv, test_covv, dim) in zip( (get_inputs(encoded_io_pairs), get_outputs(encoded_io_pairs)), (get_inputs(decoded_io_pairs), get_outputs(decoded_io_pairs)), (get_inputs(io_pairs), get_outputs(io_pairs)), - (encoded_prior_cov, encoded_obs_noise_cov), - (decoded_prior_cov, decoded_obs_noise_cov), + (encoded_input_structure_mats[:prior_cov], encoded_output_structure_mats[:obs_noise_cov]), + (decoded_input_structure_mat, decoded_output_structure_mat), (prior_cov, obs_noise_cov), (in_dim, out_dim), ) @@ -275,10 +298,10 @@ 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 + @test_throws ArgumentError initialize_and_encode_with_schedule!(schedule2a, io_pairs; obs_noise_cov = obs_noise_cov) sch2b = (decorrelate_structure_mat(), "out") schedule2b = create_encoder_schedule(sch2b) - @test_throws ArgumentError initialize_and_encode_with_schedule!(schedule2b, io_pairs, prior_cov, nothing) + @test_throws ArgumentError initialize_and_encode_with_schedule!(schedule2b, io_pairs; prior_cov = prior_cov) @@ -304,7 +327,7 @@ end bad_encoder_schedule = [(canonical_correlation(), "bad")] @test_throws ArgumentError initialize_and_encode_with_schedule!( bad_encoder_schedule, - io_pairs, + io_pairs; prior_cov, obs_noise_cov, ) @@ -314,8 +337,10 @@ end encoder_schedule = create_encoder_schedule(schedule_builder) # encode the data using the schedule - (encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) = - initialize_and_encode_with_schedule!(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + (encoded_io_pairs, encoded_input_structure_mats, encoded_output_structure_mats) = + initialize_and_encode_with_schedule!(encoder_schedule, io_pairs; prior_cov, obs_noise_cov) + encoded_prior_cov = encoded_input_structure_mats[:prior_cov] + encoded_obs_noise_cov = encoded_output_structure_mats[:obs_noise_cov] # decode the data using the schedule (decoded_io_pairs, decoded_prior_cov, decoded_obs_noise_cov) =