README.md

CalibrateEmulateSample.jl

Implements a derivative-free machine-learning-accelerated pipeline for uncertainty quantification.

Documentation		Version
License		Docs Build
Tests		Code Coverage
Downloads		JOSS
DOI

CalibrateEmulateSample.jl solves parameter estimation problems using accelerated (and approximate) Bayesian inversion. Given an observation y, a computer model (forward map) G, observational noise covariance Γ, and broad prior information on the parameters θ, it returns the joint data-informed distribution of "θ given y" for y = G(θ) + η, η ~ N(0, Γ). The framework is not inherently restricted to unimodal distributions and can be applied to computer models that are noisy or chaotic, non-differentiable, or that can only be treated as a black box (interfaced only through parameter files). As the name suggests, the package breaks the problem into three stages — calibration (via EnsembleKalmanProcesses.jl), emulation of the forward map, and sampling of the emulated posterior — following the approach described in Cleary et al. (2020).

Quick Links

• I'm getting errors about the python dependency
• How do I build prior distributions?
• How do I build good observational noise covariances
• What ensemble size should I take? Which process should I use? What is the recommended configuration?
• Where can I walk through the simple example?
• What is the EnsembleKalmanProcesses.jl package?
• What are the recommendations/defaults for dimension reduction or data processing?
• How do I plot or interpret the posterior distribution?

Features

• Derivative-free Bayesian inversion: learns the joint distribution over parameters θ given noisy observations y = G(θ) + η, without requiring gradients of the forward model G.
• Works with black-box forward models: compatible with computer models that are noisy, chaotic, non-differentiable, or only interfaced through parameter files.
• Three-stage pipeline: Calibrate with EnsembleKalmanProcesses.jl, Emulate the forward map with a Gaussian Process or Random Feature interface, and Sample the emulated posterior with Markov chain Monte Carlo.
• Flexible emulators: Gaussian Process and Random Feature emulator interfaces, with configurable kernel structures and input/output data encoding and decorrelation schedules.
• Dimension reduction and data processing: recommended defaults for encoding, decorrelating, and reducing high-dimensional input-output pairs before emulation.
• Not restricted to unimodal posteriors: supports flexible numbers of parameters and multimodal joint distributions.

Quick Example

Having installed CES and Plots, you can copy-paste the snippets below to recreate this random experiment (up to random number generation).

Calibrate

Below we will outline the current user experience for using EnsembleKalmanProcesses.jl. We solve the classic inverse problem where we learn y = G(u), noisy forward map G distributed as N(0,Γ). For example,

using LinearAlgebra
sdd=1.0
Γ = (sdd)^2*I
G(u) = [
    1/abs(u[1]),
    sum(u[2:5]),
    prod(u[3:4]),
    u[1]^2-u[2]-u[3],
    u[4],
    u[5]^3,
    ] .+ sdd*randn(6)
true_u = [3, 1, 2,-3,-4]
y = G(true_u)

We assume some prior knowledge of the parameters u in the problem (such as approximate scales, and the first parameter being positive), then we learn which parameters best fit the data with EKP.

using CalibrateEmulateSample.EnsembleKalmanProcesses
using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions

prior_u1 = constrained_gaussian("positive_with_mean_2", 2, 1, 0, Inf)
prior_u2 = constrained_gaussian("four_with_spread_5", 0, 3, -Inf, Inf, repeats=4)
prior = combine_distributions([prior_u1, prior_u2])

N_ensemble = 50
initial_ensemble = construct_initial_ensemble(prior, N_ensemble)
ensemble_kalman_process = EnsembleKalmanProcess(
    initial_ensemble, y, Γ, Inversion(), verbose=true)

N_iterations = 10
for i in 1:N_iterations
    params_i = get_ϕ_final(prior, ensemble_kalman_process)

    G_matrix = hcat(
        [G(params_i[:, i]) for i in 1:N_ensemble]... # Parallelize here!
    )

    update_ensemble!(ensemble_kalman_process, G_matrix)
end

final_solution = get_ϕ_mean_final(prior, ensemble_kalman_process)

# Let's see what's going on!
using Plots
p = plot(prior)
for (i,sp) in enumerate(p.subplots)
    vline!(sp, [true_u[i]], lc="black", lw=4)
    vline!(sp, [final_solution[i]], lc="magenta", lw=4)
end
display(p)

We now have a quick and cheap estimate of the mode of the distribution here, but no quantified uncertainty. To postprocess in a way that gives us back uncertainty, we will emulate G using our current EKP evaluations, and apply a proper sampling algorithm to the emulator. No more evaluations of G are required!

