Kernel Specification

Given the data (and context of the application), we as modelers need to specify a kernel function that captures the structure of the data. The kernel function is a measure of similarity (covariance) between two points in the input space. The most common kernel function is the squared exponential kernel (also known as Gaussian kernel or RBF kernel):

\[ \text{cov}(f_{p}, f_{q}) = k_{a, \ell}(x_p, x_q) = a^{2} \: \exp\left(-\frac{1}{2\ell^2} ||x_p - x_q||^2\right) \]

This kernel, which depends on two parameters \(a\) and \(\ell\), will produce smooth functions. The parameter \(a\) is the amplitude and \(\ell\) is the length-scale. The length-scale controls how wiggly the function is. If \(\ell\) is large, the function will be very flat , and if \(\ell\) is small, the function will be wiggly.

There are many more kernels used in applications. For example:

• We can also use the Matérn kernel, which is a generalization of the squared exponential kernel. The Matérn kernel has an additional parameter \(\nu\) that controls the smoothness of the function. The squared exponential kernel is a special case of the Matérn kernel when \(\nu \to \infty\).
• Periodic kernel, which is useful for modeling periodic functions.

For a great introduction to kernels, I recommend the example notebook “Mean and Covariance Functions” from the PyMC example gallery. If you want to go deeper into the subject then I strongly recommend “Gaussian Processes for Machine Learning - Chapter 4: Covariance Functions”

Remark: The sum of kernels is also a valid kernel. This is useful to model functions that have different scales or periodic components. For example, the kernel \(k(x, x') = k_1(x, x') + k_2(x, x')\) is a valid kernel. Also, the product of kernels is a valid kernel. For example, the kernel \(k(x, x') = k_1(x, x') \cdot k_2(x, x')\) is a valid kernel.

Let’s write the squared exponential kernel in JAX:

# https://num.pyro.ai/en/stable/examples/gp.html
def squared_exponential_kernel(
    xp: ArrayImpl, xq: ArrayImpl, amplitude: float, length_scale: float
) -> ArrayImpl:
    r = xp[:, None] - xq[None, :]
    delta = (r / length_scale) ** 2
    return amplitude**2 * jnp.exp(-0.5 * delta)

We can compute the associated kernel matrix for the training data and fixed values for the amplitude and length-scale. The kernel matrix is a matrix \(K\) such that \(K_{ij} = k(x_i, x_j)\).

kernel = squared_exponential_kernel(x_train, x_train, amplitude=1.0, length_scale=0.1)

One can visualize the kernel matrix as a heatmap to understand the structure of the covariance:

fig, ax = plt.subplots()
sns.heatmap(
    data=kernel,
    square=True,
    cmap="viridis",
    cbar=True,
    xticklabels=False,
    yticklabels=False,
    ax=ax,
)
ax.set_title("Squared Exponential Kernel", fontsize=18, fontweight="bold");

As expected, the kernel matrix is symmetric and the diagonal elements are all equal to \(a^2\).

assert (kernel == kernel.T).all().item()
assert (jnp.diag(kernel) == 1).all().item()

The kernel matrices have an additional property: they are positive semi-definite. This means that their eigenvalues are non-negative.

Remark: In this notebook we will be talking a lot about the spectrum (eigenvalues and eigenvector) of a matrix / operator. Recall that given a matrix \(A\) (or a linear operator) the eigenvalues and eigenvectors are the solutions to the equation \(A v = \lambda v\) where \(v \neq \vec{0}\) is the eigenvector and \(\lambda\) is the eigenvalue. The spectrum of a matrix is the set of its eigenvalues. For more details see the blog post “The Spectral Theorem for Matrices”. We will discuss more about eigenvalues and eigenvectors later.

We can verify this for this specific example by computing the eigenvalues of the kernel matrix on the training set:

spectrum_kernel = jnp.linalg.eigh(kernel)

fig, ax = plt.subplots(figsize=(10, 7))
ax.plot(
    spectrum_kernel.eigenvalues, marker="o", markersize=5, linestyle="--", color="black"
)
ax.set(xlabel="Eigenvalue Index (sorted)", ylabel="Eigenvalue")
ax.set_title("Training Set")
fig.suptitle(
    "Spectrum of the Squared-Exponential Kernel ($a=1$ and $\\ell=0.1$)",
    fontsize=18,
    fontweight="bold",
);

We indeed see that all of the eigenvalues are non-negative (up to some numerical error).

Gaussian Process Model

In this subsection we describe the Gaussian process model. We summarize the main concepts from Gaussian Processes for Machine Learning, Chapter 2: Regression. Recall we used the kernel to construct the kernel matrix \(K := K(X,X)\) from the training data \(X\). We can use this matrix to set a prior distribution over functions. Namely, we can consider a multivariate normal distribution with mean zero and covariance matrix \(K\) so that the sample functions can be realized by sampling from it (in practice we sample a finite number of points and not a function itself). If we denote the test set by \(X_{*}\) and all related quantities regarding this set by lower star, then the joint distribution of \(y\) (observations of the training sets) and \(f_{*}\) (latent GP realization on the test set) is given by (see Equation (2.18) in Gaussian Processes for Machine Learning, Chapter 2: Regression)

