Bugfix for MCMC with encoding (#378)

* fix additional averaging of likelihood

* weighting bug for many-data

* clearer error message

* format

* new tests and all pass

* rm accidental feature push

* added mv input testing at different input size

* remove unnec. prints

Author: odunbar

src/Emulator.jl

@@ -125,9 +125,9 @@ function Emulator(
 
     if !isnothing(obs_noise_cov)
         if haskey(encoder_kwargs, :obs_noise_cov)
-            @warn "Keyword argument `obs_noise_cov=` is deprecated and will be ignored in favor of `encoder_kwargs[:obs_noise_cov]`."
+            @warn "Keyword argument `obs_noise_cov=` is deprecated and will be ignored in favor of `encoder_kwargs=(obs_noise_cov=...)`."
         else
-            @warn "Keyword argument `obs_noise_cov=` is deprecated. Please use `encoder_kwargs[:obs_noise_cov]` instead."
+            @warn "Keyword argument `obs_noise_cov=` is deprecated. Please use `encoder_kwargs=(obs_noise_cov=...)` instead."
         end
     end
 

src/GaussianProcess.jl

@@ -167,7 +167,6 @@ function build_models!(
     end
     N_models = size(output_values, 1) #size(transformed_data)[1]
 
-
     # Use a default kernel unless a kernel was supplied to GaussianProcess
     if gp.kernel === nothing
         println("Using default squared exponential kernel, learning length scale and variance parameters")
@@ -184,7 +183,6 @@ function build_models!(
         println("Using user-defined kernel", kern)
     end
 
-
     if gp.noise_learn
         # Add white noise to kernel
         white_logstd = log(1.0)
@@ -202,9 +200,10 @@ function build_models!(
         else
             diag(output_structure_mat)
         end
-    end
 
-    logstd_regularization_noise = log.(sqrt.(regularization .* gp.alg_reg_noise))
+    end
+    regularization_noise = regularization .* gp.alg_reg_noise
+    logstd_regularization_noise = log.(sqrt.(regularization_noise))
 
     for i in 1:N_models
         logstd_regularization_i = logstd_regularization_noise[i]
@@ -228,7 +227,7 @@ function build_models!(
         push!(models, m)
 
     end
-    append!(gp.regularization, logstd_regularization_noise)
+    append!(gp.regularization, regularization_noise)
 
 end
 
@@ -262,6 +261,7 @@ function _predict(
     # Predicts columns of inputs: input_dim × N_samples
     μ = zeros(M, N_samples)
     σ2 = zeros(M, N_samples)
+    # predict method ::YType will add gp.regularization back in here, ::FType will not
     for i in 1:M
         μ[i, :], σ2[i, :] = predict_method(gp.models[i], new_inputs)
     end
@@ -430,7 +430,6 @@ AbstractGP currently does not (yet) learn hyperparameters internally. The follow
         """))
     end
 
-
     N_models = size(output_values, 1) #size(transformed_data)[1]
     # use the output_structure_matrix to scale regularization scale
     regularization = if isempty(output_structure_mats)

src/MarkovChainMonteCarlo.jl

@@ -233,7 +233,14 @@ autodiff_hessian(model::AdvancedMH.DensityModel, params, sampler::MH) where {MH
 """
 $(DocStringExtensions.TYPEDSIGNATURES)
 
-Defines the internal log-density function over a vector of observation samples using an assumed conditionally indepedent likelihood, that is with a log-likelihood of `ℓ(y,θ) = sum^n_i log( p(y_i|θ) )`. 
+Defines the internal log-density function over a vector of observation samples using an assumed conditionally indepedent likelihood, that is with a log-likelihood of `ℓ(y,θ) = sum^n_i log( p(y_i|θ) )`.
+
+Inputs:
+=======
+- θ: Parameters in physical (constrained) coordinates.
+- prior: Prior distribution as a ParameterDistribution
+- em: Emulator object with predict(.) method, evaluations returned in encoded space 
+- obs_vec: encoded data vector sample(s)
 """
 function emulator_log_density_model(
     θ,
@@ -242,21 +249,17 @@ function emulator_log_density_model(
     obs_vec::AV,
 ) where {FT <: AbstractFloat, AV <: AbstractVector}
 
-    # θ: model params we evaluate at; in original coords.
-    # transform_to_real = false means g, g_cov, obs_sample are in decorrelated coords.
+    # 
+    # transform_to_real = false means g, g_cov, obs_sample are in encoded coords.
 
-    # Recall predict() written to return multiple N_samples: expects input to be a 
-    # Matrix with N_samples columns. Returned g is likewise a Matrix, and g_cov is a
-    # 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.
+    # 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, 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, θ)
+        return sum([logpdf(MvNormal(obs, g_cov[1] * I), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
     else
-        return 1.0 / length(obs_vec) * sum([logpdf(MvNormal(obs, g_cov[1]), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
+        return sum([logpdf(MvNormal(obs, g_cov[1]), vec(g)) for obs in obs_vec]) + logpdf(prior, θ)
     end
 
 end
@@ -556,7 +559,6 @@ function MCMCWrapper(
     # 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)
 

src/ScalarRandomFeature.jl

@@ -178,7 +178,7 @@ function ScalarRandomFeatureInterface(
         "localization" => EKP.Localizers.NoLocalization(), # localization / sample error correction for small ensembles
         "cov_correction" => "nice", # type of conditioning to improve estimated covariance
         "n_cross_val_sets" => 2, # if >1 do cross validation, else if 0 do no data splitting and no training fraction
-        "overfit" => 1.0, # if >1 this forcibly overfits to the data 
+        "overfit" => 1.0, # if >1 this forcibly overfits to the data
     )
 
     if !isnothing(optimizer_options)
@@ -436,12 +436,10 @@ function build_models!(
                 ),
             )
         end
-
         # scale up the prior so that default priors are always "reasonable"
         prior_in_scale = 1.0 ./ std(input_values, dims = 2)
         prior_out_scale = std(output_values[i, :])
 
-
         prior = build_default_prior(input_dim, kernel_structure)
 
         # where prior space has changed we need to rebuild the priors
@@ -453,16 +451,17 @@ function build_models!(
             prior = build_default_prior(input_dim, kernel_structure)
 
         end
-        # [2.] Estimate covariance at mean value
+        # [2a.] Estimate the covariance at prior mean
+        n_ensemble = optimizer_options["n_ensemble"]
         μ_hp = transform_unconstrained_to_constrained(prior, mean(prior))
 
         cov_sample_multiplier = optimizer_options["cov_sample_multiplier"]
         cov_correction = optimizer_options["cov_correction"]
         overfit = max(optimizer_options["overfit"], 1e-4)
         n_cov_samples_min = n_test + 2
         n_cov_samples = Int(floor(n_cov_samples_min * max(cov_sample_multiplier, 0.0)))
-        println("estimating covariances with " * string(n_cov_samples) * " iterations...")
 
+        println("estimating covariances with " * string(n_cov_samples) * " iterations...")
         observation_vec = []
         for cv_idx in 1:n_cross_val_sets
 
@@ -504,7 +503,6 @@ function build_models!(
         end
         observation = combine_observations(observation_vec)
         # [3.] set up EKP optimization
-        n_ensemble = optimizer_options["n_ensemble"]
         n_iteration = optimizer_options["n_iteration"]
         scheduler = optimizer_options["scheduler"]
         accelerator = optimizer_options["accelerator"]
@@ -594,6 +592,7 @@ function build_models!(
         end
 
         io_pairs_i = PairedDataContainer(input_values, reshape(output_values[i, :], 1, size(output_values, 2)))
+
         # Now, fit new RF model with the optimized hyperparameters
 
         rfm_i = RFM_from_hyperparameters(

src/VectorRandomFeature.jl

@@ -160,7 +160,6 @@ Constructs a `VectorRandomFeatureInterface <: MachineLearningTool` interface for
      - "cov_correction" => "nice": type of conditioning to improve estimated covariance. "shrinkage", "shrinkage_corr" (Ledoit Wolfe 03), "nice" for (Vishny, Morzfeld et al. 2024)
       - "overfit" => 1.0: if > 1.0 forcibly overfit/under-regularize the optimizer cost, (vice versa for < 1.0).
       - "n_cross_val_sets" => 2, train fraction creates (default 5) train-test data subsets, then use 'n_cross_val_sets' of these stacked in the loss function. If set to 0, train=test on the full data provided ignoring "train_fraction".
-
 """
 function VectorRandomFeatureInterface(
     n_features::Int,
@@ -203,7 +202,7 @@ function VectorRandomFeatureInterface(
         "accelerator" => EKP.NesterovAccelerator(), # acceleration with momentum
         "cov_correction" => "nice", # type of conditioning to improve estimated covariance
         "n_cross_val_sets" => 2, # if set to 0, removes data split. i.e takes train & test to be the same data set
-        "overfit" => 1.0, # if >1 this forcibly overfits to the data 
+        "overfit" => 1.0, # if >1 this forcibly overfits to the data
     )
 
     if !isnothing(optimizer_options)
@@ -471,11 +470,9 @@ function build_models!(
         end
     end
 
-    # [2.] Estimate covariance at mean value
+    # [2a.] Estimate the cov at prior mean
+    n_ensemble = optimizer_options["n_ensemble"] # minimal ensemble size n_hp,
     μ_hp = transform_unconstrained_to_constrained(prior, mean(prior))
-    cov_sample_multiplier = optimizer_options["cov_sample_multiplier"]
-    cov_correction = optimizer_options["cov_correction"]
-    overfit = max(optimizer_options["overfit"], 1e-4)
 
     if nameof(typeof(kernel_structure)) == :SeparableKernel
         if nameof(typeof(get_output_cov_structure(kernel_structure))) == :DiagonalFactor
@@ -486,7 +483,13 @@ function build_models!(
     else
         n_cov_samples_min = (n_test * output_dim + 2)
     end
+
+    cov_sample_multiplier = optimizer_options["cov_sample_multiplier"]
+    cov_correction = optimizer_options["cov_correction"]
+    overfit = max(optimizer_options["overfit"], 1e-4)
     n_cov_samples = Int(floor(n_cov_samples_min * max(cov_sample_multiplier, 0.0)))
+
+    println("estimating covariances with " * string(n_cov_samples) * " iterations...")
     observation_vec = []
     for cv_idx in 1:n_cross_val_sets
         internal_Γ, approx_σ2 = estimate_mean_and_coeffnorm_covariance(
@@ -525,7 +528,6 @@ function build_models!(
     observation = combine_observations(observation_vec)
 
     # [3.] set up EKP optimization
-    n_ensemble = optimizer_options["n_ensemble"] # minimal ensemble size n_hp,
     n_iteration = optimizer_options["n_iteration"]
     opt_verbose_flag = optimizer_options["verbose"]
     scheduler = optimizer_options["scheduler"]