Add Numpyro MWE

Author: thomaspinder

examples/heteroscedastic_inference.py

@@ -357,8 +357,22 @@
     alpha=0.3,
     label="One std. dev.",
 )
-ax.plot(xtest.squeeze(), predictive_mean - predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
-ax.plot(xtest.squeeze(), predictive_mean + predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
+ax.plot(
+    xtest.squeeze(),
+    predictive_mean - predictive_std,
+    "--",
+    color=cols[1],
+    alpha=0.5,
+    linewidth=0.75,
+)
+ax.plot(
+    xtest.squeeze(),
+    predictive_mean + predictive_std,
+    "--",
+    color=cols[1],
+    alpha=0.5,
+    linewidth=0.75,
+)
 ax.fill_between(
     xtest.squeeze(),
     predictive_mean - 2 * predictive_std,
@@ -367,8 +381,22 @@
     alpha=0.1,
     label="Two std. dev.",
 )
-ax.plot(xtest.squeeze(), predictive_mean - 2 * predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
-ax.plot(xtest.squeeze(), predictive_mean + 2 * predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
+ax.plot(
+    xtest.squeeze(),
+    predictive_mean - 2 * predictive_std,
+    "--",
+    color=cols[1],
+    alpha=0.5,
+    linewidth=0.75,
+)
+ax.plot(
+    xtest.squeeze(),
+    predictive_mean + 2 * predictive_std,
+    "--",
+    color=cols[1],
+    alpha=0.5,
+    linewidth=0.75,
+)
 
 ax.set_title("Sparse Heteroscedastic Regression")
 ax.legend(loc="best", fontsize="small")

examples/lgcp_numpyro.py

@@ -0,0 +1,119 @@
+# %%
+import jax.numpy as jnp
+from jax import random
+from jax import config
+import numpy as np
+
+import gpjax as gpx
+from gpjax import numpyro_extras
+import numpyro
+import numpyro.distributions as dist
+from numpyro.infer import MCMC, NUTS
+import arviz as az
+
+import matplotlib.pyplot as plt
+
+# Enable x64 support for JAX
+config.update("jax_enable_x64", True)
+
+# Set random seed
+key = random.PRNGKey(42)
+
+# Configure MCMC
+num_warmup = 1000
+num_samples = 1000
+num_chains = 4
+
+# Set device count for numpyro for parallel chains
+numpyro.set_host_device_count(num_chains)
+
+# %%
+# 1. Data: Coal Mining Disasters (1851-1962)
+# Counts of disasters per year
+counts = jnp.array([
+    4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6, 3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5, 2, 2, 3, 4, 2, 1, 3, 2, 2, 1, 1, 1, 1, 3, 0, 0, 1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1, 0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0, 0, 1, 1, 0, 2, 3, 3, 1, 1, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 1, 4, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1
+], dtype=jnp.float64)
+
+years = jnp.arange(1851, 1851 + len(counts), dtype=jnp.float64).reshape(-1, 1)
+# Normalize years for better numerical stability in GP
+years_norm = (years - years.min()) / (years.max() - years.min())
+
+# %%
+# 2. Model Definition
+# We model the log-intensity log(lambda(t)) as a Gaussian Process.
+# lambda(t) = exp(f(t))
+# y_i ~ Poisson(lambda(t_i))
+
+# Mean function: Constant mean
+mean_f = gpx.mean_functions.Constant(constant=jnp.array([0.0]))
+
+# Kernel: Matern52
+# We expect changes over decades, so lengthscale should be non-trivial.
+# Since x is normalized to [0, 1], a lengthscale of 0.1 corresponds to ~11 years.
+kernel = gpx.kernels.Matern52(lengthscale=0.2, variance=0.5)
+
+prior = gpx.gps.Prior(mean_function=mean_f, kernel=kernel)
+
+def model(x, y):
+    # Register GPJax parameters (lengthscale, variance, mean_constant) with Numpyro
+    gp = numpyro_extras.register_parameters(prior)
+
+    # Sample the latent function f at the input locations x
+    f = numpyro.sample("f", gp(x))
+
+    # The intensity is exp(f)
+    rate = jnp.exp(f)
+
+    # Observation model: Poisson
+    numpyro.sample("y", dist.Poisson(rate), obs=y)
+
+# %%
+# 3. Inference
+rng_key, rng_key_ = random.split(key)
+
+kernel_nuts = NUTS(model, target_accept_prob=0.9)
+mcmc = MCMC(
+    kernel_nuts,
+    num_warmup=num_warmup,
+    num_samples=num_samples,
+    num_chains=num_chains,
+    progress_bar=True,
+    jit_model_args=True,
+)
+
+# Run MCMC
+# Note: We pass years_norm for stability, but we'll plot against original years
+mcmc.run(rng_key_, x=years_norm, y=counts)
+
+# %%
+# 4. Analysis & Plotting
+mcmc.print_summary()
+
+# Extract samples
+samples = mcmc.get_samples()
+f_samples = samples["f"]
+intensity_samples = jnp.exp(f_samples)
+
+# Compute statistics
+mean_intensity = jnp.mean(intensity_samples, axis=0)
+lower_ci = jnp.percentile(intensity_samples, 2.5, axis=0)
+upper_ci = jnp.percentile(intensity_samples, 97.5, axis=0)
+
+# Plot
+plt.figure(figsize=(12, 6))
+plt.bar(years.flatten(), counts, color="gray", alpha=0.5, label="Observed Counts", width=1.0)
+plt.plot(years.flatten(), mean_intensity, color="C0", label="Posterior Mean Intensity", linewidth=2)
+plt.fill_between(years.flatten(), lower_ci, upper_ci, color="C0", alpha=0.3, label="95% CI")
+
+plt.xlabel("Year")
+plt.ylabel("Number of Disasters")
+plt.title("Coal Mining Disasters: Log-Gaussian Cox Process (GPJax + Numpyro)")
+plt.legend()
+plt.grid(True, alpha=0.3)
+plt.tight_layout()
+plt.savefig("lgcp_coal_mining.png")
+# plt.show()
+
+# Trace plot for diagnostics
+az.plot_trace(mcmc, var_names=["kernel.lengthscale", "kernel.variance"])
+plt.tight_layout()

examples/numpyro_integration.py

@@ -0,0 +1,216 @@
+# ---
+# jupyter:
+#   jupytext:
+#     cell_metadata_filter: -all
+#     custom_cell_magics: kql
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.17.3
+#   kernelspec:
+#     display_name: python3
+#     language: python
+#     name: python3
+# ---
+
+# %% [markdown]
+# # Joint Inference with Numpyro
+#
+# In this notebook, we demonstrate how to use [Numpyro](https://num.pyro.ai/) to perform fully Bayesian inference over the hyperparameters of a Gaussian process model.
+# We will look at a scenario where we have a structured mean function (a linear model) and a GP capturing the residuals. We will infer the parameters of both the linear model and the GP jointly.
+
+# %%
+from jax import config
+import jax.numpy as jnp
+import jax.random as jr
+import matplotlib.pyplot as plt
+import numpyro
+import numpyro.distributions as dist
+from numpyro.infer import (
+    MCMC,
+    NUTS,
+)
+
+import gpjax as gpx
+from gpjax.numpyro_extras import register_parameters
+
+config.update("jax_enable_x64", True)
+
+key = jr.key(42)
+
+# %% [markdown]
+# ## Data Generation
+#
+# We generate a synthetic dataset that consists of a linear trend, a periodic component, and some noise.
+
+# %%
+N = 100
+x = jnp.sort(jr.uniform(key, shape=(N, 1), minval=0.0, maxval=10.0), axis=0)
+
+# True parameters
+true_slope = 0.5
+true_intercept = 2.0
+true_period = 2.0
+true_lengthscale = 1.0
+true_noise = 0.1
+
+# Signal
+linear_trend = true_slope * x + true_intercept
+periodic_signal = jnp.sin(2 * jnp.pi * x / true_period)
+y_clean = linear_trend + periodic_signal
+
+# Observations
+y = y_clean + true_noise * jr.normal(key, shape=x.shape)
+
+plt.figure(figsize=(10, 5))
+plt.scatter(x, y, label="Data", alpha=0.6)
+plt.plot(x, y_clean, "k--", label="True Signal")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# ## Model Definition
+#
+# We define a GP model with a generic mean function (zero for now, as we will handle the linear trend explicitly in the Numpyro model) and a kernel that is the product of a periodic kernel and an RBF kernel. This choice reflects our prior knowledge that the signal is locally periodic.
+
+# %%
+kernel = gpx.kernels.RBF() * gpx.kernels.Periodic()
+meanf = gpx.mean_functions.Zero()
+prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
+
+# We will use a ConjugatePosterior since we assume Gaussian noise
+likelihood = gpx.likelihoods.Gaussian(num_datapoints=N)
+posterior = prior * likelihood
+
+# We initialise the model parameters.
+# Note: These values will be overwritten by Numpyro samples during inference.
+D = gpx.Dataset(X=x, y=y)
+
+# %% [markdown]
+# ## Joint Inference Loop
+#
+# We define a Numpyro model function that:
+# 1. Samples the parameters for the linear trend.
+# 2. Computes the residuals (Data - Linear Trend).
+# 3. Samples the GP hyperparameters using `register_parameters`.
+# 4. Computes the GP marginal log-likelihood on the residuals.
+# 5. Adds the GP log-likelihood to the joint density.
+
+
+# %%
+def model(X, Y):
+    # 1. Sample linear model parameters
+    slope = numpyro.sample("slope", dist.Normal(0.0, 2.0))
+    intercept = numpyro.sample("intercept", dist.Normal(0.0, 2.0))
+
+    # Calculate residuals
+    trend = slope * X + intercept
+    residuals = Y - trend
+
+    # 2. Register GP parameters
+    # This automatically samples parameters from the GPJax model
+    # and returns a model with updated values.
+    # We can specify custom priors if needed, but we'll rely on defaults here.
+    # register_parameters modifies the model in-place (and returns it).
+    # Since Numpyro re-runs this function, we are overwriting the parameters
+    # of the same object repeatedly, which is fine as they are completely determined
+    # by the sample sites.
+    p_posterior = register_parameters(posterior)
+
+    # Create dataset for residuals
+    D_resid = gpx.Dataset(X=X, y=residuals)
+
+    # 3. Compute MLL
+    # We use conjugate_mll which computes log p(y | X, theta) analytically for Gaussian likelihoods.
+    mll = gpx.objectives.conjugate_mll(p_posterior, D_resid)
+
+    # 4. Add to potential
+    numpyro.factor("gp_log_lik", mll)
+
+
+# %% [markdown]
+# ## Running MCMC
+#
+# We use the NUTS sampler to draw samples from the posterior.
+
+# %%
+nuts_kernel = NUTS(model)
+mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000, num_chains=1)
+mcmc.run(jr.key(0), x, y)
+
+mcmc.print_summary()
+
+# %% [markdown]
+# ## Analysis and Plotting
+#
+# We extract the samples and plot the predictions.
+
+# %%
+samples = mcmc.get_samples()
+
+
+# Helper to get predictions
+def predict(rng_key, sample_idx):
+    # Reconstruct model with sampled values
+
+    # Linear part
+    slope = samples["slope"][sample_idx]
+    intercept = samples["intercept"][sample_idx]
+    trend = slope * x + intercept
+
+    # GP part
+    # We use numpyro.handlers.substitute to inject the sampled values into register_parameters
+    # to reconstruct the GP model state for this sample.
+    sample_dict = {k: v[sample_idx] for k, v in samples.items()}
+
+    with numpyro.handlers.substitute(data=sample_dict):
+        # We call register_parameters again to update the posterior object with this sample's values
+        p_posterior = register_parameters(posterior)
+
+    # Now predict on residuals
+    residuals = y - trend
+    D_resid = gpx.Dataset(X=x, y=residuals)
+
+    latent_dist = p_posterior.predict(x, train_data=D_resid)
+    predictive_mean = latent_dist.mean
+    predictive_std = latent_dist.stddev()
+
+    return trend + predictive_mean, predictive_std
+
+
+# Plot
+plt.figure(figsize=(12, 6))
+plt.scatter(x, y, alpha=0.5, label="Data", color="gray")
+plt.plot(x, y_clean, "k--", label="True Signal")
+
+# Compute mean prediction (using mean of samples for efficiency)
+mean_slope = jnp.mean(samples["slope"])
+mean_intercept = jnp.mean(samples["intercept"])
+mean_trend = mean_slope * x + mean_intercept
+
+mean_samples = {k: jnp.mean(v, axis=0) for k, v in samples.items()}
+with numpyro.handlers.substitute(data=mean_samples):
+    p_posterior_mean = register_parameters(posterior)
+
+residuals_mean = y - mean_trend
+D_resid_mean = gpx.Dataset(X=x, y=residuals_mean)
+latent_dist = p_posterior_mean.predict(x, train_data=D_resid_mean)
+pred_mean = latent_dist.mean
+pred_std = latent_dist.stddev()
+
+total_mean = mean_trend.flatten() + pred_mean.flatten()
+std_flat = pred_std.flatten()
+
+plt.plot(x, total_mean, "b-", label="Posterior Mean")
+plt.fill_between(
+    x.flatten(),
+    total_mean - 2 * std_flat,
+    total_mean + 2 * std_flat,
+    color="b",
+    alpha=0.2,
+    label="95% CI (GP Uncertainty)",
+)
+
+plt.legend()
+plt.show()