\[ \left( \begin{array}{c} y \\ f_* \end{array} \right) \sim \text{MultivariateNormal}(0, \mathcal{K}) \]

where

\[ \mathcal{K} = \left( \begin{array}{cc} K(X, X) + \sigma^2_n I & K(X, X_*) \\ K(X_*, X) & K(X_*, X_*) \end{array} \right) \]

Observe that we need to add the term \(\sigma^2_n I\) to the upper left component to account for noise.

As described in our main reference, to get the posterior distribution over functions we need to restrict this joint prior distribution to contain only those functions which agree with the observed data points, that is, we are interested in computing \(f_*|X, y, X_*\). Using the results of Gaussian Processes for Machine Learning, Appendinx A.2, one can show that (see Equation (2.19) in Gaussian Processes for Machine Learning, Chapter 2: Regression)

\[ f_*|X, y, X_* \sim \text{MultivariateNormal}(\bar{f}_*, \text{cov}(f_*)) \]

where

\[ \bar{f}_* = K(X_*, X)(K(X, X) + \sigma^2_n I)^{-1} y \in \mathbb{R}^{n_*} \]

and

\[ \text{cov}(f_*) = K(X_*, X_*) - K(X_*, X)(K(X, X) + \sigma^2_n I)^{-1} K(X, X_*) \in M_{n_*}(\mathbb{R}) \]

For the NumPyro implementation, we will use these expressions to sample from the posterior distribution over functions on the train and test sets. Note that we are assuming we know the parameters of the kernel function. In practice, we need to infer these parameters from the data. We will discuss this in the next sections.

Parameter Priors

After having a better understanding of the theory behind the Gaussian process model, we can now proceed to the computational aspects. As we are Bayesians (right?), we need to specify priors for the parameters of the kernel function and the noise. This bit itself is a huge topic! (see the great blog post “Priors for the parameters in a Gaussian process” by Dan Simpson). A common choice for the length-scale and amplitude parameters is to use an inverse-gamma distribution as it is positive and drops to zero for values close to zero (this help us to avoid dividing by zero). We expect different ranges for length-scale and amplitude parameters since the former controls the smoothness of the function and the latter the amplitude of the covariance. We can use the function pm.find_constrained_prior to find optimal parameters of the inverse-gamma distribution that align to our believes of the parameters given the data (e.g. plot). In addition, we can use preliz to visualize the resulting prior distributions.

inverse_gamma_params_1 = pm.find_constrained_prior(
    distribution=pm.InverseGamma,
    lower=0.01,
    upper=0.4,
    init_guess={"alpha": 5, "beta": 3},
    mass=0.94,
)
inverse_gamma_params_2 = pm.find_constrained_prior(
    distribution=pm.InverseGamma,
    lower=0.5,
    upper=1.5,
    init_guess={"alpha": 5, "beta": 6},
    mass=0.94,
)

fig, ax = plt.subplots()
pz.InverseGamma(**inverse_gamma_params_1).plot_pdf(color="C0", ax=ax)
pz.InverseGamma(**inverse_gamma_params_2).plot_pdf(color="C1", ax=ax)
ax.set(xlim=(0, 2))
ax.set_title("Prior Distributions", fontsize=18, fontweight="bold");

Remark: In this example we use the same priors for the amplitude and noise parameters as it is well suited for the synthetic data. In practice, we would need to use different priors for these parameters.

Model Specification

We now specify the model as described above: we use the kernel function to compute the kernel matrix and then use this matrix to set the prior distribution over functions as a multivariate normal distribution.

def gp_model(x, y=None, jitter=1.0e-6) -> None:
    # --- Priors ---
    kernel_amplitude = numpyro.sample(
        "kernel_amplitude",
        dist.InverseGamma(
            concentration=inverse_gamma_params_2["alpha"],
            rate=inverse_gamma_params_2["beta"],
        ),
    )
    kernel_length_scale = numpyro.sample(
        "kernel_length_scale",
        dist.InverseGamma(
            concentration=inverse_gamma_params_1["alpha"],
            rate=inverse_gamma_params_1["beta"],
        ),
    )
    noise = numpyro.sample(
        "noise",
        dist.InverseGamma(
            concentration=inverse_gamma_params_2["alpha"],
            rate=inverse_gamma_params_2["beta"],
        ),
    )
    # --- Parametrization ---
    mean = jnp.zeros(x.shape[0])
    kernel = squared_exponential_kernel(x, x, kernel_amplitude, kernel_length_scale)
    covariance_matrix = kernel + (noise + jitter) * jnp.eye(x.shape[0])
    # --- Likelihood ---
    numpyro.sample(
        "likelihood",
        dist.MultivariateNormal(loc=mean, covariance_matrix=covariance_matrix),
        obs=y,
    )

