From ba2a3b54530b4da9bb016c6fdb93c417559db2da Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 25 Jun 2025 16:35:57 -0700
Subject: [PATCH 01/75] initial gutting of Emulator transformation and replace
 with new encoders

---
 src/Emulator.jl | 252 ++++++++++++++++--------------------------------
 1 file changed, 83 insertions(+), 169 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 1904fd9be..e23c0473d 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -45,8 +45,6 @@ function predict(mlt, new_inputs; mlt_kwargs...)
     throw_define_mlt()
 end
 
-
-
 # We will define the different emulator types after the general statements
 
 """
@@ -57,28 +55,21 @@ Structure used to represent a general emulator, independently of the algorithm u
 # Fields
 $(DocStringExtensions.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
+get_io_pairs(emulator::Emulator) = emulator.io_pairs
+get_encoded_io_pairs(emulator::Emulator) = emulator.encoded_io_pairs
+get_encoder_schedule(emulator::Emulator) = emulator.encoder_schedule
 
 # Constructor for the Emulator Object
 """
@@ -88,90 +79,61 @@ 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)
+ 
 """
 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,
+    user_encoder_schedule = nothing,
+    input_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing,
+    output_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing,
     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.
+    input_structure_mat = if isnothing(input_structure_matrix)
+        Diagonal(FT.(ones(input_dim)))
+    elseif isa(input_structure_matrix, UniformScaling)
+        Diagonal(input_structure_matrix(input_dim))
     else
-        training_inputs = get_inputs(input_output_pairs)
+        input_structure_matrix
     end
-
-    # [2.] Standardize the outputs?
-    if standardize_outputs
-        training_outputs, obs_noise_cov =
-            standardize(get_outputs(input_output_pairs), obs_noise_cov, standardize_outputs_factors)
+    
+    output_structure_mat = if isnothing(output_structure_matrix)
+        Diagonal(FT.(ones(output_dim)))
+    elseif isa(output_structure_matrix, UniformScaling)
+        Diagonal(output_structure_matrix(output_dim))
     else
-        training_outputs = get_outputs(input_output_pairs)
+        output_structure_matrix
     end
 
-    # [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)
-
-        training_pairs = PairedDataContainer(training_inputs, decorrelated_training_outputs)
-        # [4.] build an emulator
-
-        build_models!(machine_learning_tool, training_pairs; mlt_kwargs...)
+    # [1.] Initializes and performs data encoding schedule
+    if !isnothing(encoder_schedule)
+        
+        encoder_schedule = create_encoder_schedule(user_encoder_schedule)
+        (encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) =
+            encode_with_schedule(
+                encoder_schedule,
+                input_output_pairs,
+                input_structure_matrix,
+                output_structure_matrix,
+            )
     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...)
+        encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) = (input_output_pairs, input_structure_matrix,  output_structure_matrix)
     end
 
+    # build_models!(machine_learning_tool, encoded_io_pairs; mlt_kwargs...)
+    # build_models!(machine_learning_tool, training_pairs; regularization_matrix = obs_noise_cov, mlt_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; mlt_kwargs...)
     return Emulator{FT}(
-        machine_learning_tool,
-        training_pairs,
-        input_mean,
-        normalize_inputs,
-        normalization,
-        standardize_outputs,
-        standardize_outputs_factors,
-        decomposition,
-        retained_svd_frac,
+    machine_learning_tool,
+        input_output_pairs,
+        encoded_io_pairs,
+        encoder_schedule,
     )
 end
 
@@ -195,23 +157,18 @@ function predict(
     emulator::Emulator{FT},
     new_inputs::AM;
     transform_to_real = false,
-    vector_rf_unstandardize = true,
     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)
-
+    # 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_io_pairs(emulator), 1)
+    
     N_samples = size(new_inputs, 2)
-
+    
     # 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))",
-            ),
-        )
-    else
+    if
         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))",
@@ -219,83 +176,40 @@ function predict(
         )
     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)
-        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
-        end
+    # encode the new input data
+    encoder_schedule = get_encoder_schedule(emulator)
+    encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_in_data), "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), get_data(encoded_inputs), mlt_kwargs...)
+
+    var_or_cov = (ndims(encoded_uncertainties[1]) == 2) ? "var" : "cov"
+
+    # return decoded or encoded?
+    if transform_to_real
+        decoded_outputs = decode_with_schedule(encoder_schedule, DataContainer(encoded_outputs), "out")
+        
+        decoded_covariances = (var_or_cov == "var") ?
+            [decode_with_schedule(encoder_schedule, Diagonal(col), "out") for col in eachcol(encoded_uncertainties) ] :
+            [decode_with_schedule(encoder_schedule, mat, "out") for mat in eachslice(encoded_uncertainties, dims=3) ] 
+        return decoded_outputs, decoded_covariances
     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 encoded_output_dim > 1
+            encoded_covariances = [Diagonal(col) for col in eachcol(encoded_uncertainties)]
+        else
+            encoded_covariances = encoded_uncertainties
         end
-
+        return encoded_outputs, encoded_covariances
     end
-
+    
 end
 
+#=
+
 # Normalization, Standardization, and Decorrelation
 """
 $(DocStringExtensions.TYPEDSIGNATURES)
@@ -540,7 +454,7 @@ function svd_reverse_transform_mean_cov(
 
     return svd_reverse_transform_mean_cov(μ, σ2_as_cov, decomposition)
 end
-
+=#
 
 
 

From 9b1eeddf59553bfaa0d3f5bb7930f3223a880089 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 26 Jun 2025 13:17:34 -0700
Subject: [PATCH 02/75] emulator bug-fixes

---
 src/Emulator.jl | 33 +++++++++++++++------------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index e23c0473d..cbd7f5cf7 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -110,27 +110,25 @@ function Emulator(
     end
 
     # [1.] Initializes and performs data encoding schedule
-    if !isnothing(encoder_schedule)
+    if !isnothing(user_encoder_schedule)
         
         encoder_schedule = create_encoder_schedule(user_encoder_schedule)
-        (encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) =
+        (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
             encode_with_schedule(
                 encoder_schedule,
                 input_output_pairs,
-                input_structure_matrix,
-                output_structure_matrix,
+                input_structure_mat,
+                output_structure_mat,
             )
     else
-        encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) = (input_output_pairs, input_structure_matrix,  output_structure_matrix)
+        (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) = (input_output_pairs, input_structure_mat,  output_structure_mat)
+        encoder_schedule = user_encoder_schedule
     end
 
-    # build_models!(machine_learning_tool, encoded_io_pairs; mlt_kwargs...)
-    # build_models!(machine_learning_tool, training_pairs; regularization_matrix = obs_noise_cov, mlt_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; mlt_kwargs...)
-    return Emulator{FT}(
-    machine_learning_tool,
+    build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
+    return Emulator{FT, typeof(encoder_schedule)}(
+        machine_learning_tool,
         input_output_pairs,
         encoded_io_pairs,
         encoder_schedule,
@@ -167,9 +165,8 @@ function predict(
     
     N_samples = size(new_inputs, 2)
     
-    # check sizing against normalization
-    if
-        size(new_inputs, 1) == input_dim || throw(
+    if !(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))",
             ),
@@ -178,7 +175,7 @@ function predict(
 
     # encode the new input data
     encoder_schedule = get_encoder_schedule(emulator)
-    encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_in_data), "in")
+    encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_inputs), "in")
 
     # predict in encoding space
     # returns outputs: [enc_out_dim x n_samples]
@@ -186,8 +183,8 @@ function predict(
     # Vector-methods uncertainties=covariances: [enc_out_dim x enc_out_dim x n_samples)
     encoded_outputs, encoded_uncertainties = predict(get_machine_learning_tool(emulator), get_data(encoded_inputs), mlt_kwargs...)
 
-    var_or_cov = (ndims(encoded_uncertainties[1]) == 2) ? "var" : "cov"
-
+    var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
+    
     # return decoded or encoded?
     if transform_to_real
         decoded_outputs = decode_with_schedule(encoder_schedule, DataContainer(encoded_outputs), "out")
@@ -195,7 +192,7 @@ function predict(
         decoded_covariances = (var_or_cov == "var") ?
             [decode_with_schedule(encoder_schedule, Diagonal(col), "out") for col in eachcol(encoded_uncertainties) ] :
             [decode_with_schedule(encoder_schedule, mat, "out") for mat in eachslice(encoded_uncertainties, dims=3) ] 
-        return decoded_outputs, decoded_covariances
+        return get_data(decoded_outputs), decoded_covariances
     else
         
         if encoded_output_dim > 1

From d231fb56ef1d9f2f09d0d229b07c758f47b0109e Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 26 Jun 2025 13:18:36 -0700
Subject: [PATCH 03/75] add new interface

---
 src/GaussianProcess.jl     | 12 +++++++++---
 src/ScalarRandomFeature.jl |  4 +++-
 src/VectorRandomFeature.jl | 27 ++++++++++-----------------
 3 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 6a59669ba..d26f9a95a 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -138,7 +138,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 
@@ -258,7 +260,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 
@@ -336,7 +340,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}
diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 2333a0226..146e5f6ec 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -306,7 +306,9 @@ 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 
diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl
index cad00bd6d..4a35cbd90 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,14 @@ 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 !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"]

From ab204d0500c0f6db034973c817f69a07d95538f0 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 26 Jun 2025 15:01:16 -0700
Subject: [PATCH 04/75] remove old normalization/standardizations code

---
 src/Emulator.jl | 275 ++++--------------------------------------------
 1 file changed, 19 insertions(+), 256 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index cbd7f5cf7..8005f21f8 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -184,19 +184,31 @@ function predict(
     encoded_outputs, encoded_uncertainties = predict(get_machine_learning_tool(emulator), get_data(encoded_inputs), mlt_kwargs...)
 
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
-    
+
     # return decoded or encoded?
     if transform_to_real
         decoded_outputs = decode_with_schedule(encoder_schedule, DataContainer(encoded_outputs), "out")
-        
-        decoded_covariances = (var_or_cov == "var") ?
-            [decode_with_schedule(encoder_schedule, Diagonal(col), "out") for col in eachcol(encoded_uncertainties) ] :
-            [decode_with_schedule(encoder_schedule, mat, "out") for mat in eachslice(encoded_uncertainties, dims=3) ] 
-        return get_data(decoded_outputs), decoded_covariances
+   
+        decoded_covariances = zeros(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_with_schedule(encoder_schedule, Diagonal(col), "out")
+            end
+        else # == "cov"
+            for (i,mat) in enumerate(eachslice(encoded_uncertainties, dims=3))
+               decoded_covariances[:,:,i] .= decode_with_schedule(encoder_schedule, mat, "out")
+            end
+        end
+            
+        return get_data(decoded_outputs), eachslice(decoded_covariances,dims=3)
     else
         
         if encoded_output_dim > 1
-            encoded_covariances = [Diagonal(col) for col in eachcol(encoded_uncertainties)]
+            encoded_covariances_mat = zeros(output_dim, output_dim, size(encoded_uncertainties)[end])
+            for (i,col) in enumerate(eachcol(encoded_uncertainties))
+                encoded_covariances_mat[:,:,i] = Diagonal(col)
+            end
+            encoded_covariances = eachslice(encoded_covariances_mat,dims=3)
         else
             encoded_covariances = encoded_uncertainties
         end
@@ -205,254 +217,5 @@ function predict(
     
 end
 
-#=
-
-# Normalization, Standardization, and Decorrelation
-"""
-$(DocStringExtensions.TYPEDSIGNATURES)
-
-Calculate the normalization of inputs.
-"""
-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
-end
-
-"""
-$(DocStringExtensions.TYPEDSIGNATURES)
-
-Normalize the input data, with a normalizing function.
-"""
-function normalize(emulator::Emulator, inputs::VOrM) where {VOrM <: AbstractVecOrMat}
-    if emulator.normalize_inputs
-        return normalize(inputs, emulator.input_mean, emulator.normalization)
-    else
-        return inputs
-    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)
-
-Standardize with a vector of factors (size equal to output dimension).
-"""
-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
-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)
-
-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).
-"""
-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)
-    else
-        return outputs, output_covs
-    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.
-"""
-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
-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
-
-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
-
-"""
-$(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}
-
-    N_predicted_points = length(σ2)
-    if !(eltype(σ2) <: AbstractMatrix)
-        throw(
-            ArgumentError(
-                "input σ2 must be a vector of eltype <: AbstractMatrix, instead  eltype(σ2) = ",
-                string(eltype(σ2)),
-            ),
-        )
-    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 * μ
-
-    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
-
-    return transformed_μ, transformed_σ2
-
-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])
-
-    return svd_reverse_transform_mean_cov(μ, σ2_as_cov, decomposition)
-end
-=#
-
-
 
 end

From 70d4a4ffce1c87b7a4db88455cd74639b0461f0b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 10:21:06 -0700
Subject: [PATCH 05/75] works when no scheduler provided

---
 src/Emulator.jl | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 8005f21f8..75179f3a7 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -111,7 +111,6 @@ function Emulator(
 
     # [1.] Initializes and performs data encoding schedule
     if !isnothing(user_encoder_schedule)
-        
         encoder_schedule = create_encoder_schedule(user_encoder_schedule)
         (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
             encode_with_schedule(
@@ -120,9 +119,9 @@ function Emulator(
                 input_structure_mat,
                 output_structure_mat,
             )
-    else
+    else 
         (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) = (input_output_pairs, input_structure_mat,  output_structure_mat)
-        encoder_schedule = user_encoder_schedule
+        encoder_schedule = []
     end
 
     # build the machine learning tool in the encoded space
@@ -176,19 +175,16 @@ function predict(
     # encode the new input data
     encoder_schedule = get_encoder_schedule(emulator)
     encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(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), get_data(encoded_inputs), mlt_kwargs...)
-
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
-
     # return decoded or encoded?
     if transform_to_real
         decoded_outputs = decode_with_schedule(encoder_schedule, DataContainer(encoded_outputs), "out")
-   
+        
         decoded_covariances = zeros(output_dim, output_dim, size(encoded_uncertainties)[end])
         if var_or_cov == "var"
             for (i,col) in enumerate(eachcol(encoded_uncertainties))

From ecb57fa957d1b09073b32661da7141d970989a7b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 11:01:31 -0700
Subject: [PATCH 06/75] trunc -> trunc_val

---
 src/Utilities.jl | 46 ++++++++++++++++++++++------------------------
 1 file changed, 22 insertions(+), 24 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 7a85a811d..1bbba08b2 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -439,17 +439,17 @@ function initialize_processor!(
         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"
+            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 = rk
+            trunc_val = rk
         end
-        sqrt_inv_sv = Diagonal(1.0 ./ sqrt.(svdA.S[1:trunc]))
-        sqrt_sv = Diagonal(sqrt.(svdA.S[1:trunc]))
+        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, :]
-        decoder_mat = svdA.Vt[1:trunc, :]' * sqrt_sv
+        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)
@@ -628,40 +628,39 @@ function initialize_processor!(
         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')
-
+        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)^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"
+            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 = min(rank(in_data), rank(out_data))
+            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"
-            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]
+            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)
-            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]
+            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 - mean(in))
-        # v = out_encoder * (out - mean(out))
+        # 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])
+        # v * u' = u * v' = Diagonal(svdio.S[1:trunc_val])
     end
 end
 
@@ -882,7 +881,6 @@ 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(

From cb0ad487a18c9e64880eb75f638805f1632281ce Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 14:44:54 -0700
Subject: [PATCH 07/75] update and cleanup example and plots

---
 .../Regression_2d_2d/compare_regression.jl    | 254 +++++++-----------
 1 file changed, 96 insertions(+), 158 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index df9282d32..be4dd83cc 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,34 @@ 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-sep-scalar",
+        "rf-sep-diag",
+        "rf-sep-sep",
+        "rf-nonsep",
     ]
-    case_mask = [1, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite)
-
+    case_mask = [3]#, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite)
+
+    #choose your encoders: 
+    encoder_schedule = [
+        #            (decorrelate_sample_cov(), "in_and_out"),
+        #            (decorrelate_structure_mat(), "in_and_out"),
+        #            (decorrelate(), "in_and_out"),
+        #            (canonical_correlation(), "in_and_out")           
+        ]
+    
+    
+    
     #problem
-    n = 200  # number of training points
+    n = 100  # 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 +78,7 @@ function main()
 
     # Add noise η
     μ = zeros(d)
-    Σ = 0.05 * [[0.8, 0.1] [0.1, 0.5]] # d x d
+    Σ = [[0.4, -0.3] [-0.3, 0.5]] # d x d
     noise_samples = rand(MvNormal(μ, Σ), n)
     # y = G(x) + η
     Y = gx .+ noise_samples
@@ -78,61 +89,51 @@ function main()
 
     #plot training data with and without noise
     if plot_flag