Emulate

We then set up the emulation framework, by getting input-output pairs from the EKP, defining which emulator to use, and defining what space we will train our emulator in. Here, we take a Random Feature emulator, and we decorrelate the training space, using prior (input) and noise (output) covariances. This stage can take several minutes.

using CalibrateEmulateSample
using CalibrateEmulateSample.Emulators
using CalibrateEmulateSample.Utilities
using CalibrateEmulateSample.DataContainers
using Statistics

# get training and test data from EKI iterations
input_output_pairs = get_training_points(ensemble_kalman_process, collect(1:N_iterations-2))
input_output_test = get_training_points(ensemble_kalman_process, collect(N_iterations-1:N_iterations))

#choose an ML tool
input_dim = ndims(prior)
n_features = 200
kernel_structure=SeparableKernel(DiagonalFactor(1e-10),OneDimFactor())

mlt = ScalarRandomFeatureInterface(
    n_features,
    input_dim,
    kernel_structure=kernel_structure,
)


# Define a data encoding recipe for the input-output data
encoder_schedule = [(decorrelate_structure_mat(), "in_and_out")]
encoder_kwargs = encoder_kwargs_from(ensemble_kalman_process, prior)

# create emulator and train
emulator = Emulator(
    mlt,
    input_output_pairs;
    encoder_schedule = deepcopy(encoder_schedule),
    encoder_kwargs = deepcopy(encoder_kwargs),
)

optimize_hyperparameters!(emulator)

It is essential to validate the emulator. Here we just predict on train, test and the true parameters to validate the mean prediction lies near 1:1. We recommend you do much more than just this to check your emulator is working as expected!

# make some predictions, and plot
pred_true_mean, pred_true_cov = predict(emulator, reshape(true_u,:,1))
pred_train_mean, pred_train_cov = predict(emulator, get_inputs(input_output_pairs))
pred_test_mean, pred_test_cov = predict(emulator, get_inputs(input_output_test))

pp = scatter(
    vec(get_outputs(input_output_pairs)),
    vec(pred_train_mean),
    xlabel="True outputs",
    ylabel="Predicted outputs",
    label="Training data (EKI iter. 1:N-2)",
    color=:lightgrey,
    markersize=3,
    markerstrokewidth = 0,
    size=(1.6*500,500),
)
scatter!(
    pp,
    vec(get_outputs(input_output_test)),
    vec(pred_test_mean),
    label="Test data (EKI iter. N-1:N)",
    color=:red,
    markersize=3,
    markerstrokewidth = 0,
)
scatter!(
    pp,
    vec(y),
    vec(pred_true_mean),
    color=:black,
    label="True parameters",
    markersize=6,
    markerstrokewidth = 4,
    marker=:+,
)

lo = minimum([minimum(pred_train_mean), minimum(pred_test_mean), minimum(y)])
hi = maximum([maximum(pred_train_mean), maximum(pred_test_mean), minimum(y)])
plot!(pp, [lo, hi], [lo, hi]; lc = :grey, ls = :dash, label = "1:1")
display(pp)

Sample

Having created a suitable emulator, we can pass everything into the sampler, here a random walk metropolis algorithm (MCMC). We find a suitable step size and then draw 100,000 samples from the emulator-based posterior.

using CalibrateEmulateSample.MarkovChainMonteCarlo
mcmc = MCMCWrapper(RWMHSampling(), y, prior, emulator; init_params = get_u_mean_final(ensemble_kalman_process))

new_step = optimize_stepsize(mcmc; init_stepsize = 1.0, N = 2000, discard_initial = 0) # find stepsize

chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) # sample

posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain)

We can plot the prior-posterior, and observe (at least) how the marginals capture the EKI optimum and the true parameter values.

# plot marginals
ppp = plot(prior, fill = :lightgray)
plot!(ppp, posterior, fill = :darkblue, alpha = 0.5)
for (i,sp) in enumerate(ppp.subplots)
    vline!(sp, [true_u[i]], lc="black", lw=4)
    vline!(sp, [final_solution[i]], lc="magenta", lw=4)
end
display(ppp)

We see that the approximate posterior contains both EKI and the true parameter, and that the parameter four_with_spread_5(dim 4) is highly constrained, while four_with_spread_5(dims 1) and (dim 4) are least constrained by the observations, coinciding with the EKI performance. More detail on correlation structure can be extracted by pair-plotting, and posterior analysis.

More runnable examples (Sinusoid, Lorenz, Cloudy, Darcy, EDMF, GCM, and dimension-reduction configurations) are in the examples/ directory.