Multi-Fidelity Gaussian Process Surrogates#
Example 1: Regression of a synthetic function#
Suppose we have a high-fidelity model \(f_h\) and a low-fidelity model \(f_\ell\) of some phenomenon, given by
\begin{align*} f_h(\mathbf{x}) &= \frac{1}{2} \sin^2\left( \frac{5}{2} x_1 + \frac{2}{3} x_2 \right) + \frac{2}{3} e^{-x_1 (x_2 - \frac{1}{2})^2} \cos^2(4x_1 + x_2) \ f_\ell(\mathbf{x}) &= 2.5 f_h(\mathbf{x}) + \frac{1}{3}\left[\sin(x_1 + x_2) + \frac{1}{2} e^{-x_1} \sin(x_1 + 7x_2)\right]. \end{align*}
These function definitions are modified from Perdikaris et al. (2015).
We want to create a surrogate for \(f_h\).
Low-fidelity GP#
Suppose we have some low-fidelity data \((\mathbf{X}_\ell, \mathbf{y}_\ell)\) and some high-fidelity data \((\mathbf{X}_h, \mathbf{y}_h)\).
N_LOW_FIDELITY = 50
N_HIGH_FIDELITY = 8
N_TEST = 100
key, subkey = jr.split(key)
X_train_l, X_train_h, X_test, y_train_l, y_train_h, y_test_l, y_test_h = generate_synthetic_data(N_LOW_FIDELITY, N_HIGH_FIDELITY, N_TEST, key=subkey)
array([<Axes: xlabel='$x_1$', ylabel='$x_2$'>,
<Axes3D: xlabel='$x_1$', ylabel='$x_2$', zlabel='$y$'>],
dtype=object)
We’ll start by fitting a Gaussian process to the low-fidelity data:
params_l = {
'log_amplitude': jnp.log(1.0),
'log_lengthscales': jnp.log(jnp.array([1.0, 1.0]))
}
key, subkey = jr.split(key)
params_l, losses_l = train_gp(
build_gp=build_gp,
init_params=params_l,
X=X_train_l,
y=y_train_l,
num_iters=5000,
learning_rate=1e-3,
batch_size=10,
key=subkey
)
Let’s visualize the low-fidelity data and the function \(f_\ell\).
X1, X2 = jnp.meshgrid(jnp.linspace(0, 1, 50), jnp.linspace(0, 1, 50))
Xq_plt = jnp.stack([X1.ravel(), X2.ravel()], axis=1)
cond_gp_plt = eval_gp(build_gp, Xq_plt, X_train_l, y_train_l, params_l)
Y_mean = cond_gp_plt.mean.reshape(*X1.shape)
array([<Axes: xlabel='$x_1$', ylabel='$x_2$'>], dtype=object)
Let’s check the fit of the low-fidelity GP to the low-fidelity data with a pairplot.
array([<Axes: xlabel='True', ylabel='Predicted'>], dtype=object)
As expected, the low-fidelity GP matches the low-fidelity function \(f_\ell\) well, but it completely misses the high-fidelity function \(f_h\).
Multi-fidelity GP#
For the multi-fidelity GP, we’ll use the nonlinear autoregressive kernel of Perdikaris et al. (2017),
where \(\tilde{m}_l\) is the mean of the low-fidelity GP.
In tinygp, this is best implemented with a tinygp.transforms.Transform object as below:
def generate_build_multi_fidelity_gp(build_gp, X_train_l, y_train_l, params_l):
"""Factory function for creating multi-fidelity GP builders."""
def build_multi_fidelity_gp(params, X):
"""Build a multi-fidelity Gaussian process with RBF kernel."""
sigma = 1e-3 # Fixing the measurement noise
amp = jnp.exp(params['log_amplitude'])
ell = jnp.exp(params['log_lengthscales'])
ell_aug = jnp.exp(params['log_lengthscale_l']) # Lengthscale for the low-fidelity GP mean
# This is the mean, m_l, of the low-fidelity GP.
# It is wrapped so that one simply needs to pass in a single input x vector.
# It is also jitted.
eval_low_fidelity_gp_mean = jit(lambda x: partial(eval_gp, build_gp)(x[None], X_train_l, y_train_l, params_l).mean)
# This is the function that lifts the input x from e.g., 2 dimensions to 3 dimensions,
# where the 3rd dimension represents the low-fidelity GP mean.
# It also applies the length scaling.
lift_and_scale = lambda x: jnp.hstack([x/ell, eval_low_fidelity_gp_mean(x)/ell_aug])
# The kernel is constructed in the lifted input space via a Transform.
# The way Transform works is that for some transformation T, Transform(T, k1) produces
# the kernel k(x, x') = k1(T(x), T(x')).
k = amp*transforms.Transform(lift_and_scale, kernels.ExpSquared())
return GaussianProcess(k, X, diag=sigma**2)
return build_multi_fidelity_gp
Here’s how we can train the multi-fidelity GP:
params_m = {
'log_amplitude': jnp.log(1.0),
'log_lengthscales': jnp.log(jnp.array([1.0, 1.0])),
'log_lengthscale_l': jnp.log(1.0)
}
build_multi_fidelity_gp = generate_build_multi_fidelity_gp(build_gp, X_train_l, y_train_l, params_l)
params_m, losses_m = train_gp(
build_gp=build_multi_fidelity_gp,
init_params=params_m,
X=X_train_h,
y=y_train_h,
num_iters=5000,
learning_rate=1e-3,
batch_size=5,
key=subkey
)
This is what the multi-fidelity GP looks like:
X1, X2 = jnp.meshgrid(jnp.linspace(0, 1, 50), jnp.linspace(0, 1, 50))
Xq_plt = jnp.stack([X1.ravel(), X2.ravel()], axis=1)
cond_gp_plt = eval_gp(build_multi_fidelity_gp, Xq_plt, X_train_h, y_train_h, params_m)
Y_mean = cond_gp_plt.mean.reshape(*X1.shape)
array([<Axes: xlabel='$x_1$', ylabel='$x_2$'>], dtype=object)
Let’s compare the accuracy of the multi-fidelity GP to that of a GP fit only to the high-fidelity data.
array([<Axes: title={'center': 'High-fidelity-only GP'}, xlabel='True', ylabel='Predicted'>,
<Axes: title={'center': 'Multi-fidelity GP'}, xlabel='True', ylabel='Predicted'>],
dtype=object)
The multi-fidelity GP has the more accurate predictions. Let’s visualize the mean predictive surface against the ground truth:
The multi-fidelity GP \(\hat{f}_m\) approximates the high-fidelity model \(f_h\) fairly well! This is a significant improvement over naively fitting to the high-fidelity data alone.
Exercises#
Decrease
N_LOW_FIDELITY. How does the multi-fidelity GP \(\hat{f}_m\) perform with less low-fidelity data?Decrease
N_HIGH_FIDELITY. How does multi-fidelity GP \(\hat{f}_m\) perform with less high-fidelity data?Increase
N_HIGH_FIDELITY. At what point is the high-fidelity-only GP \(\hat{f}_h\) just as good as the multi-fidelity GP \(\hat{f}_m\)?Add more terms (sine/cosine, exponential, quadratic, or another form) to
low_fidelity_model. How different can the low-fidelity model \(f_\ell\) be from the high-fidelity model \(f_h\) and still yield a good surrogate \(\hat{f}_m\)?
Example 2: Stochastic incompressible flow past a cylinder#
This example is adapted from Perdikaris et al. (2015). Suppose you have a flow past a cylinder, subject to random inflow boundary conditions of the form
Let \(C_\text{BP}\) be the base pressure coefficient at the rear of the cylinder (see Figure 9 of Perdikaris et al. (2015)).
Fig. 6 Finite-element mesh and flow field used to evaluate the base-pressure coefficient. Adapted from Perdikaris et al. (2015).#
The quantity of interest is the mean of the upper 40% distribution for \(C_\text{BP}\), i.e. the superquantile risk \(f(x) \equiv \mathcal{R}_{0.6}[C_\text{BP}](x)\). We have two different-fidelity models that compute \(f\). To train the surrogate, we have 8 simulations from the high-fidelity model \(f_h\) and 99 simulations from the low-fidelity model \(f_\ell\).
Here are the data:
array([<Axes: xlabel='$\\sigma_1$', ylabel='$\\sigma_2$'>,
<Axes3D: xlabel='$\\sigma_1$', ylabel='$\\sigma_2$', zlabel='$y$'>],
dtype=object)
Multi-fidelity Gaussian process for superquantile risk#
As before, we first construct the low-fidelity GP surrogate:
params_l_cyl = {
'log_amplitude': jnp.log(1.0),
'log_lengthscales': jnp.log(jnp.array([1.0, 1.0]))
}
key, subkey = jr.split(key)
params_l_cyl, losses_l_cyl = train_gp(
build_gp=build_gp,
init_params=params_l_cyl,
X=Xl_cyl,
y=yl_cyl,
num_iters=5000,
learning_rate=1e-3,
batch_size=10,
key=subkey
)
Next the multi-fidelity GP surrogate:
params_m_cyl = {
'log_amplitude': jnp.log(1.0),
'log_lengthscales': jnp.log(jnp.array([1.0, 1.0])),
'log_lengthscale_l': jnp.log(1.0)
}
build_multi_fidelity_gp = generate_build_multi_fidelity_gp(build_gp, Xl_cyl, yl_cyl, params_l_cyl)
params_m_cyl, losses_m_cyl = train_gp(
build_gp=build_multi_fidelity_gp,
init_params=params_m_cyl,
X=Xh_cyl,
y=yh_cyl,
num_iters=5000,
learning_rate=1e-3,
batch_size=5,
key=subkey
)
And let’s also construct a surrogate on just the high-fidelity data, for comparison:
params_h_cyl = {
'log_amplitude': jnp.log(1.0),
'log_lengthscales': jnp.log(jnp.array([1.0, 1.0]))
}
key, subkey = jr.split(key)
params_h_cyl, losses_h_cyl = train_gp(
build_gp=build_gp,
init_params=params_h_cyl,
X=Xh_cyl,
y=yh_cyl,
num_iters=5000,
learning_rate=1e-3,
batch_size=10,
key=subkey
)
As before, let’s visualize the predictive accuracy with some parity plots:
array([<Axes: xlabel='True model', ylabel='Surrogate'>,
<Axes: xlabel='True model', ylabel='Surrogate'>,
<Axes: xlabel='True model', ylabel='Surrogate'>], dtype=object)
The multi-fidelity GP has the best predictive accuracy. Let’s visualize the response surface of the surrogate vs. the true high-fidelity model:
The surfaces are almost right on top of each other.