Update Emulator and tests

Author: ArneBouillon

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}
 ```

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}
 ```
 

docs/src/emulate.md

@@ -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}
 ```

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)
 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)
 
 

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)
@@ -109,71 +108,60 @@ Constructor of the Emulator object,
 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`
+ - `encoder_kwargs`[=`NamedTuple()`]: a Dict or NamedTuple with keyword arguments to be passed to `initialize_and_encode_with_schedule!`
 
 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.
 Other keywords are passed to the machine learning tool initialization
 """
 function Emulator(
     machine_learning_tool::MachineLearningTool,
-    input_output_pairs::PairedDataContainer{FT};
+    input_output_pairs::PairedDataContainer{FT},
+    encoder_kwargs = NamedTuple();
     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 
+    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) =
+    encoder_schedule = create_encoder_schedule(encoder_schedule)
+    (encoded_io_pairs, input_structure_mats, output_structure_mats, _, _) =
         initialize_and_encode_with_schedule!(
-            enc_schedule,
-            input_output_pairs,
-            input_structure_matrix,
-            output_structure_matrix,
+            encoder_schedule,
+            input_output_pairs;
+            obs_noise_cov,
+            encoder_kwargs...,
         )
 
     # 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;
+        input_structure_mats,
+        output_structure_mats;
         mlt_kwargs...,
     )
-    return Emulator{FT, typeof(enc_schedule)}(machine_learning_tool, input_output_pairs, encoded_io_pairs, enc_schedule)
+    return Emulator{FT, typeof(encoder_schedule)}(machine_learning_tool, input_output_pairs, encoded_io_pairs, encoder_schedule)
 end
 
 """
@@ -209,9 +197,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 +228,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
 

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
 

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 (