first update to predict

deprecate transform_to_real

replace transform_to_real kwarg, and add add_obs_noise_cov

Author: odunbar

src/Emulator.jl

@@ -318,9 +318,10 @@ Return type of N inputs: (in the output space)
 function predict(
     emulator::Emulator{FT},
     new_inputs::AM;
-    transform_to_real = false,
+    encode=nothing, # maps decoded inputs to decoded outputs
     mlt_kwargs...,
 ) where {FT <: AbstractFloat, AM <: AbstractMatrix}
+                    
     # Check if the size of new_inputs is consistent with the training data input
     input_dim, output_dim = size(get_io_pairs(emulator), 1)
     encoded_input_dim, encoded_output_dim = size(get_encoded_io_pairs(emulator), 1)
@@ -335,8 +336,17 @@ function predict(
         )
     end
 
+    # note the logic below
+    in_already_encoded = encode ∈ ["in", "in_and_out"]
+    out_to_be_decoded = encode ∈ ["out","in_and_out"] 
+    
+    
     # encode the new input data
-    encoded_inputs = encode_data(emulator, new_inputs, "in")
+    if !in_already_encoded
+        encoded_inputs = encode_data(emulator, new_inputs, "in")
+    else
+        encoded_inputs = new_inputs
+    end
     # predict in encoding space
     # returns outputs: [enc_out_dim x n_samples]
     # Scalar-methods uncertainties=variances: [enc_out_dim x n_samples]
@@ -346,7 +356,7 @@ function predict(
     var_or_cov = (ndims(encoded_uncertainties) == 2) ? "var" : "cov"
 
     # return decoded or encoded?
-    if transform_to_real
+    if out_to_be_decoded
         decoded_outputs = decode_data(emulator, encoded_outputs, "out")
 
         decoded_covariances = zeros(eltype(encoded_outputs), output_dim, output_dim, size(encoded_uncertainties)[end])
@@ -433,7 +443,8 @@ end
 function predict(
     fmw::FMW,
     new_inputs::AM;
-    transform_to_real = false,
+    encode=nothing, # maps decoded inputs to decoded outputs
+    add_obs_noise_cov=false,
 ) where {FMW <: ForwardMapWrapper, AM <: AbstractMatrix}
     # Check if the size of new_inputs is consistent with the training input data
     input_dim, output_dim = size(get_io_pairs(fmw), 1)
@@ -449,6 +460,15 @@ function predict(
         )
     end
 
+    in_already_encoded = encode ∈ ["in", "in_and_out"]
+    out_to_be_decoded = encode ∈ ["out","in_and_out"] 
+    #need to boost to decoded inputs
+    if in_already_encoded
+        # Sample from the null space
+        decoded_inputs = ...
+    else
+        decoded_inputs = new_inputs
+    end
     # Scalar-methods uncertainties=variances: [enc_out_dim x n_samples]
     # Vector-methods uncertainties=covariances: [enc_out_dim x enc_out_dim x n_samples)
 
@@ -457,10 +477,10 @@ function predict(
     forward_map = get_forward_map(fmw)
     fm_unc = x -> forward_map(transform_unconstrained_to_constrained(prior, x))
 
-    decoded_outputs = reduce(hcat, map(fm_unc, eachcol(new_inputs))) # apply map and return: [out_dim x n_samples]
+    decoded_outputs = reduce(hcat, map(fm_unc, eachcol(decoded_inputs))) # apply map and return: [out_dim x n_samples]
 
     var_or_cov = (output_dim == 1) ? "var" : "cov"
-    if transform_to_real
+    if out_to_be_decoded
         # uncertainty returned is just `I` in encoded space
         decoded_cov = Matrix(decode_structure_matrix(fmw, I(output_dim), "out"))
 
@@ -496,4 +516,38 @@ function predict(
     end
 end
 
+
+### Deprecated keywords
+
+function predict(
+    em_or_fmw::EorFMW,
+    new_inputs::AM;
+    transform_to_real = nothing,
+    kwargs...,
+) where {AM <: AbstractMatrix,  EorFMW <: Union{Emulator, ForwardMapWrapper}}
+    
+    if !isnothing(transform_to_real)
+        Base.depwarn(
+            """`transform_to_real` keyword is deprecated. Please use the `encode` and `add_obs_noise_cov` keywords instead.
+
+Recommended usage for users is now set by default as:
+ - `encode=nothing`, `add_obs_noise_cov=false`
+This behaviour takes in non-encoded inputs, and returns non-encoded outputs. It gives only the uncertainty from the Machine Learning Tool (not inflated by observational noise)
+
+This simulation will continue with the old behavior:
+ - `transform_to_real=true` replaced with `encode=nothing, add_obs_noise_cov=true`
+ - `transform_to_real=false` replaced with `encode="out", add_obs_noise_cov=true`
+    """,
+            :predict,
+        )
+        
+        # modify kwargs
+        kw = Dict(kwargs)
+        kw[:add_obs_noise_cov] = true
+        kw[:encode] = transform_to_real ? nothing : "out"
+        predict(em_or_fmw, new_inputs; kw...)
+    end
+
+    return predict(em_or_fmw, new_inputs; kwargs...)
+
 end

src/MachineLearningTools/GaussianProcess.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,12 @@ 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 +372,15 @@ 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 +487,7 @@ 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 +500,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/ScalarRandomFeature.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/VectorRandomFeature.jl

@@ -656,7 +656,7 @@ $(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 +676,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

src/MarkovChainMonteCarlo.jl

@@ -255,7 +255,7 @@ function emulator_log_density_model(
 
     # predict is written to apply to columns.
     # Returned g is a length-1, Vector{Real} or Vector{Vector}, and g_cov is length-1 Vector{Vector} or Vector{Matrix} respectively
-    g, g_cov = Emulators.predict(em_or_fmw, reshape(θ, :, 1), transform_to_real = false)
+    g, g_cov = Emulators.predict(em_or_fmw, reshape(θ, :, 1), encode="out", add_obs_noise_cov=true)
 
     if isa(g_cov[1], Real)
         return sum([logpdf(MvNormal(obs, g_cov[1] * I), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
@@ -577,8 +577,10 @@ function MCMCWrapper(
         eachcol(observation)
     end
 
-    # encoding works on columns but mcmc wants vec-of-vec
+    # encoding data works on columns but mcmc wants vec-of-vec
     encoded_obs = [vec(encode_data(em_or_fmw, reshape(obs, :, 1), "out")) for obs in obs_slice]
+    # encoding initial condition
+    #encoded_init_params = vec(encode_data(em_or_fmw, reshape(init_params,:,1), "in"))
 
     log_posterior_map = EmulatorPosteriorModel(prior, em_or_fmw, encoded_obs)
     mh_proposal_sampler = MetropolisHastingsSampler(mcmc_alg, prior)
@@ -594,6 +596,7 @@ function MCMCWrapper(
     end
 
     sample_kwargs = (; # set defaults here
+#        :initial_params => deepcopy(encoded_init_params),
         :initial_params => deepcopy(init_params),
         :param_names => param_names,
         :discard_initial => burnin,