We can visualize the model representation in plate notation:

numpyro.render_model(
    model=gp_model,
    model_args=(x_train, y_train_obs),
    render_distributions=True,
    render_params=True,
)

Prior Predictive Check

Before fitting the model to the data, we can use the model above to generate samples from the prior distribution over functions. This will help us to understand the behavior of the model before seeing the data.

gp_numpyro_prior_predictive = Predictive(gp_model, num_samples=1_000)
rng_key, rng_subkey = random.split(rng_key)
gp_numpyro_prior_samples = gp_numpyro_prior_predictive(rng_subkey, x_train)

gp_numpyro_prior_idata = az.from_numpyro(prior=gp_numpyro_prior_samples)

fig, ax = plt.subplots()
az.plot_hdi(
    x_train,
    gp_numpyro_prior_idata.prior["likelihood"],
    hdi_prob=0.94,
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_train,
    gp_numpyro_prior_idata.prior["likelihood"],
    hdi_prob=0.5,
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$50\\%$ HDI"},
    ax=ax,
)
ax.plot(
    x_train,
    gp_numpyro_prior_idata.prior["likelihood"].mean(dim=("chain", "draw")),
    color="C0",
    linewidth=3,
    label="posterior predictive mean",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
for i in range(5):
    label = "prior samples" if i == 0 else None
    ax.plot(
        x_train,
        gp_numpyro_prior_idata.prior["likelihood"].sel(chain=0, draw=i),
        color="C0",
        alpha=0.3,
        label=label,
    )
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set(xlabel="x", ylabel="y")
ax.set_title("GP NumPyro Model - Prior Predictive", fontsize=18, fontweight="bold");

We see that the generated functions (range and smoothness) are reasonable candidates for the true function. Hence, we can proceed to fit the model to the data.

Model Fitting

We define some helper objects to fit the model to the data so that we can reuse them later.

class InferenceParams(BaseModel):
    num_warmup: int = Field(2_000, ge=1)
    num_samples: int = Field(2_000, ge=1)
    num_chains: int = Field(4, ge=1)


def run_inference(
    rng_key: ArrayImpl,
    model: Callable,
    args: InferenceParams,
    *model_args,
    **nuts_kwargs,
) -> MCMC:
    sampler = NUTS(model, **nuts_kwargs)
    mcmc = MCMC(
        sampler=sampler,
        num_warmup=args.num_warmup,
        num_samples=args.num_samples,
        num_chains=args.num_chains,
    )
    mcmc.run(rng_key, *model_args)
    return mcmc


inference_params = InferenceParams()
rng_key, rng_subkey = random.split(key=rng_key)
gp_mcmc = run_inference(rng_subkey, gp_model, inference_params, x_train, y_train_obs)

Model Diagnostics

The model samples quite fast! We can now look into the posterior distributions and model diagnostics:

gp_numpyro_idata = az.from_numpyro(posterior=gp_mcmc)

az.summary(data=gp_numpyro_idata)

axes = az.plot_trace(
    data=gp_numpyro_idata,
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (12, 7), "layout": "constrained"},
)
plt.gcf().suptitle("GP NumPyro - Trace", fontsize=18, fontweight="bold");

Everything looks good! We are ready to generate predictions for the test set.

Out of Sample Prediction

Recall from the Gaussian model description that we can generate samples on the test set as:

\[ f_*|X, y, X_* \sim \text{MultivariateNormal}(\bar{f}_*, \text{cov}(f_*)) \]

where

\[ \bar{f}_* = K(X_*, X)(K(X, X) + \sigma^2_n I)^{-1} y \]

and

\[ \text{cov}(f_*) = K(X_*, X_*) - K(X_*, X)(K(X, X) + \sigma^2_n I)^{-1} K(X, X_*) \]

We now simply write these expressions in JAX code:

def get_kernel_matrices_train_test(
    kernel, x, x_test, noise, jitter=1.0e-6, **kernel_kwargs
) -> tuple[ArrayImpl, ArrayImpl, ArrayImpl]:
    k = kernel(x, x, **kernel_kwargs) + (noise + jitter) * jnp.eye(x.shape[0])
    k_star = kernel(x_test, x, **kernel_kwargs)
    k_star_star = kernel(x_test, x_test, **kernel_kwargs) + (noise + jitter) * jnp.eye(
        x_test.shape[0]
    )
    return k, k_star, k_star_star


def sample_test(rng_key, kernel, x, y, x_test, noise, jitter=1.0e-6, **kernel_kwargs):
    k, k_star, k_star_star = get_kernel_matrices_train_test(
        kernel, x, x_test, noise, jitter, **kernel_kwargs
    )
    k_inv = jnp.linalg.inv(k)
    mean_star = k_star @ (k_inv @ y)
    cov_star = k_star_star - (k_star @ k_inv @ k_star.T)
    return random.multivariate_normal(rng_key, mean=mean_star, cov=cov_star)

We now vmap (i.e. paralellize) these functions to generate samples from the posterior distribution over functions on the test set.

