diff --git a/src/Emulator.jl b/src/Emulator.jl index 56bfe83a6..b796ef508 100644 --- a/src/Emulator.jl +++ b/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 diff --git a/src/GaussianProcess.jl b/src/GaussianProcess.jl index 9d6d403df..7f267216f 100644 --- a/src/GaussianProcess.jl +++ b/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) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 31b10ab43..c65655a6a 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/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) diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl index aca965e65..0599b09a0 100644 --- a/src/ScalarRandomFeature.jl +++ b/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,7 +451,8 @@ 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"] @@ -461,8 +460,8 @@ function build_models!( 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( diff --git a/src/VectorRandomFeature.jl b/src/VectorRandomFeature.jl index 1e029c5ea..6d5d4f5a8 100644 --- a/src/VectorRandomFeature.jl +++ b/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"] diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index 975dbec36..7546c8718 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -11,52 +11,118 @@ using CalibrateEmulateSample.Emulators using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.Utilities -function test_data(; rng_seed = 41, n = 20, var_y = 0.05, rest...) +# range 0->2, with lengthscale of transition 0.5, and y=1 at x=2 +G(x) = 5 * (tanh.((x .- 2) ./ 0.5) .+ 1) + +function test_data(prior; rng_seed = 41, n = 80, var_y = 0.05, rest...) # Seed for pseudo-random number generator rng = Random.MersenneTwister(rng_seed) # We need a GaussianProcess to run MarkovChainMonteCarlo, so let's reconstruct the one # that's tested in test/GaussianProcesses/runtests.jl Case 1 - n = 40 # number of training points - x = 2π * rand(rng, Float64, (1, n)) # predictors/features: 1 × n - σ2_y = reshape([var_y], 1, 1) - y = sin.(x) + rand(rng, Normal(0, σ2_y[1]), (1, n)) # predictands/targets: 1 × n + x = 5 * rand(rng, Float64, (1, n)) # predictors/features: 1 × n + σ2_y = var_y * I #reshape([var_y], 1, 1) + y = G(x) + sqrt(σ2_y.λ) * randn(rng, 1, n) # predictands/targets: 1 × n + + return y, + σ2_y, + PairedDataContainer(transform_constrained_to_unconstrained(prior, x), y, data_are_columns = true), + rng +end + +function validate_emulator(em, data_name, prior, test_data_kwargs; exp_name = "test") + tpk = merge(test_data_kwargs, (; rng_seed = 235412)) + + if data_name == "test_data" + _, _, validation_data, _ = test_data(prior; test_data_kwargs...) + elseif data_name == "test_data_2d" + _, _, validation_data, _ = test_data_2d(prior; test_data_kwargs...) + elseif data_name == "test_data_mv" + _, _, validation_data, _ = test_data_mv(prior; test_data_kwargs...) + end + + + μ_valid, σ2_valid = Emulators.predict(em, get_inputs(validation_data), transform_to_real = true) + rmse_valid = norm(μ_valid - get_outputs(validation_data)) ./ sqrt(size(get_outputs(validation_data), 2)) + avg_σ2_valid = mean([norm(ss) for ss in σ2_valid]) + @info "$(exp_name) \n per-point - Emulator RMSE = $rmse_valid \n Emulator STD = $(sqrt(avg_σ2_valid))\n (These numbers should be similar sized to indicate no overfitting)" + + if TEST_PLOT_OUTPUT + if data_name == "test_data" #(1D->1D) + sorted_idx = sortperm(get_inputs(validation_data)[1, :]) + sorted_in = get_inputs(validation_data)[:, sorted_idx][:] + sorted_out = get_outputs(validation_data)[:, sorted_idx][:] + μ_plot = μ_valid[:, sorted_idx][:] + σ2_plot = σ2_valid[:, sorted_idx][:] + + plot(sorted_in, sorted_out, label = "data", title = exp_name) + plot!(sorted_in, μ_plot, yerror = 2 * sqrt.(σ2_plot), label = "mean, 95% interval") + + savefig(joinpath(@__DIR__, "validate_emulator_$(exp_name).png")) + end + end +end + +function test_data_2d(prior; rng_seed = 4141, n = 100, cov_y = 0.05 * [[0.5, 0.2] [0.2, 0.5]], rest...) + rng = Random.MersenneTwister(rng_seed) + n = 100 # number of training points + + input_dim = 2 # input dim + output_dim = 2 # output dim + X = 5 * rand(rng, Float64, (input_dim, n)) # bounds from prior on input space are [-1,6] - return y, σ2_y, PairedDataContainer(x, y, data_are_columns = true), rng + # G(x1, x2) + g1x = sin.(X[1, :]) .+ cos.(X[2, :]) + g2x = sin.(X[1, :]) .- cos.(X[2, :]) + gx = zeros(2, n) + gx[1, :] = g1x + gx[2, :] = g2x + + # Add noise η + μ = zeros(output_dim) + noise_samples = rand(MvNormal(μ, cov_y), n) + + # y = G(x) + η + Y = gx .+ noise_samples + + iopairs = PairedDataContainer(transform_constrained_to_unconstrained(prior, X), Y, data_are_columns = true) + return Y, cov_y, iopairs, rng end function test_prior() - ### Define prior - umin = -1.0 - umax = 6.0 - #prior = [Priors.Prior(Uniform(umin, umax), "u")] # prior on u - prior_dist = Parameterized(Normal(0, 1)) - prior_constraint = bounded(umin, umax) - prior_name = "u" - return ParameterDistribution(prior_dist, prior_constraint, prior_name) + return constrained_gaussian("u", 2.0, 1.0, -1.0, 6.0) end -function test_prior_mv() + +function test_prior_mv(input_dim = 10) ### Define prior - return constrained_gaussian("u_mv", -1.0, 6.0, -Inf, Inf, repeats = 10) + return constrained_gaussian("u_mv", 2.0, 1.0, -1.0, 6.0, repeats = input_dim) end -function test_data_mv(; rng_seed = 41, n = 20, var_y = 0.05, input_dim = 10, rest...) +test_prior_2d() = test_prior_mv(2) + +function test_data_mv(prior; rng_seed = 41, n = 500, var_y = 0.01, input_dim = 4, rest...) + # Seed for pseudo-random number generator rng = Random.MersenneTwister(rng_seed) - n = 40 # number of training points - x = 1 / input_dim * π * rand(rng, Float64, (input_dim, n)) # predictors/features: 1 × n - σ2_y = reshape([var_y], 1, 1) - y = sin.(norm(x)) + rand(rng, Normal(0, σ2_y[1]), (1, n)) # predictands/targets: 1 × n + # number of training points + x = 5.0 * rand(rng, Float64, (input_dim, n)) # predictors/features: 1 × n + σ2_y = var_y * I #reshape([var_y], 1, 1) + # scaled by input_dim, so that the prior/training inputs remain reasonable in different input dimensions. + y = reshape(G([norm(xx) / input_dim for xx in eachcol(x)]), 1, n) + rand(rng, MvNormal(zeros(1), σ2_y), n) # predictands/targets: 1 × n - return y, σ2_y, PairedDataContainer(x, y, data_are_columns = true), rng + return y, + σ2_y, + PairedDataContainer(transform_constrained_to_unconstrained(prior, x), y, data_are_columns = true), + rng end + 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) + gp = GaussianProcess(gppackage; noise_learn = false, prediction_type = pred_type) em = Emulator(gp, iopairs; encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) return em @@ -68,8 +134,8 @@ function test_gp_1(y, σ2_y, iopairs::PairedDataContainer) # Construct kernel: # Squared exponential kernel (note that hyperparameters are on log scale) # with observational noise - GPkernel = SE(log(1.0), log(1.0)) - gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) + #GPkernel = SE(log(1.0), log(1.0)) #kernel = GPkernel + gp = GaussianProcess(gppackage; noise_learn = false, prediction_type = pred_type) em = Emulator(gp, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) return em @@ -81,24 +147,35 @@ function test_gp_and_agp_1(y, σ2_y, iopairs::PairedDataContainer) # Construct kernel: # Squared exponential kernel (note that hyperparameters are on log scale) # with observational noise - GPkernel = SE(log(1.0), log(1.0)) - gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) + #GPkernel = SE(log(1.0), log(1.0)) #kernel = GPkernel + noise_learn_gp = false + gp = GaussianProcess(gppackage; noise_learn = false, prediction_type = pred_type) em = Emulator(gp, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) Emulators.optimize_hyperparameters!(em) # now make agp from gp - agp = GaussianProcess(AGPJL(); noise_learn = true, prediction_type = pred_type) + agp = GaussianProcess(AGPJL(); noise_learn = false, prediction_type = pred_type) # NB it won't learn the noise here... gp_opt_params = Emulators.get_params(gp) gp_opt_param_names = get_param_names(gp) - kernel_params = [ - Dict( - "log_rbf_len" => model_params[1:(end - 2)], - "log_std_sqexp" => model_params[end - 1], - "log_std_noise" => model_params[end], - ) for model_params in gp_opt_params - ] + if noise_learn_gp + kernel_params = [ + Dict( + "log_rbf_len" => model_params[1:(end - 2)], + "log_std_sqexp" => model_params[end - 1], + "log_std_noise" => model_params[end], + ) for model_params in gp_opt_params + ] + else + kernel_params = [ + Dict( + "log_rbf_len" => model_params[1:(end - 1)], + "log_std_sqexp" => model_params[end], + "log_std_noise" => -10, # large+neg. will make this term ~0. (noise will come from the regularization matrix in encoder). + ) for model_params in gp_opt_params + ] + end em_agp = Emulator( agp, @@ -119,7 +196,7 @@ function test_gp_2(y, σ2_y, iopairs::PairedDataContainer) # Squared exponential kernel (note that hyperparameters are on log scale) # with observational noise GPkernel = SE(log(1.0), log(1.0)) - gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = true, prediction_type = pred_type) + gp = GaussianProcess(gppackage; kernel = GPkernel, noise_learn = false, 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; encoder_schedule = encoder_schedule, encoder_kwargs = (; obs_noise_cov = σ2_y)) @@ -128,22 +205,57 @@ function test_gp_2(y, σ2_y, iopairs::PairedDataContainer) return em end +function test_srfi(y, σ2_y, iopairs::PairedDataContainer) + # Construct kernel: + # Squared exponential kernel (note that hyperparameters are on log scale) + # with observational noise + input_dim = size(get_inputs(iopairs), 1) + n_features = 100 + kernel_structure = SeparableKernel(DiagonalFactor(), OneDimFactor()) + srfi = ScalarRandomFeatureInterface(n_features, input_dim, kernel_structure = kernel_structure) + + em = Emulator(srfi, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) + + Emulators.optimize_hyperparameters!(em) + + return em +end + +function test_vrfi(y, σ2_y, iopairs::PairedDataContainer) + # Construct kernel: + # Squared exponential kernel (note that hyperparameters are on log scale) + # with observational noise + input_dim = size(get_inputs(iopairs), 1) + output_dim = size(get_outputs(iopairs), 1) + n_features = 100 + kernel_structure = NonseparableKernel(LowRankFactor(2)) + vrfi = VectorRandomFeatureInterface(n_features, input_dim, output_dim, kernel_structure = kernel_structure) + + em = Emulator(vrfi, iopairs; encoder_schedule = [], encoder_kwargs = (; obs_noise_cov = σ2_y)) + Emulators.optimize_hyperparameters!(em) + + return em +end + + function mcmc_test_template( prior::ParameterDistribution, σ2_y, em::Emulator; + exp_name = "test", mcmc_alg = RWMHSampling(), obs_sample = [1.0], - init_params = 3.0, + init_params = transform_constrained_to_unconstrained(prior, [2.0]), step = 0.25, rng = Random.GLOBAL_RNG, target_acc = 0.25, + return_samples = false, ) if !isa(obs_sample, AbstractVecOrMat) obs_sample = reshape(collect(obs_sample), 1) # scalar -> Vector end + init_params = vec(collect(init_params)) # scalar or Vector -> Vector - init_params = reshape(collect(init_params), 1) # scalar or Vector -> Vector mcmc = MCMCWrapper(mcmc_alg, obs_sample, prior, em; init_params = init_params) # First let's run a short chain to determine a good step size new_step = optimize_stepsize(mcmc; init_stepsize = step, N = 5000, target_acc = target_acc) @@ -151,124 +263,244 @@ function mcmc_test_template( # Now begin the actual MCMC, sample is multiply exported so we qualify chain = MCMC.sample(rng, mcmc, 50_000; stepsize = new_step, discard_initial = 10000) posterior_distribution = get_posterior(mcmc, chain) - #post_mean = mean(posterior, dims=1)[1] - posterior_mean = mean(posterior_distribution) + if TEST_PLOT_OUTPUT + # plot: + pp = plot(prior, title = exp_name, color = :grey) + plot!(pp, posterior_distribution, color = :blue) + savefig(pp, joinpath(@__DIR__, "posterior_$(exp_name).png")) + end + + + if return_samples # return sanples not mean + constrained_posterior_samples = + transform_unconstrained_to_constrained(prior, get_distribution(posterior_distribution)) # still a Dict (param_name => samples) + return new_step, constrained_posterior_samples, chain + else + constrained_posterior_mean = transform_unconstrained_to_constrained(prior, mean(posterior_distribution)) + if length(constrained_posterior_mean) == 1 + return new_step, constrained_posterior_mean[1], chain + else + return new_step, constrained_posterior_mean, chain + end + end - return new_step, posterior_mean[1], chain end @testset "MarkovChainMonteCarlo" begin - obs_sample = [1.0] + + # Problem and emulator setups: then MCMC tests come after + + # [1.] 1D -> 1D + # setup + obs_sample = [5.0] + mle = 2.0 prior = test_prior() - y, σ2_y, iopairs, rng = test_data(; rng_seed = 42, n = 40, var_y = 0.05) - mcmc_params = - Dict(:mcmc_alg => RWMHSampling(), :obs_sample => obs_sample, :init_params => [3.0], :step => 0.5, :rng => rng) + test_data_kwargs = (; rng_seed = 42, n = 80, var_y = 0.05) + y, σ2_y, iopairs, rng = test_data(prior; test_data_kwargs...) # iopairs unconstrained inputs + + # mcmc setup + mcmc_params = Dict( + :mcmc_alg => RWMHSampling(), + :obs_sample => obs_sample, + :init_params => transform_constrained_to_unconstrained(prior, [2.0]), + :step => 0.25, + :rng => rng, + :target_acc => 0.25, + ) + mcmc_params_pcn = deepcopy(mcmc_params) + mcmc_params_pcn[:mcmc_alg] = pCNMHSampling() + mcmc_params_pcn[:target_acc] = 0.6 + mcmc_params_pcn[:step] = 0.025 - @testset "MV priors" begin - # 10D dist with 1 name, just build the wrapper for test - prior_mv = test_prior_mv() - y_mv, σ2_y_mv, iopairs_mv, rng_mv = test_data(input_dim = 10) - init_params = repeat([0.0], 10) - obs_sample = obs_sample # scalar or Vector -> Vector - em_mv = test_gp_mv(y_mv, σ2_y_mv, iopairs_mv) - mcmc = MCMCWrapper(RWMHSampling(), obs_sample, prior_mv, em_mv; init_params = init_params) + # build and validate the emulators + em_1, em_1b = test_gp_and_agp_1(y, σ2_y, iopairs) + em_2 = test_gp_2(y, σ2_y, iopairs) + em_s = test_srfi(y, σ2_y, iopairs) + + validate_emulator(em_s, "test_data", prior, test_data_kwargs, exp_name = "srfi_1d") + validate_emulator(em_1, "test_data", prior, test_data_kwargs, exp_name = "gpjl_1d") + validate_emulator(em_1b, "test_data", prior, test_data_kwargs, exp_name = "agpjl_1d") + validate_emulator(em_2, "test_data", prior, test_data_kwargs, exp_name = "gpjl-svd_1d") + + # [2.] 10D -> 1D + # setup + obs_sample_mv = [5.0] + input_dim = 5 + mle_norm_mv = 2.0 * input_dim + prior_mv = test_prior_mv(input_dim) + test_data_mv_kwargs = (; input_dim = input_dim) + y_mv, σ2_y_mv, iopairs_mv, rng_mv = test_data_mv(prior_mv; test_data_mv_kwargs...) # iopairs unconstrained inputs + + # build and validate the emulators + em_mv = test_gp_mv(y_mv, σ2_y_mv, iopairs_mv) + validate_emulator(em_mv, "test_data_mv", prior_mv, test_data_mv_kwargs, exp_name = "gp_mv") + + # mcmc setup + mcmc_params_mv = Dict( + :mcmc_alg => RWMHSampling(), + :obs_sample => obs_sample_mv, + :init_params => transform_constrained_to_unconstrained(prior_mv, repeat([2.1], input_dim)), + :step => 0.25, + :rng => rng_mv, + ) - end + mcmc_params_mv_pcn = deepcopy(mcmc_params_mv) + mcmc_params_mv_pcn[:mcmc_alg] = pCNMHSampling() + mcmc_params_mv_pcn[:target_acc] = 0.6 + mcmc_params_mv_pcn[:step] = 0.025 + + # [3.] 2D -> 2D + + # setup + obs_sample_2d = [1.0, 1.0] + mle_2d = [π / 2, π / 2] + prior_2d = test_prior_2d() + test_data_2d_kwargs = (; rng_seed = 4141, n = 100, cov_y = 0.05 * [[0.5, 0.2] [0.2, 0.5]]) + y_2d, σ2_y_2d, iopairs_2d, rng_2d = test_data_2d(prior_2d; test_data_2d_kwargs...) + em_v = test_vrfi(y_2d, σ2_y_2d, iopairs_2d) + validate_emulator(em_v, "test_data_2d", prior_2d, (; test_data_2d_kwargs)) + + mcmc_params_2d = Dict( + :mcmc_alg => RWMHSampling(), + :obs_sample => obs_sample_2d, + :init_params => transform_constrained_to_unconstrained(prior_2d, [2.0, 2.0]), + :step => 0.025, + :rng => rng_2d, + ) + @testset "1D-1D GP/RF & RW Metropolis" begin - @testset "Sine GP & RW Metropolis" begin - em_1, em_1b = test_gp_and_agp_1(y, σ2_y, iopairs) - new_step_1, posterior_mean_1, chain_1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...) + # test various MCMC methods + new_step_1, posterior_mean_1, chain_1 = + mcmc_test_template(prior, σ2_y, em_1; exp_name = "gpjl_1d", mcmc_params...) esjd1 = esjd(chain_1) - @info "ESJD = $esjd1" - @test isapprox(new_step_1, 0.5; atol = 0.5) + @info "ESJD [GPJL,RW] = $esjd1" + @test isapprox(new_step_1, 0.125; atol = 0.125) # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling - @test isapprox(posterior_mean_1, π / 2; atol = 4e-1) + @test isapprox(posterior_mean_1, mle; atol = 1e-1) + @info "Posterior mean: $(posterior_mean_1) ≈ $(mle)" # test the agp setup on without derivatives - new_step_1b, posterior_mean_1b, chain_1b = mcmc_test_template(prior, σ2_y, em_1b; mcmc_params...) - @test isapprox(new_step_1b, new_step_1; atol = 1e-1) - tol_small = 5e-2 - @test isapprox(posterior_mean_1b, posterior_mean_1; atol = tol_small) + new_step_1b, posterior_mean_1b, chain_1b = + mcmc_test_template(prior, σ2_y, em_1b; exp_name = "agpjl_1d", mcmc_params...) + @test isapprox(new_step_1b, new_step_1; atol = 0.125) + tol_small = 0.2 esjd1b = esjd(chain_1b) - @info "ESJD = $esjd1b" + @info "ESJD [AGPJL,RW] = $esjd1b" @test all(isapprox.(esjd1, esjd1b, rtol = 0.2)) + @test isapprox(posterior_mean_1b, posterior_mean_1; atol = tol_small) + @info "Posterior mean: $(posterior_mean_1b) ≈ $(mle)" + + # test with random features + + new_step_s, posterior_mean_s, chain_s = + mcmc_test_template(prior, σ2_y, em_s; exp_name = "srfi_1d", mcmc_params...) + esjds = esjd(chain_s) + @info "ESJD [SRFI,RW] = $esjds" + # approx [0.0002, 0.26, 0.4]" + @test all(isapprox.(esjd1, esjds, rtol = 0.5)) + @info "Posterior mean: $(posterior_mean_s) ≈ $(mle)" + @test isapprox(posterior_mean_s, mle; atol = 2e-1) # now test SVD normalization - em_2 = test_gp_2(y, σ2_y, iopairs) - _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...) + _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; exp_name = "gpjl-svd_1d", mcmc_params...) # difference between mean_1 and mean_2 only from MCMC convergence - @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1) # test diagnostic functions on the chain esjd2 = esjd(chain_2) - @info "ESJD = $esjd2" - # approx [0.04190683285347798, 0.1685296224916364, 0.4129400000002722] - @test all(isapprox.(esjd1, esjd2, rtol = 0.2)) + @info "ESJD (-with-svd)= $esjd2" + # approx [0.00015, 0.26, 0.35] + @test all(isapprox.(esjd1, esjd2, rtol = 0.5)) + @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1) + @info "Posterior mean: $(posterior_mean_2) ≈ $(mle)" # test with many slightly different samples # as vec of vec obs_sample2 = [obs_sample + 0.01 * randn(length(obs_sample)) for i in 1:100] - mcmc_params2 = mcmc_params + mcmc_params2 = deepcopy(mcmc_params) mcmc_params2[:obs_sample] = obs_sample2 - em_1 = test_gp_1(y, σ2_y, iopairs) - new_step, posterior_mean_1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params2...) - @test isapprox(new_step, 0.5; atol = 0.5) + mcmc_params2[:step] = 0.025 # less uncertainty -> smaller step + new_step, posterior_mean_1 = mcmc_test_template(prior, σ2_y, em_1; exp_name = "gpjl-samples", mcmc_params2...) + @test isapprox(new_step, 0.025; atol = 0.025) # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling - @test isapprox(posterior_mean_1, π / 2; atol = 4e-1) + esjd2 = esjd(chain_2) + @info "ESJD (-many-samples)= $esjd2" + @info "Posterior mean: $(posterior_mean_1) ≈ $(mle)" + @test isapprox(posterior_mean_1, mle; atol = 2e-1) + # as column matrix obs_sample2mat = reduce(hcat, obs_sample2) - mcmc_params2mat = mcmc_params + mcmc_params2mat = deepcopy(mcmc_params) mcmc_params2mat[:obs_sample] = obs_sample2mat - new_step, posterior_mean_1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params2mat...) - @test isapprox(new_step, 0.5; atol = 0.5) + new_step, posterior_mean_1 = + mcmc_test_template(prior, σ2_y, em_1; exp_name = "gpjl-samples-mat", mcmc_params2mat...) + @test isapprox(new_step, 0.025; atol = 0.025) # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling - @test isapprox(posterior_mean_1, π / 2; atol = 4e-1) - + @info "Posterior mean: $(posterior_mean_1) ≈ $(mle)" + @test isapprox(posterior_mean_1, mle; atol = 2e-1) - # test with int data - obs_sample3 = [1] - mcmc_params3 = mcmc_params + # test with integer data + obs_sample3 = [4] + mcmc_params3 = deepcopy(mcmc_params) mcmc_params3[:obs_sample] = obs_sample3 - em_1 = test_gp_1(y, σ2_y, iopairs) - new_step, posterior_mean_1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params3...) + new_step, posterior_mean_1 = mcmc_test_template(prior, σ2_y, em_1; exp_name = "gpjl_int_1d", mcmc_params3...) @test isapprox(new_step, 0.5; atol = 0.5) # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling - @test isapprox(posterior_mean_1, π / 2; atol = 4e-1) + @info "Posterior mean: $(posterior_mean_1) ≈ $(mle)" + @test isapprox(posterior_mean_1, mle; atol = 2e-1) + + + end + @testset "2D-2D RF & RW" begin + + new_step_v, posterior_mean_v, chain_v = mcmc_test_template(prior_2d, σ2_y_2d, em_v; mcmc_params_2d...) + @test isapprox(new_step_v, 0.5; atol = 0.5) + esjdv = esjd(chain_v) + @info "ESJD [VRFI, 2d]= $esjdv" + # [0.025, 0.003, 0.2, 0.27] + @info "Posterior mean: $(posterior_mean_v) ≈ $(mle_2d)" + @test all(isapprox.(posterior_mean_v - mle_2d, 0.0; atol = 4e-1)) + end - @testset "Sine GP & pCN" begin - mcmc_params = Dict( - :mcmc_alg => pCNMHSampling(), - :obs_sample => obs_sample, - :init_params => [3.0], - :step => 0.5, - :rng => rng, - ) - - em_1, em_1b = test_gp_and_agp_1(y, σ2_y, iopairs) - new_step_1, posterior_mean_1, chain_1 = mcmc_test_template(prior, σ2_y, em_1; mcmc_params...) + + @testset "1D-1D pCN" begin + + new_step_1, posterior_mean_1, chain_1 = + mcmc_test_template(prior, σ2_y, em_1; exp_name = "gpjl_pcn", mcmc_params_pcn...) esjd1 = esjd(chain_1) - @info "ESJD = $esjd1" - @test isapprox(new_step_1, 0.75; atol = 0.6) + @info "ESJD [GPJL,pCN] = $esjd1" + @test isapprox(new_step_1, 0.125; atol = 0.125) # difference between mean_1 and ground truth comes from MCMC convergence and GP sampling - @test isapprox(posterior_mean_1, π / 2; atol = 4e-1) + @info "Posterior mean: $(posterior_mean_1) ≈ $mle" + @test isapprox(posterior_mean_1, mle; atol = 2e-1) # test the agp setup on without derivatives - new_step_1b, posterior_mean_1b, chain_1b = mcmc_test_template(prior, σ2_y, em_1b; mcmc_params...) - @test isapprox(new_step_1b, new_step_1; atol = 1e-1) - tol_small = 5e-2 + new_step_1b, posterior_mean_1b, chain_1b = + mcmc_test_template(prior, σ2_y, em_1b; exp_name = "agpjl_pcn", mcmc_params_pcn...) + @test isapprox(new_step_1b, new_step_1; atol = 0.125) + tol_small = 1e-1 @test isapprox(posterior_mean_1b, posterior_mean_1; atol = tol_small) esjd1b = esjd(chain_1b) - @info "ESJD = $esjd1b" + @info "ESJD [AGPJL,pCN] = $esjd1b" + @info "Posterior mean: $(posterior_mean_1b) ≈ $mle" @test all(isapprox.(esjd1, esjd1b, rtol = 0.2)) + # test with random features + new_step_s, posterior_mean_s, chain_s = + mcmc_test_template(prior, σ2_y, em_s; exp_name = "srfi_1d_pcn", mcmc_params_pcn...) + @test isapprox(posterior_mean_s, mle; atol = 2e-1) + esjds = esjd(chain_s) + @info "ESJD [SRFI,pCN] = $esjds" + # approx [0.0002, 0.26, 0.4]" + @info "Posterior mean: $(posterior_mean_s) ≈ $mle" + @test all(isapprox.(esjd1, esjds, rtol = 0.5)) + # now test SVD normalization - em_2 = test_gp_2(y, σ2_y, iopairs) - _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params...) + _, posterior_mean_2, chain_2 = mcmc_test_template(prior, σ2_y, em_2; mcmc_params_pcn...) # difference between mean_1 and mean_2 only from MCMC convergence @test isapprox(posterior_mean_2, posterior_mean_1; atol = 0.1) @@ -280,37 +512,56 @@ end end - - - - mcmc_params = Dict(:obs_sample => obs_sample, :init_params => [3.0], :step => 0.25, :rng => rng, :target_acc => 0.6) # the target is usually higher in grad-based MCMC - - em_1, em_1b = test_gp_and_agp_1(y, σ2_y, iopairs) - # em_1 cannot be used here - - mcmc_algs = [ - RWMHSampling(), # sanity-check - BarkerSampling(), # ForwardDiffProtocol by default - BarkerSampling{ReverseDiffProtocol}(), # scales to high dim better, but slow. - ] - - bad_mcmc_alg = BarkerSampling{GradFreeProtocol}() - @test_throws ArgumentError mcmc_test_template(prior, σ2_y, em_1; mcmc_alg = bad_mcmc_alg, mcmc_params...) - - - # GPJL doesnt support ForwardDiff - @test_throws ErrorException mcmc_test_template(prior, σ2_y, em_1; mcmc_alg = mcmc_algs[2], mcmc_params...) - - for alg in mcmc_algs - @info "testing algorithm: $(typeof(alg))" - new_step, posterior_mean, chain = mcmc_test_template(prior, σ2_y, em_1b; mcmc_alg = alg, mcmc_params...) - esjd_tmp = esjd(chain) - @info "ESJD = $esjd_tmp" - @info posterior_mean - @testset "Sine GP & ForwardDiff variant:$(typeof(alg))" begin - @test isapprox(posterior_mean, π / 2; atol = 4e-1) + @testset "ND-1D" begin + @info "Input dimension: $(input_dim)" + for params in (mcmc_params_mv, mcmc_params_mv_pcn) + # 10D dist with 1 name, just build the wrapper for test + new_step_mv, constrained_posterior_samples_mv, chain_mv = + mcmc_test_template(prior_mv, σ2_y_mv, em_mv; exp_name = "gpjl_10d", return_samples = true, params...) + esjd_mv = esjd(chain_mv) + @info "ESJD (ND-1D) [GPJL,RW] = $esjd_mv" + # as function performs f(norm(x)), posterior should have good properties of the mean(norm(x)) + samples_mat_mv = reduce(vcat, [v for v in values(constrained_posterior_samples_mv)]) + norm_samples_mv = norm.([c for c in eachcol(samples_mat_mv)]) + posterior_norm_mean_mv = mean(norm_samples_mv) + G_mean_mv = mean(G(norm_samples_mv / input_dim)) + @info "Mean of posterior sample norm: $(posterior_norm_mean_mv) ≈ $(mle_norm_mv)" + @info "Mean of G sample norm: $(G_mean_mv) ≈ $(obs_sample_mv[1])" end + end + @testset "Autodiff MCMC variants" begin + mcmc_algs = [ + RWMHSampling(), # sanity-check + BarkerSampling(), # ForwardDiffProtocol by default + BarkerSampling{ReverseDiffProtocol}(), # scales to high dim better, but slow. + ] + + bad_mcmc_alg = BarkerSampling{GradFreeProtocol}() + bad_mcmc_params = deepcopy(mcmc_params) + bad_mcmc_params[:mcmc_alg] = bad_mcmc_alg + + @test_throws ArgumentError mcmc_test_template(prior, σ2_y, em_1; bad_mcmc_params...) + + # GPJL doesnt support ForwardDiff + bad_mcmc_params = deepcopy(mcmc_params) + bad_mcmc_params[:mcmc_alg] = mcmc_algs[2] + @test_throws ErrorException mcmc_test_template(prior, σ2_y, em_1; bad_mcmc_params...) + + for alg in mcmc_algs + mcmc_params_ad = deepcopy(mcmc_params) + mcmc_params_ad[:mcmc_alg] = alg + mcmc_params_ad[:target_acc] = 0.6 # should be > 0.5 + + @info "testing algorithm: $(typeof(alg))" + new_step, posterior_mean, chain = mcmc_test_template(prior, σ2_y, em_1b; mcmc_alg = alg, mcmc_params_ad...) + esjd_tmp = esjd(chain) + @info "ESJD = $esjd_tmp" + @info "Posterior mean: $(posterior_mean) ≈ $mle" + @testset "Sine GP & ForwardDiff variant:$(typeof(alg))" begin + @test isapprox(posterior_mean, mle; atol = 1e-1) + end + end end end