Merge pull request #373 from CliMA/ab/struct-mat-dict

Store structure matrices and vectors with names as dictionaries

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

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

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)

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

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

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

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

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