rng_key, rng_subkey = random.split(rng_key)

vmap_args = (
    random.split(
        rng_subkey, inference_params.num_samples * inference_params.num_chains
    ),
    gp_mcmc.get_samples()["noise"],
    gp_mcmc.get_samples()["kernel_amplitude"],
    gp_mcmc.get_samples()["kernel_length_scale"],
)

posterior_predictive_test = vmap(
    lambda rng_key, noise, amplitude, length_scale: sample_test(
        rng_key,
        squared_exponential_kernel,
        x_train,
        y_train_obs,
        x_test,
        noise,
        amplitude=amplitude,
        length_scale=length_scale,
    ),
)(*vmap_args)

gp_numpyro_idata.extend(
    az.from_dict(
        posterior_predictive={
            "y_pred_test": posterior_predictive_test[None, ...],
        },
        coords={"x": x_test},
        dims={"y_pred_test": ["x"]},
    )
)

Here is the final result:

fig, ax = plt.subplots()
az.plot_hdi(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"],
    hdi_prob=0.94,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"],
    hdi_prob=0.5,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$50\\%$ HDI (test)"},
    ax=ax,
)
ax.plot(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"].mean(dim=("chain", "draw")),
    color="C1",
    linewidth=3,
    label="posterior predictive mean (test)",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
ax.scatter(x_test, y_test_obs, c="C1", label="observed (test)")
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set(xlabel="x", ylabel="y")
ax.set_title("GP NumPyro Model - Posterior Predictive", fontsize=18, fontweight="bold");

It looks very good 🚀! The in-sample fit captures very well the latent Gaussian process mean (black line). The out of sample also works well close to the boundaries (note how the uncertainty grows as we get far away from the training space, as expected). In the following plot we add indicators of one length scale the posterior distribution shift from the boundaries, until where we expect the model to work relatively well.

kernel_length_scale_hdi = az.hdi(gp_numpyro_idata.posterior)[
    "kernel_length_scale"
].to_numpy()


fig, ax = plt.subplots()
plt.axvspan(
    xmin=x_train.min().item() - kernel_length_scale_hdi[1],
    xmax=x_train.min().item() - kernel_length_scale_hdi[0],
    color="C7",
    alpha=0.3,
    label="x_train_min - length scale $94\\%$ HDI",
)
plt.axvspan(
    xmin=x_train.max().item() + kernel_length_scale_hdi[1],
    xmax=x_train.max().item() + kernel_length_scale_hdi[0],
    color="C8",
    alpha=0.3,
    label="x_train_max + length scale $94\\%$ HDI",
)
az.plot_hdi(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"],
    hdi_prob=0.94,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"],
    hdi_prob=0.5,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$50\\%$ HDI (test)"},
    ax=ax,
)
ax.plot(
    x_test,
    gp_numpyro_idata.posterior_predictive["y_pred_test"].mean(dim=("chain", "draw")),
    color="C1",
    linewidth=3,
    label="posterior predictive mean (test)",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
ax.scatter(x_test, y_test_obs, c="C1", label="observed (test)")
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=2)
ax.set(xlabel="x", ylabel="y")
ax.set_title("GP NumPyro Model - Posterior Predictive", fontsize=18, fontweight="bold");

Hilbert Space Approximation Step by Step

In the last part we already implemented the spectral density for the squared exponential kernel in one dimension. Next, we need to write the expressions of the eigenvalues and eigenvector of the Laplacian in the interval \([-L, L]\) (we use the l_max for \(L\) as we do not like upper-case for function arguments 😅).

def laplacian_sqrt_eigenvalues(l_max: float, m: int):
    sqrt_base_eigenvalue = jnp.pi / (2 * l_max)
    return sqrt_base_eigenvalue * jnp.arange(1, 1 + m)


def diag_squared_exponential_spectral_density(
    amplitude: float, length_scale: float, l_max: float, m: int
):
    sqrt_eigenvalues = laplacian_sqrt_eigenvalues(l_max, m)
    return squared_exponential_spectral_density(
        sqrt_eigenvalues, amplitude, length_scale
    )

Recall that the weights in the approximation formula are given by the spectral density evaluated on the square root of the eigenvalues of the Laplacian. Let’s see how these look for various length scales and fixed amplitude \(a=1\) and box length \(L=1.5\).

m = 8
l_max = 1.5


fig, ax = plt.subplots()
for i, length_scale in enumerate(length_scales):
    ax.plot(
        laplacian_sqrt_eigenvalues(l_max=l_max, m=m),
        diag_squared_exponential_spectral_density(
            amplitude=1.0, length_scale=length_scale, l_max=l_max, m=m
        ),
        marker="o",
        c=f"C{i}",
        linewidth=3,
        markersize=8,
        label=f"length_scale={length_scale}",
    )
ax.legend(loc="upper right")
ax.set(
    xlabel="Laplacian Sqrt Eigenvalue $\\sqrt{\\lambda}$",
    ylabel="Spectral Density $S(\\sqrt{\\lambda})$",
)
ax.set_title(
    """Squared Exponential Spectral Density
    as a Function of the Laplacian Eigenvalues (sqrt)
    """,
    fontsize=18,
    fontweight="bold",
);

From this plot we see how the concentration of the spectral density varies with the langth-scale parameter. The larger the length-scale, the more concentrated the spectral density is around zero.

We continue by looking into the Laplacian eigenfunctions

\[ \phi_j(x) = \sqrt{\frac{1}{L}} \sin\left(\frac{\pi j (x + L)}{2L}\right) \]

We can compute them in a vectorized way:

def laplace_eigenfunctions(x, l_max, m):
    sqrt_eigenvalues = laplacian_sqrt_eigenvalues(l_max, m)
    rep_x = jnp.tile(l_max + x[:, None], reps=m)
    diag_sqrt_eigenvalues = jnp.diag(sqrt_eigenvalues)
    num = jnp.sin(rep_x @ diag_sqrt_eigenvalues)
    den = jnp.sqrt(l_max)
    return num / den


phi = laplace_eigenfunctions(x_train, l_max, m)

assert phi.shape == (n_train, m)

We can visualize the first few eigenfunctions on the domain of the training data:

fig, ax = plt.subplots()
for i in range(m):
    ax.plot(
        x_train,
        phi[:, i],
        linewidth=3,
        c=f"C{i}",
        label=f"eigenfunction {i + 1}",
    )
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=4)
ax.set(xlabel="x", ylabel="eigenfunction value")
ax.set_title("Laplacian Eigenfunctions", fontsize=18, fontweight="bold");

