replace transform_to_real kwarg, and add add_obs_noise_cov

Author: odunbar

src/Emulator.jl

@@ -319,7 +319,6 @@ function predict(
     emulator::Emulator{FT},
     new_inputs::AM;
     encode=nothing, # maps decoded inputs to decoded outputs
-    add_obs_noise_cov=false,
     mlt_kwargs...,
 ) where {FT <: AbstractFloat, AM <: AbstractMatrix}
 

src/MachineLearningTools/GaussianProcess.jl

@@ -76,7 +76,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
@@ -258,7 +258,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)
@@ -282,11 +282,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!(
@@ -369,13 +370,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
@@ -482,7 +485,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)
@@ -495,8 +498,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,