Add explicit prior specification

Author: thomaspinder

.github/workflows/test_docs.yml

@@ -3,6 +3,11 @@ name: Test documentation
 on:
   pull_request:
 
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
 jobs:
   test-docs:
     # Functionality for testing documentation builds on multiple OSes and Python versions
@@ -42,3 +47,26 @@ jobs:
           uv sync --extra docs
           uv run python docs/scripts/gen_examples.py --execute
           uv run mkdocs build
+
+      - name: Upload built docs artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: docs-site-html
+          path: site
+
+      - name: Upload Pages artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: site
+
+  deploy-docs-preview:
+    needs: test-docs
+    if: github.event_name == 'pull_request'
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy MkDocs preview
+        id: deployment
+        uses: actions/deploy-pages@v4

examples/lgcp_numpyro.py

@@ -30,6 +30,7 @@
 from numpyro.infer import (
     MCMC,
     NUTS,
+    Predictive,
 )
 
 import gpjax as gpx
@@ -67,20 +68,38 @@
 plt.scatter(x, y, label="Data", alpha=0.6)
 plt.plot(x, y_clean, "k--", label="True Signal")
 plt.legend()
-plt.show()
+# 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()
+# Define priors
+lengthscale_prior = dist.LogNormal(0.0, 1.0)
+variance_prior = dist.LogNormal(0.0, 1.0)
+period_prior = dist.LogNormal(0.0, 0.5)
+noise_prior = dist.LogNormal(0.0, 1.0)
+
+# Define Kernel with priors
+# We can explicitly attach priors to the parameters
+kernel = gpx.kernels.RBF(
+    lengthscale=gpx.parameters.PositiveReal(1.0, prior=lengthscale_prior),
+    variance=gpx.parameters.PositiveReal(1.0, prior=variance_prior),
+) * gpx.kernels.Periodic(
+    lengthscale=gpx.parameters.PositiveReal(1.0, prior=lengthscale_prior),
+    period=gpx.parameters.PositiveReal(1.0, prior=period_prior),
+)
+
 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)
+likelihood = gpx.likelihoods.Gaussian(
+    num_datapoints=N,
+    obs_stddev=gpx.parameters.NonNegativeReal(1.0, prior=noise_prior),
+)
 posterior = prior * likelihood
 
 # We initialise the model parameters.
@@ -111,7 +130,8 @@ def model(X, Y):
     # 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.
+    # We attached priors to the parameters during model definition,
+    # so register_parameters will use those.
     # 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
@@ -150,67 +170,59 @@ def model(X, Y):
 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
+def predict_fn(X_new, Y_train):
+    # 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))
 
-    # 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()}
+    # Calculate residuals
+    trend_train = slope * x + intercept
+    residuals = Y_train - trend_train
 
-    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)
+    # 2. Register GP parameters
+    p_posterior = register_parameters(posterior)
 
-    # Now predict on residuals
-    residuals = y - trend
+    # Create dataset for residuals
     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()
+    # 3. Compute latent GP distribution
+    latent_dist = p_posterior.predict(X_new, train_data=D_resid)
 
-    return trend + predictive_mean, predictive_std
+    # 4. Sample latent function values
+    f = numpyro.sample("f", latent_dist)
+    f = f.reshape((-1, 1))
 
+    # 5. Compute and return total prediction
+    total_prediction = slope * X_new + intercept + f
+    numpyro.deterministic("y_pred", total_prediction)
+    return total_prediction
 
-# 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
+# Create predictive utility
+predictive = Predictive(predict_fn, posterior_samples=samples)
 
-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)
+# Generate predictions
+predictions = predictive(jr.key(1), X_new=x, Y_train=y)
+y_pred = predictions["y_pred"]
 
-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()
+# Compute statistics
+mean_prediction = jnp.mean(y_pred, axis=0)
+std_prediction = jnp.std(y_pred, axis=0)
 
-total_mean = mean_trend.flatten() + pred_mean.flatten()
-std_flat = pred_std.flatten()
+# 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")
 
-plt.plot(x, total_mean, "b-", label="Posterior Mean")
+plt.plot(x, mean_prediction, "b-", label="Posterior Mean")
 plt.fill_between(
     x.flatten(),
-    total_mean - 2 * std_flat,
-    total_mean + 2 * std_flat,
+    mean_prediction.flatten() - 2 * std_prediction.flatten(),
+    mean_prediction.flatten() + 2 * std_prediction.flatten(),
     color="b",
     alpha=0.2,
     label="95% CI (GP Uncertainty)",
 )
 
 plt.legend()
-plt.show()
+# plt.show()

examples/numpyro_integration.py

@@ -68,6 +68,9 @@ def sample(self, key, sample_shape=()):
         def affine_transformation(_x):
             return self.loc + covariance_root @ _x
 
+        if not sample_shape:
+            return affine_transformation(white_noise)
+
         return vmap(affine_transformation)(white_noise)
 
     @property

gpjax/distributions.py

@@ -6,29 +6,10 @@
 import numpyro.distributions as dist
 
 from gpjax.parameters import (
-    FillTriangularTransform,
     Parameter,
 )
 
 
-def _get_default_prior(tag, shape, ndim):
-    if tag in ("positive", "non_negative"):
-        return dist.LogNormal(0.0, 1.0).expand(shape).to_event(ndim)
-    if tag == "real":
-        return dist.Normal(0.0, 1.0).expand(shape).to_event(ndim)
-    if tag == "sigmoid":
-        return dist.Uniform(0.0, 1.0).expand(shape).to_event(ndim)
-    if tag == "lower_triangular":
-        N = shape[-1]
-        K = N * (N + 1) // 2
-        batch_shape = shape[:-2]
-        base_shape = batch_shape + (K,)
-        base_dist = dist.Normal(0.0, 1.0).expand(base_shape).to_event(1)
-        td = dist.TransformedDistribution(base_dist, FillTriangularTransform())
-        return td.to_event(len(batch_shape))
-    return dist.Normal(0.0, 1.0).expand(shape).to_event(ndim)
-
-
 def register_parameters(
     model: nnx.Module,
     priors: tp.Dict[str, dist.Distribution] | None = None,
@@ -71,7 +52,12 @@ def _param_callback(path, param):
         # Determine prior
         prior = priors.get(name)
         if prior is None:
-            prior = _get_default_prior(param.tag, param.value.shape, param.value.ndim)
+            # Check for attached prior
+            numpyro_props = getattr(param, "numpyro_properties", {})
+            prior = numpyro_props.get("prior")
+
+        if prior is None:
+            return param
 
         # Sample
         value = numpyro.sample(name, prior)