Remark: Although it is not strictly necessary, it is recommended to center the training data x_train so that it has mean zero and it is symmetrically located in a bounded box \([-L, L]\). This is automatically handled in the PyMC implementation as we will see below (see “Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming”).

Now we can put everything together as a function that computes the Hilbert space approximation of the squared exponential kernel in one dimension:

def hs_approx_squared_exponential_ncp(x, amplitude, length_scale, l_max, m):
    phi = laplace_eigenfunctions(x, l_max, m)
    spd = jnp.sqrt(
        diag_squared_exponential_spectral_density(amplitude, length_scale, l_max, m)
    )
    with numpyro.plate("basis", m):
        beta = numpyro.sample("beta", dist.Normal(0, 1))

    return numpyro.deterministic("f", phi @ (spd * beta))

Remark: This implementation is using the so called non-centered parametrization. It is the same mathematical model but just written in a way that it samples more efficiently. This is a common practice in hierarchical models. The main idea is that sampling from a normal distribution

\[x \sim \text{Normal}(\mu, \sigma)\]

is equivalent to sampling from the standard normal distribution and then transforming the samples as

\[x = \mu + \sigma z \quad \text{with} \quad z \sim \text{Normal}(0, 1).\]

For details, see the example “A Primer on Bayesian Methods for Multilevel Modeling”.

In order to get a feeling on the implementation, let’s sample from this component using the mean of the parameters (amplitude and length-scale) from the model fitted in the first part. We use take \(40\) samples with \(m=8\) eigenfunctions to approximate the kernel and we keep the box size \(L=1.5\).

# prior sample from the HSGP component
def prior_sample_f_approx_se_ncp(rng_key, *args, **kwargs):
    model = seed(hs_approx_squared_exponential_ncp, rng_key)
    model_trace = trace(model).get_trace(*args, **kwargs)
    return model_trace["f"]["value"]


# set parameters
m = 8
l_max = 1.5
amplitude = gp_mcmc.get_samples()["kernel_amplitude"].mean()
length_scale = gp_mcmc.get_samples()["kernel_length_scale"].mean()

# vectorize the computation
# see https://num.pyro.ai/en/stable/tutorials/bayesian_regression.html#Predictive-Utility-With-Effect-Handlers
vmap_prior_sample_f_approx_se_ncp = vmap(
    lambda rng_key: prior_sample_f_approx_se_ncp(
        rng_key=rng_key,
        x=x_train,
        amplitude=amplitude,
        length_scale=length_scale,
        l_max=l_max,
        m=m,
    ),
    in_axes=0,
    out_axes=-1,
)

n_prior_samples = 40
rng_key, rng_subkey = random.split(rng_key)
prior_samples_f_approx_se_ncp = vmap_prior_sample_f_approx_se_ncp(
    rng_key=random.split(rng_subkey, n_prior_samples)
)

fig, ax = plt.subplots()
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.plot(x_train, prior_samples_f_approx_se_ncp, c="C3", alpha=0.5, linewidth=0.5)
ax.plot(
    x_train,
    prior_samples_f_approx_se_ncp[:, 0],
    c="C3",
    alpha=1,
    linewidth=2,
    label="prior sample (given amplitude and length_scale)",
)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set(xlabel="x", ylabel="y")
ax.set_title(
    "HSGP Prior Samples (Given Amplitude and Length Scale)",
    fontsize=18,
    fontweight="bold",
);