-        p1 = plot(
-            X[1, :],
-            X[2, :],
-            g1x,
-            st = :surface,
-            camera = (30, 60),
+        n_pts = 200
+        x1 = range(-2*pi, stop = 2*π, length = n_pts)
+        x2 = range(-2*pi, stop = 2*π, length = n_pts)
+        
+        p1 = contourf(
+            x1,
+            x2,
+            g1.(x1',x2),
             c = :cividis,
             xlabel = "x1",
             ylabel = "x2",
-            zguidefontrotation = 90,
-        )
+            aspect_ratio = :equal            
 
-        figpath = joinpath(output_directory, "observed_y1nonoise.png")
-        savefig(figpath)
-
-        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)
-
-        p3 = plot(
-            X[1, :],
-            X[2, :],
-            Y[1, :],
-            st = :surface,
-            camera = (30, 60),
-            c = :cividis,
-            xlabel = "x1",
-            ylabel = "x2",
-            zguidefontrotation = 90,
+        scatter!(
+            p1,
+            X[1,:],
+            X[2,:],
+            c=:black,
+            ms=3,
+            label="train point"
         )
-        figpath = joinpath(output_directory, "observed_y1.png")
+        
+        figpath = joinpath(output_directory, "g1_true.png")
         savefig(figpath)
-
-        p4 = plot(
-            X[1, :],
-            X[2, :],
-            Y[2, :],
-            st = :surface,
-            camera = (30, 60),
+        p2 = contourf(
+            x1,
+            x2,
+            g2.(x1',x2),
             c = :cividis,
             xlabel = "x1",
             ylabel = "x2",
-            zguidefontrotation = 90,
+            aspect_ratio = :equal            
+
         )
-        figpath = joinpath(output_directory, "observed_y2.png")
+        scatter!(
+            p2,
+            X[1,:],
+            X[2,:],
+            c=:black,
+            ms=3,
+            label="train point"
+        )
+        
+        figpath = joinpath(output_directory, "g2_true.png")
         savefig(figpath)
 
     end
@@ -148,98 +149,74 @@ function main()
         # 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
-
-
+            Dict("n_iteration" => 10, "n_features_opt" => 150, "n_ensemble" => 80, "cov_sample_multiplier" => 5.0)
+        nugget = 1e-4
 
+        # data processing schedule
+   
         if case == "gp-skljl"
             gppackage = SKLJL()
-            gaussian_process = GaussianProcess(gppackage, noise_learn = true)
-            emulator = Emulator(gaussian_process, iopairs, obs_noise_cov = Σ, normalize_inputs = true)
+            mlt = GaussianProcess(gppackage, noise_learn = 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(
+            mlt = GaussianProcess(gppackage, noise_learn = true)
+           
+        elseif case == "rf-sep-scalar"
+            mlt = 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(
+        elseif case == "rf-sep-diag"
+            mlt = 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(
+         elseif case == "rf-sep-sep"
+            mlt = 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(
+        elseif case == "rf-nonsep"
+            mlt = 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.")
 
+        emulator = Emulator(
+            mlt,
+            iopairs,
+            user_encoder_schedule = encoder_schedule,
+            input_structure_matrix = prior_cov,
+            output_structure_matrix = Σ,
+        )
+
+        
         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 = 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")
 
@@ -252,32 +229,26 @@ function main()
 
             mean_grid = reshape(em_mean[y_i, :], n_pts, n_pts) # 2 x 40000
             if plot_flag
-                p5 = plot(
+                p5 = contourf(
                     x1,
                     x2,
                     mean_grid,
-                    st = :surface,
-                    camera = (30, 60),
                     c = :cividis,
                     xlabel = "x1",
                     ylabel = "x2",
-                    zlabel = "mean of y" * string(y_i),
-                    zguidefontrotation = 90,
+                    title="mean output $y_i",
                 )
             end
             var_grid = reshape(em_var[y_i, :], n_pts, n_pts)
             if plot_flag
-                p6 = plot(
+                p6 = contourf(
                     x1,
                     x2,
                     var_grid,
-                    st = :surface,
-                    camera = (30, 60),
                     c = :cividis,
                     xlabel = "x1",
                     ylabel = "x2",
-                    zlabel = "var of y" * string(y_i),
-                    zguidefontrotation = 90,
+                    title="var output $y_i",
                 )
 
                 plot(p5, p6, layout = (1, 2), legend = false)
@@ -289,42 +260,11 @@ function main()
         # 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,
-            )
-            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"))
-
-        end
+        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)
 
@@ -342,17 +282,15 @@ 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"))

From b070a6c03c8ddaccc2b70b7ef69530a0b6896964 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 14:57:49 -0700
Subject: [PATCH 08/75] gp case

---
 .../Emulator/Regression_2d_2d/compare_regression.jl   |  6 +++---
 src/Emulator.jl                                       | 11 +++++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index be4dd83cc..0c2014e82 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -47,7 +47,7 @@ function main()
         "rf-sep-sep",
         "rf-nonsep",
     ]
-    case_mask = [3]#, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite)
+    case_mask = [1]#, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite)
 
     #choose your encoders: 
     encoder_schedule = [
@@ -149,8 +149,8 @@ function main()
         # common random feature setup
         n_features = 300
         optimizer_options =
-            Dict("n_iteration" => 10, "n_features_opt" => 150, "n_ensemble" => 80, "cov_sample_multiplier" => 5.0)
-        nugget = 1e-4
+            Dict("n_iteration" => 10, "n_features_opt" => 60, "n_ensemble" => 30, "cov_sample_multiplier" => 5.0, "scheduler" => DataMisfitController(terminate_at=10))
+        nugget = 1.0
 
         # data processing schedule
    
diff --git a/src/Emulator.jl b/src/Emulator.jl
index 75179f3a7..e83f26221 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -124,6 +124,15 @@ function Emulator(
         encoder_schedule = []
     end
 
+    @info input_structure_mat
+    @info output_structure_mat
+    @info encoded_input_structure_mat
+    @info encoded_output_structure_mat
+    @info cov(get_inputs(encoded_io_pairs), dims=2)
+    @info cov(get_outputs(encoded_io_pairs), dims=2)
+    @info get_inputs(encoded_io_pairs) * get_outputs(encoded_io_pairs)'
+    @info get_outputs(encoded_io_pairs)* get_inputs(encoded_io_pairs)'
+    
     # build the machine learning tool in the encoded space
     build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
     return Emulator{FT, typeof(encoder_schedule)}(
@@ -175,11 +184,13 @@ function predict(
     # encode the new input data
     encoder_schedule = get_encoder_schedule(emulator)
     encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_inputs), "in")
+    @info cov(get_data(encoded_inputs),dims=2)
     # 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), get_data(encoded_inputs), mlt_kwargs...)
+    @info cov(encoded_outputs,dims=2)
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
     # return decoded or encoded?
     if transform_to_real

From 7f85bcc4848bf477b3983f818e3c04d641efe3b2 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 15:25:50 -0700
Subject: [PATCH 09/75] cycle over cases and encoders

---
 .../Regression_2d_2d/compare_regression.jl    | 52 ++++++++++++-------
 src/Emulator.jl                               | 14 +----
 2 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index 0c2014e82..084a11b97 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -47,16 +47,25 @@ function main()
         "rf-sep-sep",
         "rf-nonsep",
     ]
-    case_mask = [1]#, 3:length(cases)...] # (KEEP set to 1:length(cases) when pushing for buildkite)
-
-    #choose your encoders: 
-    encoder_schedule = [
-        #            (decorrelate_sample_cov(), "in_and_out"),
-        #            (decorrelate_structure_mat(), "in_and_out"),
-        #            (decorrelate(), "in_and_out"),
-        #            (canonical_correlation(), "in_and_out")           
-        ]
-    
+
+    encoder_names = [
+        "no-proc",
+        "sample-proc",
+        "struct-mat-proc",
+        "combined-proc",
+        "cca-proc",
+    ]
+    encoders = [
+        [],
+        (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 = [1:5...]
+    case_mask = [1] # (KEEP set to 1:length(cases) when pushing for buildkite)
     
     
     #problem
@@ -138,10 +147,16 @@ function main()
 
     end
 
-    for case in cases[case_mask]
+    for case in cases[case_mask] 
+        for (enc_name, enc) in zip(encoder_names[encoder_mask], encoders[encoder_mask])
 
         println(" ")
-        @info "running case $case"
+        @info """
+
+        -------------
+        running case $case with encoder $enc_name
+        -------------
+        """
 
         # common Gaussian feature setup
         pred_type = YType()
@@ -149,8 +164,8 @@ function main()
         # common random feature setup
         n_features = 300
         optimizer_options =
-            Dict("n_iteration" => 10, "n_features_opt" => 60, "n_ensemble" => 30, "cov_sample_multiplier" => 5.0, "scheduler" => DataMisfitController(terminate_at=10))
-        nugget = 1.0
+            Dict("n_iteration" => 10, "n_features_opt" => 60, "n_ensemble" => 30, "cov_sample_multiplier" => 5.0, "verbose" => true)
+        nugget = 1e-7
 
         # data processing schedule
    
@@ -167,7 +182,7 @@ function main()
             mlt = ScalarRandomFeatureInterface(
                 n_features,
                 p,
-                kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()),
+                kernel_structure = SeparableKernel(LowRankFactor(1, nugget), OneDimFactor()),
                 optimizer_options = optimizer_options,
             )
         elseif case == "rf-sep-diag"
@@ -200,7 +215,7 @@ function main()
         emulator = Emulator(
             mlt,
             iopairs,
-            user_encoder_schedule = encoder_schedule,
+            user_encoder_schedule = enc,
             input_structure_matrix = prior_cov,
             output_structure_matrix = Σ,
         )
@@ -253,7 +268,7 @@ function main()
 
                 plot(p5, p6, layout = (1, 2), legend = false)
 
-                savefig(joinpath(output_directory, case * "_y" * string(y_i) * "_predictions.png"))
+                savefig(joinpath(output_directory, case *"_"* enc_name * "_y" * string(y_i) * "_predictions.png"))
             end
         end
 
@@ -293,9 +308,10 @@ function main()
                     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/src/Emulator.jl b/src/Emulator.jl
index e83f26221..f41a948da 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -122,17 +122,7 @@ function Emulator(
     else 
         (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) = (input_output_pairs, input_structure_mat,  output_structure_mat)
         encoder_schedule = []
-    end
-
-    @info input_structure_mat
-    @info output_structure_mat
-    @info encoded_input_structure_mat
-    @info encoded_output_structure_mat
-    @info cov(get_inputs(encoded_io_pairs), dims=2)
-    @info cov(get_outputs(encoded_io_pairs), dims=2)
-    @info get_inputs(encoded_io_pairs) * get_outputs(encoded_io_pairs)'
-    @info get_outputs(encoded_io_pairs)* get_inputs(encoded_io_pairs)'
-    
+    end    
     # build the machine learning tool in the encoded space
     build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
     return Emulator{FT, typeof(encoder_schedule)}(
@@ -184,13 +174,11 @@ function predict(
     # encode the new input data
     encoder_schedule = get_encoder_schedule(emulator)
     encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_inputs), "in")
-    @info cov(get_data(encoded_inputs),dims=2)
     # 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), get_data(encoded_inputs), mlt_kwargs...)
-    @info cov(encoded_outputs,dims=2)
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
     # return decoded or encoded?
     if transform_to_real

From 2f9fd402f76b7977e9f1cb7d5f70d4125240affa Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 15:30:59 -0700
Subject: [PATCH 10/75] n=200

---
 examples/Emulator/Regression_2d_2d/compare_regression.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index 084a11b97..c0e03e940 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -69,7 +69,7 @@ function main()
     
     
     #problem
-    n = 100  # number of training points
+    n = 200  # number of training points
     p = 2   # input dim 
     d = 2   # output dim
     prior_cov = (pi^2)*I(p)

From 7271dbef7442e056e4a685a6092d3c81e29a96b4 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 27 Jun 2025 15:36:48 -0700
Subject: [PATCH 11/75] smaller noise and new schedule train

---
 .../Emulator/Regression_2d_2d/compare_regression.jl  | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index c0e03e940..50d0e6a3d 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -54,17 +54,23 @@ function main()
         "struct-mat-proc",
         "combined-proc",
         "cca-proc",
+        "whiten-then-cca",
     ]
     encoders = [
         [],
         (decorrelate_sample_cov(), "in_and_out"),
         (decorrelate_structure_mat(), "in_and_out"),
         (decorrelate(), "in_and_out"),
-        (canonical_correlation(), "in_and_out")           
+        (canonical_correlation(), "in_and_out"),
+        [
+            (decorrelate_sample_cov(), "in"),
+            (decorrelate_structure_mat(), "out"),
+            (canonical_correlation(), "in_and_out"),
+        ],
     ]
 
     # USER CHOICES 
-    encoder_mask = [1:5...]
+    encoder_mask = [1:6...]
     case_mask = [1] # (KEEP set to 1:length(cases) when pushing for buildkite)
     
     
@@ -87,7 +93,7 @@ function main()
 
     # Add noise η
     μ = zeros(d)
-    Σ = [[0.4, -0.3] [-0.3, 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

From 116b142de72574a126d847003c097fe280d24b69 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 12:13:33 -0700
Subject: [PATCH 12/75] more flexible regularizer (now it is not always going
 to be "I"

---
 src/ScalarRandomFeature.jl | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 146e5f6ec..c636e0fe2 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)
@@ -315,7 +327,6 @@ function build_models!(
     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)
 
@@ -352,7 +363,7 @@ function build_models!(
         n_cross_val_sets = 1 # now just pretend there is one partition for looping purposes
         n_train = n_data
         n_test = n_data
-    else
+   else
         train_fraction = optimizer_options["train_fraction"]
         n_train = Int(floor(train_fraction * n_data))
         n_test = n_data - n_train
@@ -375,9 +386,6 @@ function build_models!(
 
 
 
-    #regularization = I = 1.0 in scalar case
-    regularization = I
-
     @info (
         "hyperparameter learning for $n_rfms models using $n_train training points, $n_test validation points and $n_features_opt features"
     )
@@ -385,6 +393,9 @@ function build_models!(
     diagnostics = zeros(n_iteration, n_rfms)
     for i in 1:n_rfms
 
+        #regularization = I #
+        regularization = output_structure_matrix[i,i]*I
+
         io_pairs_opt = PairedDataContainer(input_values, reshape(output_values[i, :], 1, size(output_values, 2)))
 
         multithread = optimizer_options["multithread"]
@@ -548,6 +559,7 @@ function build_models!(
 
         push!(rfms, rfm_i)
         push!(fitted_features, fitted_features_i)
+        push!(get_regularization(srfi), regularization)
 
     end
     push!(optimizer, diagnostics)
@@ -592,9 +604,13 @@ function predict(
             tullio_threading = tullio_threading,
         )
     end
-
+    
     # add the noise contribution from the regularization
-    σ2[:, :] = σ2[:, :] .+ 1.0
-
+    reg = get_regularization(srfi)[1]
+    reg_diag  = isa(reg, UniformScaling) ? reg.λ*ones(M) : diag(reg)
+    for i = 1:M
+        σ2[i, :] .+= reg_diag[i]
+    end
+    
     return μ, σ2
 end

From 4e768e7b235b1804df2344355108fb54850ae5bb Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 13:29:48 -0700
Subject: [PATCH 13/75] rm unnecessary plots from compare_regression

---
 .../Regression_2d_2d/compare_regression.jl    | 29 +++++++------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index 50d0e6a3d..cecb7085a 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -43,7 +43,6 @@ function main()
         "gp-skljl",
         "gp-gpjl", # Very slow prediction...
         "rf-sep-scalar",
-        "rf-sep-diag",
         "rf-sep-sep",
         "rf-nonsep",
     ]
@@ -54,7 +53,7 @@ function main()
         "struct-mat-proc",
         "combined-proc",
         "cca-proc",
-        "whiten-then-cca",
+        "sample-struct-proc",
     ]
     encoders = [
         [],
@@ -65,13 +64,12 @@ function main()
         [
             (decorrelate_sample_cov(), "in"),
             (decorrelate_structure_mat(), "out"),
-            (canonical_correlation(), "in_and_out"),
         ],
     ]
 
     # USER CHOICES 
-    encoder_mask = [1:6...]
-    case_mask = [1] # (KEEP set to 1:length(cases) when pushing for buildkite)
+    encoder_mask = [3,4,6] # best performing
+    case_mask = [1,3,4,5] # (KEEP set to 1:length(cases) when pushing for buildkite)
     
     
     #problem
@@ -168,10 +166,10 @@ function main()
         pred_type = YType()
 
         # common random feature setup
-        n_features = 300
+        n_features = 200
         optimizer_options =
-            Dict("n_iteration" => 10, "n_features_opt" => 60, "n_ensemble" => 30, "cov_sample_multiplier" => 5.0, "verbose" => true)
-        nugget = 1e-7
+            Dict("n_iteration" => 5, "n_features_opt" => 200, "n_ensemble" => 80, "cov_sample_multiplier" => 5.0)
+        nugget = 1e-8
 
         # data processing schedule
    
@@ -188,17 +186,9 @@ function main()
             mlt = ScalarRandomFeatureInterface(
                 n_features,
                 p,
-                kernel_structure = SeparableKernel(LowRankFactor(1, nugget), OneDimFactor()),
+                kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()),
                 optimizer_options = optimizer_options,
             )
-        elseif case == "rf-sep-diag"
-            mlt = VectorRandomFeatureInterface(
-                n_features,
-                p,
-                d,
-                kernel_structure = SeparableKernel(LowRankFactor(2, nugget), DiagonalFactor()),
-                optimizer_options = deepcopy(optimizer_options),
-            )
          elseif case == "rf-sep-sep"
             mlt = VectorRandomFeatureInterface(
                 n_features,
@@ -221,7 +211,7 @@ function main()
         emulator = Emulator(
             mlt,
             iopairs,
-            user_encoder_schedule = enc,
+            encoder_schedule = enc,
             input_structure_matrix = prior_cov,
             output_structure_matrix = Σ,
         )
@@ -288,7 +278,7 @@ function main()
         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
 
@@ -317,6 +307,7 @@ function main()
                 savefig(joinpath(output_directory, case*"_"* enc_name * "_y" * string(y_i) * "_difference_truth_prediction.png"))
             end
         end
+            =#
         end
     end
 end

From e82bf367e10937e023a2cfa63f659fa4e754a09b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 14:14:42 -0700
Subject: [PATCH 14/75] new message

---
 src/Utilities.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 1bbba08b2..75b649645 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -443,10 +443,13 @@ function initialize_processor!(
             @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

From 8215bd844e023bcd19b4193d13b5f6634d343512 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 14:15:15 -0700
Subject: [PATCH 15/75] converted Ishigami

---
 examples/Emulator/Ishigami/emulate.jl | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl
index 45e04acd1..c3034d911 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,18 @@ function main()
         output[i] = y[ind[i]] + noise[i]
     end
     iopairs = PairedDataContainer(input, output)
-
+    encoder_schedule = (decorrelate_structure_mat(), "in_and_out") # a copy taken for each repeat
+    
     cases = ["Prior", "GP", "RF-scalar"]
-    case = cases[3]
-    decorrelate = true
-    nugget = Float64(1e-12)
+    case = cases[2]
+    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
@@ -102,7 +103,7 @@ function main()
     result_preds = []
     opt_diagnostics = []
     for rep_idx in 1:n_repeats
-
+        
         # Build ML tools
         if case == "GP"
             gppackage = Emulators.SKLJL()
@@ -123,7 +124,7 @@ function main()
         end
 
         # Emulate
-        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate)
+        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, copy(encoder_schedule))
         optimize_hyperparameters!(emulator)
 
         # get EKP errors - just stored in "optimizer" box for now

From dfbaa15faa5061cf0ca8a835161c096f0712937e Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 14:16:41 -0700
Subject: [PATCH 16/75] user_... to ...

---
 src/Emulator.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index f41a948da..c89aea8ae 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -84,7 +84,7 @@ Keyword Arguments
 function Emulator(
     machine_learning_tool::MachineLearningTool,
     input_output_pairs::PairedDataContainer{FT};
-    user_encoder_schedule = nothing,
+    encoder_schedule = nothing,
     input_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing,
     output_structure_matrix::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing,
     mlt_kwargs...,
@@ -110,26 +110,26 @@ function Emulator(
     end
 
     # [1.] Initializes and performs data encoding schedule
-    if !isnothing(user_encoder_schedule)
-        encoder_schedule = create_encoder_schedule(user_encoder_schedule)
+    if !isnothing(encoder_schedule)
+        enc_schedule = create_encoder_schedule(encoder_schedule)
         (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
             encode_with_schedule(
-                encoder_schedule,
+                enc_schedule,
                 input_output_pairs,
                 input_structure_mat,
                 output_structure_mat,
             )
     else 
         (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) = (input_output_pairs, input_structure_mat,  output_structure_mat)
-        encoder_schedule = []
+        enc_schedule = []
     end    
     # build the machine learning tool in the encoded space
     build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
-    return Emulator{FT, typeof(encoder_schedule)}(
+    return Emulator{FT, typeof(enc_schedule)}(
         machine_learning_tool,
         input_output_pairs,
         encoded_io_pairs,
-        encoder_schedule,
+        enc_schedule,
     )
 end
 

From 0972ad212d45f69a38c42ef17e56f4c2e3132a87 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 14:20:10 -0700
Subject: [PATCH 17/75] emulate

---
 examples/Emulator/Ishigami/emulate.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl
index c3034d911..92b6be8f0 100644
--- a/examples/Emulator/Ishigami/emulate.jl
+++ b/examples/Emulator/Ishigami/emulate.jl
@@ -124,7 +124,7 @@ function main()
         end
 
         # Emulate
-        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, copy(encoder_schedule))
+        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = copy(encoder_schedule))
         optimize_hyperparameters!(emulator)
 
         # get EKP errors - just stored in "optimizer" box for now

From fcca4b91f1afa2900f6a69240c7f37767d303d92 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 14:27:09 -0700
Subject: [PATCH 18/75] copy -> deepcopy

---
 examples/Emulator/Ishigami/emulate.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl
index 92b6be8f0..165ea4103 100644
--- a/examples/Emulator/Ishigami/emulate.jl
+++ b/examples/Emulator/Ishigami/emulate.jl
@@ -124,7 +124,7 @@ function main()
         end
 
         # Emulate
-        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = copy(encoder_schedule))
+        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
         optimize_hyperparameters!(emulator)
 
         # get EKP errors - just stored in "optimizer" box for now

From 506dac32004f597f6ead1ef7dcb8d5d218316c10 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 15:01:19 -0700
Subject: [PATCH 19/75] G-function update

---
 .../Emulator/G-function/emulate-test-n-features.jl |  8 ++++----
 examples/Emulator/G-function/emulate.jl            | 14 ++++++--------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/examples/Emulator/G-function/emulate-test-n-features.jl b/examples/Emulator/G-function/emulate-test-n-features.jl
index 042694fbe..6a214f203 100644
--- a/examples/Emulator/G-function/emulate-test-n-features.jl
+++ b/examples/Emulator/G-function/emulate-test-n-features.jl
@@ -82,7 +82,8 @@ function main()
         output[i] = y[ind[i]] + noise[i]
     end
     iopairs = PairedDataContainer(input, output)
-
+    
+    
     # analytic sobol indices taken from
     # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8989694/pdf/main.pdf
     a = [(i - 1.0) / 2.0 for i in 1:n_dimensions]  # a_i < a_j => a_i more sensitive
@@ -95,7 +96,7 @@ function main()
 
     cases = ["Prior", "GP", "RF-scalar"]
     case = cases[2]
-    decorrelate = true
+    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
     nugget = Float64(1e-12)
 
     n_features_vec = [25, 50, 100, 200, 400, 800] # < data points ideally as output dim = 1
@@ -116,7 +117,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 +152,7 @@ function main()
 
             # Emulate
             ttt[f_idx, rep_idx] = @elapsed begin
-                emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate)
+                emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * 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..dcc541c53 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 = 10# repeat exp with same data.
+    n_dimensions = 6
     # To create the sampling
     n_data_gen = 800
 
@@ -90,21 +91,18 @@ 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
+    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
     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,7 @@ function main()
 
         # Emulate
         times[rep_idx] = @elapsed begin
-            emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, decorrelate = decorrelate)
+            emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
             optimize_hyperparameters!(emulator)
         end
 

From 77b1f2006bea2a106315e997b77b5ba61011ce68 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 15:04:20 -0700
Subject: [PATCH 20/75] utils

---
 examples/Emulator/G-function/emulate-test-n-features.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/Emulator/G-function/emulate-test-n-features.jl b/examples/Emulator/G-function/emulate-test-n-features.jl
index 6a214f203..22c558bac 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 CalibrateEmuatleSample.Utilities
 
 using CairoMakie, ColorSchemes #for plots
 seed = 2589436

From 9b24e1fdb2c3b31b6dabb4c7be175b39ae92a877 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 17:05:15 -0700
Subject: [PATCH 21/75] L63

---
 examples/Emulator/L63/emulate.jl                | 16 +++++++---------
 examples/Emulator/L63/emulate_diff-rank-test.jl |  7 ++++---
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/examples/Emulator/L63/emulate.jl b/examples/Emulator/L63/emulate.jl
index 64ee4b6c3..5ca1f9a7b 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,8 +107,8 @@ function main()
         "RF-svd-sep",
     ]
 
-    case = cases[7]
-
+    case = cases[3] # 7
+    
     nugget = Float64(1e-8)
     u_test = []
     u_hist = []
@@ -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..05bbe1952 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,12 @@ 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
 

From 8584b19dfbd45ab0d1d52a905c5dce7e0cd21744 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 17:07:15 -0700
Subject: [PATCH 22/75] L63

---
 .../Emulator/G-function/emulate-test-n-features.jl     |  8 +++-----
 examples/Emulator/G-function/emulate.jl                | 10 +++++-----
 examples/Emulator/Ishigami/emulate.jl                  |  5 ++---
 examples/Emulator/Regression_2d_2d/Project.toml        |  4 ++--
 4 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/examples/Emulator/G-function/emulate-test-n-features.jl b/examples/Emulator/G-function/emulate-test-n-features.jl
index 22c558bac..0b339ac40 100644
--- a/examples/Emulator/G-function/emulate-test-n-features.jl
+++ b/examples/Emulator/G-function/emulate-test-n-features.jl
@@ -12,7 +12,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses
 using CalibrateEmulateSample.Emulators
 using CalibrateEmulateSample.DataContainers
 using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers
-using CalibrateEmuatleSample.Utilities
+using CalibrateEmulateSample.Utilities
 
 using CairoMakie, ColorSchemes #for plots
 seed = 2589436
@@ -84,7 +84,6 @@ function main()
     end
     iopairs = PairedDataContainer(input, output)
     
-    
     # analytic sobol indices taken from
     # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8989694/pdf/main.pdf
     a = [(i - 1.0) / 2.0 for i in 1:n_dimensions]  # a_i < a_j => a_i more sensitive
@@ -94,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]
-    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+    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
@@ -153,7 +151,7 @@ function main()
 
             # Emulate
             ttt[f_idx, rep_idx] = @elapsed begin
-                emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
+                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 dcc541c53..11f86be28 100644
--- a/examples/Emulator/G-function/emulate.jl
+++ b/examples/Emulator/G-function/emulate.jl
@@ -40,8 +40,8 @@ function main()
 
     rng = MersenneTwister(seed)
 
-    n_repeats = 10# repeat exp with same data.
-    n_dimensions = 6
+    n_repeats = 5# repeat exp with same data.
+    n_dimensions = 10
     # To create the sampling
     n_data_gen = 800
 
@@ -81,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
@@ -93,7 +94,6 @@ function main()
 
     cases = ["Prior", "GP", "RF-scalar"]
     case = cases[3]
-    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
     nugget = Float64(1e-12)
 
     overrides = Dict(
@@ -142,7 +142,7 @@ function main()
 
         # Emulate
         times[rep_idx] = @elapsed begin
-            emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
+            emulator = Emulator(mlt, iopairs; output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
             optimize_hyperparameters!(emulator)
         end
 
@@ -305,7 +305,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 165ea4103..d3b980e7a 100644
--- a/examples/Emulator/Ishigami/emulate.jl
+++ b/examples/Emulator/Ishigami/emulate.jl
@@ -81,10 +81,9 @@ function main()
         output[i] = y[ind[i]] + noise[i]
     end
     iopairs = PairedDataContainer(input, output)
-    encoder_schedule = (decorrelate_structure_mat(), "in_and_out") # a copy taken for each repeat
-    
     cases = ["Prior", "GP", "RF-scalar"]
     case = cases[2]
+    encoder_schedule = (decorrelate_structure_mat(), "out")
     nugget = Float64(1e-4)
     overrides = Dict(
         "scheduler" => DataMisfitController(terminate_at = 1e4),
@@ -124,7 +123,7 @@ function main()
         end
 
         # Emulate
-        emulator = Emulator(mlt, iopairs; obs_noise_cov = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
+        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/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"

From b07b02e51b666ae338aaea46661e89477ddb1e5c Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 17:09:46 -0700
Subject: [PATCH 23/75] sensible default for EKI-produced data

---
 .../Regression_2d_2d/compare_regression.jl    |  9 ++--
 src/Emulator.jl                               | 45 +++++++++++++------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index cecb7085a..5a2bd426d 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -49,22 +49,19 @@ function main()
 
     encoder_names = [
         "no-proc",
+        "sample-struct-proc",
         "sample-proc",
         "struct-mat-proc",
         "combined-proc",
         "cca-proc",
-        "sample-struct-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"),
-        [
-            (decorrelate_sample_cov(), "in"),
-            (decorrelate_structure_mat(), "out"),
-        ],
     ]
 
     # USER CHOICES 
diff --git a/src/Emulator.jl b/src/Emulator.jl
index c89aea8ae..e549e6bad 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -87,12 +87,20 @@ function Emulator(
     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}
 
     # 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=`")
+    end
+
     input_structure_mat = if isnothing(input_structure_matrix)
         Diagonal(FT.(ones(input_dim)))
     elseif isa(input_structure_matrix, UniformScaling)
@@ -110,19 +118,30 @@ function Emulator(
     end
 
     # [1.] Initializes and performs data encoding schedule
-    if !isnothing(encoder_schedule)
-        enc_schedule = create_encoder_schedule(encoder_schedule)
-        (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
-            encode_with_schedule(
-                enc_schedule,
-                input_output_pairs,
-                input_structure_mat,
-                output_structure_mat,
-            )
-    else 
-        (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) = (input_output_pairs, input_structure_mat,  output_structure_mat)
-        enc_schedule = []
-    end    
+    # 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
+            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_mat, encoded_output_structure_mat) =
+        encode_with_schedule(
+            enc_schedule,
+            input_output_pairs,
+            input_structure_mat,
+            output_structure_mat,
+        )
+    
     # build the machine learning tool in the encoded space
     build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
     return Emulator{FT, typeof(enc_schedule)}(

From 8549029bf7ec9b8391d6f1bb230745d62e303f52 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 18:16:57 -0700
Subject: [PATCH 24/75] new encode/decode works with sampling too

---
 src/Emulator.jl              | 32 ++++++++++++++++---
 src/MarkovChainMonteCarlo.jl | 60 ++++--------------------------------
 2 files changed, 34 insertions(+), 58 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index e549e6bad..50b241a00 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -1,7 +1,9 @@
 module Emulators
 
 using ..DataContainers
-
+import ..Utilities.encode_with_schedule
+import ..Utilities.decode_with_schedule
+ 
 using DocStringExtensions
 using Statistics
 using Distributions
@@ -14,7 +16,7 @@ export Emulator
 export calculate_normalization
 export build_models!
 export optimize_hyperparameters!
-export predict
+export predict, encode_with_schedule, decode_with_schedule
 
 """
 $(DocStringExtensions.TYPEDEF)
@@ -161,6 +163,22 @@ function optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) w
     optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...)
 end
 
+function encode_with_schedule(emulator::Emulator, data::AM, in_or_out::AS) where {AS <: AbstractString, AM <: AbstractMatrix}
+    return get_data(encode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out))
+end
+    
+function encode_with_schedule(emulator::Emulator, data::DC, in_or_out::AS) where {AS <: AbstractString, DC <: DataContainer}
+    return encode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+end
+
+function decode_with_schedule(emulator::Emulator, data::AM, in_or_out::AS) where {AS <: AbstractString, AM <: AbstractMatrix}
+    return get_data(decode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out))
+end
+    
+function decode_with_schedule(emulator::Emulator, data::DC, in_or_out::AS) where {AS <: AbstractString, DC <: DataContainer}
+    return decode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+end
+
 
 """
 $(DocStringExtensions.TYPEDSIGNATURES)
@@ -219,8 +237,14 @@ function predict(
         
         if encoded_output_dim > 1
             encoded_covariances_mat = zeros(output_dim, output_dim, size(encoded_uncertainties)[end])
-            for (i,col) in enumerate(eachcol(encoded_uncertainties))
-                encoded_covariances_mat[:,:,i] = Diagonal(col)
+            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] = Diagonal(mat)
+                end
             end
             encoded_covariances = eachslice(encoded_covariances_mat,dims=3)
         else
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index 46c30051f..9a8452de4 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,9 +250,8 @@ 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)
 
         return 1.0 / length(obs_vec) * sum([logpdf(MvNormal(obs, g_cov[1] * I), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
@@ -594,10 +547,9 @@ function MCMCWrapper(
     kwargs...,
 ) where {AV <: AbstractVector, AMorAV <: Union{AbstractVector, AbstractMatrix}}
 
-    # decorrelate observations into a vector
-    decorrelated_obs = to_decorrelated(observation, em)
-
-    log_posterior_map = EmulatorPosteriorModel(prior, em, decorrelated_obs)
+    # encoding works on columns but mcmc wants vec-of-vec
+    encoded_obs = [vec(encode_with_schedule(em, reshape(obs, :, 1), "out")) for obs in eachcol(observation)]    
+    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 +569,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, observation, encoded_obs, log_posterior_map, mh_proposal_sampler, sample_kwargs)
 end
 
 """

From f24dc1b9d59bf6491e65dbbfc424fc17ca5025d5 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 1 Jul 2025 18:17:26 -0700
Subject: [PATCH 25/75] sinusoid

---
 examples/Sinusoid/emulate.jl | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/examples/Sinusoid/emulate.jl b/examples/Sinusoid/emulate.jl
index dcf0027c0..02ddb6663 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)
 
 

From 26dc1ae5d35857c58d95a4f8fc430fd6587982be Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 15:57:54 -0700
Subject: [PATCH 26/75] new encoder methods and exports

---
 src/Emulator.jl | 50 ++++++++++++++++++++++++++++---------------------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 50b241a00..46015c2f5 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -16,8 +16,8 @@ export Emulator
 export calculate_normalization
 export build_models!
 export optimize_hyperparameters!
-export predict, encode_with_schedule, decode_with_schedule
-
+export predict, encode_data, decode_data, encode_structure_mat, decode_structure_mat
+export get_machine_learning_tool, get_io_pairs, get_encoded_io_pairs, get_encoder_schedule
 """
 $(DocStringExtensions.TYPEDEF)
 
@@ -163,22 +163,31 @@ function optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) w
     optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...)
 end
 
-function encode_with_schedule(emulator::Emulator, data::AM, in_or_out::AS) where {AS <: AbstractString, AM <: AbstractMatrix}
-    return get_data(encode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out))
-end
     
-function encode_with_schedule(emulator::Emulator, data::DC, in_or_out::AS) where {AS <: AbstractString, DC <: DataContainer}
-    return encode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+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 encode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+    end
 end
 
-function decode_with_schedule(emulator::Emulator, data::AM, in_or_out::AS) where {AS <: AbstractString, AM <: AbstractMatrix}
-    return get_data(decode_with_schedule(get_encoder_schedule(emulator), DataContainer(data), in_or_out))
+function encode_structure_mat(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+    return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
+
     
-function decode_with_schedule(emulator::Emulator, data::DC, in_or_out::AS) where {AS <: AbstractString, DC <: DataContainer}
-    return decode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+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 decode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
+    end
 end
 
+function decode_structure_mat(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+    return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
+end
 
 """
 $(DocStringExtensions.TYPEDSIGNATURES)
@@ -196,7 +205,7 @@ function predict(
     # 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_io_pairs(emulator), 1)
+    encoded_input_dim, encoded_output_dim = size(get_encoded_io_pairs(emulator), 1)
     
     N_samples = size(new_inputs, 2)
     
@@ -209,41 +218,40 @@ function predict(
     end
 
     # encode the new input data
-    encoder_schedule = get_encoder_schedule(emulator)
-    encoded_inputs = encode_with_schedule(encoder_schedule, DataContainer(new_inputs), "in")
+    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), get_data(encoded_inputs), mlt_kwargs...)
+    encoded_outputs, encoded_uncertainties = predict(get_machine_learning_tool(emulator), encoded_inputs, mlt_kwargs...)
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
     # return decoded or encoded?
     if transform_to_real
-        decoded_outputs = decode_with_schedule(encoder_schedule, DataContainer(encoded_outputs), "out")
+        decoded_outputs = decode_data(emulator, encoded_outputs, "out")
         
         decoded_covariances = zeros(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_with_schedule(encoder_schedule, Diagonal(col), "out")
+               decoded_covariances[:,:,i] .= decode_structure_mat(emulator, Diagonal(col), "out")
             end
         else # == "cov"
             for (i,mat) in enumerate(eachslice(encoded_uncertainties, dims=3))
-               decoded_covariances[:,:,i] .= decode_with_schedule(encoder_schedule, mat, "out")
+               decoded_covariances[:,:,i] .= decode_structure_mat(emulator, mat, "out")
             end
         end
             
-        return get_data(decoded_outputs), eachslice(decoded_covariances,dims=3)
+        return decoded_outputs, eachslice(decoded_covariances,dims=3)
     else
         
         if encoded_output_dim > 1
-            encoded_covariances_mat = zeros(output_dim, output_dim, size(encoded_uncertainties)[end])
+            encoded_covariances_mat = zeros(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] = Diagonal(mat)
+                    encoded_covariances_mat[:,:,i] = mat
                 end
             end
             encoded_covariances = eachslice(encoded_covariances_mat,dims=3)

From dea9670ee691c64dac544c9e6561a37d82f90a29 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 15:58:13 -0700
Subject: [PATCH 27/75] bugfix truncation of variance

---
 src/Utilities.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 75b649645..c3f53703a 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -438,7 +438,7 @@ function initialize_processor!(
         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
+            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
@@ -636,7 +636,7 @@ function initialize_processor!(
         # 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
+            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

From ec0aaeb9e79962df722b614810e0d1db4d1d7f8a Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 15:58:45 -0700
Subject: [PATCH 28/75] use new method to encode

---
 src/MarkovChainMonteCarlo.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index 9a8452de4..4ea545e59 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -548,7 +548,8 @@ function MCMCWrapper(
 ) where {AV <: AbstractVector, AMorAV <: Union{AbstractVector, AbstractMatrix}}
 
     # encoding works on columns but mcmc wants vec-of-vec
-    encoded_obs = [vec(encode_with_schedule(em, reshape(obs, :, 1), "out")) for obs in eachcol(observation)]    
+    encoded_obs = [vec(encode_data(em, reshape(obs, :, 1), "out")) for obs in eachcol(observation)]
+
     log_posterior_map = EmulatorPosteriorModel(prior, em, encoded_obs)
     mh_proposal_sampler = MetropolisHastingsSampler(mcmc_alg, prior)
 
@@ -729,6 +730,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 = [

From 3534ebb4e14f0a1f6c0b02d39259a732645d662b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 16:00:21 -0700
Subject: [PATCH 29/75] fix regularization for non-I matrices

---
 src/ScalarRandomFeature.jl | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index c636e0fe2..547857b2b 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -384,6 +384,8 @@ function build_models!(
         end
     end
 
+    #regularization = I #
+    regularization = Diagonal(output_structure_matrix) # creates diag matrix of diagonal
 
 
     @info (
@@ -392,10 +394,8 @@ function build_models!(
     n_iteration = optimizer_options["n_iteration"]
     diagnostics = zeros(n_iteration, n_rfms)
     for i in 1:n_rfms
-
-        #regularization = I #
-        regularization = output_structure_matrix[i,i]*I
-
+        regularization_i = regularization[i,i]*I
+        
         io_pairs_opt = PairedDataContainer(input_values, reshape(output_values[i, :], 1, size(output_values, 2)))
 
         multithread = optimizer_options["multithread"]
@@ -435,7 +435,7 @@ function build_models!(
                 srfi,
                 rng,
                 μ_hp,
-                regularization,
+                regularization_i,
                 n_features_opt,
                 train_idx[cv_idx],
                 test_idx[cv_idx],
@@ -447,7 +447,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(Γ)
@@ -493,7 +493,7 @@ function build_models!(
                     srfi,
                     rng,
                     lvec,
-                    regularization,
+                    regularization_i,
                     n_features_opt,
                     train_idx[cv_idx],
                     test_idx[cv_idx],
@@ -549,7 +549,7 @@ function build_models!(
             srfi,
             rng,
             hp_optimal,
-            regularization,
+            regularization_i,
             n_features,
             batch_sizes,
             input_dim,
@@ -559,9 +559,9 @@ function build_models!(
 
         push!(rfms, rfm_i)
         push!(fitted_features, fitted_features_i)
-        push!(get_regularization(srfi), regularization)
 
     end
+    push!(get_regularization(srfi), regularization)
     push!(optimizer, diagnostics)
 
 end
@@ -578,7 +578,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,
@@ -605,9 +605,10 @@ function predict(
         )
     end
     
-    # add the noise contribution from the regularization
+    # 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 = 1:M
         σ2[i, :] .+= reg_diag[i]
     end

From 77628e53232f6730a39133dc04b31e246121434e Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 16:01:45 -0700
Subject: [PATCH 30/75] update Lorenz and Lorenz spatial dep examples

---
 examples/Lorenz/calibrate_spatial_dep.jl      | 18 ++--
 examples/Lorenz/emulate_sample.jl             | 82 ++++++------------
 examples/Lorenz/emulate_sample_spatial_dep.jl | 84 +++++++------------
 examples/Lorenz/plot_spatial_dep.jl           | 10 +--
 4 files changed, 72 insertions(+), 122 deletions(-)

diff --git a/examples/Lorenz/calibrate_spatial_dep.jl b/examples/Lorenz/calibrate_spatial_dep.jl
index d39ab9f6a..ff7e988e5 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_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..4e28d5add 100644
--- a/examples/Lorenz/emulate_sample.jl
+++ b/examples/Lorenz/emulate_sample.jl
@@ -5,7 +5,6 @@ include(joinpath(@__DIR__, "..", "ci", "linkfig.jl"))
 using Distributions  # probability distributions and associated functions
 using LinearAlgebra
 ENV["GKSwstype"] = "100"
-using StatsPlots
 using Plots
 using Random
 using JLD2
@@ -18,36 +17,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 +35,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 +85,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 +94,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 +104,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 +117,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 +132,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,22 +165,19 @@ 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
-
-
+        # 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,
             obs_noise_cov = Γy,
-            normalize_inputs = normalized,
-            standardize_outputs = standardize,
-            standardize_outputs_factors = norm_factor,
-            retained_svd_frac = retained_svd_frac,
-            decorrelate = decorrelate,
         )
         optimize_hyperparameters!(emulator)
 
@@ -266,13 +232,19 @@ 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..de99067d0 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_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,47 @@ 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)
@@ -206,7 +180,7 @@ function main()
         # First let's run a short chain to determine a good step size
         mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0)
         new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0)
-
+        
         # Now begin the actual MCMC
         println("Begin MCMC - with step size ", new_step)
         chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000)
@@ -256,6 +230,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..12136c740 100644
--- a/examples/Lorenz/plot_spatial_dep.jl
+++ b/examples/Lorenz/plot_spatial_dep.jl
@@ -4,9 +4,6 @@
 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]
 
@@ -35,10 +32,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 +63,7 @@ p1 = plot(
 )
 
 p2 = plot(
-    1:length(y),
+    1:ny,
     y,
     ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))),
     label = "data",

From 6b78e0b86c80ff54bfda23c601b25e94d32585eb Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 2 Jul 2025 17:57:54 -0700
Subject: [PATCH 31/75] updated darcy

---
 examples/Darcy/calibrate.jl      |  2 +-
 examples/Darcy/emulate_sample.jl | 17 +++++++----------
 2 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/examples/Darcy/calibrate.jl b/examples/Darcy/calibrate.jl
index f79d5dda7..62ceb1df9 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)),
diff --git a/examples/Darcy/emulate_sample.jl b/examples/Darcy/emulate_sample.jl
index 5ba96b8fe..b4df4fb9a 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)])
 

From c8661dc91b206575dfc4862a9fcd56f0603d3282 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 09:51:44 -0700
Subject: [PATCH 32/75] updated cloudy example

---
 examples/Cloudy/Cloudy_calibrate.jl      | 12 ++++--
 examples/Cloudy/Cloudy_emulate_sample.jl | 49 +++++++++++-------------
 examples/Cloudy/Project.toml             |  2 +-
 3 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/examples/Cloudy/Cloudy_calibrate.jl b/examples/Cloudy/Cloudy_calibrate.jl
index 037de2d65..31e430e94 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 = [0]
 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[1] = 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] - 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] - 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..c74927fc2 100644
--- a/examples/Cloudy/Cloudy_emulate_sample.jl
+++ b/examples/Cloudy/Cloudy_emulate_sample.jl
@@ -16,9 +16,8 @@ using CairoMakie, PairPlots
 using CalibrateEmulateSample.Emulators
 using CalibrateEmulateSample.MarkovChainMonteCarlo
 using CalibrateEmulateSample.Utilities
-using EnsembleKalmanProcesses
-using EnsembleKalmanProcesses.ParameterDistributions
-using EnsembleKalmanProcesses.DataContainers
+using CalibrateEmulateSample.EnsembleKalmanProcesses
+#using EnsembleKalmanProcesses.ParameterDistributions
 
 function get_standardizing_factors(data::Array{FT, 2}) where {FT}
     # Input: data size: N_data x N_ensembles
@@ -100,7 +99,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 +115,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 +139,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 +152,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 +164,34 @@ 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
+        
+        # Data processing
+        if case == "rf-nonsep"
+            encoder_schedule = []
+        else
+            encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+        end
+        
         # output data with information from Γy, if required
         # Note: The `standardize_outputs_factors` are only used under the
         # condition that `standardize_outputs` is true.
         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 +200,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"

From 7f5bef3689716acad46c0d3350e7a277edde803b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 09:53:11 -0700
Subject: [PATCH 33/75] remove standardization

---
 examples/Cloudy/Cloudy_emulate_sample.jl | 16 +++-------------
 1 file changed, 3 insertions(+), 13 deletions(-)

diff --git a/examples/Cloudy/Cloudy_emulate_sample.jl b/examples/Cloudy/Cloudy_emulate_sample.jl
index c74927fc2..cee2ba915 100644
--- a/examples/Cloudy/Cloudy_emulate_sample.jl
+++ b/examples/Cloudy/Cloudy_emulate_sample.jl
@@ -17,14 +17,6 @@ using CalibrateEmulateSample.Emulators
 using CalibrateEmulateSample.MarkovChainMonteCarlo
 using CalibrateEmulateSample.Utilities
 using CalibrateEmulateSample.EnsembleKalmanProcesses
-#using EnsembleKalmanProcesses.ParameterDistributions
-
-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
 
 ################################################################################
 #                                                                              #
@@ -173,12 +165,10 @@ function main()
         if case == "rf-nonsep"
             encoder_schedule = []
         else
-            encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+            encoder_schedule = (decorrelate_structure_mat(), "in_and_out") 
         end
-        
-        # output data with information from Γy, if required
-        # Note: The `standardize_outputs_factors` are only used under the
-        # condition that `standardize_outputs` is true.
+
+        # build emulator
         emulator = Emulator(
             mlt,
             input_output_pairs;

From 23d764a88c3586a0bba86a2b89459049e135dd0e Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 10:31:18 -0700
Subject: [PATCH 34/75] updated EDMF-data

---
 examples/EDMF_data/emulator-rank-test.jl |  53 ++++++------
 examples/EDMF_data/rebuild_priors.jl     |  39 +++++++++
 examples/EDMF_data/uq_for_edmf.jl        | 102 +++--------------------
 3 files changed, 77 insertions(+), 117 deletions(-)
 create mode 100644 examples/EDMF_data/rebuild_priors.jl

diff --git a/examples/EDMF_data/emulator-rank-test.jl b/examples/EDMF_data/emulator-rank-test.jl
index f26a07ef8..cfdadd6b2 100644
--- a/examples/EDMF_data/emulator-rank-test.jl
+++ b/examples/EDMF_data/emulator-rank-test.jl
@@ -32,19 +32,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 = 4
     case = cases[case_id]
     rank_test = rank_cases[case_id]
     n_repeats = 1
@@ -157,6 +155,19 @@ function main()
         train_pairs = DataContainers.PairedDataContainer(train_inputs, train_outputs)
     end
 
+    include("rebuild_priors.jl")
+    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 +199,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 +220,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,19 +233,7 @@ function main()
                     optimizer_options = overrides,
                 )
 
-            elseif case ∈ ["RF-vector-svd-nonsep", "RF-prior"]
-                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,
-                )
-            elseif case ∈ ["RF-vector-nosvd-nonsep"]
+            elseif case ∈ ["RF-nonsep", "RF-prior"]
                 kernel_structure = NonseparableKernel(LowRankFactor(rank_val, nugget))
                 n_features = 500
 
@@ -246,19 +245,19 @@ function main()
                     kernel_structure = kernel_structure,
                     optimizer_options = overrides,
                 )
-                decorrelate = false
             end
-
+            
+            # data processing:
+            retain_var = 0.95
+            encoder_schedule = (decorrelate_structure_mat(retain_var=retain_var), "in_and_out")
+            
             # Fit an emulator to the data
-            normalized = true
-
             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,
                 )
 
                 # Optimize the GP hyperparameters for better fit
diff --git a/examples/EDMF_data/rebuild_priors.jl b/examples/EDMF_data/rebuild_priors.jl
new file mode 100644
index 000000000..af23d8d74
--- /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..ae5a3f253 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_σ"]
@@ -235,7 +196,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(kernel_rank, nugget))
         n_features = 500
 
@@ -248,7 +209,7 @@ function main()
             optimizer_options = overrides,
         )
 
-    elseif case ∈ ["RF-vector-svd-nonsep", "RF-prior"]
+    elseif case ∈ ["RF-nonsep", "RF-prior"]
         kernel_structure = NonseparableKernel(LowRankFactor(kernel_rank, nugget))
         n_features = 500
 
@@ -260,64 +221,25 @@ function main()
             kernel_structure = kernel_structure,
             optimizer_options = overrides,
         )
-    elseif case ∈ ["RF-vector-nosvd-nonsep"]
-        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,
-        )
-        decorrelate = false
     end
 
+    # data processing:
+    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+    
     # Fit an emulator to the data
-    normalized = true
-
     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,
     )
 
     # 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)
 

From 21143937be12239c897c8b8758118ddd53ff04ae Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 10:40:43 -0700
Subject: [PATCH 35/75] updated GCM

---
 examples/GCM/emulate_sample_script.jl | 40 +++++++++++++++------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/examples/GCM/emulate_sample_script.jl b/examples/GCM/emulate_sample_script.jl
index 4c0f9ca9c..25115f47a 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,29 @@ 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 +273,6 @@ function main()
 
     ### MCMC
 
-    priorfile = "priors.jld2"
-    prior = load(priorfile)["prior"]
 
     ##
     ###  Sample: Markov Chain Monte Carlo

From b16fb4b6e44bde3cc6e69d2527dc35862572446a Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 11:06:35 -0700
Subject: [PATCH 36/75] typo

---
 examples/EDMF_data/emulator-rank-test.jl | 10 ++++++----
 examples/EDMF_data/uq_for_edmf.jl        |  4 +++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/examples/EDMF_data/emulator-rank-test.jl b/examples/EDMF_data/emulator-rank-test.jl
index cfdadd6b2..92ddd2c3d 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()
 
@@ -42,7 +44,7 @@ function main()
         1:5, # input rank always 5, output rank 1:5
     ]
 
-    case_id = 4
+    case_id = 3
     case = cases[case_id]
     rank_test = rank_cases[case_id]
     n_repeats = 1
@@ -155,7 +157,6 @@ function main()
         train_pairs = DataContainers.PairedDataContainer(train_inputs, train_outputs)
     end
 
-    include("rebuild_priors.jl")
     prior_config = get_prior_config(exp_name)
     means = prior_config["prior_mean"]
     std = prior_config["unconstrained_σ"]
@@ -197,7 +198,7 @@ function main()
                 "scheduler" => DataMisfitController(terminate_at = 1000),
                 "cov_sample_multiplier" => 0.2,
                 "n_iteration" => n_iteration,
-                "n_features_opt" => Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works
+                "n_features_opt" => 50, #Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works
                 "localization" => SEC(0.05),
                 "n_ensemble" => 100,
                 "n_cross_val_sets" => n_cross_val_sets,
@@ -248,7 +249,7 @@ function main()
             end
             
             # data processing:
-            retain_var = 0.95
+            retain_var = 0.9
             encoder_schedule = (decorrelate_structure_mat(retain_var=retain_var), "in_and_out")
             
             # Fit an emulator to the data
@@ -258,6 +259,7 @@ function main()
                     train_pairs;
                     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/uq_for_edmf.jl b/examples/EDMF_data/uq_for_edmf.jl
index ae5a3f253..002ee3d65 100644
--- a/examples/EDMF_data/uq_for_edmf.jl
+++ b/examples/EDMF_data/uq_for_edmf.jl
@@ -224,7 +224,8 @@ function main()
     end
 
     # data processing:
-    encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
+    retain_var = 0.9
+    encoder_schedule = (decorrelate_structure_mat(retain_var=retain_var), "in_and_out")
     
     # Fit an emulator to the data
     emulator = Emulator(
@@ -232,6 +233,7 @@ function main()
         input_output_pairs;
         input_structure_matrix = cov(prior),
         output_structure_matrix = truth_cov,
+        encoder_schedule=encoder_schedule,
     )
 
     # Optimize the GP hyperparameters for better fit

From 9b06db2fb641f856cdbf231c18da341248d94eb7 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 11:10:20 -0700
Subject: [PATCH 37/75] typo

---
 examples/EDMF_data/emulator-rank-test.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/EDMF_data/emulator-rank-test.jl b/examples/EDMF_data/emulator-rank-test.jl
index 92ddd2c3d..e240c8422 100644
--- a/examples/EDMF_data/emulator-rank-test.jl
+++ b/examples/EDMF_data/emulator-rank-test.jl
@@ -198,7 +198,7 @@ function main()
                 "scheduler" => DataMisfitController(terminate_at = 1000),
                 "cov_sample_multiplier" => 0.2,
                 "n_iteration" => n_iteration,
-                "n_features_opt" => 50, #Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works
+                "n_features_opt" => Int(floor((max_feature_size / 5))),# here: /5 with rank <= 3 works
                 "localization" => SEC(0.05),
                 "n_ensemble" => 100,
                 "n_cross_val_sets" => n_cross_val_sets,

From 80b49e5692be2e7d351b4d6a285f931114cbb1b3 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 14:47:42 -0700
Subject: [PATCH 38/75] added more API and ==

---
 src/Emulator.jl  | 16 ++++++++++------
 src/Utilities.jl |  4 ++++
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 46015c2f5..9b028ac9d 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -3,7 +3,11 @@ 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
 using Distributions
@@ -16,7 +20,7 @@ export Emulator
 export calculate_normalization
 export build_models!
 export optimize_hyperparameters!
-export predict, encode_data, decode_data, encode_structure_mat, decode_structure_mat
+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)
@@ -172,7 +176,7 @@ function encode_data(emulator::Emulator, data::MorDC, in_or_out::AS) where {AS <
     end
 end
 
-function encode_structure_mat(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+function encode_structure_matrix(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
     return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
@@ -185,7 +189,7 @@ function decode_data(emulator::Emulator, data::MorDC, in_or_out::AS) where {AS <
     end
 end
 
-function decode_structure_mat(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+function decode_structure_matrix(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
     return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
@@ -232,11 +236,11 @@ function predict(
         decoded_covariances = zeros(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_mat(emulator, Diagonal(col), "out")
+               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_mat(emulator, mat, "out")
+               decoded_covariances[:,:,i] .= decode_structure_matrix(emulator, mat, "out")
             end
         end
             
diff --git a/src/Utilities.jl b/src/Utilities.jl
index c3f53703a..e3bbfc5a4 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -119,6 +119,10 @@ abstract type QuartileScaling <: UnivariateAffineScaling end
 abstract type MinMaxScaling <: UnivariateAffineScaling end
 abstract type ZScoreScaling <: UnivariateAffineScaling end
 
+# 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))
+
 # Processors
 """
 $(TYPEDEF)

From c2893c6373c0b0ddd7950e956cabb65ae220ee15 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 17:22:44 -0700
Subject: [PATCH 39/75] add ! and mor eunit test coverage for emulators

---
 src/Emulator.jl            |   2 +-
 src/Utilities.jl           |   4 +-
 test/Emulator/runtests.jl  | 146 ++++++++++++++++++++++++++++++-------
 test/Utilities/runtests.jl |  31 ++++++--
 4 files changed, 147 insertions(+), 36 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 9b028ac9d..cf607ce32 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -141,7 +141,7 @@ function Emulator(
         
     enc_schedule = create_encoder_schedule(encoder_schedule)
     (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
-        encode_with_schedule(
+        encode_with_schedule!(
             enc_schedule,
             input_output_pairs,
             input_structure_mat,
diff --git a/src/Utilities.jl b/src/Utilities.jl
index e3bbfc5a4..d3f0f6bcf 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -21,7 +21,7 @@ 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 create_encoder_schedule, encode_with_schedule, encode_with_schedule!, decode_with_schedule
 export initialize_processor!,
     initialize_and_encode_data!, encode_data, decode_data, encode_structure_matrix, decode_structure_matrix
 
@@ -902,7 +902,7 @@ $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 encode_with_schedule!(
     encoder_schedule::VV,
     io_pairs::PDC,
     input_structure_mat::USorM1,
diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl
index 370e419bd..e7016713d 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -9,26 +9,115 @@ using LinearAlgebra
 
 using CalibrateEmulateSample.Emulators
 using CalibrateEmulateSample.DataContainers
+using CalibrateEmulateSample.Utilities
 
 #build an unknown type
 struct MLTester <: Emulators.MachineLearningTool end
 
+@testset "Emulators" begin
+    #build some quick data + noise
+    m = 50
+    d = 6
+    p = 10
+    x = rand(p, m) #R^3
+    y = rand(d, m) #R^6
+
+    # "noise"
+    μ = zeros(d)
+    Σ = rand(d, d)
+    Σ = Σ' * Σ
+    noise_samples = rand(MvNormal(μ, Σ), m)
+    y += noise_samples
+
+    io_pairs = PairedDataContainer(x, y, data_are_columns = true)
+    @test get_inputs(io_pairs) == x
+    @test get_outputs(io_pairs) == y
+
+    # Unknown ML type
+    mlt = MLTester()
+    @test_throws ErrorException emulator = Emulator(
+        mlt,
+        io_pairs,
+    )
+
+    # 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 = encode_with_schedule!(enc_sch, io_pairs, 1.0*I(p), 1.0*I(d))
+    @test get_encoder_schedule(em)[1][1] == enc_sch[1][1] # inputs: proc
+    @test get_encoder_schedule(em)[1][3] == enc_sch[1][3] # inputs: apply_to
+    @test get_encoder_schedule(em)[2][1] == enc_sch[2][1] # outputs...
+    @test get_encoder_schedule(em)[2][3] == enc_sch[2][3]
+    @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) 
+
+    #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_I = encode_structure_matrix(em, I(d), "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 isapprox(norm(encoded_I-enc_I_out), 0, atol = tol*d*d)
+    decoded_dc = decode_data(em, encoded_dc, "in")
+    decoded_mat = decode_data(em, encoded_mat, "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 isapprox(norm(decoded_I - I(d)), 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=3.0*I)
+    @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)
+    em1 = Emulator(gp, io_pairs, obs_noise_cov=3*I)
+    em2 = Emulator(gp, io_pairs, obs_noise_cov=3.0*I, output_structure_matrix=2.0*I)
+    enc_sch1 = create_encoder_schedule([(decorrelate_sample_cov(), "in"), (decorrelate_structure_mat(), "out")])
+    enc_sch2 = deepcopy(enc_sch1)
+    (_, _, _) =
+        encode_with_schedule!(
+            enc_sch1,
+            io_pairs,
+            1.0*I(p),
+            3.0*I(d), # obs noise cov becomes the output structure matrix
+        )
+    @test get_encoder_schedule(em1) == enc_sch1
+
+    (_, _, _) =
+        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
+   
+    
+    
+   
+end
+
+
+#=
 function constructor_tests(
-    iopairs::PairedDataContainer,
+    io_pairs::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)
+    input_mean = vec(mean(get_inputs(io_pairs), dims = 2)) #column vector
+    normalization = sqrt(inv(Symmetric(cov(get_inputs(io_pairs), dims = 2))))
+    normalization = Emulators.calculate_normalization(get_inputs(io_pairs))
+    norm_inputs = Emulators.normalize(get_inputs(io_pairs), input_mean, normalization)
     # [4.] test emulator preserves the structures
     mlt = MLTester()
     @test_throws ErrorException emulator = Emulator(
         mlt,
-        iopairs,
+        io_pairs,
         obs_noise_cov = Σ,
         normalize_inputs = true,
         standardize_outputs = false,
@@ -39,7 +128,7 @@ function constructor_tests(
     gp = GaussianProcess(GPJL())
     emulator = Emulator(
         gp,
-        iopairs,
+        io_pairs,
         obs_noise_cov = Σ,
         normalize_inputs = false,
         standardize_outputs = false,
@@ -55,7 +144,7 @@ function constructor_tests(
     gp = GaussianProcess(GPJL())
     emulator2 = Emulator(
         gp,
-        iopairs,
+        io_pairs,
         obs_noise_cov = Σ,
         normalize_inputs = true,
         standardize_outputs = false,
@@ -63,14 +152,14 @@ function constructor_tests(
     )
     train_inputs = get_inputs(emulator2.training_pairs)
     @test train_inputs == norm_inputs
-    train_inputs2 = Emulators.normalize(emulator2, get_inputs(iopairs))
+    train_inputs2 = Emulators.normalize(emulator2, get_inputs(io_pairs))
     @test train_inputs2 == norm_inputs
 
     # reverse standardise
     gp = GaussianProcess(GPJL())
     emulator3 = Emulator(
         gp,
-        iopairs,
+        io_pairs,
         obs_noise_cov = Σ,
         normalize_inputs = false,
         standardize_outputs = true,
@@ -79,7 +168,7 @@ function constructor_tests(
     )
 
     #standardized and decorrelated (sd) data
-    s_y, s_y_cov = Emulators.standardize(get_outputs(iopairs), Σ, norm_factors)
+    s_y, s_y_cov = Emulators.standardize(get_outputs(io_pairs), Σ, 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
@@ -101,9 +190,9 @@ 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
+    io_pairs = PairedDataContainer(x, y, data_are_columns = true)
+    @test get_inputs(io_pairs) == x
+    @test get_outputs(io_pairs) == y
 
     # [1.] test SVD (2D y version)
     test_SVD = svd(Σ)
@@ -141,17 +230,17 @@ end
 
     # [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))
+    input_mean = vec(mean(get_inputs(io_pairs), dims = 2)) #column vector
+    normalization = sqrt(inv(Symmetric(cov(get_inputs(io_pairs), dims = 2))))
+    @test all(isapprox.(Emulators.calculate_normalization(get_inputs(io_pairs)), 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))
+    norm_inputs = Emulators.normalize(get_inputs(io_pairs), input_mean, normalization)
+    @test all(isapprox.(norm_inputs, normalization * (get_inputs(io_pairs) .- 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)]
+    reduced_inputs = get_inputs(io_pairs)[:, 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 
@@ -168,18 +257,18 @@ end
     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
+    s_y, s_y_cov = Emulators.standardize(get_outputs(io_pairs), Σ, norm_factors)
+    @test s_y == get_outputs(io_pairs) ./ norm_factors
     @test s_y_cov == Σ ./ (norm_factors .* norm_factors')
 
     @testset "Emulators - dense Array Σ" begin
-        constructor_tests(iopairs, Σ, norm_factors, decomposition)
+        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
 
         # truncation - only test here, where singular value of Σ differ
         gp = GaussianProcess(GPJL())
         emulator4 = Emulator(
             gp,
-            iopairs,
+            io_pairs,
             obs_noise_cov = Σ,
             normalize_inputs = false,
             standardize_outputs = false,
@@ -196,10 +285,10 @@ end
         Σ = Diagonal(rand(d) .+ 1.0)
         noise_samples = rand(MvNormal(zeros(d), Σ), m)
         y += noise_samples
-        iopairs = PairedDataContainer(x, y, data_are_columns = true)
+        io_pairs = PairedDataContainer(x, y, data_are_columns = true)
         _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0)
 
-        constructor_tests(iopairs, Σ, norm_factors, decomposition)
+        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
     end
 
     @testset "Emulators - UniformScaling Σ" begin
@@ -209,9 +298,10 @@ end
         Σ = UniformScaling(1.3)
         noise_samples = rand(MvNormal(zeros(d), Σ), m)
         y += noise_samples
-        iopairs = PairedDataContainer(x, y, data_are_columns = true)
+        io_pairs = PairedDataContainer(x, y, data_are_columns = true)
         _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0)
 
-        constructor_tests(iopairs, Σ, norm_factors, decomposition)
+        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
     end
 end
+=#
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 687532b22..74cffceff 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -104,8 +104,29 @@ end
     @test get_data_mean(DD) == [1]
     @test get_encoder_mat(DD) == [2]
     @test get_decoder_mat(DD) == [3]
-    # get some data as IO pairs for functional tests
 
+
+    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
     out_dim = 50
     samples = 120
@@ -150,7 +171,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)
+            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 +258,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))
@@ -302,7 +323,7 @@ end
     @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)
+    @test_throws ArgumentError 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")
@@ -312,7 +333,7 @@ end
 
     # 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)
+        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) =

From 71b077ed8bb5e3a01fc363df1588d35143d7df53 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 17:23:10 -0700
Subject: [PATCH 40/75] remove large comment

---
 test/Emulator/runtests.jl | 205 --------------------------------------
 1 file changed, 205 deletions(-)

diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl
index e7016713d..17956a57b 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -100,208 +100,3 @@ struct MLTester <: Emulators.MachineLearningTool end
    
 end
 
-
-#=
-function constructor_tests(
-    io_pairs::PairedDataContainer,
-    Σ::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing},
-    norm_factors,
-    decomposition,
-) where {FT <: AbstractFloat}
-
-    input_mean = vec(mean(get_inputs(io_pairs), dims = 2)) #column vector
-    normalization = sqrt(inv(Symmetric(cov(get_inputs(io_pairs), dims = 2))))
-    normalization = Emulators.calculate_normalization(get_inputs(io_pairs))
-    norm_inputs = Emulators.normalize(get_inputs(io_pairs), input_mean, normalization)
-    # [4.] test emulator preserves the structures
-    mlt = MLTester()
-    @test_throws ErrorException emulator = Emulator(
-        mlt,
-        io_pairs,
-        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,
-        io_pairs,
-        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,
-        io_pairs,
-        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(io_pairs))
-    @test train_inputs2 == norm_inputs
-
-    # reverse standardise
-    gp = GaussianProcess(GPJL())
-    emulator3 = Emulator(
-        gp,
-        io_pairs,
-        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(io_pairs), Σ, 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
-    d = 6
-    p = 10
-    x = rand(p, m) #R^3
-    y = rand(d, m) #R^6
-
-    # "noise"
-    μ = zeros(d)
-    Σ = rand(d, d)
-    Σ = Σ' * Σ
-    noise_samples = rand(MvNormal(μ, Σ), m)
-    y += noise_samples
-
-    io_pairs = PairedDataContainer(x, y, data_are_columns = true)
-    @test get_inputs(io_pairs) == x
-    @test get_outputs(io_pairs) == 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)
-
-    # 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])
-
-    # 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 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(io_pairs), dims = 2)) #column vector
-    normalization = sqrt(inv(Symmetric(cov(get_inputs(io_pairs), dims = 2))))
-    @test all(isapprox.(Emulators.calculate_normalization(get_inputs(io_pairs)), normalization, atol = 1e-12))
-
-
-    norm_inputs = Emulators.normalize(get_inputs(io_pairs), input_mean, normalization)
-    @test all(isapprox.(norm_inputs, normalization * (get_inputs(io_pairs) .- 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(io_pairs)[:, 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))
-
-    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(io_pairs), Σ, norm_factors)
-    @test s_y == get_outputs(io_pairs) ./ norm_factors
-    @test s_y_cov == Σ ./ (norm_factors .* norm_factors')
-
-    @testset "Emulators - dense Array Σ" begin
-        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
-
-        # truncation - only test here, where singular value of Σ differ
-        gp = GaussianProcess(GPJL())
-        emulator4 = Emulator(
-            gp,
-            io_pairs,
-            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
-        io_pairs = PairedDataContainer(x, y, data_are_columns = true)
-        _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0)
-
-        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
-    end
-
-    @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
-        io_pairs = PairedDataContainer(x, y, data_are_columns = true)
-        _, decomposition = Emulators.svd_transform(y, Σ, retained_svd_frac = 1.0)
-
-        constructor_tests(io_pairs, Σ, norm_factors, decomposition)
-    end
-end
-=#

From fc4ced9b1269a373e676b74424452228ee6c0353 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 3 Jul 2025 17:33:01 -0700
Subject: [PATCH 41/75] more cases

---
 test/Emulator/runtests.jl | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl
index 17956a57b..a8fc26fc8 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -71,18 +71,23 @@ struct MLTester <: Emulators.MachineLearningTool end
     @test isapprox(norm(decoded_I - I(d)), 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=3.0*I)
+    @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)
-    em1 = Emulator(gp, io_pairs, obs_noise_cov=3*I)
-    em2 = Emulator(gp, io_pairs, obs_noise_cov=3.0*I, output_structure_matrix=2.0*I)
+    em1 = Emulator(gp, io_pairs, 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 = deepcopy(enc_sch1)
+    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")])
     (_, _, _) =
         encode_with_schedule!(
             enc_sch1,
             io_pairs,
             1.0*I(p),
-            3.0*I(d), # obs noise cov becomes the output structure matrix
+            Σ, # obs noise cov becomes the output structure matrix
         )
     @test get_encoder_schedule(em1) == enc_sch1
 
@@ -94,7 +99,15 @@ struct MLTester <: Emulators.MachineLearningTool end
             2.0*I(d), # out_struct_mat overrides obs noise cov
         )
     @test get_encoder_schedule(em2) == enc_sch2
-   
+    
+    (_, _, _) =
+        encode_with_schedule!(
+            enc_sch3,
+            io_pairs,
+            Γ, 
+            I(d), 
+        )
+    @test get_encoder_schedule(em3) == enc_sch3
     
     
    

From f9e1f8f9236272abddf9660832f51bd12707c615 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 10:56:15 -0700
Subject: [PATCH 42/75] GP test updated

---
 test/GaussianProcess/runtests.jl | 82 +++++++++-----------------------
 1 file changed, 22 insertions(+), 60 deletions(-)

diff --git a/test/GaussianProcess/runtests.jl b/test/GaussianProcess/runtests.jl
index f92f24393..1228c4ca7 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,32 +63,22 @@ using CalibrateEmulateSample.DataContainers
     @test gp1.prediction_type == pred_type
     @test gp1.alg_reg_noise == 1e-4
 
-
     em1 = Emulator(
         gp1,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     )
 
-    @test_logs (:warn,) (:warn,) Emulator(
+    @test_logs (:warn,) Emulator(
         gp1,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     ) # 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,
+        encoder_schedule = [],
     )
     @test length(gp1.models) == 1
 
@@ -104,10 +95,7 @@ using CalibrateEmulateSample.DataContainers
     @test_throws ArgumentError Emulator(
         agp,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     )
 
     gp1_opt_params = Emulators.get_params(gp1)[1] # one model only
@@ -122,21 +110,15 @@ using CalibrateEmulateSample.DataContainers
     em_agp_from_gp1 = Emulator(
         agp,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
         kernel_params = kernel_params,
     )
     optimize_hyperparameters!(em_agp_from_gp1)
     # skip rebuild:
-    @test_logs (:warn,) (:warn,) Emulator(
+    @test_logs (:warn,) Emulator(
         agp,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
         kernel_params = kernel_params,
     )
 
@@ -158,10 +140,7 @@ using CalibrateEmulateSample.DataContainers
     em2 = Emulator(
         gp2,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     )
 
     Emulators.optimize_hyperparameters!(em2)
@@ -182,36 +161,22 @@ using CalibrateEmulateSample.DataContainers
     em3 = Emulator(
         gp3,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     )
-    @test_logs (:warn,) (:warn,) Emulator(
+    @test_logs (:warn,) Emulator(
         gp3,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = [],
     )
     Emulator(
         gp3,
         iopairs,
-        obs_noise_cov = nothing,
-        normalize_inputs = false,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        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,16 +215,16 @@ using CalibrateEmulateSample.DataContainers
     @test get_inputs(iopairs2) == X
     @test get_outputs(iopairs2) == Y
 
+    encoder_schedule = (decorrelate_structure_mat(), "out")
+    
     # 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,
+        encoder_schedule = deepcopy(encoder_schedule),
+        output_structure_matrix = Σ,
     )
 
     gp4 = GaussianProcess(gppackage; kernel = nothing, noise_learn = false, prediction_type = pred_type)
@@ -267,10 +232,8 @@ using CalibrateEmulateSample.DataContainers
     em4_noise_from_Σ = Emulator(
         gp4,
         iopairs2,
-        obs_noise_cov = Σ,
-        normalize_inputs = true,
-        standardize_outputs = false,
-        retained_svd_frac = 1.0,
+        encoder_schedule = deepcopy(encoder_schedule),
+        output_structure_matrix = Σ,
     )
 
     Emulators.optimize_hyperparameters!(em4_noise_learnt)
@@ -319,9 +282,8 @@ using CalibrateEmulateSample.DataContainers
     em_agp_from_gp4 = Emulator(
         agp4,
         iopairs2,
-        obs_noise_cov = Σ,
-        normalize_inputs = true,
-        retained_svd_frac = 1.0,
+        encoder_schedule = deepcopy(encoder_schedule),
+        output_structure_matrix = Σ,
         kernel_params = kernel_params,
     )
 

From bf1a8c378d982e376cc1ef2bc2896173c39ffdc7 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 11:00:51 -0700
Subject: [PATCH 43/75] use def. enc. for GP tests

---
 test/GaussianProcess/runtests.jl | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/test/GaussianProcess/runtests.jl b/test/GaussianProcess/runtests.jl
index 1228c4ca7..4e8434256 100644
--- a/test/GaussianProcess/runtests.jl
+++ b/test/GaussianProcess/runtests.jl
@@ -215,7 +215,6 @@ using CalibrateEmulateSample.Utilities
     @test get_inputs(iopairs2) == X
     @test get_outputs(iopairs2) == Y
 
-    encoder_schedule = (decorrelate_structure_mat(), "out")
     
     # 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)
@@ -223,7 +222,6 @@ using CalibrateEmulateSample.Utilities
     em4_noise_learnt = Emulator(
         gp4_noise_learnt,
         iopairs2,
-        encoder_schedule = deepcopy(encoder_schedule),
         output_structure_matrix = Σ,
     )
 
@@ -232,7 +230,6 @@ using CalibrateEmulateSample.Utilities
     em4_noise_from_Σ = Emulator(
         gp4,
         iopairs2,
-        encoder_schedule = deepcopy(encoder_schedule),
         output_structure_matrix = Σ,
     )
 
@@ -282,7 +279,6 @@ using CalibrateEmulateSample.Utilities
     em_agp_from_gp4 = Emulator(
         agp4,
         iopairs2,
-        encoder_schedule = deepcopy(encoder_schedule),
         output_structure_matrix = Σ,
         kernel_params = kernel_params,
     )

From 1396b421d882c8039f7a413f80109e5643c00663 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 11:57:04 -0700
Subject: [PATCH 44/75] all return types

---
 src/Emulator.jl | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index cf607ce32..be6b3a108 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -228,23 +228,31 @@ function predict(
     # 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...)
+    
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
+
     # return decoded or encoded?
     if transform_to_real
         decoded_outputs = decode_data(emulator, encoded_outputs, "out")
-        
+            
         decoded_covariances = zeros(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")
+                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")
+                decoded_covariances[:,:,i] .= decode_structure_matrix(emulator, mat, "out")
             end
         end
-            
-        return decoded_outputs, eachslice(decoded_covariances,dims=3)
+
+        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
+        
     else
         
         if encoded_output_dim > 1
@@ -267,5 +275,4 @@ function predict(
     
 end
 
-
 end

From ee83d041ed1c106c12cc0189a48625f6ebe60068 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 12:16:45 -0700
Subject: [PATCH 45/75] output cases

---
 src/Emulator.jl | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index be6b3a108..33b969095 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -198,6 +198,10 @@ $(DocStringExtensions.TYPEDSIGNATURES)
 
 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:
+  - 1-D: mean [1 x N], cov [1 x N]
+  - p-D: mean [p x N], cov N x [p x p] 
 """
 function predict(
     emulator::Emulator{FT},
@@ -255,22 +259,23 @@ function predict(
         
     else
         
-        if encoded_output_dim > 1
-            encoded_covariances_mat = zeros(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
+        encoded_covariances_mat = zeros(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
-            encoded_covariances = eachslice(encoded_covariances_mat,dims=3)
-        else
-            encoded_covariances = encoded_uncertainties
+        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
-        return encoded_outputs, encoded_covariances
     end
     
 end

From fc65322d5bf818f3d3aba9ff8235ae5b93046363 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 12:17:22 -0700
Subject: [PATCH 46/75] updated RF tests

---
 test/RandomFeature/runtests.jl | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/test/RandomFeature/runtests.jl b/test/RandomFeature/runtests.jl
index baa23f674..d170eab28 100644
--- a/test/RandomFeature/runtests.jl
+++ b/test/RandomFeature/runtests.jl
@@ -270,7 +270,7 @@ rng = Random.MersenneTwister(seed)
         output_dim = 1
         n = 50                                       # number of training points
         x = reshape(2.0 * π * rand(n), 1, n)         # unif(0,2π) predictors/features: 1 x n
-        obs_noise_cov = 0.05^2 * I
+        obs_noise_cov= 0.05^2 * I
         y = reshape(sin.(x) + 0.05 * randn(n)', 1, n) # predictands/targets: 1 x n
         iopairs = PairedDataContainer(x, y, data_are_columns = true)
 
@@ -304,11 +304,11 @@ rng = Random.MersenneTwister(seed)
             rng = rng,
             optimizer_options = Dict("n_cross_val_sets" => 0),
         )
-
+        
         # 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_μ)
 

From a35fcc7cebc9a1c6e675577edbb365f644631b14 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 13:38:32 -0700
Subject: [PATCH 47/75] update interfcace to abstractGP

---
 src/GaussianProcess.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index d26f9a95a..9c130d684 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -410,7 +410,7 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow
             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; obsdim=2)
 
         data_i = output_values[i, :]
         opt_post_fx = posterior(opt_fx, data_i)
@@ -434,7 +434,7 @@ 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

From 14e522e5ce9f5557b7d9d7207dc11372894087a1 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 13:39:32 -0700
Subject: [PATCH 48/75] updated run tests for MCMC

---
 src/Emulator.jl                        |  4 +-
 src/MarkovChainMonteCarlo.jl           | 23 ++++++----
 test/MarkovChainMonteCarlo/runtests.jl | 62 +++++++++-----------------
 3 files changed, 39 insertions(+), 50 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 33b969095..972d0d8c4 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -239,7 +239,7 @@ function predict(
     if transform_to_real
         decoded_outputs = decode_data(emulator, encoded_outputs, "out")
             
-        decoded_covariances = zeros(output_dim, output_dim, size(encoded_uncertainties)[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")
@@ -259,7 +259,7 @@ function predict(
         
     else
         
-        encoded_covariances_mat = zeros(encoded_output_dim, encoded_output_dim, size(encoded_uncertainties)[end])
+        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)
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index 4ea545e59..f3caa5e69 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -495,13 +495,13 @@ AbstractMCMC's terminology).
 # Fields
 $(DocStringExtensions.TYPEDFIELDS)
 """
-struct MCMCWrapper{AMorAV <: Union{AbstractVector, AbstractMatrix}, AV <: AbstractVector}
+struct MCMCWrapper{AVV, 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."
@@ -528,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.
@@ -547,8 +546,16 @@ function MCMCWrapper(
     kwargs...,
 ) where {AV <: AbstractVector, AMorAV <: Union{AbstractVector, AbstractMatrix}}
 
+    # make into iterable over vectors
+    obs_slice = isa(observation, AbstractMatrix) ? eachcol(observation) : observation
+    if eltype(obs_slice) <: Number # just one observation
+        obs_slice = [obs_slice]
+    end
+    
     # encoding works on columns but mcmc wants vec-of-vec
-    encoded_obs = [vec(encode_data(em, reshape(obs, :, 1), "out")) for obs in eachcol(observation)]
+    
+    encoded_obs = [vec(encode_data(em, reshape(obs, :, 1), "out")) for obs in obs_slice]
+        
 
     log_posterior_map = EmulatorPosteriorModel(prior, em, encoded_obs)
     mh_proposal_sampler = MetropolisHastingsSampler(mcmc_alg, prior)
@@ -570,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, encoded_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
 
 """
diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl
index e0e4c3793..b44605a6c 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:
@@ -72,17 +73,14 @@ function test_gp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor = nothing
     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,
+        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:
@@ -93,11 +91,8 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor =
     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,
+        output_structure_matrix = σ2_y,
+        encoder_schedule = [],
     )
     Emulators.optimize_hyperparameters!(em)
 
@@ -118,11 +113,8 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor =
     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,
+        output_structure_matrix = σ2_y,
+        encoder_schedule = [],
         kernel_params = kernel_params,
     )
 
@@ -130,7 +122,7 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer; norm_factor =
 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,14 +130,16 @@ 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)
+    retain_var = 0.95 
+    encoder_schedule = [
+        (quartile_scale(), "out"),
+        (decorrelate_structure_mat(retain_var=retain_var), "out")
+    ]
     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,
+        output_structure_matrix = σ2_y,
+        encoder_schedule = encoder_schedule
     )
     Emulators.optimize_hyperparameters!(em)
 
@@ -157,7 +151,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 +183,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 +214,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 +285,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)

From 0469c7ccd5a3c6bcb7c6651febd7b0129b3c0966 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 13:46:22 -0700
Subject: [PATCH 49/75] format

---
 examples/Cloudy/Cloudy_calibrate.jl           |   4 +-
 examples/Cloudy/Cloudy_emulate_sample.jl      |   6 +-
 examples/Darcy/emulate_sample.jl              |   6 +-
 examples/EDMF_data/emulator-rank-test.jl      |   8 +-
 examples/EDMF_data/plot_posterior.jl          |  13 +-
 examples/EDMF_data/rebuild_priors.jl          |  22 +-
 examples/EDMF_data/uq_for_edmf.jl             |  14 +-
 .../G-function/emulate-test-n-features.jl     |   9 +-
 examples/Emulator/G-function/emulate.jl       |   7 +-
 examples/Emulator/Ishigami/emulate.jl         |   7 +-
 examples/Emulator/L63/emulate.jl              |   4 +-
 .../Emulator/L63/emulate_diff-rank-test.jl    |   7 +-
 examples/Emulator/L63/plot_results.jl         |   2 +-
 .../Regression_2d_2d/compare_regression.jl    | 297 ++++++++----------
 examples/GCM/emulate_sample_script.jl         |  17 +-
 examples/Lorenz/calibrate_spatial_dep.jl      |   4 +-
 examples/Lorenz/emulate_sample.jl             |  23 +-
 examples/Lorenz/emulate_sample_spatial_dep.jl |  16 +-
 examples/Lorenz/plot_spatial_dep.jl           |   7 +-
 examples/Sinusoid/emulate.jl                  |   4 +-
 src/Emulator.jl                               | 111 ++++---
 src/GaussianProcess.jl                        |   4 +-
 src/MarkovChainMonteCarlo.jl                  |  10 +-
 src/RandomFeature.jl                          |   4 +-
 src/ScalarRandomFeature.jl                    |  16 +-
 src/Utilities.jl                              |  14 +-
 src/VectorRandomFeature.jl                    |   2 +-
 test/Emulator/runtests.jl                     |  94 +++---
 test/GaussianProcess/runtests.jl              |  83 +----
 test/MarkovChainMonteCarlo/runtests.jl        |  37 +--
 test/RandomFeature/runtests.jl                |   4 +-
 test/Utilities/runtests.jl                    |  12 +-
 32 files changed, 396 insertions(+), 472 deletions(-)

diff --git a/examples/Cloudy/Cloudy_calibrate.jl b/examples/Cloudy/Cloudy_calibrate.jl
index 31e430e94..70fa8a2b2 100644
--- a/examples/Cloudy/Cloudy_calibrate.jl
+++ b/examples/Cloudy/Cloudy_calibrate.jl
@@ -169,10 +169,10 @@ for n in 1:N_iter
     G_ens = hcat(G_n...)  # reformat
     terminate = EnsembleKalmanProcesses.update_ensemble!(ekiobj, G_ens)
     if !isnothing(terminate)
-        n_iter[1] = n-1
+        n_iter[1] = n - 1
         break
     end
-        
+
 end
 
 
diff --git a/examples/Cloudy/Cloudy_emulate_sample.jl b/examples/Cloudy/Cloudy_emulate_sample.jl
index cee2ba915..030f503be 100644
--- a/examples/Cloudy/Cloudy_emulate_sample.jl
+++ b/examples/Cloudy/Cloudy_emulate_sample.jl
@@ -113,7 +113,7 @@ function main()
     # 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)
-    test_pairs =  get_training_points(ekiobj, length(get_u(ekiobj)) - 1: length(get_u(ekiobj)) - 1)
+    test_pairs = get_training_points(ekiobj, (length(get_u(ekiobj)) - 1):(length(get_u(ekiobj)) - 1))
     for case in cases[case_mask]
 
         println(" ")
@@ -160,12 +160,12 @@ function main()
             error("Case $case is not implemented yet.")
 
         end
-        
+
         # Data processing
         if case == "rf-nonsep"
             encoder_schedule = []
         else
-            encoder_schedule = (decorrelate_structure_mat(), "in_and_out") 
+            encoder_schedule = (decorrelate_structure_mat(), "in_and_out")
         end
 
         # build emulator
diff --git a/examples/Darcy/emulate_sample.jl b/examples/Darcy/emulate_sample.jl
index b4df4fb9a..ca046dbb8 100644
--- a/examples/Darcy/emulate_sample.jl
+++ b/examples/Darcy/emulate_sample.jl
@@ -18,7 +18,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers
 using CalibrateEmulateSample.ParameterDistributions
 using CalibrateEmulateSample.DataContainers
 
-    
+
 include("GModel.jl")
 
 function main()
@@ -93,11 +93,11 @@ function main()
 
         # data processing
         encoding_schedule = (decorrelate_structure_mat(), "in_and_out")
-        
+
         emulator = Emulator(
             mlt,
             input_output_pairs;
-            input_structure_matrix= cov(prior),
+            input_structure_matrix = cov(prior),
             output_structure_matrix = Γy,
             encoding_schedule = encoding_schedule,
         )
diff --git a/examples/EDMF_data/emulator-rank-test.jl b/examples/EDMF_data/emulator-rank-test.jl
index e240c8422..4bf51fd4d 100644
--- a/examples/EDMF_data/emulator-rank-test.jl
+++ b/examples/EDMF_data/emulator-rank-test.jl
@@ -247,11 +247,11 @@ function main()
                     optimizer_options = overrides,
                 )
             end
-            
+
             # data processing:
             retain_var = 0.9
-            encoder_schedule = (decorrelate_structure_mat(retain_var=retain_var), "in_and_out")
-            
+            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(
@@ -259,7 +259,7 @@ function main()
                     train_pairs;
                     input_structure_matrix = cov(prior),
                     output_structure_matrix = truth_cov,
-                    encoder_schedule=encoder_schedule,
+                    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
index af23d8d74..8d652c496 100644
--- a/examples/EDMF_data/rebuild_priors.jl
+++ b/examples/EDMF_data/rebuild_priors.jl
@@ -1,13 +1,13 @@
- # 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
-    =#
+# 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()
@@ -34,6 +34,6 @@ function get_prior_config(s::SS) where {SS <: AbstractString}
         config["unconstrained_σ"] = 1.0
     else
         throw(ArgumentError("prior for experiment $s not found, please implement in uq_for_edmf.jl"))
-        end
+    end
     return config
 end
diff --git a/examples/EDMF_data/uq_for_edmf.jl b/examples/EDMF_data/uq_for_edmf.jl
index 002ee3d65..a708b00f5 100644
--- a/examples/EDMF_data/uq_for_edmf.jl
+++ b/examples/EDMF_data/uq_for_edmf.jl
@@ -160,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 =
@@ -171,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"
@@ -197,7 +197,7 @@ function main()
             noise_learn = false,
         )
     elseif case ∈ ["RF-lr-lr"]
-        kernel_structure = SeparableKernel(LowRankFactor(5, nugget), LowRankFactor(kernel_rank, nugget))
+        kernel_structure = SeparableKernel(LowRankFactor(1, nugget), LowRankFactor(kernel_rank, nugget))
         n_features = 500
 
         mlt = VectorRandomFeatureInterface(
@@ -225,15 +225,15 @@ function main()
 
     # data processing:
     retain_var = 0.9
-    encoder_schedule = (decorrelate_structure_mat(retain_var=retain_var), "in_and_out")
-    
+    encoder_schedule = (decorrelate_structure_mat(retain_var = retain_var), "in_and_out")
+
     # Fit an emulator to the data
     emulator = Emulator(
         mlt,
         input_output_pairs;
         input_structure_matrix = cov(prior),
         output_structure_matrix = truth_cov,
-        encoder_schedule=encoder_schedule,
+        encoder_schedule = encoder_schedule,
     )
 
     # 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 0b339ac40..2dc1c5776 100644
--- a/examples/Emulator/G-function/emulate-test-n-features.jl
+++ b/examples/Emulator/G-function/emulate-test-n-features.jl
@@ -83,7 +83,7 @@ function main()
         output[i] = y[ind[i]] + noise[i]
     end
     iopairs = PairedDataContainer(input, output)
-    
+
     # analytic sobol indices taken from
     # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC8989694/pdf/main.pdf
     a = [(i - 1.0) / 2.0 for i in 1:n_dimensions]  # a_i < a_j => a_i more sensitive
@@ -151,7 +151,12 @@ function main()
 
             # Emulate
             ttt[f_idx, rep_idx] = @elapsed begin
-                emulator = Emulator(mlt, iopairs; output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
+                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 11f86be28..3cc3ab123 100644
--- a/examples/Emulator/G-function/emulate.jl
+++ b/examples/Emulator/G-function/emulate.jl
@@ -142,7 +142,12 @@ function main()
 
         # Emulate
         times[rep_idx] = @elapsed begin
-            emulator = Emulator(mlt, iopairs; output_structure_matrix = Γ * I, encoder_schedule = deepcopy(encoder_schedule))
+            emulator = Emulator(
+                mlt,
+                iopairs;
+                output_structure_matrix = Γ * I,
+                encoder_schedule = deepcopy(encoder_schedule),
+            )
             optimize_hyperparameters!(emulator)
         end
 
diff --git a/examples/Emulator/Ishigami/emulate.jl b/examples/Emulator/Ishigami/emulate.jl
index d3b980e7a..74421011d 100644
--- a/examples/Emulator/Ishigami/emulate.jl
+++ b/examples/Emulator/Ishigami/emulate.jl
@@ -91,7 +91,7 @@ function main()
         "n_ensemble" => 50,
         "n_iteration" => 20,
     )
-    
+
     if case == "Prior"
         # don't do anything
         overrides["n_iteration"] = 0
@@ -102,7 +102,7 @@ function main()
     result_preds = []
     opt_diagnostics = []
     for rep_idx in 1:n_repeats
-        
+
         # Build ML tools
         if case == "GP"
             gppackage = Emulators.SKLJL()
@@ -123,7 +123,8 @@ function main()
         end
 
         # Emulate
-        emulator = Emulator(mlt, iopairs; output_structure_matrix = Γ*I, encoder_schedule = deepcopy(encoder_schedule))
+        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 5ca1f9a7b..b623a89b5 100644
--- a/examples/Emulator/L63/emulate.jl
+++ b/examples/Emulator/L63/emulate.jl
@@ -108,7 +108,7 @@ function main()
     ]
 
     case = cases[3] # 7
-    
+
     nugget = Float64(1e-8)
     u_test = []
     u_hist = []
@@ -215,7 +215,7 @@ 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; 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 05bbe1952..e60b7ebcc 100644
--- a/examples/Emulator/L63/emulate_diff-rank-test.jl
+++ b/examples/Emulator/L63/emulate_diff-rank-test.jl
@@ -191,7 +191,12 @@ function main()
                 encoder_schedule = nothing
             end
             ttt[rank_id, rep_idx] = @elapsed begin
-                emulator = Emulator(mlt, iopairs; output_structure_matrix = Γy, encoder_schedule=deepcopy(encoder_schedule))
+                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/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index 5a2bd426d..3b8b11f8b 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -42,19 +42,12 @@ function main()
     cases = [
         "gp-skljl",
         "gp-gpjl", # Very slow prediction...
-        "rf-sep-scalar",
-        "rf-sep-sep",
+        "rf-lr-scalar",
+        "rf-lr-lr",
         "rf-nonsep",
     ]
 
-    encoder_names = [
-        "no-proc",
-        "sample-struct-proc",
-        "sample-proc",
-        "struct-mat-proc",
-        "combined-proc",
-        "cca-proc",
-    ]
+    encoder_names = ["no-proc", "sample-struct-proc", "sample-proc", "struct-mat-proc", "combined-proc", "cca-proc"]
     encoders = [
         [], # no proc
         nothing, # default proc
@@ -65,21 +58,21 @@ function main()
     ]
 
     # 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)
-    
-    
+    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
-    prior_cov = (pi^2)*I(p)
-    X = rand(MvNormal(zeros(p),prior_cov), n)
+    prior_cov = (pi^2) * I(p)
+    X = rand(MvNormal(zeros(p), prior_cov), n)
     # G(x1, x2)
     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)
+    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)
@@ -88,7 +81,7 @@ function main()
 
     # Add noise η
     μ = zeros(d)
-    Σ = 0.1*[[0.4, -0.3] [-0.3, 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
@@ -100,184 +93,168 @@ function main()
     #plot training data with and without noise
     if plot_flag
         n_pts = 200
-        x1 = range(-2*pi, stop = 2*π, length = n_pts)
-        x2 = range(-2*pi, stop = 2*π, length = n_pts)
-        
+        x1 = range(-2 * pi, stop = 2 * π, length = n_pts)
+        x2 = range(-2 * pi, stop = 2 * π, length = n_pts)
+
         p1 = contourf(
             x1,
             x2,
-            g1.(x1',x2),
+            g1.(x1', x2),
             c = :cividis,
             xlabel = "x1",
             ylabel = "x2",
-            aspect_ratio = :equal            
-
+            aspect_ratio = :equal,
         )
-        scatter!(
-            p1,
-            X[1,:],
-            X[2,:],
-            c=:black,
-            ms=3,
-            label="train point"
-        )
-        
+        scatter!(p1, X[1, :], X[2, :], c = :black, ms = 3, label = "train point")
+
         figpath = joinpath(output_directory, "g1_true.png")
         savefig(figpath)
         p2 = contourf(
             x1,
             x2,
-            g2.(x1',x2),
+            g2.(x1', x2),
             c = :cividis,
             xlabel = "x1",
             ylabel = "x2",
-            aspect_ratio = :equal            
-
+            aspect_ratio = :equal,
         )
-        scatter!(
-            p2,
-            X[1,:],
-            X[2,:],
-            c=:black,
-            ms=3,
-            label="train point"
-        )
-        
+        scatter!(p2, X[1, :], X[2, :], c = :black, ms = 3, label = "train point")
+
         figpath = joinpath(output_directory, "g2_true.png")
         savefig(figpath)
 
     end
 
-    for case in cases[case_mask] 
+    for case in cases[case_mask]
         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-sep-scalar"
-            mlt = ScalarRandomFeatureInterface(
-                n_features,
-                p,
-                kernel_structure = SeparableKernel(LowRankFactor(2, nugget), OneDimFactor()),
-                optimizer_options = optimizer_options,
-            )
-         elseif case == "rf-sep-sep"
-            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),
-            )
- 
-        end
+            println(" ")
+            @info """
 
-        emulator = Emulator(
-            mlt,
-            iopairs,
-            encoder_schedule = enc,
-            input_structure_matrix = prior_cov,
-            output_structure_matrix = Σ,
-        )
+            -------------
+            running case $case with encoder $enc_name
+            -------------
+            """
 
-        
-        optimize_hyperparameters!(emulator) # although RF already optimized
+            # common Gaussian feature setup
+            pred_type = YType()
 
-        # 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")
+            # 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
 
-        #plot predictions
-        for y_i in 1:d
+            if case == "gp-skljl"
+                gppackage = SKLJL()
+                mlt = GaussianProcess(gppackage, noise_learn = true)
 
-            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
+            elseif case == "gp-gpjl"
+                @warn "gp-gpjl case is very slow at prediction"
+                gppackage = GPJL()
+                mlt = GaussianProcess(gppackage, noise_learn = true)
 
-            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",
+            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 = contourf(
-                    x1,
-                    x2,
-                    var_grid,
-                    c = :cividis,
-                    xlabel = "x1",
-                    ylabel = "x2",
-                    title="var output $y_i",
+            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),
+                )
+
+            end
+
+            emulator = Emulator(
+                mlt,
+                iopairs,
+                encoder_schedule = enc,
+                input_structure_matrix = prior_cov,
+                output_structure_matrix = Σ,
+            )
 
-                plot(p5, p6, layout = (1, 2), legend = false)
 
-                savefig(joinpath(output_directory, case *"_"* enc_name * "_y" * string(y_i) * "_predictions.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
 
-        # Plot the true components of G(x1, x2)
-        g1_true = g1(inputs)
-        g1_true_grid = reshape(g1_true, n_pts, n_pts)
-       
+            # Plot the true components of G(x1, x2)
+            g1_true = g1(inputs)
+            g1_true_grid = reshape(g1_true, n_pts, n_pts)
+
 
-        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)
+            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
+            # 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,)
@@ -303,7 +280,7 @@ function main()
 
                 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 25115f47a..07098b641 100644
--- a/examples/GCM/emulate_sample_script.jl
+++ b/examples/GCM/emulate_sample_script.jl
@@ -55,7 +55,7 @@ function main()
     priorfile = "priors.jld2"
     prior = load(priorfile)["prior"]
 
-    
+
     #take only first 400 points
     iter_mask = [1, 2, 4]
     data_mask = 1:96
@@ -158,28 +158,23 @@ function main()
     end
 
     if emulator_type == "VectorRFR-nondiag"
-        
-        encoder_schedule = [
-            (quartile_scale(), "in"),
-            (decorrelate_structure_mat(), "out"),
-        ]
+
+        encoder_schedule = [(quartile_scale(), "in"), (decorrelate_structure_mat(), "out")]
         emulator = Emulator(
             mlt,
             input_output_pairs;
             output_structure_matrix = obs_noise_cov,
-            encoder_schedule=encoder_schedule,
+            encoder_schedule = encoder_schedule,
         )
     else
-        encoder_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 = obs_noise_cov,
-            encoder_schedule=encoder_schedule,
+            encoder_schedule = encoder_schedule,
         )
 
     end
diff --git a/examples/Lorenz/calibrate_spatial_dep.jl b/examples/Lorenz/calibrate_spatial_dep.jl
index ff7e988e5..82fab7267 100644
--- a/examples/Lorenz/calibrate_spatial_dep.jl
+++ b/examples/Lorenz/calibrate_spatial_dep.jl
@@ -193,7 +193,7 @@ y_ens = hcat(
 )
 
 # estimate noise as the variability of these pushed-forward samples on the attractor
-obs_noise_cov = cov(y_ens, dims = 2) + 1e-2*I
+obs_noise_cov = cov(y_ens, dims = 2) + 1e-2 * I
 y = y_ens[:, 1]
 
 pl = 4.0
@@ -249,7 +249,7 @@ for i in 1:N_iter
         [
             lorenz_forward(
                 EnsembleMemberConfig(params_i[:, j]),
-                x_on_attractor[:,(i-1) * N_ens + j],
+                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 4e28d5add..338620169 100644
--- a/examples/Lorenz/emulate_sample.jl
+++ b/examples/Lorenz/emulate_sample.jl
@@ -167,11 +167,8 @@ function main()
 
         # Process data
         retain_var = 0.95
-        encoder_schedule = [
-            (quartile_scale(), "in"),
-            (decorrelate_structure_mat(retain_var = retain_var), "out"),
-        ]
-        
+        encoder_schedule = [(quartile_scale(), "in"), (decorrelate_structure_mat(retain_var = retain_var), "out")]
+
         emulator = Emulator(
             mlt,
             input_output_pairs;
@@ -232,19 +229,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,:])
-        
+        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, xlims = xlims, ylims=ylims)
+        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 de99067d0..50b634ddc 100644
--- a/examples/Lorenz/emulate_sample_spatial_dep.jl
+++ b/examples/Lorenz/emulate_sample_spatial_dep.jl
@@ -23,7 +23,7 @@ using CalibrateEmulateSample.DataContainers
 function main()
 
     cases = [
-        "GP", 
+        "GP",
         "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs
     ]
 
@@ -83,11 +83,11 @@ function main()
             )
         elseif case == "RF-scalar"
             overrides = Dict(
-         #       "verbose" => true,
+                #       "verbose" => true,
                 "scheduler" => DataMisfitController(terminate_at = 1000.0),
                 "cov_sample_multiplier" => 1.0,
                 "n_iteration" => 8,
-                "n_features_opt" => 40,                
+                "n_features_opt" => 40,
             )
             n_features = 80
             kernel_structure = SeparableKernel(LowRankFactor(1, nugget), OneDimFactor())
@@ -113,8 +113,8 @@ function main()
 
         # plot training points in constrained space
         if case == cases[mask[1]]
-            i1=1
-            i2=10
+            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)
@@ -139,9 +139,7 @@ function main()
 
         # data processing configuration
         retain_var = 0.95
-        encoder_schedule = [
-            (decorrelate_structure_mat(retain_var = retain_var), "in_and_out"),
-        ]
+        encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")]
 
         emulator = Emulator(
             mlt,
@@ -180,7 +178,7 @@ function main()
         # First let's run a short chain to determine a good step size
         mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0)
         new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0)
-        
+
         # Now begin the actual MCMC
         println("Begin MCMC - with step size ", new_step)
         chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000)
diff --git a/examples/Lorenz/plot_spatial_dep.jl b/examples/Lorenz/plot_spatial_dep.jl
index 12136c740..4265469b7 100644
--- a/examples/Lorenz/plot_spatial_dep.jl
+++ b/examples/Lorenz/plot_spatial_dep.jl
@@ -1,11 +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
-]
-case = cases[2]
+cases = ["GP", "RF-scalar"]
+case = cases[1]
 
 # load packages
 # CES 
diff --git a/examples/Sinusoid/emulate.jl b/examples/Sinusoid/emulate.jl
index 02ddb6663..0b02a6259 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, 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 
@@ -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, encoder_schedule = encoder_schedule_rf, output_structure_matrix = Γ)
 optimize_hyperparameters!(emulator_random_features)
 
 
diff --git a/src/Emulator.jl b/src/Emulator.jl
index 972d0d8c4..66f67cbf2 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -93,7 +93,7 @@ function Emulator(
     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}
 
@@ -101,10 +101,14 @@ function Emulator(
     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`.")
+        @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=`")
+        @warn(
+            "Keyword `obs_noise_cov=` is now deprecated and will be ignored. \n Continuing with value of `output_structure_matrix=`"
+        )
     end
 
     input_structure_mat = if isnothing(input_structure_matrix)
@@ -114,7 +118,7 @@ function Emulator(
     else
         input_structure_matrix
     end
-    
+
     output_structure_mat = if isnothing(output_structure_matrix)
         Diagonal(FT.(ones(output_dim)))
     elseif isa(output_structure_matrix, UniformScaling)
@@ -137,25 +141,21 @@ function Emulator(
         else
             push!(encoder_schedule, (decorrelate_structure_mat(), "out"))
         end
-    end        
-        
+    end
+
     enc_schedule = create_encoder_schedule(encoder_schedule)
     (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
-        encode_with_schedule!(
-            enc_schedule,
-            input_output_pairs,
-            input_structure_mat,
-            output_structure_mat,
-        )
-    
+        encode_with_schedule!(enc_schedule, input_output_pairs, input_structure_mat, output_structure_mat)
+
     # build the machine learning tool in the encoded space
-    build_models!(machine_learning_tool, encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat; mlt_kwargs...)
-    return Emulator{FT, typeof(enc_schedule)}(
+    build_models!(
         machine_learning_tool,
-        input_output_pairs,
         encoded_io_pairs,
-        enc_schedule,
+        encoded_input_structure_mat,
+        encoded_output_structure_mat;
+        mlt_kwargs...,
     )
+    return Emulator{FT, typeof(enc_schedule)}(machine_learning_tool, input_output_pairs, encoded_io_pairs, enc_schedule)
 end
 
 """
@@ -167,29 +167,45 @@ function optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) w
     optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...)
 end
 
-    
-function encode_data(emulator::Emulator, data::MorDC, in_or_out::AS) where {AS <: AbstractString, MorDC <: Union{AbstractMatrix,DataContainer}}
-    if isa(data, AbstractMatrix) 
+
+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 encode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
     end
 end
 
-function encode_structure_matrix(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+function encode_structure_matrix(
+    emulator::Emulator,
+    structure_mat::USorM,
+    in_or_out::AS,
+) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
     return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
-    
-function decode_data(emulator::Emulator, data::MorDC, in_or_out::AS) where {AS <: AbstractString, MorDC <: Union{AbstractMatrix,DataContainer}}
-    if isa(data, AbstractMatrix) 
+
+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 decode_with_schedule(get_encoder_schedule(emulator), data, in_or_out)
     end
 end
 
-function decode_structure_matrix(emulator::Emulator, structure_mat::USorM, in_or_out::AS) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+function decode_structure_matrix(
+    emulator::Emulator,
+    structure_mat::USorM,
+    in_or_out::AS,
+) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
     return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
@@ -214,9 +230,9 @@ function predict(
     # 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)
-    
+
     if !(size(new_inputs, 1) == input_dim)
         throw(
             ArgumentError(
@@ -232,52 +248,53 @@ function predict(
     # 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...)
-    
+
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
 
     # return decoded or encoded?
     if transform_to_real
         decoded_outputs = decode_data(emulator, encoded_outputs, "out")
-            
+
         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")
+            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")
+            for (i, mat) in enumerate(eachslice(encoded_uncertainties, dims = 3))
+                decoded_covariances[:, :, i] .= decode_structure_matrix(emulator, mat, "out")
             end
         end
 
         if output_dim > 1
-            return decoded_outputs, eachslice(decoded_covariances,dims=3)
+            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,:,:]
+            return decoded_outputs, decoded_covariances[1, :, :]
         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)
+
+        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
+            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    
+            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,:,:]
+            return encoded_outputs, encoded_covariances_mat[1, :, :]
         end
     end
-    
+
 end
 
 end
diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 9c130d684..7a2414e31 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -410,7 +410,7 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow
             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; obsdim=2)
+        opt_fx = opt_f(input_values', regularization_noise; obsdim = 2)
 
         data_i = output_values[i, :]
         opt_post_fx = posterior(opt_fx, data_i)
@@ -434,7 +434,7 @@ 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; obsdim=2)
+        pred = pred_gp(new_inputs; obsdim = 2)
         μ[i, :] = mean(pred)
         σ2[i, :] = var(pred)
     end
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index f3caa5e69..fb3048fdc 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -251,7 +251,7 @@ function emulator_log_density_model(
     # 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)
-    
+
     if isa(g_cov[1], Real)
 
         return 1.0 / length(obs_vec) * sum([logpdf(MvNormal(obs, g_cov[1] * I), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
@@ -551,11 +551,11 @@ function MCMCWrapper(
     if eltype(obs_slice) <: Number # just one observation
         obs_slice = [obs_slice]
     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, encoded_obs)
     mh_proposal_sampler = MetropolisHastingsSampler(mcmc_alg, prior)
@@ -737,7 +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 547857b2b..8a3321f17 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -363,7 +363,7 @@ function build_models!(
         n_cross_val_sets = 1 # now just pretend there is one partition for looping purposes
         n_train = n_data
         n_test = n_data
-   else
+    else
         train_fraction = optimizer_options["train_fraction"]
         n_train = Int(floor(train_fraction * n_data))
         n_test = n_data - n_train
@@ -394,8 +394,8 @@ 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
-        
+        regularization_i = regularization[i, i] * I
+
         io_pairs_opt = PairedDataContainer(input_values, reshape(output_values[i, :], 1, size(output_values, 2)))
 
         multithread = optimizer_options["multithread"]
@@ -604,14 +604,14 @@ function predict(
             tullio_threading = tullio_threading,
         )
     end
-    
+
     # 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 = 1:M
+    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 d3f0f6bcf..47afa03cf 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -120,8 +120,10 @@ abstract type MinMaxScaling <: UnivariateAffineScaling end
 abstract type ZScoreScaling <: UnivariateAffineScaling end
 
 # 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))
+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))
 
 # Processors
 """
@@ -447,11 +449,11 @@ function initialize_processor!(
             @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)
+            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
@@ -636,7 +638,7 @@ function initialize_processor!(
         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
@@ -652,7 +654,7 @@ function initialize_processor!(
             # 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
diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl
index 4a35cbd90..d8ff493d3 100644
--- a/src/VectorRandomFeature.jl
+++ b/src/VectorRandomFeature.jl
@@ -451,7 +451,7 @@ function build_models!(
     else
         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 a8fc26fc8..5c97c0033 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -35,10 +35,7 @@ struct MLTester <: Emulators.MachineLearningTool end
 
     # Unknown ML type
     mlt = MLTester()
-    @test_throws ErrorException emulator = Emulator(
-        mlt,
-        io_pairs,
-    )
+    @test_throws ErrorException emulator = Emulator(mlt, io_pairs)
 
     # test getters & defaults
     gp = GaussianProcess(GPJL())
@@ -46,70 +43,71 @@ struct MLTester <: Emulators.MachineLearningTool end
     @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 = encode_with_schedule!(enc_sch, io_pairs, 1.0*I(p), 1.0*I(d))
+    enc_sch = create_encoder_schedule(default_encoder)
+    enc_io_pairs, enc_I_in, enc_I_out = encode_with_schedule!(enc_sch, io_pairs, 1.0 * I(p), 1.0 * I(d))
     @test get_encoder_schedule(em)[1][1] == enc_sch[1][1] # inputs: proc
     @test get_encoder_schedule(em)[1][3] == enc_sch[1][3] # inputs: apply_to
     @test get_encoder_schedule(em)[2][1] == enc_sch[2][1] # outputs...
     @test get_encoder_schedule(em)[2][3] == enc_sch[2][3]
-    @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 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)
 
     #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_I = encode_structure_matrix(em, I(d), "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 isapprox(norm(encoded_I-enc_I_out), 0, atol = tol*d*d)
+    @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 isapprox(norm(encoded_I - enc_I_out), 0, atol = tol * d * d)
     decoded_dc = decode_data(em, encoded_dc, "in")
     decoded_mat = decode_data(em, encoded_mat, "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 isapprox(norm(decoded_I - I(d)), 0, atol = tol*d*d)
+    @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 isapprox(norm(decoded_I - I(d)), 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)
-    em1 = Emulator(gp, io_pairs, 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=Γ)
+    @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,
+    )
+    em1 = Emulator(gp, io_pairs, 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")])
-    (_, _, _) =
-        encode_with_schedule!(
-            enc_sch1,
-            io_pairs,
-            1.0*I(p),
-            Σ, # obs noise cov becomes the output structure matrix
-        )
+    (_, _, _) = 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
 
-    (_, _, _) =
-        encode_with_schedule!(
-            enc_sch2,
-            io_pairs,
-            4.0*I(p),
-            2.0*I(d), # out_struct_mat overrides obs noise cov
-        )
+    (_, _, _) = 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
-    
-    (_, _, _) =
-        encode_with_schedule!(
-            enc_sch3,
-            io_pairs,
-            Γ, 
-            I(d), 
-        )
+
+    (_, _, _) = encode_with_schedule!(enc_sch3, io_pairs, Γ, I(d))
     @test get_encoder_schedule(em3) == enc_sch3
-    
-    
-   
-end
 
+end
diff --git a/test/GaussianProcess/runtests.jl b/test/GaussianProcess/runtests.jl
index 4e8434256..13b43d9fc 100644
--- a/test/GaussianProcess/runtests.jl
+++ b/test/GaussianProcess/runtests.jl
@@ -63,23 +63,11 @@ using CalibrateEmulateSample.Utilities
     @test gp1.prediction_type == pred_type
     @test gp1.alg_reg_noise == 1e-4
 
-    em1 = Emulator(
-        gp1,
-        iopairs,
-        encoder_schedule = [],
-    )
+    em1 = Emulator(gp1, iopairs, encoder_schedule = [])
 
-    @test_logs (:warn,) Emulator(
-        gp1,
-        iopairs,
-        encoder_schedule = [],
-    ) # check that gp1 does not get more models added under second call
+    @test_logs (:warn,) Emulator(gp1, iopairs, encoder_schedule = []) # check that gp1 does not get more models added under second call
 
-    Emulator(
-        gp1,
-        iopairs,
-        encoder_schedule = [],
-    )
+    Emulator(gp1, iopairs, encoder_schedule = [])
     @test length(gp1.models) == 1
 
     Emulators.optimize_hyperparameters!(em1)
@@ -92,11 +80,7 @@ using CalibrateEmulateSample.Utilities
 
     # 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,
-        encoder_schedule = [],
-    )
+    @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
@@ -107,20 +91,10 @@ using CalibrateEmulateSample.Utilities
         "log_std_noise" => gp1_opt_params[end],
     )
 
-    em_agp_from_gp1 = Emulator(
-        agp,
-        iopairs,
-        encoder_schedule = [],
-        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,) Emulator(
-        agp,
-        iopairs,
-        encoder_schedule = [],
-        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)
@@ -137,11 +111,7 @@ using CalibrateEmulateSample.Utilities
 
     gp2 = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type)
 
-    em2 = Emulator(
-        gp2,
-        iopairs,
-        encoder_schedule = [],
-    )
+    em2 = Emulator(gp2, iopairs, encoder_schedule = [])
 
     Emulators.optimize_hyperparameters!(em2)
 
@@ -158,21 +128,9 @@ using CalibrateEmulateSample.Utilities
     GPkernel = var * se
 
     gp3 = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type)
-    em3 = Emulator(
-        gp3,
-        iopairs,
-        encoder_schedule = [],
-    )
-    @test_logs (:warn,) Emulator(
-        gp3,
-        iopairs,
-        encoder_schedule = [],
-    )
-    Emulator(
-        gp3,
-        iopairs,
-        encoder_schedule = [],
-    )
+    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)
@@ -215,23 +173,15 @@ using CalibrateEmulateSample.Utilities
     @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,
-        output_structure_matrix = Σ,
-    )
+    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,
-        output_structure_matrix = Σ,
-    )
+    em4_noise_from_Σ = Emulator(gp4, iopairs2, output_structure_matrix = Σ)
 
     Emulators.optimize_hyperparameters!(em4_noise_learnt)
     Emulators.optimize_hyperparameters!(em4_noise_from_Σ)
@@ -276,12 +226,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, 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 b44605a6c..27621f09d 100644
--- a/test/MarkovChainMonteCarlo/runtests.jl
+++ b/test/MarkovChainMonteCarlo/runtests.jl
@@ -70,12 +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; output_structure_matrix = σ2_y, encoder_schedule = [])
     Emulators.optimize_hyperparameters!(em)
     return em
 end
@@ -88,12 +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; output_structure_matrix = σ2_y, encoder_schedule = [])
     Emulators.optimize_hyperparameters!(em)
 
     # now make agp from gp
@@ -110,13 +100,8 @@ 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, output_structure_matrix = σ2_y, encoder_schedule = [], kernel_params = kernel_params)
 
     return em, em_agp
 end
@@ -130,17 +115,9 @@ function test_gp_2(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)
-    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
-    )
+    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
diff --git a/test/RandomFeature/runtests.jl b/test/RandomFeature/runtests.jl
index d170eab28..1201faa8e 100644
--- a/test/RandomFeature/runtests.jl
+++ b/test/RandomFeature/runtests.jl
@@ -270,7 +270,7 @@ rng = Random.MersenneTwister(seed)
         output_dim = 1
         n = 50                                       # number of training points
         x = reshape(2.0 * π * rand(n), 1, n)         # unif(0,2π) predictors/features: 1 x n
-        obs_noise_cov= 0.05^2 * I
+        obs_noise_cov = 0.05^2 * I
         y = reshape(sin.(x) + 0.05 * randn(n)', 1, n) # predictands/targets: 1 x n
         iopairs = PairedDataContainer(x, y, data_are_columns = true)
 
@@ -304,7 +304,7 @@ rng = Random.MersenneTwister(seed)
             rng = rng,
             optimizer_options = Dict("n_cross_val_sets" => 0),
         )
-        
+
         # build emulators
         em_srfi = Emulator(srfi, iopairs, output_structure_matrix = obs_noise_cov)
         n_srfi = length(get_rfms(srfi))
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 74cffceff..eeb84d297 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -108,15 +108,15 @@ end
 
     cc = canonical_correlation()
     @test get_retain_var(cc) == 1.0
-    cc2 = canonical_correlation(retain_var=0.7)
+    cc2 = canonical_correlation(retain_var = 0.7)
     @test get_retain_var(cc2) == 0.7
-    cc3 = CanonicalCorrelation([1],[2],[3],1.0,"test")
+    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()
@@ -124,9 +124,9 @@ end
     dd_copy = decorrelate()
     @test cc == cc_copy
     @test dd == dd_copy
-    
+
     # get some data as IO pairs for functional tests
-    
+
     in_dim = 10
     out_dim = 50
     samples = 120

From 333609fff064fc10f7e9851a6bb3190e275f4118 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 13:53:38 -0700
Subject: [PATCH 50/75] missed by foramtter

---
 .../Regression_2d_2d/compare_regression.jl    | 20 ++-----------------
 1 file changed, 2 insertions(+), 18 deletions(-)

diff --git a/examples/Emulator/Regression_2d_2d/compare_regression.jl b/examples/Emulator/Regression_2d_2d/compare_regression.jl
index 3b8b11f8b..fd55ca02d 100644
--- a/examples/Emulator/Regression_2d_2d/compare_regression.jl
+++ b/examples/Emulator/Regression_2d_2d/compare_regression.jl
@@ -96,28 +96,12 @@ function main()
         x1 = range(-2 * pi, stop = 2 * π, length = n_pts)
         x2 = range(-2 * pi, stop = 2 * π, length = n_pts)
 
-        p1 = contourf(
-            x1,
-            x2,
-            g1.(x1', x2),
-            c = :cividis,
-            xlabel = "x1",
-            ylabel = "x2",
-            aspect_ratio = :equal,
-        )
+        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")
 
         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,
-        )
+        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")
 
         figpath = joinpath(output_directory, "g2_true.png")

From a5502027876f193e7cf31090ffeb73f1e6db3f5c Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 14:23:09 -0700
Subject: [PATCH 51/75] docs

---
 docs/src/API/Emulators.md             | 9 ++++-----
 docs/src/API/MarkovChainMonteCarlo.md | 6 ------
 docs/src/API/RandomFeatures.md        | 4 ++--
 3 files changed, 6 insertions(+), 13 deletions(-)

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

From 0cff1e592e67cfb46a59270025e45e2e6fa7f334 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 4 Jul 2025 14:38:09 -0700
Subject: [PATCH 52/75] docs

---
 docs/src/API/GaussianProcess.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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

From e258bfde7e0bf6396ef56c106f2950a7c585831c Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 8 Jul 2025 13:57:10 -0700
Subject: [PATCH 53/75] docstrings and comments addressed

---
 src/Emulator.jl              | 76 +++++++++++++++++++++++++++++-------
 src/MarkovChainMonteCarlo.jl |  2 +-
 src/Utilities.jl             | 42 --------------------
 test/Utilities/runtests.jl   | 28 -------------
 4 files changed, 63 insertions(+), 85 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 66f67cbf2..cdd098580 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -23,7 +23,7 @@ export optimize_hyperparameters!
 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:
 
@@ -54,12 +54,12 @@ 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, VV <: AbstractVector}
     "Machine learning tool, defined as a struct of type MachineLearningTool."
@@ -72,20 +72,49 @@ struct Emulator{FT <: AbstractFloat, VV <: AbstractVector}
     encoder_schedule::VV
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Gets the `machine_learning_tool` field of the `Emulator`
+"""
 get_machine_learning_tool(emulator::Emulator) = emulator.machine_learning_tool
+
+"""
+$(TYPEDSIGNATURES)
+
+Gets the `io_pairs` field of the `Emulator`
+"""
 get_io_pairs(emulator::Emulator) = emulator.io_pairs
+
+"""
+$(TYPEDSIGNATURES)
+
+Gets the `encoded_io_pairs` field of the `Emulator`
+"""
 get_encoded_io_pairs(emulator::Emulator) = emulator.encoded_io_pairs
+
+"""
+$(TYPEDSIGNATURES)
+
+Gets the `encoder_schedul` field of the `Emulator`
+"""
 get_encoder_schedule(emulator::Emulator) = emulator.encoder_schedule
 
 # Constructor for the Emulator Object
 """
-$(DocStringExtensions.TYPEDSIGNATURES)
+$(TYPEDSIGNATURES)
+
+Constructor of the Emulator object,
 
 Positional Arguments
- - `machine_learning_tool` ::MachineLearningTool,
- - `input_output_pairs` ::PairedDataContainer
+ - `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`
+
 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. `nothing` sets a default `I(input_dim)`
+ - `output_structure_matrix` [=`nothing`] Some encoders make use of an input structure (e.g., the prior covariance matrix). Particularly useful for few samples. `nothing` sets a default `I(input_dim)`
+Other keywords are passed to the machine learning tool initialization
 """
 function Emulator(
     machine_learning_tool::MachineLearningTool,
@@ -114,7 +143,7 @@ function Emulator(
     input_structure_mat = if isnothing(input_structure_matrix)
         Diagonal(FT.(ones(input_dim)))
     elseif isa(input_structure_matrix, UniformScaling)
-        Diagonal(input_structure_matrix(input_dim))
+       input_structure_matrix(input_dim)
     else
         input_structure_matrix
     end
@@ -122,7 +151,7 @@ function Emulator(
     output_structure_mat = if isnothing(output_structure_matrix)
         Diagonal(FT.(ones(output_dim)))
     elseif isa(output_structure_matrix, UniformScaling)
-        Diagonal(output_structure_matrix(output_dim))
+        output_structure_matrix(output_dim)
     else
         output_structure_matrix
     end
@@ -159,15 +188,19 @@ function Emulator(
 end
 
 """
-$(DocStringExtensions.TYPEDSIGNATURES)
+$(TYPEDSIGNATURES)
 
-Optimizes the hyperparameters in the machine learning tool.
+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 optimize_hyperparameters!(emulator::Emulator{FT}, args...; kwargs...) where {FT <: AbstractFloat}
     optimize_hyperparameters!(emulator.machine_learning_tool, args...; kwargs...)
 end
 
+"""
+$(TYPEDSIGNATURES)
 
+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 encode_data(
     emulator::Emulator,
     data::MorDC,
@@ -180,6 +213,11 @@ function encode_data(
     end
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Encode a new structure matrix in the input space (`"in"`) or output space (`"out"`). with the stored and initialized encoder schedule. 
+"""
 function encode_structure_matrix(
     emulator::Emulator,
     structure_mat::USorM,
@@ -189,6 +227,11 @@ function encode_structure_matrix(
 end
 
 
+"""
+$(TYPEDSIGNATURES)
+
+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 decode_data(
     emulator::Emulator,
     data::MorDC,
@@ -201,6 +244,11 @@ function decode_data(
     end
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Decode a new structure matrix in the input space (`"in"`) or output space (`"out"`). with the stored and initialized encoder schedule. 
+"""
 function decode_structure_matrix(
     emulator::Emulator,
     structure_mat::USorM,
@@ -210,12 +258,12 @@ function decode_structure_matrix(
 end
 
 """
-$(DocStringExtensions.TYPEDSIGNATURES)
+$(TYPEDSIGNATURES)
 
 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:
+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] 
 """
@@ -233,7 +281,7 @@ function predict(
 
     N_samples = size(new_inputs, 2)
 
-    if !(size(new_inputs, 1) == input_dim)
+    if 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))",
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index fb3048fdc..ce1a9996a 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -495,7 +495,7 @@ AbstractMCMC's terminology).
 # Fields
 $(DocStringExtensions.TYPEDFIELDS)
 """
-struct MCMCWrapper{AVV, 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
     "[output_dim x N_samples] matrix, of given observation data."
diff --git a/src/Utilities.jl b/src/Utilities.jl
index 47afa03cf..e222fdb4b 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -67,48 +67,6 @@ 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 
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index eeb84d297..9e560ac6b 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

From 955bdae3bd3e4af7c35f234a56268f32bfc8400b Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 8 Jul 2025 14:27:29 -0700
Subject: [PATCH 54/75] format

---
 src/Emulator.jl              | 2 +-
 src/MarkovChainMonteCarlo.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index cdd098580..652f984f8 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -143,7 +143,7 @@ function Emulator(
     input_structure_mat = if isnothing(input_structure_matrix)
         Diagonal(FT.(ones(input_dim)))
     elseif isa(input_structure_matrix, UniformScaling)
-       input_structure_matrix(input_dim)
+        input_structure_matrix(input_dim)
     else
         input_structure_matrix
     end
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index ce1a9996a..6091a88f6 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -495,7 +495,7 @@ AbstractMCMC's terminology).
 # Fields
 $(DocStringExtensions.TYPEDFIELDS)
 """
-struct MCMCWrapper{AVV <: AbstractVector , 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
     "[output_dim x N_samples] matrix, of given observation data."

From bffca73d73ea88017a912ee014b31596d78a8611 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 9 Jul 2025 10:45:08 -0700
Subject: [PATCH 55/75] add docs page

---
 docs/src/data_processing.md | 69 +++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)
 create mode 100644 docs/src/data_processing.md

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

From 7ad963ca340c2f0388a158405816ad249c18de07 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 9 Jul 2025 11:23:16 -0700
Subject: [PATCH 56/75] add docs into index

---
 docs/make.jl | 1 +
 1 file changed, 1 insertion(+)

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

From 5c176e59c01fbd7911cfd9fd81b6a13112df2ed1 Mon Sep 17 00:00:00 2001
From: Arne Bouillon <arne.bouillon@kuleuven.be>
Date: Mon, 7 Jul 2025 18:15:55 -0700
Subject: [PATCH 57/75] Reduce number of methods; give data processors separate
 files

---
 src/Utilities.jl                       | 841 +++----------------------
 src/Utilities/canonical_correlation.jl | 218 +++++++
 src/Utilities/decorrelator.jl          | 216 +++++++
 src/Utilities/elementwise_scaler.jl    | 179 ++++++
 test/Utilities/runtests.jl             |   3 +-
 5 files changed, 689 insertions(+), 768 deletions(-)
 create mode 100644 src/Utilities/canonical_correlation.jl
 create mode 100644 src/Utilities/decorrelator.jl
 create mode 100644 src/Utilities/elementwise_scaler.jl

diff --git a/src/Utilities.jl b/src/Utilities.jl
index e222fdb4b..1272b21ae 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -12,22 +12,9 @@ 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, 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
 
 
 
@@ -67,15 +54,12 @@ function get_training_points(
     return training_points
 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
+# Data processing tooling:
 
-abstract type UnivariateAffineScaling end
-abstract type QuartileScaling <: UnivariateAffineScaling end
-abstract type MinMaxScaling <: UnivariateAffineScaling end
-abstract type ZScoreScaling <: UnivariateAffineScaling end
+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
 
 # define how to have equality
 Base.:(==)(a::DCP, b::DCP) where {DCP <: DataContainerProcessor} =
@@ -83,693 +67,48 @@ Base.:(==)(a::DCP, b::DCP) where {DCP <: DataContainerProcessor} =
 Base.:(==)(a::PDCP, b::PDCP) where {PDCP <: PairedDataContainerProcessor} =
     all(getfield(a, f) == getfield(b, f) for f in fieldnames(PDCP))
 
-# 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) / 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 `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
-
-"""
-$(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"]
+function _encode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
+    input_data, output_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
-
-    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,
-    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::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,
-    data,
-    structure_mat::USorM,
-    apply_to::AS,
-) where {PDCP <: PairedDataContainerProcessor, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString}
+function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
     input_data, output_data = data
-    initialize_processor!(dcp, input_data, output_data, structure_mat, apply_to)
 
     if apply_to == "in"
-        return encode_data(dcp, input_data)
+        return decode_data(proc, input_data)
     elseif apply_to == "out"
-        return encode_data(dcp, output_data)
+        return decode_data(proc, output_data)
     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}
+function _initialize_and_encode_data!(proc::P, data, structure_mats, apply_to::AS) where {P <: DataProcessor, 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)
+    input_structure_mat, output_structure_mat = structure_mats
+
+    if P isa PairedDataContainerProcessor
+        initialize_processor!(proc, input_data, output_data, apply_to == "in" ? input_structure_mat : output_structure_mat, apply_to)
     else
-        bad_apply_to(apply_to)
+        if apply_to == "in"
+            initialize_processor!(proc, input_data, input_structure_mat)
+        else
+            initialize_processor!(proc, output_data, output_structure_mat)
+        end
     end
+
+    _encode_data(proc, data, apply_to)
 end
 
 """
@@ -799,44 +138,18 @@ enc_schedule = [
 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
 
@@ -848,21 +161,13 @@ 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,
@@ -878,21 +183,16 @@ function encode_with_schedule!(
     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
@@ -900,6 +200,7 @@ function encode_with_schedule!(
     return processed_io_pairs, processed_input_structure_mat, processed_output_structure_mat
 end
 
+
 # Functions to encode/decode with initialized schedule
 """
 $TYPEDSIGNATURES
@@ -911,33 +212,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
 
@@ -948,14 +238,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"])
         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
@@ -980,15 +269,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)
@@ -1014,7 +302,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
@@ -1022,11 +309,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
 
@@ -1043,7 +329,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
@@ -1051,7 +336,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
@@ -1061,5 +346,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..5ee1426c0
--- /dev/null
+++ b/src/Utilities/canonical_correlation.jl
@@ -0,0 +1,218 @@
+# 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,
+    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::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
diff --git a/src/Utilities/decorrelator.jl b/src/Utilities/decorrelator.jl
new file mode 100644
index 000000000..69906a9b6
--- /dev/null
+++ b/src/Utilities/decorrelator.jl
@@ -0,0 +1,216 @@
+# 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::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) / 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::MM) where {MM <: 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::MM) where {MM <: 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..c8b84454b
--- /dev/null
+++ b/src/Utilities/elementwise_scaler.jl
@@ -0,0 +1,179 @@
+# 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::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
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 9e560ac6b..aebcc63b6 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -72,7 +72,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 = 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]
@@ -297,7 +297,6 @@ end
     bad_encoder_schedule = [(canonical_correlation(), func, "bad")]
     @test_throws ArgumentError 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")
 
 

From 2110fea9ff42ee59a9d7b6c67093b2022b8fd9bb Mon Sep 17 00:00:00 2001
From: Arne Bouillon <arne.bouillon@kuleuven.be>
Date: Mon, 7 Jul 2025 18:19:00 -0700
Subject: [PATCH 58/75] Pass both structure matrices to PDCPs

---
 src/Utilities.jl                       | 12 +++++++-----
 src/Utilities/canonical_correlation.jl |  3 ++-
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 1272b21ae..7156af1cb 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -95,16 +95,18 @@ function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS
 end
 
 function _initialize_and_encode_data!(proc::P, data, structure_mats, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
-    input_data, output_data = data
-    input_structure_mat, output_structure_mat = structure_mats
-
     if P isa PairedDataContainerProcessor
-        initialize_processor!(proc, input_data, output_data, apply_to == "in" ? input_structure_mat : output_structure_mat, apply_to)
+        initialize_processor!(proc, data..., structure_mats..., apply_to)
     else
+        input_data, output_data = data
+        input_structure_mat, output_structure_mat = structure_mats
+
         if apply_to == "in"
             initialize_processor!(proc, input_data, input_structure_mat)
-        else
+        elseif apply_to == "out"
             initialize_processor!(proc, output_data, output_structure_mat)
+        else
+            bad_apply_to(apply_to)
         end
     end
 
diff --git a/src/Utilities/canonical_correlation.jl b/src/Utilities/canonical_correlation.jl
index 5ee1426c0..de08b4db2 100644
--- a/src/Utilities/canonical_correlation.jl
+++ b/src/Utilities/canonical_correlation.jl
@@ -170,7 +170,8 @@ initialize_processor!(
     cc::CanonicalCorrelation,
     in_data::MM,
     out_data::MM,
-    structure_matrix,
+    input_structure_matrix,
+    output_structure_matrix,
     apply_to::AS,
 ) where {MM <: AbstractMatrix, AS <: AbstractString} = initialize_processor!(cc, in_data, out_data, apply_to)
 

From 01112fea757624da0018740f39ea10ec41160974 Mon Sep 17 00:00:00 2001
From: Arne Bouillon <arne.bouillon@kuleuven.be>
Date: Mon, 7 Jul 2025 18:31:33 -0700
Subject: [PATCH 59/75] Format code

---
 src/Utilities.jl | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 7156af1cb..53a23cfc7 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -94,7 +94,12 @@ function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS
     end
 end
 
-function _initialize_and_encode_data!(proc::P, data, structure_mats, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
+function _initialize_and_encode_data!(
+    proc::P,
+    data,
+    structure_mats,
+    apply_to::AS,
+) where {P <: DataProcessor, AS <: AbstractString}
     if P isa PairedDataContainerProcessor
         initialize_processor!(proc, data..., structure_mats..., apply_to)
     else
@@ -188,7 +193,12 @@ function initialize_and_encode_with_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)
+        processed = _initialize_and_encode_data!(
+            processor,
+            processed_io_pairs,
+            (processed_input_structure_mat, processed_output_structure_mat),
+            apply_to,
+        )
 
         if apply_to == "in"
             processed_input_structure_mat = encode_structure_matrix(processor, processed_input_structure_mat)

From 90e26e30b44c250e58c60657ba8c7c39846c5660 Mon Sep 17 00:00:00 2001
From: Arne Bouillon <arne.bouillon@kuleuven.be>
Date: Wed, 9 Jul 2025 11:33:59 -0700
Subject: [PATCH 60/75] Fix bugs

---
 src/Utilities.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 53a23cfc7..a6cfc1657 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -72,7 +72,7 @@ Base.:(==)(a::PDCP, b::PDCP) where {PDCP <: PairedDataContainerProcessor} =
 
 
 function _encode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
-    input_data, output_data = data
+    input_data, output_data = get_data(data)
     if apply_to == "in"
         return encode_data(proc, input_data)
     elseif apply_to == "out"
@@ -83,7 +83,7 @@ function _encode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS
 end
 
 function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS <: AbstractString}
-    input_data, output_data = data
+    input_data, output_data = get_data(data)
 
     if apply_to == "in"
         return decode_data(proc, input_data)
@@ -100,10 +100,10 @@ function _initialize_and_encode_data!(
     structure_mats,
     apply_to::AS,
 ) where {P <: DataProcessor, AS <: AbstractString}
-    if P isa PairedDataContainerProcessor
+    if proc isa PairedDataContainerProcessor
         initialize_processor!(proc, data..., structure_mats..., apply_to)
     else
-        input_data, output_data = data
+        input_data, output_data = get_data(data)
         input_structure_mat, output_structure_mat = structure_mats
 
         if apply_to == "in"

From 8d33f8fd31812c0060974b4ec5f9e1f6153225cf Mon Sep 17 00:00:00 2001
From: Arne Bouillon <arne.bouillon@kuleuven.be>
Date: Wed, 9 Jul 2025 11:34:26 -0700
Subject: [PATCH 61/75] Adapt tests to new functions

---
 src/Emulator.jl            |  7 ++++++-
 src/Utilities.jl           | 13 ++++++++++---
 test/Emulator/runtests.jl  | 12 ++++++------
 test/Utilities/runtests.jl | 15 +++++++++------
 4 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 652f984f8..288ed472f 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -174,7 +174,12 @@ function Emulator(
 
     enc_schedule = create_encoder_schedule(encoder_schedule)
     (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
-        encode_with_schedule!(enc_schedule, input_output_pairs, input_structure_mat, output_structure_mat)
+        initialize_and_encode_with_schedule!(
+            enc_schedule,
+            input_output_pairs,
+            input_structure_mat,
+            output_structure_mat,
+        )
 
     # build the machine learning tool in the encoded space
     build_models!(
diff --git a/src/Utilities.jl b/src/Utilities.jl
index a6cfc1657..f8c1ae154 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -14,7 +14,14 @@ using ..DataContainers
 export get_training_points
 
 export PairedDataContainerProcessor, DataContainerProcessor
-export create_encoder_schedule, initialize_and_encode_with_schedule!, encode_with_schedule, decode_with_schedule
+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
 
 
 
@@ -101,7 +108,7 @@ function _initialize_and_encode_data!(
     apply_to::AS,
 ) where {P <: DataProcessor, AS <: AbstractString}
     if proc isa PairedDataContainerProcessor
-        initialize_processor!(proc, data..., structure_mats..., apply_to)
+        initialize_processor!(proc, get_data(data)..., structure_mats..., apply_to)
     else
         input_data, output_data = get_data(data)
         input_structure_mat, output_structure_mat = structure_mats
@@ -250,7 +257,7 @@ 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)
diff --git a/test/Emulator/runtests.jl b/test/Emulator/runtests.jl
index 5c97c0033..5fae44e86 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -44,11 +44,11 @@ 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 = encode_with_schedule!(enc_sch, io_pairs, 1.0 * I(p), 1.0 * I(d))
+    enc_io_pairs, enc_I_in, enc_I_out = initialize_and_encode_with_schedule!(enc_sch, io_pairs, 1.0 * I(p), 1.0 * I(d))
     @test get_encoder_schedule(em)[1][1] == enc_sch[1][1] # inputs: proc
-    @test get_encoder_schedule(em)[1][3] == enc_sch[1][3] # inputs: apply_to
+    @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][3] == enc_sch[2][3]
+    @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)
 
@@ -91,7 +91,7 @@ struct MLTester <: Emulators.MachineLearningTool end
     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")])
-    (_, _, _) = encode_with_schedule!(
+    (_, _, _) = initialize_and_encode_with_schedule!(
         enc_sch1,
         io_pairs,
         1.0 * I(p),
@@ -99,7 +99,7 @@ struct MLTester <: Emulators.MachineLearningTool end
     )
     @test get_encoder_schedule(em1) == enc_sch1
 
-    (_, _, _) = encode_with_schedule!(
+    (_, _, _) = initialize_and_encode_with_schedule!(
         enc_sch2,
         io_pairs,
         4.0 * I(p),
@@ -107,7 +107,7 @@ struct MLTester <: Emulators.MachineLearningTool end
     )
     @test get_encoder_schedule(em2) == enc_sch2
 
-    (_, _, _) = encode_with_schedule!(enc_sch3, io_pairs, Γ, I(d))
+    (_, _, _) = initialize_and_encode_with_schedule!(enc_sch3, io_pairs, Γ, I(d))
     @test get_encoder_schedule(em3) == enc_sch3
 
 end
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index aebcc63b6..1d5250197 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -143,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)
@@ -293,18 +293,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 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) =

From c16da9ea9e6306d8371181db90e3a2491ca912e6 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 9 Jul 2025 15:58:49 -0700
Subject: [PATCH 62/75] replace logic for input obs

---
 src/MarkovChainMonteCarlo.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index 6091a88f6..c5f9ba6a4 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -547,13 +547,13 @@ function MCMCWrapper(
 ) where {AV <: AbstractVector, AMorAV <: Union{AbstractVector, AbstractMatrix}}
 
     # make into iterable over vectors
-    obs_slice = isa(observation, AbstractMatrix) ? eachcol(observation) : observation
-    if eltype(obs_slice) <: Number # just one observation
-        obs_slice = [obs_slice]
-    end
+    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]
 
 

From fbf72491b0c3b17b2f28e721fa442a74908ea801 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 11:35:00 -0700
Subject: [PATCH 63/75] format and emulate docs page

---
 docs/src/emulate.md               | 37 +++----------------------------
 examples/Lorenz/emulate_sample.jl |  8 +------
 src/MarkovChainMonteCarlo.jl      |  2 +-
 3 files changed, 5 insertions(+), 42 deletions(-)

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/Lorenz/emulate_sample.jl b/examples/Lorenz/emulate_sample.jl
index 338620169..b38854360 100644
--- a/examples/Lorenz/emulate_sample.jl
+++ b/examples/Lorenz/emulate_sample.jl
@@ -169,13 +169,7 @@ 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,
-            obs_noise_cov = Γy,
-        )
+        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
diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
index c5f9ba6a4..31b10ab43 100644
--- a/src/MarkovChainMonteCarlo.jl
+++ b/src/MarkovChainMonteCarlo.jl
@@ -551,7 +551,7 @@ function MCMCWrapper(
         observation
     else # NB a vector is treated as a column here:
         eachcol(observation)
-    end    
+    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]

From 966bebd981dc255071d83ea545f1b95afd36c390 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 12:06:50 -0700
Subject: [PATCH 64/75] n_iter scalar

---
 examples/Cloudy/Cloudy_calibrate.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/examples/Cloudy/Cloudy_calibrate.jl b/examples/Cloudy/Cloudy_calibrate.jl
index 70fa8a2b2..620901856 100644
--- a/examples/Cloudy/Cloudy_calibrate.jl
+++ b/examples/Cloudy/Cloudy_calibrate.jl
@@ -160,7 +160,7 @@ dummy = ones(n_params)
 dist_type = ParticleDistributions.GammaPrimitiveParticleDistribution(dummy...)
 model_settings = DynamicalModel.ModelSettings(kernel, dist_type, moments, tspan)
 # EKI iterations
-n_iter = [0]
+n_iter = 0
 for n in 1:N_iter
     # Return transformed parameters in physical/constrained space
     ϕ_n = get_ϕ_final(priors, ekiobj)
@@ -169,7 +169,7 @@ for n in 1:N_iter
     G_ens = hcat(G_n...)  # reformat
     terminate = EnsembleKalmanProcesses.update_ensemble!(ekiobj, G_ens)
     if !isnothing(terminate)
-        n_iter[1] = n - 1
+        n_iter = n - 1
         break
     end
 
@@ -208,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] - 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, :]))
@@ -265,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] - 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, :]))

From 15959c01845e94124f0f89da719efed53a4eebf1 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 12:14:16 -0700
Subject: [PATCH 65/75] n_iter -> scalar

---
 examples/Cloudy/Cloudy_calibrate.jl      |  2 +-
 examples/Darcy/calibrate.jl              |  6 +++---
 examples/Lorenz/calibrate.jl             |  8 ++++----
 examples/Lorenz/calibrate_spatial_dep.jl |  2 +-
 examples/Sinusoid/calibrate.jl           | 13 +++++++++----
 5 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/examples/Cloudy/Cloudy_calibrate.jl b/examples/Cloudy/Cloudy_calibrate.jl
index 620901856..1b4fd34b5 100644
--- a/examples/Cloudy/Cloudy_calibrate.jl
+++ b/examples/Cloudy/Cloudy_calibrate.jl
@@ -160,7 +160,7 @@ dummy = ones(n_params)
 dist_type = ParticleDistributions.GammaPrimitiveParticleDistribution(dummy...)
 model_settings = DynamicalModel.ModelSettings(kernel, dist_type, moments, tspan)
 # EKI iterations
-n_iter = 0
+n_iter = N_iter
 for n in 1:N_iter
     # Return transformed parameters in physical/constrained space
     ϕ_n = get_ϕ_final(priors, ekiobj)
diff --git a/examples/Darcy/calibrate.jl b/examples/Darcy/calibrate.jl
index 62ceb1df9..b807a6766 100644
--- a/examples/Darcy/calibrate.jl
+++ b/examples/Darcy/calibrate.jl
@@ -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/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 82fab7267..25152046c 100644
--- a/examples/Lorenz/calibrate_spatial_dep.jl
+++ b/examples/Lorenz/calibrate_spatial_dep.jl
@@ -238,7 +238,7 @@ method = Inversion(prior)
 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
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,
 )

From b0911e47cabc648b97230e89c8e6889800888ba7 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 14:34:41 -0700
Subject: [PATCH 66/75] remove default structure matrix, extend framework of
 encoding a nothing, add tests

---
 src/Emulator.jl                        | 38 ++++---------
 src/Utilities.jl                       | 74 ++++++++++++++++----------
 src/Utilities/canonical_correlation.jl | 10 +++-
 src/Utilities/decorrelator.jl          | 30 ++++++++---
 src/Utilities/elementwise_scaler.jl    | 10 +++-
 test/Emulator/runtests.jl              | 14 +++--
 test/Utilities/runtests.jl             |  8 +++
 7 files changed, 116 insertions(+), 68 deletions(-)

diff --git a/src/Emulator.jl b/src/Emulator.jl
index 288ed472f..1c836a295 100644
--- a/src/Emulator.jl
+++ b/src/Emulator.jl
@@ -112,8 +112,8 @@ Positional Arguments
 
 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. `nothing` sets a default `I(input_dim)`
- - `output_structure_matrix` [=`nothing`] Some encoders make use of an input structure (e.g., the prior covariance matrix). Particularly useful for few samples. `nothing` sets a default `I(input_dim)`
+ - `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 Emulator(
@@ -140,22 +140,6 @@ function Emulator(
         )
     end
 
-    input_structure_mat = if isnothing(input_structure_matrix)
-        Diagonal(FT.(ones(input_dim)))
-    elseif isa(input_structure_matrix, UniformScaling)
-        input_structure_matrix(input_dim)
-    else
-        input_structure_matrix
-    end
-
-    output_structure_mat = if isnothing(output_structure_matrix)
-        Diagonal(FT.(ones(output_dim)))
-    elseif isa(output_structure_matrix, UniformScaling)
-        output_structure_matrix(output_dim)
-    else
-        output_structure_matrix
-    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)
@@ -173,20 +157,20 @@ function Emulator(
     end
 
     enc_schedule = create_encoder_schedule(encoder_schedule)
-    (encoded_io_pairs, encoded_input_structure_mat, encoded_output_structure_mat) =
+    (encoded_io_pairs, encoded_input_structure_matrix, encoded_output_structure_matrix) =
         initialize_and_encode_with_schedule!(
             enc_schedule,
             input_output_pairs,
-            input_structure_mat,
-            output_structure_mat,
+            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_mat,
-        encoded_output_structure_mat;
+        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)
@@ -225,9 +209,9 @@ Encode a new structure matrix in the input space (`"in"`) or output space (`"out
 """
 function encode_structure_matrix(
     emulator::Emulator,
-    structure_mat::USorM,
+    structure_mat::USorMorN,
     in_or_out::AS,
-) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}}
     return encode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
@@ -256,9 +240,9 @@ Decode a new structure matrix in the input space (`"in"`) or output space (`"out
 """
 function decode_structure_matrix(
     emulator::Emulator,
-    structure_mat::USorM,
+    structure_mat::USorMorN,
     in_or_out::AS,
-) where {AS <: AbstractString, USorM <: Union{UniformScaling, AbstractMatrix}}
+) where {AS <: AbstractString, USorMorN <: Union{UniformScaling, AbstractMatrix, Nothing}}
     return decode_with_schedule(get_encoder_schedule(emulator), structure_mat, in_or_out)
 end
 
diff --git a/src/Utilities.jl b/src/Utilities.jl
index f8c1ae154..fa36c5988 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -102,27 +102,33 @@ function _decode_data(proc::P, data, apply_to::AS) where {P <: DataProcessor, AS
 end
 
 function _initialize_and_encode_data!(
-    proc::P,
+    proc::PairedDataContainerProcessor,
     data,
     structure_mats,
     apply_to::AS,
-) where {P <: DataProcessor, AS <: AbstractString}
-    if proc isa PairedDataContainerProcessor
-        initialize_processor!(proc, get_data(data)..., structure_mats..., apply_to)
-    else
-        input_data, output_data = get_data(data)
-        input_structure_mat, output_structure_mat = structure_mats
+) where {AS <: AbstractString}
+    initialize_processor!(proc, get_data(data)..., structure_mats..., apply_to)
+    return _encode_data(proc, data, apply_to)
+end
 
-        if apply_to == "in"
-            initialize_processor!(proc, input_data, input_structure_mat)
-        elseif apply_to == "out"
-            initialize_processor!(proc, output_data, output_structure_mat)
-        else
-            bad_apply_to(apply_to)
-        end
+function _initialize_and_encode_data!(
+    proc::DataContainerProcessor,
+    data,
+    structure_mats,
+    apply_to::AS,
+) where {AS <: AbstractString}
+    input_data, output_data = get_data(data)
+    input_structure_mat, output_structure_mat = structure_mats
+
+    if apply_to == "in"
+        initialize_processor!(proc, input_data, input_structure_mat)
+    elseif apply_to == "out"
+        initialize_processor!(proc, output_data, output_structure_mat)
+    else
+        bad_apply_to(apply_to)
     end
 
-    _encode_data(proc, data, apply_to)
+    return _encode_data(proc, data, apply_to)
 end
 
 """
@@ -138,15 +144,15 @@ 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)
@@ -184,13 +190,13 @@ 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::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)
@@ -219,8 +225,22 @@ function initialize_and_encode_with_schedule!(
     return processed_io_pairs, processed_input_structure_mat, processed_output_structure_mat
 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 de08b4db2..e8be4a611 100644
--- a/src/Utilities/canonical_correlation.jl
+++ b/src/Utilities/canonical_correlation.jl
@@ -203,7 +203,10 @@ $(TYPEDSIGNATURES)
 
 Apply the `CanonicalCorrelation` encoder to a provided structure matrix
 """
-function encode_structure_matrix(cc::CanonicalCorrelation, structure_matrix::MM) where {MM <: AbstractMatrix}
+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
@@ -213,7 +216,10 @@ $(TYPEDSIGNATURES)
 
 Apply the `CanonicalCorrelation` decoder to a provided structure matrix
 """
-function decode_structure_matrix(cc::CanonicalCorrelation, enc_structure_matrix::MM) where {MM <: AbstractMatrix}
+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
index 69906a9b6..62b51da1a 100644
--- a/src/Utilities/decorrelator.jl
+++ b/src/Utilities/decorrelator.jl
@@ -121,8 +121,8 @@ Computes and populates the `data_mean` and `encoder_mat` and `decoder_mat` field
 function initialize_processor!(
     dd::Decorrelator,
     data::MM,
-    structure_matrix::USorM,
-) where {MM <: AbstractMatrix, USorM <: Union{UniformScaling, AbstractMatrix}}
+    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
@@ -132,8 +132,20 @@ function initialize_processor!(
         # 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)
+            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)
@@ -200,7 +212,10 @@ $(TYPEDSIGNATURES)
 
 Apply the `Decorrelator` encoder to a provided structure matrix
 """
-function encode_structure_matrix(dd::Decorrelator, structure_matrix::MM) where {MM <: AbstractMatrix}
+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
@@ -210,7 +225,10 @@ $(TYPEDSIGNATURES)
 
 Apply the `Decorrelator` decoder to a provided structure matrix
 """
-function decode_structure_matrix(dd::Decorrelator, enc_structure_matrix::MM) where {MM <: AbstractMatrix}
+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
index c8b84454b..e139be4b3 100644
--- a/src/Utilities/elementwise_scaler.jl
+++ b/src/Utilities/elementwise_scaler.jl
@@ -165,7 +165,10 @@ $(TYPEDSIGNATURES)
 
 Apply the `ElementwiseScaler` encoder to a provided structure matrix
 """
-function encode_structure_matrix(es::ElementwiseScaler, structure_matrix::MM) where {MM <: AbstractMatrix}
+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
 
@@ -174,6 +177,9 @@ $(TYPEDSIGNATURES)
 
 Apply the `ElementwiseScaler` decoder to a provided structure matrix
 """
-function decode_structure_matrix(es::ElementwiseScaler, enc_structure_matrix::MM) where {MM <: AbstractMatrix}
+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/test/Emulator/runtests.jl b/test/Emulator/runtests.jl
index 5fae44e86..503aedd06 100644
--- a/test/Emulator/runtests.jl
+++ b/test/Emulator/runtests.jl
@@ -44,28 +44,34 @@ 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, 1.0 * I(p), 1.0 * I(d))
+    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_I = encode_structure_matrix(em, I(d), "out")
+    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 isapprox(norm(encoded_I - enc_I_out), 0, atol = tol * d * d)
+    @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 isapprox(norm(decoded_I - I(d)), 0, atol = tol * d * d)
+    @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 = Σ)
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 1d5250197..6871b9608 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -272,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)

From 40ed75f75254ef52cb2122f34a010c0e377a8cf4 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 14:43:05 -0700
Subject: [PATCH 67/75] behaviour when passing nothing into ML tools

---
 src/ScalarRandomFeature.jl | 3 ++-
 src/VectorRandomFeature.jl | 4 +++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 8a3321f17..afba754b9 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -323,6 +323,7 @@ function build_models!(
     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)
@@ -385,7 +386,7 @@ function build_models!(
     end
 
     #regularization = I #
-    regularization = Diagonal(output_structure_matrix) # creates diag matrix of diagonal
+    regularization =  isnothing(output_structure_matrix) ? 1.0*I : Diagonal(output_structure_matrix) 
 
 
     @info (
diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl
index d8ff493d3..9864acef8 100644
--- a/src/VectorRandomFeature.jl
+++ b/src/VectorRandomFeature.jl
@@ -445,7 +445,9 @@ function build_models!(
     )
 
     # think of the output_structure_matrix as the observational noise covariance, or a related quantity
-    if !isposdef(output_structure_matrix)
+    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

From 3fd4a3b36aac80decb2469b4de3eb026c1fb5e71 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 10 Jul 2025 15:11:14 -0700
Subject: [PATCH 68/75] format

---
 examples/Lorenz/emulate_sample.jl | 1 +
 src/ScalarRandomFeature.jl        | 4 ++--
 src/VectorRandomFeature.jl        | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/examples/Lorenz/emulate_sample.jl b/examples/Lorenz/emulate_sample.jl
index b38854360..02ed90d5a 100644
--- a/examples/Lorenz/emulate_sample.jl
+++ b/examples/Lorenz/emulate_sample.jl
@@ -6,6 +6,7 @@ using Distributions  # probability distributions and associated functions
 using LinearAlgebra
 ENV["GKSwstype"] = "100"
 using Plots
+using StatsPlots
 using Random
 using JLD2
 
diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index afba754b9..a0309c546 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -323,7 +323,7 @@ function build_models!(
     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)
@@ -386,7 +386,7 @@ function build_models!(
     end
 
     #regularization = I #
-    regularization =  isnothing(output_structure_matrix) ? 1.0*I : Diagonal(output_structure_matrix) 
+    regularization = isnothing(output_structure_matrix) ? 1.0 * I : Diagonal(output_structure_matrix)
 
 
     @info (
diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl
index 9864acef8..50eedc316 100644
--- a/src/VectorRandomFeature.jl
+++ b/src/VectorRandomFeature.jl
@@ -446,7 +446,7 @@ 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
+        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")

From 28295e93c25632445019c818b367223e078fd6a7 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 15:00:52 -0700
Subject: [PATCH 69/75] when not learning noise, regularization is scaled with
 output matrix

---
 src/GaussianProcess.jl | 54 ++++++++++++++++++++++++++++++++----------
 1 file changed, 41 insertions(+), 13 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 7a2414e31..721765d7f 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
@@ -172,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)
@@ -179,9 +190,12 @@ 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 = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    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)
@@ -196,12 +210,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
 
@@ -299,12 +314,15 @@ 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 = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization_noise_vec = gp.alg_reg_noise .* regularization
+    @info regularization_noise_vec
     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,)
@@ -314,6 +332,8 @@ function build_models!(
         push!(models, m)
         @info(m.kernel)
     end
+    append!(gp.regularization, regularization_noise_vec)
+
 end
 
 
@@ -331,7 +351,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
@@ -392,7 +414,9 @@ 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 = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    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
@@ -405,12 +429,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; obsdim = 2)
+        opt_fx = opt_f(input_values', regularization_noise_i; obsdim = 2)
 
         data_i = output_values[i, :]
         opt_post_fx = posterior(opt_fx, data_i)
@@ -418,6 +443,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
 
@@ -438,7 +465,8 @@ function predict(gp::GaussianProcess{AGPJL}, new_inputs::AM) where {AM <: Abstra
         μ[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

From 0e5ee0cbf5671d94b739b9380f7f8088f2762ee0 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 15:09:18 -0700
Subject: [PATCH 70/75] remove @info

---
 src/GaussianProcess.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 721765d7f..c21286623 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -191,7 +191,8 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    
     logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise))
 
     for i in 1:N_models
@@ -315,9 +316,8 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
     regularization_noise_vec = gp.alg_reg_noise .* regularization
-    @info regularization_noise_vec
     for i in 1:N_models
         regularization_noise_i = regularization_noise_vec[i]
         kernel_i = deepcopy(kern)
@@ -415,7 +415,7 @@ 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 = isnothing(output_structure_matrix) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
     regularization_noise = gp.alg_reg_noise .* regularization
 
     # now obtain the values of the hyperparameters

From 42947adf5a566198ae96cc80ff07c2963d8ee0cb Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 17:03:32 -0700
Subject: [PATCH 71/75] typo

---
 src/GaussianProcess.jl     | 14 ++++++++++----
 src/ScalarRandomFeature.jl |  6 +++---
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index c21286623..1d5f54a14 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -191,8 +191,10 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
-    
+    regularization =
+        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
+        diag(output_structure_matrix)
+
     logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise))
 
     for i in 1:N_models
@@ -316,7 +318,9 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization =
+        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
+        diag(output_structure_matrix)
     regularization_noise_vec = gp.alg_reg_noise .* regularization
     for i in 1:N_models
         regularization_noise_i = regularization_noise_vec[i]
@@ -415,7 +419,9 @@ 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 = (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) : diag(output_structure_matrix)
+    regularization =
+        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
+        diag(output_structure_matrix)
     regularization_noise = gp.alg_reg_noise .* regularization
 
     # now obtain the values of the hyperparameters
diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index a0309c546..1e6a20f09 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -385,9 +385,9 @@ function build_models!(
         end
     end
 
-    #regularization = I #
-    regularization = isnothing(output_structure_matrix) ? 1.0 * I : Diagonal(output_structure_matrix)
-
+    regularization =
+        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * I(N_rfms) :
+        Diagonal(output_structure_matrix)
 
     @info (
         "hyperparameter learning for $n_rfms models using $n_train training points, $n_test validation points and $n_features_opt features"

From 8cf8a181207744a747c5d8cdca0f0a46e2f9feba Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 17:26:56 -0700
Subject: [PATCH 72/75] missing \lambda in uniform scaling case

---
 src/GaussianProcess.jl     | 30 +++++++++++++++++++++---------
 src/ScalarRandomFeature.jl |  8 ++++++--
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 1d5f54a14..32ab4d39b 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -191,9 +191,13 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization =
-        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
-        diag(output_structure_matrix)
+    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))
 
@@ -318,9 +322,13 @@ function build_models!(
         println("Learning additive white noise")
     end
     # use the output_structure_matrix to scale regularization scale
-    regularization =
-        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
-        diag(output_structure_matrix)
+    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]
@@ -419,9 +427,13 @@ 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 =
-        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * ones(N_models) :
-        diag(output_structure_matrix)
+    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
diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 1e6a20f09..70c4bfe35 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -385,9 +385,13 @@ function build_models!(
         end
     end
 
-    regularization =
-        (isnothing(output_structure_matrix) || isa(output_structure_matrix, UniformScaling)) ? 1.0 * I(N_rfms) :
+    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"

From 179f7fc88af2b6d2cca0ff1aa91c2dc81acb9106 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 18:36:32 -0700
Subject: [PATCH 73/75] typo

---
 src/ScalarRandomFeature.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 70c4bfe35..52f6900cf 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -386,7 +386,7 @@ function build_models!(
     end
 
     regularization = if isnothing(output_structure_matrix)
-        1.0 * I(N_rfms) 
+        1.0 * I(n_rfms) 
     elseif isa(output_structure_matrix, UniformScaling) 
         output_structure_matrix 
     else

From 6765404b36084d74077297f555a27b0452908c33 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 11 Jul 2025 18:44:17 -0700
Subject: [PATCH 74/75] format

---
 src/ScalarRandomFeature.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index 52f6900cf..b23983f17 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -386,9 +386,9 @@ function build_models!(
     end
 
     regularization = if isnothing(output_structure_matrix)
-        1.0 * I(n_rfms) 
-    elseif isa(output_structure_matrix, UniformScaling) 
-        output_structure_matrix 
+        1.0 * I(n_rfms)
+    elseif isa(output_structure_matrix, UniformScaling)
+        output_structure_matrix
     else
         Diagonal(output_structure_matrix)
     end

From b00169205d5d73ed830abe21ae77ff64422d7c31 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Mon, 14 Jul 2025 10:16:02 -0700
Subject: [PATCH 75/75] missed format file

---
 src/GaussianProcess.jl | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl
index 32ab4d39b..3fd1ae3b8 100644
--- a/src/GaussianProcess.jl
+++ b/src/GaussianProcess.jl
@@ -192,11 +192,11 @@ function build_models!(
     end
     # use the output_structure_matrix to scale regularization scale
     regularization = if isnothing(output_structure_matrix)
-        1.0 * ones(N_models) 
+        1.0 * ones(N_models)
     elseif isa(output_structure_matrix, UniformScaling)
-        output_structure_matrix.λ*ones(N_models)
+        output_structure_matrix.λ * ones(N_models)
     else
-       diag(output_structure_matrix)
+        diag(output_structure_matrix)
     end
 
     logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise))
@@ -323,11 +323,11 @@ function build_models!(
     end
     # use the output_structure_matrix to scale regularization scale
     regularization = if isnothing(output_structure_matrix)
-        1.0 * ones(N_models) 
+        1.0 * ones(N_models)
     elseif isa(output_structure_matrix, UniformScaling)
-        output_structure_matrix.λ*ones(N_models)
+        output_structure_matrix.λ * ones(N_models)
     else
-       diag(output_structure_matrix)
+        diag(output_structure_matrix)
     end
     regularization_noise_vec = gp.alg_reg_noise .* regularization
     for i in 1:N_models
@@ -428,11 +428,11 @@ 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)
-        1.0 * ones(N_models) 
+        1.0 * ones(N_models)
     elseif isa(output_structure_matrix, UniformScaling)
-        output_structure_matrix.λ*ones(N_models)
+        output_structure_matrix.λ * ones(N_models)
     else
-       diag(output_structure_matrix)
+        diag(output_structure_matrix)
     end
     regularization_noise = gp.alg_reg_noise .* regularization