new predict method API, and reduced sampling in input space (#404)

* added ForwardMapWrapper, made compatible with MCMC and encoder utilities

* format

* Update src/Emulator.jl

Co-authored-by: ArneBouillon <45404227+ArneBouillon@users.noreply.github.com>

* first update to predict

deprecate transform_to_real

replace transform_to_real kwarg, and add add_obs_noise_cov

* restore format

* initial encode/decode inputs

* add orth. compl. piece

* format

* precompiles

* corrected logic for encode/decode flags

* add maxlog to warning message

* add get_en/decoder_from_schedule

* running version or sampler with encoded inputs (just decoding for posterior currently)

* return full affine shift from get_encoder_from_schedule, and compute gaussian lifts correctly

* compute proposals correctly, resolve input dim checks

* format

* add guardrailsfor getting empty encoders

* reshape bug in test

* format

* due to reverse-diff slow downs with LinearMaps.jl,  reduced the mcmc steps during tests

* change lengths

* tests for utils

* test retrieval of encoder

* tests for the enriching

* typo

* format

* updated and added new docstrings

* larger tol as it is borderline

* Address review comments: reduced prior storage/sampling, always use decode_and_add_noise

* add a few extra tests

* suitable tols

* typos

* add scaling, rename boost_for_loss... to noise_injector_threshold

* resolve review

* remove all "transform_to_real" in tests

* add one example to test the deprecation message...

---------

Co-authored-by: ArneBouillon <45404227+ArneBouillon@users.noreply.github.com>

Author: odunbar

src/Emulator.jl

@@ -78,7 +78,7 @@ struct GaussianProcess{GPPackage, FT, VV <: AbstractVector} <: MachineLearningTo
     noise_learn::Bool
     "Additional observational or regularization noise in used in GP algorithms"
     alg_reg_noise::FT
-    "Prediction type (`y` to predict the data, `f` to predict the latent function)."
+    "[Deprecated - use `add_obs_noise_cov` kwarg when calling `predict(`] 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
@@ -260,7 +260,7 @@ end
 function _predict(
     gp::GaussianProcess,
     new_inputs::AbstractMatrix{FT},
-    predict_method::Function,
+    predict_method::Function;
 ) where {FT <: AbstractFloat}
     M = length(gp.models)
     N_samples = size(new_inputs, 2)
@@ -284,11 +284,17 @@ predict(gp::GaussianProcess{GPJL}, new_inputs::AbstractMatrix{FT}, ::FType) wher
 """
 $(DocStringExtensions.TYPEDSIGNATURES)
 
-Predict means and covariances in decorrelated output space using Gaussian process models.
+Predict means and covariances in decorrelated output space using Gaussian process models. The use of stored `FType` and `YType` to control this method is deprecated, the return covariance is now determined by the `predict(` kwarg `add_obs_noise_cov` 
 """
-predict(gp::GaussianProcess{GPJL}, new_inputs::AbstractMatrix{FT}) where {FT <: AbstractFloat} =
-    predict(gp, new_inputs, gp.prediction_type)
-
+function predict(
+    gp::GaussianProcess{GPJL},
+    new_inputs::AbstractMatrix{FT};
+    add_obs_noise_cov = false,
+    mlt_kwargs...,
+) where {FT <: AbstractFloat}
+    pred_type = add_obs_noise_cov ? YType() : FType()
+    return predict(gp, new_inputs, pred_type)
+end
 
 #now we build the SKLJL implementation
 function build_models!(
@@ -371,13 +377,20 @@ function _SKJL_predict_function(gp_model::PyObject, new_inputs::AbstractMatrix{F
     μ, σ = gp_model.predict(new_inputs', return_std = true)
     return μ, (σ .* σ)
 end
-function predict(gp::GaussianProcess{SKLJL}, new_inputs::AbstractMatrix{FT}) where {FT <: AbstractFloat}
+function predict(
+    gp::GaussianProcess{SKLJL},
+    new_inputs::AbstractMatrix{FT};
+    add_obs_noise_cov = false,
+    mlt_kwargs...,
+) where {FT <: AbstractFloat}
     μ, σ2 = _predict(gp, new_inputs, _SKJL_predict_function)
 
     # 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.
-    for i in 1:size(σ2, 2)
-        σ2[:, i] = σ2[:, i] + gp.regularization
+    if add_obs_noise_cov
+        for i in 1:size(σ2, 2)
+            σ2[:, i] = σ2[:, i] + gp.regularization
+        end
     end
 
     return μ, σ2
@@ -484,7 +497,12 @@ function optimize_hyperparameters!(gp::GaussianProcess{AGPJL}, args...; kwargs..
     @info "AbstractGP already built. Continuing..."
 end
 
-function predict(gp::GaussianProcess{AGPJL}, new_inputs::AM) where {AM <: AbstractMatrix}
+function predict(
+    gp::GaussianProcess{AGPJL},
+    new_inputs::AM;
+    add_obs_noise_cov = false,
+    mlt_kwargs...,
+) where {AM <: AbstractMatrix}
 
     N_models = length(gp.models)
     N_samples = size(new_inputs, 2)
@@ -497,8 +515,10 @@ function predict(gp::GaussianProcess{AGPJL}, new_inputs::AM) where {AM <: Abstra
         μ[i, :] = mean(pred)
         σ2[i, :] = var(pred)
     end
-    for i in 1:size(σ2, 2)
-        σ2[:, i] .= σ2[:, i] + gp.regularization
+    if add_obs_noise_cov
+        for i in 1:size(σ2, 2)
+            σ2[:, i] .= σ2[:, i] + gp.regularization
+        end
     end
     return μ, σ2
 end

src/MachineLearningTools/GaussianProcess.jl

@@ -631,6 +631,8 @@ function predict(
     srfi::ScalarRandomFeatureInterface,
     new_inputs::MM;
     multithread = "ensemble",
+    add_obs_noise_cov = false,
+    mlt_kwargs...,
 ) where {MM <: AbstractMatrix}
     M = length(get_rfms(srfi))
     N_samples = size(new_inputs, 2)
@@ -653,11 +655,13 @@ function predict(
     end
 
     # add the noise contribution stored within the regularization
-    reg = get_regularization(srfi)[1]
-    reg_diag = isa(reg, UniformScaling) ? reg.λ * ones(M) : diag(reg)
+    if add_obs_noise_cov
+        reg = get_regularization(srfi)[1]
+        reg_diag = isa(reg, UniformScaling) ? reg.λ * ones(M) : diag(reg)
 
-    for i in 1:M
-        σ2[i, :] .+= reg_diag[i]
+        for i in 1:M
+            σ2[i, :] .+= reg_diag[i]
+        end
     end
 
     return μ, σ2

src/MachineLearningTools/ScalarRandomFeature.jl

@@ -656,7 +656,12 @@ $(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.
 """
-function predict(vrfi::VectorRandomFeatureInterface, new_inputs::M) where {M <: AbstractMatrix}
+function predict(
+    vrfi::VectorRandomFeatureInterface,
+    new_inputs::M;
+    add_obs_noise_cov = false,
+    mlt_kwargs...,
+) where {M <: AbstractMatrix}
     input_dim = get_input_dim(vrfi)
     output_dim = get_output_dim(vrfi)
     rfm = get_rfms(vrfi)[1]
@@ -676,12 +681,14 @@ function predict(vrfi::VectorRandomFeatureInterface, new_inputs::M) where {M <:
     # sizes (output_dim x n_test), (output_dim x output_dim x n_test) 
     # add the noise contribution from the regularization
     # note this is because we are predicting the data here, not the latent function.
-    lambda = get_regularization(vrfi)[1]
-    for i in 1:N_samples
-        σ2[:, :, i] = 0.5 * (σ2[:, :, i] + permutedims(σ2[:, :, i], (2, 1))) + lambda
-
-        if !isposdef(σ2[:, :, i])
-            σ2[:, :, i] = posdef_correct(σ2[:, :, i])
+    if add_obs_noise_cov
+        lambda = get_regularization(vrfi)[1]
+        for i in 1:N_samples
+            σ2[:, :, i] = 0.5 * (σ2[:, :, i] + permutedims(σ2[:, :, i], (2, 1))) + lambda
+
+            if !isposdef(σ2[:, :, i])
+                σ2[:, :, i] = posdef_correct(σ2[:, :, i])
+            end
         end
     end
     return μ, σ2