These prior predictive samples look very reasonable!

Model Specification

Having implemented the Hilbert space approximation component, we are ready to define the whole model in NumPyro. We use the same priors for the amplitude and length-scale parameters as in the first part. We also use the same priors for the noise parameter.

def hsgp_model(x, l_max, m, y=None) -> None:
    """Hilbert Space Gaussian Process (HSGP) model with a
    squared-exponential kernel in one dimension.

    Parameters
    ----------
    x : feature variable
    l_max : box size L (well, the box size is in fact 2L)
    m : number of eigenfunctions to select
    y : target variable
    """
    # --- Priors ---
    kernel_amplitude = numpyro.sample(
        "kernel_amplitude",
        dist.InverseGamma(
            concentration=inverse_gamma_params_2["alpha"],
            rate=inverse_gamma_params_2["beta"],
        ),
    )
    kernel_length_scale = numpyro.sample(
        "kernel_length_scale",
        dist.InverseGamma(
            concentration=inverse_gamma_params_1["alpha"],
            rate=inverse_gamma_params_1["beta"],
        ),
    )
    noise = numpyro.sample(
        "noise",
        dist.InverseGamma(
            concentration=inverse_gamma_params_2["alpha"],
            rate=inverse_gamma_params_2["beta"],
        ),
    )
    # --- Parametrization ---
    f = hs_approx_squared_exponential_ncp(
        x, kernel_amplitude, kernel_length_scale, l_max, m
    )
    # --- Likelihood ---
    with numpyro.plate("data", x.shape[0]):
        numpyro.sample("likelihood", dist.Normal(loc=f, scale=noise), obs=y)

numpyro.render_model(
    model=hsgp_model,
    model_kwargs={"x": x_train, "l_max": 1.5, "m": 8, "y": y_train_obs},
    render_distributions=True,
    render_params=True,
)

Prior Predictive Check

For this specific example we choose:

l_max = 1.3
m = 20

Before fitting the model to the data, we can use the model above to generate samples from the prior distribution over functions.

hsgp_numpyro_prior_predictive = Predictive(hsgp_model, num_samples=1_000)
rng_key, rng_subkey = random.split(rng_key)
hsgp_numpyro_prior_samples = hsgp_numpyro_prior_predictive(
    rng_subkey, x_train, l_max, m
)

hsgp_numpyro_prior_idata = az.from_numpyro(prior=hsgp_numpyro_prior_samples)

fig, ax = plt.subplots()
az.plot_hdi(
    x_train,
    hsgp_numpyro_prior_idata.prior["likelihood"],
    hdi_prob=0.94,
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_train,
    hsgp_numpyro_prior_idata.prior["likelihood"],
    hdi_prob=0.5,
    color="C0",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$50\\%$ HDI (test)"},
    ax=ax,
)
ax.plot(
    x_train,
    hsgp_numpyro_prior_idata.prior["likelihood"].mean(dim=("chain", "draw")),
    color="C0",
    linewidth=3,
    label="posterior predictive mean (test)",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
for i in range(5):
    label = "prior samples" if i == 0 else None
    ax.plot(
        x_train,
        hsgp_numpyro_prior_idata.prior["likelihood"].sel(chain=0, draw=i),
        color="C0",
        alpha=0.3,
        label=label,
    )
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=3)
ax.set(xlabel="x", ylabel="y")
ax.set_title("HSGP NumPyro Model - Prior Predictive", fontsize=18, fontweight="bold");

The prior predictive samples look similar to the ones from the first part (vanilla Gaussian process model).

Model Fitting

We can fit the model as usual:

rng_key, rng_subkey = random.split(rng_key)
hsgp_mcmc = run_inference(
    rng_subkey,
    hsgp_model,
    inference_params,
    x_train,
    l_max,
    m,
    y_train_obs,
    target_accept_prob=0.9,
)

Observe that it runs very fast!

Model Diagnostics

The model diagnostics look good:

hsgp_numpyro_idata = az.from_numpyro(posterior=hsgp_mcmc)

az.summary(
    data=hsgp_numpyro_idata,
    var_names=["kernel_amplitude", "kernel_length_scale", "noise", "beta"],
)

axes = az.plot_trace(
    data=hsgp_numpyro_idata,
    var_names=["kernel_amplitude", "kernel_length_scale", "noise", "beta"],
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (12, 9), "layout": "constrained"},
)
plt.gcf().suptitle("HSGP NumPyro - Trace", fontsize=18, fontweight="bold");

We can compare the posterior distributions of the parameters of the Hilbert space approximation with the ones of the Gaussian process model:

axes = az.plot_forest(
    data=[gp_numpyro_idata, hsgp_numpyro_idata],
    model_names=["GP", "HSGP"],
    var_names=["kernel_amplitude", "kernel_length_scale", "noise"],
    combined=True,
    hdi_prob=0.94,
    figsize=(10, 6),
)
axes[0].figure.suptitle("NumPyro Models", fontsize=18, fontweight="bold", x=0.32);

They look very similar!

Out of Sample Prediction

Finally, we can generate samples from the posterior distribution over functions on the test set. Note that we need to ensure these values lie inside the approximation domain \([-L, L]\).

hpgp_mcmc_samples_no_f = {
    k: v for (k, v) in hsgp_mcmc.get_samples().items() if k != "f"
}

rng_key, rng_subkey = random.split(rng_key)
hsgp_numpyro_idata.extend(
    az.from_numpyro(
        posterior_predictive=Predictive(hsgp_model, hpgp_mcmc_samples_no_f)(
            rng_subkey, x_test, l_max, m
        )
    )
)

fig, ax = plt.subplots()
az.plot_hdi(
    x_test,
    hsgp_numpyro_idata.posterior_predictive["likelihood"],
    hdi_prob=0.94,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_numpyro_idata.posterior_predictive["likelihood"],
    hdi_prob=0.5,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$50\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_numpyro_idata.posterior_predictive["f"],
    hdi_prob=0.94,
    color="C3",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$f \\: 94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_numpyro_idata.posterior_predictive["f"],
    hdi_prob=0.5,
    color="C3",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$f \\: 50\\%$ HDI (test)"},
    ax=ax,
)
ax.plot(
    x_test,
    hsgp_numpyro_idata.posterior_predictive["f"].mean(dim=("chain", "draw")),
    color="C3",
    linewidth=3,
    label="posterior predictive mean (test)",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
ax.scatter(x_test, y_test_obs, c="C1", label="observed (test)")
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=4)
ax.set(xlabel="x", ylabel="y")
ax.set_title(
    "HSGP NumPyro Model - Posterior Predictive", fontsize=18, fontweight="bold"
);

We have successfully implemented the Hilbert space Gaussian process approximation in NumPyro and applied it to the same synthetic data as in the first part. The results look very good 🚀!

Next, let’s see how to leverage PyMC’s API to do the same!

HSGP.prior

This first implementation follows the PyMC’s GP module API and therefore changing from a vanilla Gaussian process to the corresponding Hilbert space approximation requires changing no more than \(3\) lines of code:

with pm.Model() as hsgp_pymc_model:
    x_data = pm.MutableData("x_data", value=x_train)
    y_data = pm.MutableData("y_data", y_train_obs)

    kernel_amplitude = pm.InverseGamma(
        "kernel_amplitude",
        alpha=inverse_gamma_params_2["alpha"],
        beta=inverse_gamma_params_2["beta"],
    )
    kernel_length_scale = pm.InverseGamma(
        "kernel_length_scale",
        alpha=inverse_gamma_params_1["alpha"],
        beta=inverse_gamma_params_1["beta"],
    )
    noise = pm.InverseGamma(
        "noise",
        alpha=inverse_gamma_params_2["alpha"],
        beta=inverse_gamma_params_2["beta"],
    )

    mean = pm.gp.mean.Zero()
    cov = kernel_amplitude**2 * pm.gp.cov.ExpQuad(input_dim=1, ls=kernel_length_scale)
    gp = pm.gp.HSGP(m=[20], L=[1.3], mean_func=mean, cov_func=cov)
    f = gp.prior("f", X=x_data[:, None])

    pm.Normal("likelihood", mu=f, sigma=noise, observed=y_data)


pm.model_to_graphviz(model=hsgp_pymc_model)

Let’s inspect the HSGP class

print(pm.gp.HSGP.__doc__)

    Hilbert Space Gaussian process approximation.

    The `gp.HSGP` class is an implementation of the Hilbert Space Gaussian process.  It is a
    reduced rank GP approximation that uses a fixed set of basis vectors whose coefficients are
    random functions of a stationary covariance function's power spectral density.  Its usage
    is largely similar to `gp.Latent`.  Like `gp.Latent`, it does not assume a Gaussian noise model
    and can be used with any likelihood, or as a component anywhere within a model.  Also like
    `gp.Latent`, it has `prior` and `conditional` methods.  It supports any sum of covariance
    functions that implement a `power_spectral_density` method. (Note, this excludes the
    `Periodic` covariance function, which uses a different set of basis functions for a
    low rank approximation, as described in `HSGPPeriodic`.).

    For information on choosing appropriate `m`, `L`, and `c`, refer to Ruitort-Mayol et al. or to
    the PyMC examples that use HSGP.

    To work with the HSGP in its "linearized" form, as a matrix of basis vectors and a vector of
    coefficients, see the method `prior_linearized`.

    Parameters
    ----------
    m: list
        The number of basis vectors to use for each active dimension (covariance parameter
        `active_dim`).
    L: list
        The boundary of the space for each `active_dim`.  It is called the boundary condition.
        Choose L such that the domain `[-L, L]` contains all points in the column of X given by the
        `active_dim`.
    c: float
        The proportion extension factor.  Used to construct L from X.  Defined as `S = max|X|` such
        that `X` is in `[-S, S]`.  `L` is calculated as `c * S`.  One of `c` or `L` must be
        provided.  Further information can be found in Ruitort-Mayol et al.
    drop_first: bool
        Default `False`. Sometimes the first basis vector is quite "flat" and very similar to
        the intercept term.  When there is an intercept in the model, ignoring the first basis
        vector may improve sampling. 
    parameterization: str
        Whether to use the `centered` or `noncentered` parameterization when multiplying the
        basis by the coefficients.
    cov_func: Covariance function, must be an instance of `Stationary` and implement a
        `power_spectral_density` method.
    mean_func: None, instance of Mean
        The mean function.  Defaults to zero.

    Examples
    --------
    .. code:: python

        # A three dimensional column vector of inputs.
        X = np.random.rand(100, 3)

        with pm.Model() as model:
            # Specify the covariance function.
            # Three input dimensions, but we only want to use the last two.
            cov_func = pm.gp.cov.ExpQuad(3, ls=0.1, active_dims=[1, 2])

            # Specify the HSGP.
            # Use 25 basis vectors across each active dimension for a total of 25 * 25 = 625.
            # The value `c = 4` means the boundary of the approximation
            # lies at four times the half width of the data.
            # In this example the data lie between zero and one,
            # so the boundaries occur at -1.5 and 2.5.  The data, both for
            # training and prediction should reside well within that boundary..
            gp = pm.gp.HSGP(m=[25, 25], c=4.0, cov_func=cov_func)

            # Place a GP prior over the function f.
            f = gp.prior("f", X=X)

        ...

        # After fitting or sampling, specify the distribution
        # at new points with .conditional
        Xnew = np.linspace(-1, 2, 50)[:, None]

        with model:
            fcond = gp.conditional("fcond", Xnew=Xnew)

    References
    ----------
    -   Ruitort-Mayol, G., and Anderson, M., and Solin, A., and Vehtari, A. (2022). Practical
        Hilbert Space Approximate Bayesian Gaussian Processes for Probabilistic Programming

    -   Solin, A., Sarkka, S. (2019) Hilbert Space Methods for Reduced-Rank Gaussian Process
        Regression.
    

Observe that the parametrization argument allow us to select between centered and non-centered parametrizations automatically 🤩!

Sampling from this model is much faster that the GP model from the fist part!

rng_key, rng_subkey = random.split(rng_key)

with hsgp_pymc_model:
    hsgp_pymc_idata = pm.sample(
        target_accept=0.9,
        draws=inference_params.num_samples,
        chains=inference_params.num_chains,
        nuts_sampler="numpyro",
        random_seed=rng_subkey[0].item(),
    )

axes = az.plot_trace(
    data=hsgp_pymc_idata,
    var_names=["kernel_amplitude", "kernel_length_scale", "noise"],
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (12, 7), "layout": "constrained"},
)
plt.gcf().suptitle("HSGP PyMC - Trace", fontsize=18, fontweight="bold");

We can generate out of sample predictions using the conditional method as before.

with hsgp_pymc_model:
    x_star_data = pm.MutableData("x_star_data", x_test)
    f_star = gp.conditional("f_star", x_star_data[:, None])
    pm.set_data({"x_data": x_test, "y_data": np.ones_like(x_test)})
    hsgp_pymc_idata.extend(
        pm.sample_posterior_predictive(
            trace=hsgp_pymc_idata,
            var_names=["f_star", "likelihood"],
            random_seed=rng_subkey[1].item(),
        )
    )

Let’s see the final fit:

fig, ax = plt.subplots()
az.plot_hdi(
    x_test,
    hsgp_pymc_idata.posterior_predictive["likelihood"],
    hdi_prob=0.94,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$\\: 94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_pymc_idata.posterior_predictive["likelihood"],
    hdi_prob=0.5,
    color="C1",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$\\: 50\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_pymc_idata.posterior_predictive["f_star"],
    hdi_prob=0.94,
    color="C3",
    smooth=False,
    fill_kwargs={"alpha": 0.2, "label": "$f \\: 94\\%$ HDI (test)"},
    ax=ax,
)
az.plot_hdi(
    x_test,
    hsgp_pymc_idata.posterior_predictive["f_star"],
    hdi_prob=0.5,
    color="C3",
    smooth=False,
    fill_kwargs={"alpha": 0.4, "label": "$f \\: 50\\%$ HDI (test)"},
    ax=ax,
)
ax.plot(
    x_test,
    hsgp_pymc_idata.posterior_predictive["f_star"].mean(dim=("chain", "draw")),
    color="C3",
    linewidth=3,
    label="posterior predictive mean (test)",
)
ax.scatter(x_train, y_train_obs, c="C0", label="observed (train)")
ax.scatter(x_test, y_test_obs, c="C1", label="observed (test)")
ax.plot(x_train, y_train, color="black", linewidth=3, alpha=0.7, label="mean (latent)")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=4)
ax.set(xlabel="x", ylabel="y")
ax.set_title("HSGP PyMC Model - Posterior Predictive", fontsize=18, fontweight="bold");

🙌 !