From 4e8f8d382b268ae5bbe832c77b0d9795b3e09e95 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:27:35 +0200 Subject: [PATCH 01/81] test(oracle): pin conjugate MLL, predict, and LOOCV to closed form conjugate_mll had no value-level test anywhere in the suite, yet serves as the oracle for the Kalman MLL and collapsed_elbo. These closed-form pins, computed through an independent jnp.linalg path, give the reference frame its ground truth ahead of the v1.0 conditioning refactor. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- tests/test_conditioning_oracle.py | 91 +++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 tests/test_conditioning_oracle.py diff --git a/tests/test_conditioning_oracle.py b/tests/test_conditioning_oracle.py new file mode 100644 index 000000000..db76c180d --- /dev/null +++ b/tests/test_conditioning_oracle.py @@ -0,0 +1,91 @@ +"""Closed-form oracles for the conjugate GP quantities. + +These tests pin the *values* of ``conjugate_mll``, ``ConjugatePosterior.predict`` +and ``conjugate_loocv`` on a tiny fixed dataset, computed through an independent +linear-algebra path (direct ``jnp.linalg.solve``/``slogdet``, never +``GaussianDistribution``). Everything downstream in the test suite — the Kalman +MLL, ``collapsed_elbo`` — is validated against ``conjugate_mll``, so these +oracles are the ground truth the rest of the suite stands on. +""" + +import jax.numpy as jnp + +import gpjax as gpx + +JITTER = 1e-6 +OBS_STDDEV = 0.3 +LENGTHSCALE = 0.7 +VARIANCE = 1.3 + +X = jnp.array([[0.0], [0.45], [1.1]]) +Y = jnp.array([[0.2], [-0.1], [0.6]]) +XTEST = jnp.array([[0.2], [0.85]]) + + +def _rbf(x1, x2): + sq_dists = (x1[:, None, 0] - x2[None, :, 0]) ** 2 + return VARIANCE * jnp.exp(-0.5 * sq_dists / LENGTHSCALE**2) + + +def _posterior(): + kernel = gpx.kernels.RBF(lengthscale=LENGTHSCALE, variance=VARIANCE) + prior = gpx.gps.Prior( + mean_function=gpx.mean_functions.Zero(), kernel=kernel, jitter=JITTER + ) + likelihood = gpx.likelihoods.Gaussian( + num_datapoints=X.shape[0], obs_stddev=OBS_STDDEV + ) + return prior * likelihood + + +def _sigma(): + n = X.shape[0] + return _rbf(X, X) + (JITTER + OBS_STDDEV**2) * jnp.eye(n) + + +def test_conjugate_mll_matches_closed_form(): + posterior = _posterior() + data = gpx.Dataset(X=X, y=Y) + + sigma = _sigma() + n = X.shape[0] + quad = Y[:, 0] @ jnp.linalg.solve(sigma, Y[:, 0]) + _, logdet = jnp.linalg.slogdet(sigma) + expected = -0.5 * (quad + logdet + n * jnp.log(2.0 * jnp.pi)) + + actual = gpx.objectives.conjugate_mll(posterior, data) + assert jnp.allclose(actual, expected, atol=1e-10) + + +def test_predict_matches_closed_form(): + posterior = _posterior() + data = gpx.Dataset(X=X, y=Y) + + sigma = _sigma() + kxt = _rbf(X, XTEST) + ktt = _rbf(XTEST, XTEST) + sigma_inv_y = jnp.linalg.solve(sigma, Y[:, 0]) + expected_mean = kxt.T @ sigma_inv_y + expected_cov = ktt - kxt.T @ jnp.linalg.solve(sigma, kxt) + JITTER * jnp.eye(2) + + predictive = posterior.predict(XTEST, data) + assert jnp.allclose(predictive.mean, expected_mean, atol=1e-8) + assert jnp.allclose(predictive.covariance_matrix, expected_cov, atol=1e-8) + + +def test_loocv_matches_closed_form(): + posterior = _posterior() + data = gpx.Dataset(X=X, y=Y) + + sigma = _sigma() + sigma_inv = jnp.linalg.inv(sigma) + resid = Y[:, 0] + diag = jnp.diag(sigma_inv) + loo_means = resid - sigma_inv @ resid / diag + loo_vars = 1.0 / diag + expected = jnp.sum( + -0.5 * (jnp.log(2 * jnp.pi * loo_vars) + (resid - loo_means) ** 2 / loo_vars) + ) + + actual = gpx.objectives.conjugate_loocv(posterior, data) + assert jnp.allclose(actual, expected, atol=1e-8) From 0ff8f2da47732d3502978ebbf84f3862a4344397 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:28:54 +0200 Subject: [PATCH 02/81] test(equivalence): pin cross-derivation agreement; xfail two-owner jitter bug MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit collapsed_elbo(z=X) vs conjugate_mll and whitened-vs-unwhitened predicts at matched parameters now guard the five independent derivations of the conjugate conditioning algebra. The strict xfail documents that at non-default jitter the derivations factorise different matrices — the bug the v1.0 conditioning module removes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- tests/test_conditioning_equivalence.py | 113 +++++++++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 tests/test_conditioning_equivalence.py diff --git a/tests/test_conditioning_equivalence.py b/tests/test_conditioning_equivalence.py new file mode 100644 index 000000000..c9bbdeb5f --- /dev/null +++ b/tests/test_conditioning_equivalence.py @@ -0,0 +1,113 @@ +"""Equivalence tests across independently-derived conditioning code paths. + +Five call sites in the library derive the conjugate conditioning algebra +independently (predict, MLL, LOOCV, collapsed ELBO, variational predicts). +These tests pin the mathematical identities that must hold across them, so +that consolidating the derivations cannot silently change behaviour. +""" + +import jax.numpy as jnp +import pytest + +import gpjax as gpx + + +def _make(jitter=1e-6): + x = jnp.linspace(0.0, 1.0, 12).reshape(-1, 1) + y = jnp.sin(3.0 * x) + data = gpx.Dataset(X=x, y=y) + prior = gpx.gps.Prior( + mean_function=gpx.mean_functions.Constant(), + kernel=gpx.kernels.RBF(), + jitter=jitter, + ) + likelihood = gpx.likelihoods.Gaussian(num_datapoints=data.n, obs_stddev=0.2) + return prior * likelihood, data + + +def test_collapsed_elbo_equals_mll_when_z_is_x(): + posterior, data = _make() + q = gpx.variational_families.CollapsedVariationalGaussian( + posterior=posterior, inducing_inputs=data.X + ) + elbo = gpx.objectives.collapsed_elbo(q, data) + mll = gpx.objectives.conjugate_mll(posterior, data) + # The identity is exact only in the jitter -> 0 limit: the family's jitter + # enters Kzz while the model's enters Sigma, so a small O(jitter/noise) + # discrepancy is expected even when both knobs are 1e-6. + assert jnp.allclose(elbo, mll, atol=2e-4) + + +@pytest.mark.xfail( + strict=True, + reason="jitter has two owners (Prior.jitter vs the variational family's " + "jitter); collapsed_elbo and conjugate_mll factorise different matrices " + "at non-default jitter. Resolved by the v1.0 conditioning stack.", +) +def test_collapsed_elbo_equals_mll_when_z_is_x_nondefault_jitter(): + posterior, data = _make(jitter=1e-3) + q = gpx.variational_families.CollapsedVariationalGaussian( + posterior=posterior, inducing_inputs=data.X + ) + elbo = gpx.objectives.collapsed_elbo(q, data) + mll = gpx.objectives.conjugate_mll(posterior, data) + assert jnp.allclose(elbo, mll, atol=2e-4) + + +def test_whitened_matches_unwhitened_at_matched_parameters(): + posterior, data = _make() + z = jnp.linspace(0.0, 1.0, 5).reshape(-1, 1) + q_white = gpx.variational_families.WhitenedVariationalGaussian( + posterior=posterior, inducing_inputs=z + ) + q_plain = gpx.variational_families.VariationalGaussian( + posterior=posterior, inducing_inputs=z + ) + # At default parameters, whitened q(u) = N(0, I) and unwhitened + # q(u) = N(0, I) describe different measures UNLESS the predictive + # reduces identically; we instead match parameters explicitly: + # unwhitened (mu, sqrt) = (m(z) + Lz mu_w, Lz sqrt_w). + import equinox as eqx + + kernel = posterior.prior.kernel + kzz = kernel.gram(z).as_matrix() + 1e-6 * jnp.eye(z.shape[0]) + lz = jnp.linalg.cholesky(kzz) + mu_white = jnp.array([[0.3], [-0.2], [0.1], [0.4], [-0.5]]) + sqrt_white = 0.1 * jnp.eye(5) + 0.05 * jnp.tril(jnp.ones((5, 5)), k=-1) + + mean_z = posterior.prior.mean_function(z) + mu_plain = mean_z + lz @ mu_white + sqrt_plain = lz @ sqrt_white + + q_white = eqx.tree_at( + lambda q: (q.variational_mean, q.variational_root_covariance), + q_white, + ( + type(q_white.variational_mean)(mu_white) + if hasattr(type(q_white.variational_mean), "unwrap") + else mu_white, + type(q_white.variational_root_covariance)(sqrt_white) + if hasattr(type(q_white.variational_root_covariance), "unwrap") + else sqrt_white, + ), + ) + q_plain = eqx.tree_at( + lambda q: (q.variational_mean, q.variational_root_covariance), + q_plain, + ( + type(q_plain.variational_mean)(mu_plain) + if hasattr(type(q_plain.variational_mean), "unwrap") + else mu_plain, + type(q_plain.variational_root_covariance)(sqrt_plain) + if hasattr(type(q_plain.variational_root_covariance), "unwrap") + else sqrt_plain, + ), + ) + + xtest = jnp.linspace(-0.2, 1.2, 7).reshape(-1, 1) + dist_white = q_white(xtest) + dist_plain = q_plain(xtest) + assert jnp.allclose(dist_white.mean, dist_plain.mean, atol=1e-5) + assert jnp.allclose( + dist_white.covariance_matrix, dist_plain.covariance_matrix, atol=1e-5 + ) From d195be8a4a3dc0f194d75d2cd20a684a932d9976 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:29:29 +0200 Subject: [PATCH 03/81] test(integration): fail loudly when golden values drift _compare previously swallowed AssertionError with a print, so the harness could never fail. Failures are now collected per-example and raised at the end of test(), making the four golden-value pins a real no-behaviour-change net for the v1.0 refactor. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- tests/integration_tests.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/integration_tests.py b/tests/integration_tests.py index aff981c72..4edb4e306 100644 --- a/tests/integration_tests.py +++ b/tests/integration_tests.py @@ -46,6 +46,7 @@ class Result: def __post_init__(self): self.name: str = self.path.split("/")[-1].split(".")[0].replace("_", "-") + self.failures: list = [] def _compare( self, @@ -56,11 +57,14 @@ def _compare( ): if variable_name == "history" and not self.compare_history: return - try: - value = operation(observed_variables[variable_name]) - assert abs(true_value - value) < self.precision - except AssertionError as e: - print(e) + value = operation(observed_variables[variable_name]) + if not abs(true_value - value) < self.precision: + message = ( + f"{self.name}: {variable_name} drifted from golden value " + f"{true_value} (got {value}, precision {self.precision})" + ) + print(message) + self.failures.append(message) def test(self): notebook = jupytext.read(self.path) @@ -101,6 +105,11 @@ def test(self): self._compare( observed_variables=loc, variable_name=k, true_value=truth, operation=op ) + if self.failures: + raise AssertionError( + f"{self.name}: golden-value drift detected:\n" + + "\n".join(self.failures) + ) # %% From 5d96ee23696e7a52f83a5a9257c04d71c878dd70 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:32:49 +0200 Subject: [PATCH 04/81] test(integration): re-pin golden values to current-main behaviour The newly-loud harness exposed pre-existing drift in all four examples: the collapsed/uncollapsed goldens predated the real-data example swap (#696) and the regression/heteroscedastic goldens predated subsequent behaviour fixes (#707/#708/#713 and dependency bumps). The toothless harness never noticed. Re-pinned so the net measures the v1.0 refactor, not history. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- tests/integration_tests.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/integration_tests.py b/tests/integration_tests.py index 4edb4e306..6dd3c7ae6 100644 --- a/tests/integration_tests.py +++ b/tests/integration_tests.py @@ -117,8 +117,8 @@ def test(self): path="docs/examples/regression.py", comparisons={ "history": (55.07405622, get_last), - "predictive_mean": (36.24383416, jnp.sum), - "predictive_std": (197.04727051, jnp.sum), + "predictive_mean": (37.91222107, jnp.sum), + "predictive_std": (202.36889441, jnp.sum), }, ) regression.test() @@ -127,9 +127,9 @@ def test(self): sparse = Result( path="docs/examples/collapsed_vi.py", comparisons={ - "history": (1924.7634809, get_last), - "predictive_mean": (-8.39869652, jnp.sum), - "predictive_std": (255.74838027, jnp.sum), + "history": (1851.11700608, get_last), + "predictive_mean": (1.37497714, jnp.sum), + "predictive_std": (248.32254630, jnp.sum), }, ) sparse.test() @@ -138,9 +138,9 @@ def test(self): stochastic = Result( path="docs/examples/uncollapsed_vi.py", comparisons={ - "history": (-2678.41302494, get_last), - "meanf": (-54.14787028, jnp.sum), - "sigma": (121.4298333, jnp.sum), + "history": (59440.08265547, get_last), + "meanf": (-55.18585235, jnp.sum), + "sigma": (555.41381240, jnp.sum), }, ) stochastic.test() @@ -149,7 +149,7 @@ def test(self): heteroscedastic = Result( path="docs/examples/heteroscedastic_inference.py", comparisons={ - "history": (-141.590, get_last), + "history": (-139.22405213, get_last), }, ) heteroscedastic.test() From 880c58f235e258a88da0b16174ecba40c85abbaf Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:33:35 +0200 Subject: [PATCH 05/81] =?UTF-8?q?feat(linalg):=20stabilised=5Fcholesky=20?= =?UTF-8?q?=E2=80=94=20single=20stabilise-and-factor=20seam?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- gpjax/linalg/__init__.py | 2 ++ gpjax/linalg/utils.py | 11 +++++++++++ tests/test_linalg.py | 16 ++++++++++++++++ 3 files changed, 29 insertions(+) diff --git a/gpjax/linalg/__init__.py b/gpjax/linalg/__init__.py index 50a605cb6..7e3854c06 100644 --- a/gpjax/linalg/__init__.py +++ b/gpjax/linalg/__init__.py @@ -6,6 +6,7 @@ cholesky_factor, logdet, logdet_from_factor, + stabilised_cholesky, ) __all__ = [ @@ -15,4 +16,5 @@ "cholesky_factor", "logdet", "logdet_from_factor", + "stabilised_cholesky", ] diff --git a/gpjax/linalg/utils.py b/gpjax/linalg/utils.py index 11fbced0e..a409876b4 100644 --- a/gpjax/linalg/utils.py +++ b/gpjax/linalg/utils.py @@ -19,6 +19,17 @@ def add_jitter(matrix: Array, jitter: float | Array = 1e-6) -> Array: return matrix + jnp.eye(matrix.shape[0]) * jitter +def stabilised_cholesky(matrix: Array, jitter: float | Array) -> Array: + """Lower Cholesky factor of ``matrix + jitter * I``. + + The single stabilise-and-factor entry point for GP conditioning: the + jitter policy is applied here and nowhere else. Structure-aware + dispatch over lineax operators arrives with the linalg deepening; the + interface will not change. + """ + return jnp.linalg.cholesky(add_jitter(matrix, jitter)) + + @functools.singledispatch def cholesky_factor(op: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator: """Cholesky factor of a PSD operator. Returns lower-triangular L s.t. A = L L^T.""" diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 42bf07777..7ce0242a6 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -264,3 +264,19 @@ def test_kronecker_structures(): kron = Kronecker(A=A, B=B) assert kron.in_structure().shape == (6,) assert kron.out_structure().shape == (6,) + + +def test_stabilised_cholesky_identity(): + from gpjax.linalg.utils import stabilised_cholesky + + factor = stabilised_cholesky(jnp.eye(3), 1e-2) + assert jnp.allclose(factor, jnp.sqrt(1.01) * jnp.eye(3), atol=1e-12) + + +def test_stabilised_cholesky_reconstructs(): + from gpjax.linalg.utils import stabilised_cholesky + + root = jnp.array([[1.0, 0.0], [0.4, 0.8]]) + psd = root @ root.T + factor = stabilised_cholesky(psd, 1e-3) + assert jnp.allclose(factor @ factor.T, psd + 1e-3 * jnp.eye(2), atol=1e-10) From 5814aa3d68e28b250b1a2fa2dfee630538a4e6bd Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:34:09 +0200 Subject: [PATCH 06/81] feat(dataset): static n_total metadata; get_batch stamps full size The full-dataset size a minibatch ELBO needs now travels on the one object that knows it, as static pytree aux_data, instead of being smuggled through likelihood constructors. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- gpjax/dataset.py | 8 ++++++-- gpjax/fit.py | 3 ++- tests/test_dataset.py | 27 +++++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/gpjax/dataset.py b/gpjax/dataset.py index 8dc1c3631..affb3b257 100644 --- a/gpjax/dataset.py +++ b/gpjax/dataset.py @@ -31,10 +31,14 @@ class Dataset: Args: X: input data. y: output data. + n_total: full-dataset size when this object is a minibatch view of a + larger dataset (stamped by ``gpjax.fit.get_batch``); ``None`` means + the dataset is self-describing (``n_total == n``). """ X: Optional[Num[Array, "N D"]] = None y: Optional[Num[Array, "N Q"]] = None + n_total: Optional[int] = None def __post_init__(self) -> None: r"""Checks that the shapes of $X$ and $y$ are compatible, @@ -93,11 +97,11 @@ def num_outputs(self) -> int: return self.y.shape[1] def tree_flatten(self): - return (self.X, self.y), None + return (self.X, self.y), self.n_total @classmethod def tree_unflatten(cls, aux_data, children): - return cls(*children) + return cls(*children, n_total=aux_data) def _check_shape( diff --git a/gpjax/fit.py b/gpjax/fit.py index c65caabbb..bd7923fc2 100644 --- a/gpjax/fit.py +++ b/gpjax/fit.py @@ -388,7 +388,8 @@ def get_batch(train_data: Dataset, batch_size: int, key: KeyArray) -> Dataset: # Subsample mini-batch indices with replacement. indices = jr.choice(key, n, (batch_size,), replace=True) - return Dataset(X=x[indices], y=y[indices]) + full_size = train_data.n_total if train_data.n_total is not None else n + return Dataset(X=x[indices], y=y[indices], n_total=full_size) def _check_model(model: tp.Any) -> None: diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 94bef84cd..db3d7d5f2 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -198,3 +198,30 @@ def test_dataset_multi_output_properties(): D = Dataset(X=X, y=y) assert D.multi_output is True assert D.num_outputs == 3 + + +def test_n_total_defaults_to_none(): + x = jnp.linspace(0.0, 1.0, 10).reshape(-1, 1) + dataset = Dataset(X=x, y=jnp.sin(x)) + assert dataset.n_total is None + + +def test_get_batch_stamps_n_total(): + import jax + import jax.random as jr + + from gpjax.fit import get_batch + + x = jnp.linspace(0.0, 1.0, 25).reshape(-1, 1) + dataset = Dataset(X=x, y=jnp.sin(x)) + batch = get_batch(dataset, batch_size=4, key=jr.key(0)) + assert batch.n == 4 + assert batch.n_total == 25 + + # The stamp is static aux_data: it survives pytree operations. + mapped = jax.tree_util.tree_map(lambda leaf: leaf, batch) + assert mapped.n_total == 25 + + # Batching a batch preserves the original full size. + rebatch = get_batch(batch, batch_size=2, key=jr.key(1)) + assert rebatch.n_total == 25 From 8c7ba10b90849140e3d0447f6889d73dcc12112c Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:35:15 +0200 Subject: [PATCH 07/81] =?UTF-8?q?feat(likelihoods)!:=20pure=20conditional?= =?UTF-8?q?=20families=20=E2=80=94=20drop=20num=5Fdatapoints=20and=20noise?= =?UTF-8?q?=5Fprior?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Of ~245 occurrences of num_datapoints, only four were real reads: the ELBO minibatch scale (now served by Dataset.n_total) and latent sizing (moves to data-contact time in the JointModel rewrite). No likelihood used the value internally, and nothing validated it — a wrong value silently mis-scaled the ELBO. noise_prior moves to the model layer, where priors live. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- gpjax/likelihoods.py | 22 ++-------------------- gpjax/objectives.py | 14 ++++++-------- 2 files changed, 8 insertions(+), 28 deletions(-) diff --git a/gpjax/likelihoods.py b/gpjax/likelihoods.py index ce34293b5..29d80e676 100644 --- a/gpjax/likelihoods.py +++ b/gpjax/likelihoods.py @@ -14,7 +14,6 @@ import abc from dataclasses import dataclass -from typing import TYPE_CHECKING import beartype.typing as tp import equinox as eqx @@ -44,9 +43,6 @@ ScalarFloat, ) -if TYPE_CHECKING: - from gpjax.gps import Prior - def _diagonal_scale(op): """Unwrap a single TaggedLinearOperator layer and return the inner diagonal, or None. @@ -85,22 +81,18 @@ class AbstractLikelihood(_SummaryMixin, eqx.Module): `link_function` methods. """ - num_datapoints: int = eqx.field(static=True) integrator: AbstractIntegrator = eqx.field(static=True) def __init__( self, - num_datapoints: int, integrator: AbstractIntegrator = GHQuadratureIntegrator(), ): """Initializes the likelihood. Args: - num_datapoints (int): the number of data points. integrator (AbstractIntegrator): The integrator to be used for computing expected log likelihoods. Must be an instance of `AbstractIntegrator`. """ - self.num_datapoints = num_datapoints self.integrator = integrator def __call__( @@ -254,21 +246,16 @@ def moments( class AbstractHeteroscedasticLikelihood(AbstractLikelihood): r"""Base class for heteroscedastic likelihoods with latent noise processes.""" - noise_prior: tp.Any noise_transform: AbstractNoiseTransform def __init__( self, - num_datapoints: int, - noise_prior: Prior, noise_transform: tp.Union[ AbstractNoiseTransform, tp.Callable[[Float[Array, ...]], Float[Array, ...]], ] = SoftplusTransform(), integrator: AbstractIntegrator = GHQuadratureIntegrator(), ): - self.noise_prior = noise_prior - if isinstance(noise_transform, AbstractNoiseTransform): self.noise_transform = noise_transform else: @@ -281,7 +268,7 @@ def __init__( # Users should implement AbstractNoiseTransform for custom transforms. self.noise_transform = SoftplusTransform() - super().__init__(num_datapoints=num_datapoints, integrator=integrator) + super().__init__(integrator=integrator) def __call__( self, @@ -331,14 +318,12 @@ class Gaussian(AbstractLikelihood): def __init__( self, - num_datapoints: int, obs_stddev: tp.Union[ScalarFloat, Float[Array, "#N"], NonNegativeReal] = 1.0, integrator: AbstractIntegrator = AnalyticalGaussianIntegrator(), ): r"""Initializes the Gaussian likelihood. Args: - num_datapoints (int): the number of data points. obs_stddev (Union[ScalarFloat, Float[Array, "#N"]]): the standard deviation of the Gaussian observation noise. integrator (AbstractIntegrator): The integrator to be used for computing expected log @@ -350,7 +335,7 @@ def __init__( self.obs_stddev = obs_stddev self.num_outputs = 1 - super().__init__(num_datapoints, integrator) + super().__init__(integrator) def link_function(self, f: Float[Array, ...]) -> npd.Normal: r"""The link function of the Gaussian likelihood. @@ -412,21 +397,18 @@ class MultiOutputGaussian(Gaussian): """Gaussian likelihood with per-output noise variance. Args: - num_datapoints: Total number of observations (N, not N*P). num_outputs: Number of output dimensions (P). obs_stddev: Per-output noise standard deviation. Scalar broadcasts to [P]. """ def __init__( self, - num_datapoints: int, num_outputs: int, obs_stddev: tp.Union[float, Float[Array, " P"]] = 1.0, ): if isinstance(obs_stddev, (int, float)): obs_stddev = jnp.full(num_outputs, float(obs_stddev)) super().__init__( - num_datapoints=num_datapoints, obs_stddev=NonNegativeReal(jnp.asarray(obs_stddev)), ) self.num_outputs = num_outputs diff --git a/gpjax/objectives.py b/gpjax/objectives.py index 54b79649e..b52fa262d 100644 --- a/gpjax/objectives.py +++ b/gpjax/objectives.py @@ -323,12 +323,8 @@ def elbo(variational_family: VF, data: Dataset) -> ScalarFloat: var_exp = variational_expectation(variational_family, data) # For batch size b, we compute n/b * sum_i[ int log(p(y|f(xi))) q(f(xi)) df(xi)] - KL[q(f(.)) || p(f(.))] - return ( - jnp.sum(var_exp) - * variational_family.posterior.likelihood.num_datapoints - / data.n - - kl - ) + full_size = data.n_total if data.n_total is not None else data.n + return jnp.sum(var_exp) * full_size / data.n - kl def variational_expectation( @@ -529,7 +525,8 @@ def heteroscedastic_elbo_conjugate( return_parts=True, ) - scale = likelihood.num_datapoints / data.n + full_size = data.n_total if data.n_total is not None else data.n + scale = full_size / data.n return scale * jnp.sum(expected_ll) - variational_family.prior_kl() @@ -550,7 +547,8 @@ def heteroscedastic_elbo_chained(variational_family: HVF, data: Dataset) -> Scal noise_stats=noise_stats, ) - scale = likelihood.num_datapoints / data.n + full_size = data.n_total if data.n_total is not None else data.n + scale = full_size / data.n return scale * jnp.sum(expected_ll) - variational_family.prior_kl() From 9b9a43c49bf8e63dc847cab578699718e31f2a87 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 01:42:44 +0200 Subject: [PATCH 08/81] =?UTF-8?q?feat!:=20v1.0=20conditioning=20architectu?= =?UTF-8?q?re=20=E2=80=94=20JointModel,=20Posterior,=20condition?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The API now mirrors the maths: prior * likelihood -> JointModel (the joint p(f,y), the trainable object); model.condition(D) — sugar: model | D — returns an immutable Posterior pytree caching the Cholesky factor and representer weights. The predictive, log_marginal_likelihood, loo, and pathwise sample_approx are views of that one factorisation, deleting the eleven independent derivations and the two-owner jitter split (prior.jitter is now the single knob, applied once inside conditioning). - gpjax/conditioning.py: deep module (Posterior, ExactPosterior, LatentPosterior); MO validation moves to condition time; sample_approx refuses multi-output loudly instead of silently broadcasting wrong. - gps.py: Prior (AbstractPrior folded in), ConjugateModel, NonConjugateModel (lazy latent, sized at data contact), HeteroscedasticModel (owns noise_prior — likelihoods are pure conditionals again, killing the likelihoods->gps circular import). Deleted: AbstractPrior, AbstractPosterior, LatentPosterior marker, ChainedPosterior marker, construct_posterior (now construct_model). - objectives: conjugate_mll/conjugate_loocv/log_posterior_density are one-line views of the conditioned posterior. - fit: _prepare_model hook sizes lazily-initialised state from data. - predict(t, D) survives as documented one-line sugar everywhere. - return_covariance_type kwarg renamed to covariance. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- gpjax/__init__.py | 16 + gpjax/conditioning.py | 417 +++++++++++++++ gpjax/fit.py | 25 +- gpjax/gps.py | 922 ++++++++++---------------------- gpjax/models/oilmm.py | 34 +- gpjax/objectives.py | 146 ++--- gpjax/state_space/fit.py | 6 +- gpjax/state_space/gps.py | 38 +- gpjax/state_space/objectives.py | 2 +- gpjax/variational_families.py | 33 +- 10 files changed, 839 insertions(+), 800 deletions(-) create mode 100644 gpjax/conditioning.py diff --git a/gpjax/__init__.py b/gpjax/__init__.py index 12a13859b..9cf22da92 100644 --- a/gpjax/__init__.py +++ b/gpjax/__init__.py @@ -32,6 +32,7 @@ variational_families, ) from gpjax.citation import cite +from gpjax.conditioning import Posterior from gpjax.dataset import Dataset from gpjax.distributions import GaussianDistribution from gpjax.fit import ( @@ -39,6 +40,14 @@ fit_lbfgs, fit_scipy, ) +from gpjax.gps import ( + ConjugateModel, + HeteroscedasticModel, + JointModel, + NonConjugateModel, + Prior, + construct_model, +) from gpjax.summary import summarise __license__ = "MIT" @@ -48,9 +57,16 @@ __version__ = "0.18.0" __all__ = [ + "ConjugateModel", "Dataset", "GaussianDistribution", + "HeteroscedasticModel", + "JointModel", + "NonConjugateModel", + "Posterior", + "Prior", "cite", + "construct_model", "fit", "fit_lbfgs", "fit_scipy", diff --git a/gpjax/conditioning.py b/gpjax/conditioning.py new file mode 100644 index 000000000..315da9148 --- /dev/null +++ b/gpjax/conditioning.py @@ -0,0 +1,417 @@ +# Copyright 2026 The GPJax Contributors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +r"""The deep conditioning module. + +Conditioning a Gaussian process on data is one computation: stabilise and +factor the training covariance once, then derive every downstream quantity — +the predictive distribution, the log marginal likelihood (evidence), +leave-one-out densities, and pathwise samples — as views of that single +factorisation. This module is the one home of that algebra. + +A :class:`Posterior` is an immutable pytree produced by +``model.condition(train_data)`` (equivalently ``model | train_data``). It +caches the Cholesky factor and representer weights, so repeated queries never +re-factorise. +""" + +from abc import abstractmethod +from typing import Literal +import warnings + +import beartype.typing as tp +import equinox as eqx +import jax.numpy as jnp +import jax.random as jr +import jax.scipy as jsp +from jaxtyping import ( + Float, + Num, +) +import lineax as lx +import numpyro.distributions as npd + +from gpjax.dataset import Dataset +from gpjax.distributions import GaussianDistribution +from gpjax.kernels import RFF +from gpjax.linalg.utils import stabilised_cholesky +from gpjax.parameters import _val +from gpjax.typing import ( + Array, + FunctionalSample, + KeyArray, + ScalarFloat, +) + + +class Posterior(eqx.Module): + r"""A conditioned Gaussian process, :math:`p(f \mid \mathcal{D})`. + + The result of conditioning a joint model on data. Immutable: the + factorisation of the training covariance is computed once at + ``condition`` time and cached on this object; every query is a view of + it. Query the process at test inputs by calling it:: + + posterior = model.condition(train_data) # or: model | train_data + predictive = posterior(test_inputs) + """ + + @abstractmethod + def __call__( + self, + test_inputs: Num[Array, "N D"], + *, + covariance: Literal["dense", "diagonal"] = "dense", + ) -> GaussianDistribution: + r"""Evaluate the conditioned process at the given test inputs. + + Args: + test_inputs: Input locations at which to query the process. + covariance: Whether to return the dense joint covariance over the + test inputs or only the marginal (diagonal) variances. + + Returns: + GaussianDistribution: The predictive distribution at the inputs. + """ + raise NotImplementedError + + def predict( + self, + test_inputs: Num[Array, "N D"], + train_data: tp.Optional[Dataset] = None, + *, + covariance: Literal["dense", "diagonal"] = "dense", + ) -> GaussianDistribution: + r"""Sugar for calling the posterior: ``predict(t) == self(t)``. + + Retained for signature compatibility with the pre-v1.0 API; + ``train_data`` is accepted and ignored — this process is already + conditioned on its training set. + """ + del train_data + return self(test_inputs, covariance=covariance) + + +class ExactPosterior(Posterior): + r"""Exactly conditioned GP: a Gaussian likelihood integrated analytically. + + Caches the lower Cholesky factor of + :math:`\Sigma = K_{xx} + \texttt{jitter}\,\mathbf{I} + \mathrm{diag}(\sigma^2)` + and the representer weights :math:`\alpha = \Sigma^{-1}(y - m(x))`. The + predictive moments, the evidence, LOO densities, and pathwise samples are + all views of these two objects. + """ + + prior: tp.Any + likelihood: tp.Any + train_data: Dataset + cholesky_factor: Float[Array, "NP NP"] + representer_weights: Float[Array, "NP 1"] + residual: Float[Array, "NP 1"] + log_marginal_likelihood: ScalarFloat + + def __init__(self, prior: tp.Any, likelihood: tp.Any, train_data: Dataset): + from gpjax.kernels.multioutput.base import MultiOutputKernel + + kernel = prior.kernel + if isinstance(kernel, MultiOutputKernel): + if not train_data.multi_output: + raise ValueError("MultiOutputKernel requires multi-output data.") + if train_data.num_outputs != kernel.num_outputs: + raise ValueError( + f"Dataset has {train_data.num_outputs} outputs " + f"but kernel expects {kernel.num_outputs}." + ) + + x, y = train_data.X, train_data.y + mean_x = prior.mean_function(x) + y_flat, mean_flat = likelihood.prepare_targets(y, mean_x) + noise = likelihood.noise_vector(train_data.n) + + gram_plus_noise = kernel.gram(x).as_matrix() + jnp.diag(noise) + factor = stabilised_cholesky(gram_plus_noise, prior.jitter) + residual = (y_flat - mean_flat).reshape(-1, 1) + weights = jsp.linalg.cho_solve((factor, True), residual) + + num_scalars = residual.shape[0] + half_logdet = jnp.sum(jnp.log(jnp.diagonal(factor))) + evidence = ( + -0.5 + * ( + jnp.sum(residual * weights) + + num_scalars * jnp.log(2.0 * jnp.pi) + ) + - half_logdet + ) + + self.prior = prior + self.likelihood = likelihood + self.train_data = train_data + self.cholesky_factor = factor + self.representer_weights = weights + self.residual = residual + self.log_marginal_likelihood = jnp.squeeze(evidence) + + def __call__( + self, + test_inputs: Num[Array, "N D"], + *, + covariance: Literal["dense", "diagonal"] = "dense", + ) -> GaussianDistribution: + kernel = self.prior.kernel + x = self.train_data.X + num_outputs = self.likelihood.num_outputs + + cross_cov = kernel.cross_covariance(x, test_inputs) + solved_cross = jsp.linalg.solve_triangular( + self.cholesky_factor, cross_cov, lower=True + ) + + mean_test_raw = self.prior.mean_function(test_inputs) + mean_test = ( + jnp.tile(mean_test_raw, (num_outputs, 1)) + if num_outputs > 1 + else mean_test_raw + ) + mean = mean_test + jnp.matmul(cross_cov.T, self.representer_weights) + + if covariance == "diagonal" and num_outputs > 1: + warnings.warn( + "Diagonal covariance is not yet supported for multi-output GPs. " + "Returning full covariance.", + stacklevel=2, + ) + covariance = "dense" + + if covariance == "dense": + test_gram = kernel.gram(test_inputs).as_matrix() + predictive_cov = test_gram - jnp.matmul(solved_cross.T, solved_cross) + predictive_cov = predictive_cov + self.prior.jitter * jnp.eye( + predictive_cov.shape[0] + ) + scale = lx.MatrixLinearOperator(predictive_cov) + else: + test_var_diag = lx.diagonal(kernel.diagonal(test_inputs)) + marginal_var = ( + test_var_diag + - jnp.einsum("ij,ji->i", solved_cross.T, solved_cross) + + self.prior.jitter + ) + scale = lx.DiagonalLinearOperator(jnp.atleast_1d(marginal_var.squeeze())) + + return GaussianDistribution(loc=jnp.atleast_1d(mean.squeeze()), scale=scale) + + def loo(self) -> Float[Array, " NP"]: + r"""Per-point leave-one-out predictive log-densities. + + Computed from the cached factor via Rasmussen & Williams eq. 5.12 — + no model is refit. Sum the result for the LOOCV objective. + """ + factor = self.cholesky_factor + num_scalars = factor.shape[0] + factor_inv = jsp.linalg.solve_triangular( + factor, jnp.eye(num_scalars), lower=True + ) + precision_diag = jnp.sum(factor_inv**2, axis=0).reshape(-1, 1) + + loo_means = ( + self.residual - self.representer_weights / precision_diag + ) + loo_vars = 1.0 / precision_diag + loo_dist = npd.Normal(loc=loo_means, scale=jnp.sqrt(loo_vars)) + return loo_dist.log_prob(self.residual).squeeze(-1) + + def sample_approx( + self, + num_samples: int, + key: KeyArray, + num_features: int | None = 100, + ) -> FunctionalSample: + r"""Draw approximate posterior samples via pathwise conditioning. + + Decomposes each sample into Fourier features of the prior plus + canonical features weighted through the cached training factor + (Wilson et al., 2020). + + Args: + num_samples: The desired number of samples. + key: The random seed used for the sample(s). + num_features: The number of Fourier features used to approximate + the prior component of each sample. + + Returns: + FunctionalSample: A function evaluating the sample draws at any + inputs; the same draw is returned for all queries. + """ + if (not isinstance(num_samples, int)) or num_samples <= 0: + raise ValueError("num_samples must be a positive integer") + if self.likelihood.num_outputs > 1: + raise ValueError( + "sample_approx does not support multi-output likelihoods yet." + ) + + freq_key, weight_key, noise_key = jr.split(key, 3) + fourier_feature_fn = _build_fourier_features_fn( + self.prior, num_features, freq_key + ) + fourier_weights = jr.normal(weight_key, [num_samples, 2 * num_features]) + + x = self.train_data.X + obs_var = _val(self.likelihood.obs_stddev) ** 2 + observation_noise = jnp.sqrt(obs_var) * jr.normal( + noise_key, [self.train_data.n, num_samples] + ) + prior_features = fourier_feature_fn(x) + perturbed_residual = ( + self.residual + + observation_noise + - jnp.inner(prior_features, fourier_weights) + ) + canonical_weights = jsp.linalg.cho_solve( + (self.cholesky_factor, True), perturbed_residual + ) + + def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]: + fourier_features = fourier_feature_fn(test_inputs) + weight_space_contribution = jnp.inner(fourier_features, fourier_weights) + canonical_features = self.prior.kernel.cross_covariance(test_inputs, x) + function_space_contribution = jnp.matmul( + canonical_features, canonical_weights + ) + return ( + self.prior.mean_function(test_inputs) + + weight_space_contribution + + function_space_contribution + ) + + return sample_fn + + +class LatentPosterior(Posterior): + r"""Approximately conditioned GP for non-Gaussian likelihoods. + + Conditioning is on the model's whitened latent vector rather than on the + observations directly: the cached factor is the prior gram's Cholesky, and + the latent plays the role of the representer weights. The joint + log-density (``log_posterior_density``) is the quantity MCMC or MAP + optimisation targets. + """ + + prior: tp.Any + likelihood: tp.Any + train_data: Dataset + cholesky_factor: Float[Array, "N N"] + latent: Float[Array, "N 1"] + + def __init__( + self, + prior: tp.Any, + likelihood: tp.Any, + latent: tp.Any, + train_data: Dataset, + ): + self.prior = prior + self.likelihood = likelihood + self.train_data = train_data + self.cholesky_factor = stabilised_cholesky( + prior.kernel.gram(train_data.X).as_matrix(), prior.jitter + ) + self.latent = _val(latent) + + def __call__( + self, + test_inputs: Num[Array, "N D"], + *, + covariance: Literal["dense", "diagonal"] = "dense", + ) -> GaussianDistribution: + kernel = self.prior.kernel + x = self.train_data.X + + cross_cov = kernel.cross_covariance(x, test_inputs) + solved_cross = jsp.linalg.solve_triangular( + self.cholesky_factor, cross_cov, lower=True + ) + + mean_test = self.prior.mean_function(test_inputs) + mean = mean_test + jnp.matmul(solved_cross.T, self.latent) + + if covariance == "dense": + test_gram = kernel.gram(test_inputs).as_matrix() + predictive_cov = test_gram - jnp.matmul(solved_cross.T, solved_cross) + predictive_cov = predictive_cov + self.prior.jitter * jnp.eye( + predictive_cov.shape[0] + ) + scale = lx.MatrixLinearOperator(predictive_cov) + else: + test_var_diag = lx.diagonal(kernel.diagonal(test_inputs)) + marginal_var = ( + test_var_diag + - jnp.einsum("ij,ji->i", solved_cross.T, solved_cross) + + self.prior.jitter + ) + scale = lx.DiagonalLinearOperator(jnp.atleast_1d(marginal_var.squeeze())) + + return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), scale) + + @property + def log_posterior_density(self) -> ScalarFloat: + r"""Unnormalised log-posterior density of the whitened latent model. + + :math:`\log p(y \mid f(x)) + \log \mathcal{N}(w_x \mid 0, I)` where + :math:`f(x) = m(x) + L_x w_x`. + """ + mean_x = self.prior.mean_function(self.train_data.X) + latent_function = mean_x + self.cholesky_factor @ self.latent + observation_density = self.likelihood.link_function(latent_function) + whitened_prior = npd.Normal(loc=0.0, scale=1.0) + return ( + observation_density.log_prob(self.train_data.y).sum() + + whitened_prior.log_prob(self.latent).sum() + ) + + +def _build_fourier_features_fn( + prior: tp.Any, num_features: int, key: KeyArray +) -> tp.Callable[[Float[Array, "N D"]], Float[Array, "N L"]]: + r"""Return a function evaluating features sampled from the Fourier feature + decomposition of the prior's kernel. + + Args: + prior (Prior): The Prior distribution. + num_features (int): The number of feature functions to be sampled. + key (KeyArray): The random seed used. + + Returns: + Callable: A callable function evaluating the sampled feature functions. + """ + if (not isinstance(num_features, int)) or num_features <= 0: + raise ValueError("num_features must be a positive integer") + + approximate_kernel = RFF( + base_kernel=prior.kernel, num_basis_fns=num_features, key=key + ) + + def eval_fourier_features(test_inputs: Float[Array, "N D"]) -> Float[Array, "N L"]: + feature_matrix = approximate_kernel.compute_features(x=test_inputs) + feature_matrix *= jnp.sqrt(_val(prior.kernel.variance) / num_features) + return feature_matrix + + return eval_fourier_features + + +__all__ = [ + "ExactPosterior", + "LatentPosterior", + "Posterior", +] diff --git a/gpjax/fit.py b/gpjax/fit.py index bd7923fc2..20a66281d 100644 --- a/gpjax/fit.py +++ b/gpjax/fit.py @@ -66,7 +66,7 @@ def fit( >>> >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood >>> @@ -109,6 +109,8 @@ def fit( _check_log_rate(log_rate) _check_verbose(verbose) + model = _prepare_model(model, train_data) + # Use paramax.unwrap for the constrained -> unconstrained -> constrained cycle. # paramax handles the bijection automatically via AbstractUnwrappable subclasses. @@ -192,7 +194,7 @@ def fit_scipy( >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood @@ -208,6 +210,8 @@ def fit_scipy( _check_num_iters(max_iters) _check_verbose(verbose) + model = _prepare_model(model, train_data) + # Split model into trainable arrays and static parts params, static = eqx.partition(model, eqx.is_array) @@ -287,7 +291,7 @@ def fit_lbfgs( >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood @@ -302,6 +306,8 @@ def fit_lbfgs( _check_train_data(train_data) _check_num_iters(max_iters) + model = _prepare_model(model, train_data) + # Split model into trainable arrays and static parts params, static = eqx.partition(model, eqx.is_array) @@ -392,6 +398,19 @@ def get_batch(train_data: Dataset, batch_size: int, key: KeyArray) -> Dataset: return Dataset(X=x[indices], y=y[indices], n_total=full_size) +def _prepare_model(model: Model, train_data: Dataset) -> Model: + """Run any data-dependent initialisation the model defines. + + JointModels use this to size lazily-initialised state (e.g. the + non-conjugate latent vector) from the training data. + """ + from gpjax.gps import JointModel + + if isinstance(model, JointModel): + return model._prepare(train_data) + return model + + def _check_model(model: tp.Any) -> None: """Check that the model is a subclass of eqx.Module.""" if not isinstance(model, eqx.Module): diff --git a/gpjax/gps.py b/gpjax/gps.py index aef48ab6f..c232a3890 100644 --- a/gpjax/gps.py +++ b/gpjax/gps.py @@ -12,15 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +r"""Gaussian process priors and joint models. + +The API mirrors the mathematics: + +.. code-block:: python + + prior = gpx.Prior(mean_function=meanf, kernel=kernel) # p(f) + likelihood = gpx.likelihoods.Gaussian() # p(y | f) + model = prior * likelihood # p(f, y) + posterior = model.condition(train_data) # p(f | D) + predictive = posterior(test_inputs) + +A :class:`JointModel` is the trainable object — ``gpx.fit`` optimises its +hyperparameters. Conditioning it on data returns an immutable +:class:`gpjax.conditioning.Posterior` that caches the factorisation. +""" -from abc import abstractmethod from typing import Literal import beartype.typing as tp import equinox as eqx import jax.numpy as jnp import jax.random as jr -import jax.scipy as jsp from jaxtyping import ( Float, Num, @@ -28,23 +42,24 @@ import lineax as lx from paramax import AbstractUnwrappable +from gpjax.conditioning import ( + ExactPosterior, + LatentPosterior, + Posterior, + _build_fourier_features_fn, +) from gpjax.dataset import Dataset from gpjax.distributions import GaussianDistribution -from gpjax.kernels import RFF from gpjax.kernels.base import AbstractKernel from gpjax.likelihoods import ( AbstractHeteroscedasticLikelihood, AbstractLikelihood, Gaussian, - HeteroscedasticGaussian, NonGaussian, ) from gpjax.linalg.utils import add_jitter from gpjax.mean_functions import AbstractMeanFunction -from gpjax.parameters import ( - Real, - _val, -) +from gpjax.parameters import Real from gpjax.summary import _SummaryMixin from gpjax.typing import ( Array, @@ -60,99 +75,12 @@ HL = tp.TypeVar("HL", bound=AbstractHeteroscedasticLikelihood) -class AbstractPrior(_SummaryMixin, eqx.Module, tp.Generic[M, K]): - r"""Abstract Gaussian process prior.""" - - kernel: K - mean_function: M - jitter: float = eqx.field(static=True, default=1e-6) - - def __init__( - self, - kernel: K, - mean_function: M, - jitter: float = 1e-6, - ): - r"""Construct a Gaussian process prior. - - Args: - kernel: kernel object inheriting from AbstractKernel. - mean_function: mean function object inheriting from AbstractMeanFunction. - """ - self.kernel = kernel - self.mean_function = mean_function - self.jitter = jitter - - def __call__( - self, - test_inputs: Num[Array, "N D"], - *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", - ) -> GaussianDistribution: - r"""Evaluate the Gaussian process at the given points. - - The output of this function is a ``GaussianDistribution`` from which - the latent function's mean and covariance can be evaluated and the - distribution can be sampled. - - Under the hood, ``__call__`` invokes the ``predict`` method. Classes - inheriting ``AbstractPrior`` should not overwrite ``__call__`` and - should instead define a ``predict`` method. - - Args: - test_inputs: Input locations where the GP should be evaluated. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. - - Returns: - GaussianDistribution: A multivariate normal random variable representation - of the Gaussian process. - """ - return self.predict( - test_inputs, - return_covariance_type=return_covariance_type, - ) - - @abstractmethod - def predict( - self, - test_inputs: Num[Array, "N D"], - *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", - ) -> GaussianDistribution: - r"""Evaluate the predictive distribution. - - Compute the latent function's multivariate normal distribution for a - given set of parameters. For any class inheriting the `AbstractPrior` class, - this method must be implemented. - - Args: - test_inputs: Input locations where the GP should be evaluated. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. - - Returns: - GaussianDistribution: A multivariate normal random variable representation - of the Gaussian process. - """ - raise NotImplementedError - - ####################### # GP Priors ####################### -class Prior(AbstractPrior[M, K]): +class Prior(_SummaryMixin, eqx.Module, tp.Generic[M, K]): r"""A Gaussian process prior object. - The GP is parameterised by a - mean - and kernel - function. - A Gaussian process prior parameterised by a mean function $m(\cdot)$ and a kernel function $k(\cdot, \cdot)$ is given by $p(f(\cdot)) = \mathcal{GP}(m(\cdot), k(\cdot, \cdot))$. @@ -171,88 +99,110 @@ class Prior(AbstractPrior[M, K]): :doc:`/examples/regression` puts one to work end to end. """ + kernel: K + mean_function: M + jitter: float = eqx.field(static=True, default=1e-6) + + def __init__( + self, + kernel: K, + mean_function: M, + jitter: float = 1e-6, + ): + r"""Construct a Gaussian process prior. + + Args: + kernel: kernel object inheriting from AbstractKernel. + mean_function: mean function object inheriting from AbstractMeanFunction. + jitter: the model's single numerical-stabilisation knob. Applied + exactly once, inside conditioning. + """ + self.kernel = kernel + self.mean_function = mean_function + self.jitter = jitter + if tp.TYPE_CHECKING: @tp.overload - def __mul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]": ... + def __mul__(self, other: GL) -> "ConjugateModel[M, K, GL]": ... @tp.overload - def __mul__(self, other: NGL) -> "NonConjugatePosterior[Prior[M, K], NGL]": ... + def __mul__(self, other: NGL) -> "NonConjugateModel[M, K, NGL]": ... @tp.overload - def __mul__(self, other: L) -> "AbstractPosterior[Prior[M, K], L]": ... + def __mul__(self, other: L) -> "JointModel[M, K, L]": ... def __mul__(self, other): - r"""Combine the prior with a likelihood to form a posterior distribution. + r"""Combine the prior with a likelihood to form a joint model. - The product of a prior and likelihood is proportional to the posterior - distribution. By computing the product of a GP prior and a likelihood - object, a posterior GP object will be returned. Mathematically, this can - be described by: + The product of a prior and likelihood is the joint distribution over + latent function and observations, .. math:: - p(f(\cdot) \mid y) \propto p(y \mid f(\cdot))p(f(\cdot)), + p(f(\cdot), y) = p(y \mid f(\cdot))\,p(f(\cdot)), - where $p(y | f(\cdot))$ is the likelihood and $p(f(\cdot))$ is the prior. + where $p(y | f(\cdot))$ is the likelihood and $p(f(\cdot))$ is the + prior. Conditioning the returned model on data yields the posterior. Example: >>> import gpjax as gpx >>> meanf = gpx.mean_functions.Zero() >>> kernel = gpx.kernels.RBF() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=100) - >>> prior * likelihood + >>> likelihood = gpx.likelihoods.Gaussian() + >>> model = prior * likelihood + Args: - other (Likelihood): The likelihood distribution of the observed dataset. + other (AbstractLikelihood): The likelihood of the observations. Returns: - Posterior: The relevant GP posterior for the given prior and - likelihood. Special cases are accounted for where the model - is conjugate. + JointModel: The joint model for the given prior and likelihood. + The concrete type reflects conjugacy. """ - return construct_posterior(prior=self, likelihood=other) + return construct_model(prior=self, likelihood=other) if tp.TYPE_CHECKING: @tp.overload - def __rmul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]": ... + def __rmul__(self, other: GL) -> "ConjugateModel[M, K, GL]": ... @tp.overload - def __rmul__(self, other: NGL) -> "NonConjugatePosterior[Prior[M, K], NGL]": ... + def __rmul__(self, other: NGL) -> "NonConjugateModel[M, K, NGL]": ... @tp.overload - def __rmul__(self, other: L) -> "AbstractPosterior[Prior[M, K], L]": ... + def __rmul__(self, other: L) -> "JointModel[M, K, L]": ... def __rmul__(self, other): - r"""Combine the prior with a likelihood to form a posterior distribution. + r"""Order-invariant product: ``likelihood * prior``.""" + return self.__mul__(other) - Reimplement the multiplication operator to allow for order-invariant - product of a likelihood and a prior i.e., likelihood * prior. + def __call__( + self, + test_inputs: Num[Array, "N D"], + *, + covariance: Literal["dense", "diagonal"] = "dense", + ) -> GaussianDistribution: + r"""Evaluate the prior process at the given points. Args: - other (Likelihood): The likelihood distribution of the observed - dataset. + test_inputs: Input locations where the GP should be evaluated. + covariance: Whether to return the dense joint covariance at the + test inputs or only the marginal (diagonal) variances. Returns: - Posterior: The relevant GP posterior for the given prior and - likelihood. Special cases are accounted for where the model - is conjugate. + GaussianDistribution: A multivariate normal random variable + representation of the Gaussian process. """ - return self.__mul__(other) + return self.predict(test_inputs, covariance=covariance) def predict( self, test_inputs: Num[Array, "N D"], *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", + covariance: Literal["dense", "diagonal"] = "dense", ) -> GaussianDistribution: - r"""Compute the predictive prior distribution for a given set of - parameters. The output of this function is a ``GaussianDistribution`` - for a given set of inputs. - - In the following example, we compute the predictive prior distribution - and then evaluate it on the interval :math:`[0, 1]`: + r"""Compute the prior predictive distribution at the test inputs. Example: >>> import gpjax as gpx @@ -263,27 +213,24 @@ def predict( >>> prior.predict(jnp.linspace(0, 1, 100)[:, None]) Args: - test_inputs (Float[Array, "N D"]): The inputs at which to evaluate the - prior distribution. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. + test_inputs (Float[Array, "N D"]): The inputs at which to evaluate + the prior distribution. + covariance: Whether to return the dense joint covariance at the + test inputs or only the marginal (diagonal) variances. Returns: - GaussianDistribution: A multivariate normal random variable representation - of the Gaussian process. + GaussianDistribution: A multivariate normal random variable + representation of the Gaussian process. """ - mean_at_test = self.mean_function(test_inputs) - if return_covariance_type == "dense": - Kxx_dense = add_jitter( + if covariance == "dense": + gram_dense = add_jitter( self.kernel.gram(test_inputs).as_matrix(), self.jitter ) - cov = lx.MatrixLinearOperator(Kxx_dense) + cov = lx.MatrixLinearOperator(gram_dense) else: - Ktt_diag = lx.diagonal(self.kernel.diagonal(test_inputs)) - var = Ktt_diag + self.jitter + gram_diag = lx.diagonal(self.kernel.diagonal(test_inputs)) + var = gram_diag + self.jitter cov = lx.DiagonalLinearOperator(jnp.atleast_1d(var.squeeze())) return GaussianDistribution( @@ -298,27 +245,14 @@ def sample_approx( ) -> FunctionalSample: r"""Approximate samples from the Gaussian process prior. - Build an approximate sample from the Gaussian process prior. This method - provides a function that returns the evaluations of a sample across any - given inputs. - - In particular, we approximate the Gaussian processes' prior as the + Build an approximate sample from the Gaussian process prior via the finite feature approximation - $\hat{f}(x) = \sum_{i=1}^m\phi_i(x)\theta_i$ where $\phi_i$ are $m$ features - sampled from the Fourier feature decomposition of the model's kernel and - $\theta_i$ are samples from a unit Gaussian. - - A key property of such functional samples is that the same sample draw is - evaluated for all queries. Consistency is a property that is prohibitively costly - to ensure when sampling exactly from the GP prior, as the cost of exact sampling - scales cubically with the size of the sample. In contrast, finite feature representations - can be evaluated with constant cost regardless of the required number of queries. - - In the following example, we build 10 such samples and then evaluate them - over the interval $[0, 1]$: + $\hat{f}(x) = \sum_{i=1}^m\phi_i(x)\theta_i$ where $\phi_i$ are $m$ + features sampled from the Fourier feature decomposition of the model's + kernel and $\theta_i$ are samples from a unit Gaussian. - For a `prior` distribution, the following code snippet will - build and evaluate an approximate sample. + The same sample draw is evaluated for all queries, at constant cost + per query. Example: >>> import gpjax as gpx @@ -336,14 +270,13 @@ def sample_approx( Args: num_samples (int): The desired number of samples. key (KeyArray): The random seed used for the sample(s). - num_features (int): The number of features used when approximating the - kernel. + num_features (int): The number of features used when approximating + the kernel. Returns: - FunctionalSample: A function representing an approximate sample from the - Gaussian process prior. + FunctionalSample: A function representing an approximate sample + from the Gaussian process prior. """ - if (not isinstance(num_samples, int)) or num_samples <= 0: raise ValueError("num_samples must be a positive integer") @@ -360,146 +293,112 @@ def sample_fn(test_inputs: Float[Array, "N D"]) -> Float[Array, "N B"]: return sample_fn -P = tp.TypeVar("P", bound=AbstractPrior) - - ####################### -# GP Posteriors +# Joint models ####################### -class AbstractPosterior(_SummaryMixin, eqx.Module, tp.Generic[P, L]): - r"""Abstract Gaussian process posterior. - - The base GP posterior object conditioned on an observed dataset. All - posterior objects should inherit from this class. +class JointModel(_SummaryMixin, eqx.Module, tp.Generic[M, K, L]): + r"""The joint distribution $p(f, y) = p(y \mid f)\,p(f)$. + + Pairs a :class:`Prior` with a likelihood. This is the *trainable* object: + ``gpx.fit`` optimises its hyperparameters. Conditioning it on data — + ``model.condition(D)`` or ``model | D`` — produces the posterior process. + + The base class carries no inference of its own; concrete subclasses + (:class:`ConjugateModel`, :class:`NonConjugateModel`, + :class:`HeteroscedasticModel`) define what conditioning means for their + likelihood. A bare ``JointModel`` is a lightweight pairing used where + inference is delegated elsewhere (e.g. variational families over a + latent noise process). """ - prior: AbstractPrior + prior: Prior likelihood: tp.Any - jitter: float = eqx.field(static=True, default=1e-6) - def __init__( - self, - prior: AbstractPrior[M, K], - likelihood: L, - jitter: float = 1e-6, - ): - r"""Construct a Gaussian process posterior. + def __init__(self, prior: Prior[M, K], likelihood: L): + r"""Construct a joint model. Args: - prior (AbstractPrior): The prior distribution. - likelihood (AbstractLikelihood): The likelihood distribution. - jitter (float): A small constant added to the diagonal of the - covariance matrix to ensure numerical stability. + prior (Prior): The prior process. + likelihood (AbstractLikelihood): The observation likelihood. """ self.prior = prior self.likelihood = likelihood - self.jitter = jitter - def __call__( - self, - test_inputs: Num[Array, "N D"], - train_data: Dataset, - *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", - ) -> GaussianDistribution: - r"""Evaluate the Gaussian process posterior at the given points. - - The output of this function is a ``GaussianDistribution`` from which - the latent function's mean and covariance can be evaluated and the - distribution can be sampled. - - Under the hood, ``__call__`` invokes the ``predict`` method. Classes - inheriting ``AbstractPosterior`` should not overwrite ``__call__`` and - should instead define a ``predict`` method. + def condition(self, train_data: Dataset) -> Posterior: + r"""Condition the joint model on data, returning the posterior process. Args: - test_inputs: Input locations where the GP should be evaluated. - train_data: Training dataset to condition on. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. + train_data: The observations to condition on. Returns: - GaussianDistribution: A multivariate normal random variable representation - of the Gaussian process. + Posterior: The conditioned process $p(f \mid \mathcal{D})$. """ - return self.predict( - test_inputs, - train_data, - return_covariance_type=return_covariance_type, + raise NotImplementedError( + f"{type(self).__name__} does not define direct conditioning; " + "use a variational family for inference." ) - @abstractmethod - def predict( + def __or__(self, train_data: Dataset) -> Posterior: + r"""Sugar for conditioning: ``model | D`` reads as $p(f \mid \mathcal{D})$.""" + return self.condition(train_data) + + def _prepare(self, train_data: Dataset) -> "JointModel": + r"""Hook for data-dependent initialisation; returns a ready-to-fit model.""" + del train_data + return self + + def __call__( self, test_inputs: Num[Array, "N D"], train_data: Dataset, *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", + covariance: Literal["dense", "diagonal"] = "dense", ) -> GaussianDistribution: - r"""Compute the latent function's multivariate normal distribution for a - given set of parameters. For any class inheriting the `AbstractPosterior` class, - this method must be implemented. - - Args: - test_inputs: Input locations where the GP should be evaluated. - train_data: Training dataset to condition on. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. + r"""Sugar: condition on ``train_data`` and query at ``test_inputs``. - Returns: - GaussianDistribution: A multivariate normal random variable representation - of the Gaussian process. + Equivalent to ``self.condition(train_data)(test_inputs)``. """ - raise NotImplementedError - - -class LatentPosterior(AbstractPosterior[P, L]): - r"""A posterior shell used to expose prior structure without inference.""" + return self.predict(test_inputs, train_data, covariance=covariance) def predict( self, test_inputs: Num[Array, "N D"], train_data: Dataset, *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", + covariance: Literal["dense", "diagonal"] = "dense", ) -> GaussianDistribution: - raise NotImplementedError( - "LatentPosteriors are a lightweight wrapper for priors and do not " - "implement predictive distributions. Use a variational family for inference." - ) - + r"""Sugar: condition on ``train_data`` and query at ``test_inputs``. -class ConjugatePosterior(AbstractPosterior[P, GL]): - r"""A Conjuate Gaussian process posterior object. + Defined as exactly ``self.condition(train_data)(test_inputs)``. When + making repeated predictions, condition once and reuse the returned + posterior — the factorisation is cached there. - A Gaussian process posterior distribution when the constituent likelihood - function is a Gaussian distribution. In such cases, the latent function values - $f$ can be analytically integrated out of the posterior distribution. - As such, many computational operations can be simplified; something we make use - of in this object. + Args: + test_inputs: A Jax array of test inputs. + train_data: A `gpx.Dataset` to condition on. + covariance: Whether to return the dense joint covariance at the + test inputs or only the marginal (diagonal) variances. - For a Gaussian process prior $p(\mathbf{f})$ and a Gaussian likelihood - $p(y | \mathbf{f}) = \mathcal{N}(y\mid \mathbf{f}, \sigma^2))$ where - $\mathbf{f} = f(\mathbf{x})$, the predictive posterior distribution at - a set of inputs $\mathbf{x}$ is given by + Returns: + GaussianDistribution: The predictive distribution. + """ + return self.condition(train_data)(test_inputs, covariance=covariance) - .. math:: - \begin{aligned} - p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star}, \mathbf{f} \mid \mathbf{y})\\ - & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}} - \end{aligned} +class ConjugateModel(JointModel[M, K, GL]): + r"""A joint model with Gaussian likelihood: conditioning is exact. - where + For a Gaussian process prior $p(\mathbf{f})$ and a Gaussian likelihood + $p(y | \mathbf{f}) = \mathcal{N}(y\mid \mathbf{f}, \sigma^2))$, the latent + function can be analytically integrated out. Conditioning returns the + closed-form posterior .. math:: \begin{aligned} - \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y} \\ + p(\mathbf{f}^{\star}\mid \mathbf{y}) & =\mathcal{N}(\mathbf{f}^{\star}; + \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}}),\\ + \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y}, \\ \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}). \end{aligned} @@ -507,120 +406,29 @@ class ConjugatePosterior(AbstractPosterior[P, GL]): >>> import gpjax as gpx >>> import jax.numpy as jnp >>> + >>> xtrain = jnp.linspace(0, 1).reshape(-1, 1) + >>> D = gpx.Dataset(X=xtrain, y=jnp.sin(xtrain)) + >>> >>> prior = gpx.gps.Prior( ... mean_function = gpx.mean_functions.Zero(), ... kernel = gpx.kernels.RBF() ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=100) - >>> - >>> posterior = prior * likelihood + >>> model = prior * gpx.likelihoods.Gaussian() + >>> posterior = model.condition(D) + >>> predictive = posterior(xtrain) + >>> evidence = posterior.log_marginal_likelihood """ - def predict( - self, - test_inputs: Num[Array, "M D"], - train_data: Dataset, - *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", - ) -> GaussianDistribution: - r"""Query the predictive posterior distribution. - - Conditional on a training data set, compute the GP's posterior - predictive distribution for a given set of parameters. The returned function - can be evaluated at a set of test inputs to compute the corresponding - predictive density. - - The predictive distribution of a conjugate GP is given by - $$ - p(\mathbf{f}^{\star}\mid \mathbf{y}) & = \int p(\mathbf{f}^{\star} \mathbf{f} \mid \mathbf{y})\\ - & =\mathcal{N}(\mathbf{f}^{\star} \boldsymbol{\mu}_{\mid \mathbf{y}}, \boldsymbol{\Sigma}_{\mid \mathbf{y}} - $$ - where - $$ - \boldsymbol{\mu}_{\mid \mathbf{y}} & = k(\mathbf{x}^{\star}, \mathbf{x})\left(k(\mathbf{x}, \mathbf{x}')+\sigma^2\mathbf{I}_n\right)^{-1}\mathbf{y} \\ - \boldsymbol{\Sigma}_{\mid \mathbf{y}} & =k(\mathbf{x}^{\star}, \mathbf{x}^{\star\prime}) -k(\mathbf{x}^{\star}, \mathbf{x})\left( k(\mathbf{x}, \mathbf{x}') + \sigma^2\mathbf{I}_n \right)^{-1}k(\mathbf{x}, \mathbf{x}^{\star}). - $$ - - The conditioning set is a GPJax `Dataset` object, whilst predictions - are made on a regular Jax array. - - Example: - >>> import gpjax as gpx - >>> import jax.numpy as jnp - >>> - >>> xtrain = jnp.linspace(0, 1).reshape(-1, 1) - >>> ytrain = jnp.sin(xtrain) - >>> D = gpx.Dataset(X=xtrain, y=ytrain) - >>> xtest = jnp.linspace(0, 1).reshape(-1, 1) - >>> - >>> prior = gpx.gps.Prior(mean_function = gpx.mean_functions.Zero(), kernel = gpx.kernels.RBF()) - >>> posterior = prior * gpx.likelihoods.Gaussian(num_datapoints = D.n) - >>> predictive_dist = posterior(xtest, D) - - Args: - test_inputs (Num[Array, "N D"]): A Jax array of test inputs at which the - predictive distribution is evaluated. - train_data (Dataset): A `gpx.Dataset` object that contains the input and - output data used for training dataset. - return_covariance_type: Literal denoting whether to return the full covariance - of the joint predictive distribution at the test_inputs (dense) - or just the the standard-deviation of the predictive distribution at - the test_inputs. + def condition(self, train_data: Dataset) -> ExactPosterior: + r"""Condition on data exactly. Returns: - GaussianDistribution: A function that accepts an input array and - returns the predictive distribution as a `GaussianDistribution`. + ExactPosterior: The closed-form posterior process, with the + training-covariance factorisation cached. Exposes the + predictive (via ``__call__``), ``log_marginal_likelihood``, + ``loo`` and ``sample_approx``. """ - import warnings - - kernel = self.prior.kernel - x, y = train_data.X, train_data.y - P = self.likelihood.num_outputs - - # Prepare targets via likelihood protocol (identity for single-output, - # output-major reshape for multi-output) - mx = self.prior.mean_function(x) - y_flat, mx_flat = self.likelihood.prepare_targets(y, mx) - noise = self.likelihood.noise_vector(train_data.n) - - Kxx = kernel.gram(x) - Kxx_dense = add_jitter(Kxx.as_matrix(), self.jitter) - Sigma_dense = Kxx_dense + jnp.diag(noise) - L_sigma = jnp.linalg.cholesky(Sigma_dense) - - Kxt = kernel.cross_covariance(x, test_inputs) - L_inv_Kxt = jsp.linalg.solve_triangular(L_sigma, Kxt, lower=True) - L_inv_y_diff = jsp.linalg.solve_triangular( - L_sigma, y_flat - mx_flat, lower=True - ) - - mean_t_raw = self.prior.mean_function(test_inputs) - mean_t = jnp.tile(mean_t_raw, (P, 1)) if P > 1 else mean_t_raw - mean = mean_t + jnp.matmul(L_inv_Kxt.T, L_inv_y_diff) - - # Diagonal covariance not yet supported for multi-output - if return_covariance_type == "diagonal" and P > 1: - warnings.warn( - "Diagonal covariance is not yet supported for multi-output GPs. " - "Returning full covariance.", - stacklevel=2, - ) - return_covariance_type = "dense" - - if return_covariance_type == "dense": - Ktt = kernel.gram(test_inputs).as_matrix() - covariance = Ktt - jnp.matmul(L_inv_Kxt.T, L_inv_Kxt) - covariance = add_jitter(covariance, self.prior.jitter) - cov = lx.MatrixLinearOperator(covariance) - else: - Ktt_diag = lx.diagonal(kernel.diagonal(test_inputs)) - var = ( - Ktt_diag - - jnp.einsum("ij,ji->i", L_inv_Kxt.T, L_inv_Kxt) - + self.prior.jitter - ) - cov = lx.DiagonalLinearOperator(jnp.atleast_1d(var.squeeze())) - return GaussianDistribution(loc=jnp.atleast_1d(mean.squeeze()), scale=cov) + return ExactPosterior(self.prior, self.likelihood, train_data) def sample_approx( self, @@ -629,228 +437,132 @@ def sample_approx( key: KeyArray, num_features: int | None = 100, ) -> FunctionalSample: - r"""Draw approximate samples from the Gaussian process posterior. - - Build an approximate sample from the Gaussian process posterior. This method - provides a function that returns the evaluations of a sample across any given - inputs. - - Unlike when building approximate samples from a Gaussian process prior, decompositions - based on Fourier features alone rarely give accurate samples. Therefore, we must also - include an additional set of features (known as canonical features) to better model the - transition from Gaussian process prior to Gaussian process posterior. For more details - see [Wilson et. al. (2020)](https://arxiv.org/abs/2002.09309). - - In particular, we approximate the Gaussian processes' posterior as the finite - feature approximation - $\hat{f}(x) = \sum_{i=1}^m \phi_i(x)\theta_i + \sum{j=1}^N v_jk(.,x_j)$ - where $\phi_i$ are m features sampled from the Fourier feature decomposition of - the model's kernel and $k(., x_j)$ are N canonical features. The Fourier - weights $\theta_i$ are samples from a unit Gaussian. See - [Wilson et. al. (2020)](https://arxiv.org/abs/2002.09309) for expressions - for the canonical weights $v_j$. - - A key property of such functional samples is that the same sample draw is - evaluated for all queries. Consistency is a property that is prohibitively costly - to ensure when sampling exactly from the GP prior, as the cost of exact sampling - scales cubically with the size of the sample. In contrast, finite feature representations - can be evaluated with constant cost regardless of the required number of queries. - - Args: - num_samples (int): The desired number of samples. - key (KeyArray): The random seed used for the sample(s). - num_features (int): The number of features used when approximating the - kernel. + r"""Sugar: ``self.condition(train_data).sample_approx(...)``. - Returns: - FunctionalSample: A function representing an approximate sample from the Gaussian - process prior. + Draw approximate posterior samples via pathwise conditioning + (Wilson et al., 2020). """ - if (not isinstance(num_samples, int)) or num_samples <= 0: - raise ValueError("num_samples must be a positive integer") - - # sample fourier features - freq_key, weight_key, noise_key = jr.split(key, 3) - fourier_feature_fn = _build_fourier_features_fn( - self.prior, num_features, freq_key + return self.condition(train_data).sample_approx( + num_samples, key, num_features ) - fourier_weights = jr.normal(weight_key, [num_samples, 2 * num_features]) - - obs_var = _val(self.likelihood.obs_stddev) ** 2 - Kxx = self.prior.kernel.gram(train_data.X) - Sigma_dense = add_jitter(Kxx.as_matrix(), obs_var + self.jitter) - L_sigma = jnp.linalg.cholesky(Sigma_dense) - eps = jnp.sqrt(obs_var) * jr.normal(noise_key, [train_data.n, num_samples]) - y = train_data.y - self.prior.mean_function(train_data.X) - Phi = fourier_feature_fn(train_data.X) - # Solve L_sigma @ canonical_weights = rhs - rhs = y + eps - jnp.inner(Phi, fourier_weights) - canonical_weights = jsp.linalg.cho_solve((L_sigma, True), rhs) # [N, B] - - def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]: - fourier_features = fourier_feature_fn(test_inputs) - weight_space_contribution = jnp.inner(fourier_features, fourier_weights) - canonical_features = self.prior.kernel.cross_covariance( - test_inputs, train_data.X - ) - function_space_contribution = jnp.matmul( - canonical_features, canonical_weights - ) - - return ( - self.prior.mean_function(test_inputs) - + weight_space_contribution - + function_space_contribution - ) - - return sample_fn +class NonConjugateModel(JointModel[M, K, NGL]): + r"""A joint model with non-Gaussian likelihood. -class NonConjugatePosterior(AbstractPosterior[P, NGL]): - r"""A non-conjugate Gaussian process posterior object. + Exact conditioning is intractable; the model instead carries a whitened + latent vector $w_x$ as a trainable parameter, and conditioning produces + the approximate posterior implied by its current value. Markov chain Monte + Carlo, variational inference, or MAP optimisation (via + ``gpx.objectives.log_posterior_density``) refine it. - A Gaussian process posterior object for models where the likelihood is - non-Gaussian. Unlike the `ConjugatePosterior` object, the - `NonConjugatePosterior` object does not provide an exact marginal - log-likelihood function. Instead, the `NonConjugatePosterior` object - represents the posterior distributions as a function of the model's - hyperparameters and the latent function. Markov chain Monte Carlo, - variational inference, or Laplace approximations can then be used to sample - from, or optimise an approximation to, the posterior distribution. + The latent is sized by the training data, so it is initialised lazily on + first contact with data — ``gpx.fit`` does this automatically, or call + :meth:`init_latent` explicitly. """ latent: tp.Any def __init__( self, - prior: P, + prior: Prior[M, K], likelihood: NGL, latent: tp.Union[Float[Array, "N 1"], AbstractUnwrappable, None] = None, - jitter: float = 1e-6, - key: KeyArray = jr.key(42), ): - r"""Construct a non-conjugate Gaussian process posterior. + r"""Construct a non-conjugate joint model. Args: - prior (AbstractPrior): The prior distribution. - likelihood (AbstractLikelihood): The likelihood distribution. - jitter (float): A small constant added to the diagonal of the - covariance matrix to ensure numerical stability. + prior (Prior): The prior process. + likelihood (AbstractLikelihood): The observation likelihood. + latent: Whitened latent function values at the training inputs. + ``None`` (the default) defers initialisation to first data + contact. """ - super().__init__(prior=prior, likelihood=likelihood, jitter=jitter) + super().__init__(prior=prior, likelihood=likelihood) + if latent is None or isinstance(latent, AbstractUnwrappable): + self.latent = latent + else: + self.latent = Real(latent) - if latent is None: - latent = jr.normal(key, shape=(self.likelihood.num_datapoints, 1)) + def init_latent( + self, num_datapoints: int, key: KeyArray = jr.key(42) + ) -> "NonConjugateModel[M, K, NGL]": + r"""Return a copy of this model with the latent vector initialised. - self.latent = ( - latent if isinstance(latent, AbstractUnwrappable) else Real(latent) + Args: + num_datapoints: The number of training observations the latent + must cover. + key: The random seed for the initial values. + """ + latent = jr.normal(key, shape=(num_datapoints, 1)) + return NonConjugateModel( + prior=self.prior, likelihood=self.likelihood, latent=latent ) - def predict( - self, - test_inputs: Num[Array, "M D"], - train_data: Dataset, - *, - return_covariance_type: Literal["dense", "diagonal"] = "dense", - ) -> GaussianDistribution: - r"""Query the predictive posterior distribution. + def _prepare(self, train_data: Dataset) -> "NonConjugateModel[M, K, NGL]": + if self.latent is not None: + return self + return self.init_latent(train_data.n) - Conditional on a set of training data, compute the GP's posterior - predictive distribution for a given set of parameters. The returned - function can be evaluated at a set of test inputs to compute the - corresponding predictive density. Note, to gain predictions on the scale - of the original data, the returned distribution will need to be - transformed through the likelihood function's inverse link function. + def condition(self, train_data: Dataset) -> LatentPosterior: + r"""Return the approximate posterior implied by the current latent. - Args: - test_inputs (Num[Array, "N D"]): A Jax array of test inputs at which the - predictive distribution is evaluated. - train_data (Dataset): A `gpx.Dataset` object that contains the input - and output data used for training dataset. - return_covariance_type: Literal denoting whether to return the full - covariance of the joint predictive distribution at the test_inputs - (dense) or just the the standard-deviation of the predictive - distribution at the test_inputs. + A ``None`` latent conditions at the prior mean (zeros in whitened + space). Returns: - GaussianDistribution: A function that accepts an - input array and returns the predictive distribution as - a `dx.Distribution`. + LatentPosterior: The conditioned process. Exposes the predictive + (via ``__call__``) and ``log_posterior_density``. """ - x = train_data.X - t = test_inputs - mean_function = self.prior.mean_function - kernel = self.prior.kernel - - # Precompute lower triangular of Gram matrix - Kxx = kernel.gram(x) - Kxx_dense = add_jitter(Kxx.as_matrix(), self.prior.jitter) - Lx = jnp.linalg.cholesky(Kxx_dense) - - Kxt = kernel.cross_covariance(x, t) - # Lx^{-1} Kxt - Lx_inv_Kxt = jsp.linalg.solve_triangular(Lx, Kxt, lower=True) - - mean_t = mean_function(t) - # Whitened function values, wx, corresponding to the inputs, x - wx = _val(self.latent) - - # mut + Ktx Lx^{-1} wx - mean = mean_t + jnp.matmul(Lx_inv_Kxt.T, wx) - - if return_covariance_type == "dense": - Ktt = kernel.gram(test_inputs).as_matrix() - covariance = Ktt - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt) - covariance = add_jitter(covariance, self.prior.jitter) - cov = lx.MatrixLinearOperator(covariance) - else: - Ktt_diag = lx.diagonal(kernel.diagonal(test_inputs)) - var = ( - Ktt_diag - - jnp.einsum("ij,ji->i", Lx_inv_Kxt.T, Lx_inv_Kxt) - + self.prior.jitter - ) - cov = lx.DiagonalLinearOperator(jnp.atleast_1d(var.squeeze())) - - return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), cov) + latent = self.latent + if latent is None: + latent = jnp.zeros((train_data.n, 1)) + return LatentPosterior(self.prior, self.likelihood, latent, train_data) -class HeteroscedasticPosterior(LatentPosterior[P, HL]): - r"""Posterior shell for heteroscedastic likelihoods. +class HeteroscedasticModel(JointModel[M, K, HL]): + r"""A joint model with input-dependent (heteroscedastic) noise. - The posterior retains both the signal and noise priors; inference is delegated - to variational families and specialised objectives. - """ + The joint holds *two* priors — one over the signal process and one over + the latent noise process — which is why it is constructed directly rather + than via ``prior * likelihood``: - noise_prior: tp.Any - noise_posterior: tp.Any + .. code-block:: python - def __init__( - self, - prior: AbstractPrior[M, K], - likelihood: HL, - jitter: float = 1e-6, - ): - if likelihood.noise_prior is None: - raise ValueError("Heteroscedastic likelihoods require a noise_prior.") - super().__init__(prior=prior, likelihood=likelihood, jitter=jitter) - self.noise_prior = likelihood.noise_prior - self.noise_posterior = LatentPosterior( - prior=self.noise_prior, likelihood=likelihood, jitter=jitter + model = gpx.gps.HeteroscedasticModel( + prior=signal_prior, + likelihood=gpx.likelihoods.HeteroscedasticGaussian(), + noise_prior=noise_prior, ) + Inference is delegated to + :class:`gpjax.variational_families.HeteroscedasticVariationalFamily` and + the ``heteroscedastic_elbo`` objective; the noise process is exposed as + the nested joint model :attr:`noise_model`. + """ -class ChainedPosterior(HeteroscedasticPosterior[P, HL]): - r"""Posterior routed for heteroscedastic likelihoods using chained bounds.""" + noise_prior: tp.Any + noise_model: tp.Any def __init__( self, - prior: AbstractPrior[M, K], + prior: Prior[M, K], likelihood: HL, - jitter: float = 1e-6, + noise_prior: Prior, ): - super().__init__(prior=prior, likelihood=likelihood, jitter=jitter) + r"""Construct a heteroscedastic joint model. + + Args: + prior (Prior): The prior over the signal process. + likelihood (AbstractHeteroscedasticLikelihood): The observation + likelihood. + noise_prior (Prior): The prior over the latent noise process. + """ + if noise_prior is None: + raise ValueError("Heteroscedastic models require a noise_prior.") + super().__init__(prior=prior, likelihood=likelihood) + self.noise_prior = noise_prior + self.noise_model = JointModel(prior=noise_prior, likelihood=likelihood) ####################### @@ -859,41 +571,33 @@ def __init__( @tp.overload -def construct_posterior(prior: P, likelihood: GL) -> ConjugatePosterior[P, GL]: ... - - -@tp.overload -def construct_posterior(prior: P, likelihood: NGL) -> NonConjugatePosterior[P, NGL]: ... +def construct_model(prior: Prior, likelihood: GL) -> ConjugateModel: ... @tp.overload -def construct_posterior( - prior: P, likelihood: HeteroscedasticGaussian -) -> HeteroscedasticPosterior[P, HeteroscedasticGaussian]: ... - +def construct_model(prior: Prior, likelihood: NGL) -> NonConjugateModel: ... -@tp.overload -def construct_posterior( - prior: P, likelihood: AbstractHeteroscedasticLikelihood -) -> ChainedPosterior[P, AbstractHeteroscedasticLikelihood]: ... +def construct_model( + prior: Prior, likelihood: AbstractLikelihood +) -> "JointModel": + r"""Construct the joint model for a prior/likelihood pair. -def construct_posterior( - prior: AbstractPrior, likelihood: AbstractLikelihood -) -> "AbstractPosterior": - r"""Utility function for constructing a posterior object from a prior and - likelihood. The function will automatically select the correct posterior - object based on the likelihood. + Selects the concrete :class:`JointModel` subclass from the likelihood's + conjugacy. This is what ``prior * likelihood`` calls. Args: - prior (Prior): The Prior distribution. - likelihood (AbstractLikelihood): The likelihood that represents our - beliefs around the distribution of the data. + prior (Prior): The prior process. + likelihood (AbstractLikelihood): The observation likelihood. Returns: - AbstractPosterior: A posterior distribution. If the likelihood is - Gaussian, then a `ConjugatePosterior` will be returned. Otherwise, a - `NonConjugatePosterior` will be returned. + JointModel: A ``ConjugateModel`` for Gaussian likelihoods, a + ``NonConjugateModel`` otherwise. + + Raises: + ValueError: For heteroscedastic likelihoods, which carry a second + prior and must be constructed directly via + ``HeteroscedasticModel(prior, likelihood, noise_prior=...)``. """ # Multi-output validation from gpjax.kernels.multioutput.base import MultiOutputKernel @@ -912,59 +616,25 @@ def construct_posterior( "Multi-output kernels require a MultiOutputGaussian likelihood." ) - if isinstance(likelihood, Gaussian): - return ConjugatePosterior(prior=prior, likelihood=likelihood) - - if ( - isinstance(likelihood, HeteroscedasticGaussian) - and likelihood.supports_tight_bound() - ): - return HeteroscedasticPosterior(prior=prior, likelihood=likelihood) - if isinstance(likelihood, AbstractHeteroscedasticLikelihood): - return ChainedPosterior(prior=prior, likelihood=likelihood) - - return NonConjugatePosterior(prior=prior, likelihood=likelihood) - - -def _build_fourier_features_fn( - prior: Prior, num_features: int, key: KeyArray -) -> tp.Callable[[Float[Array, "N D"]], Float[Array, "N L"]]: - r"""Return a function that evaluates features sampled from the Fourier feature - decomposition of the prior's kernel. - - Args: - prior (Prior): The Prior distribution. - num_features (int): The number of feature functions to be sampled. - key (KeyArray): The random seed used. - - Returns: - Callable: A callable function evaluating the sampled feature functions. - """ - if (not isinstance(num_features, int)) or num_features <= 0: - raise ValueError("num_features must be a positive integer") - - # Approximate kernel with feature decomposition - approximate_kernel = RFF( - base_kernel=prior.kernel, num_basis_fns=num_features, key=key - ) + raise ValueError( + "Heteroscedastic likelihoods carry a second (noise) prior, which " + "the two-operand product cannot express. Construct the model " + "directly: HeteroscedasticModel(prior, likelihood, " + "noise_prior=...)." + ) - def eval_fourier_features(test_inputs: Float[Array, "N D"]) -> Float[Array, "N L"]: - Phi = approximate_kernel.compute_features(x=test_inputs) - Phi *= jnp.sqrt(_val(prior.kernel.variance) / num_features) - return Phi + if isinstance(likelihood, Gaussian): + return ConjugateModel(prior=prior, likelihood=likelihood) - return eval_fourier_features + return NonConjugateModel(prior=prior, likelihood=likelihood) __all__ = [ - "AbstractPosterior", - "AbstractPrior", - "ChainedPosterior", - "ConjugatePosterior", - "HeteroscedasticPosterior", - "LatentPosterior", - "NonConjugatePosterior", + "ConjugateModel", + "HeteroscedasticModel", + "JointModel", + "NonConjugateModel", "Prior", - "construct_posterior", + "construct_model", ] diff --git a/gpjax/models/oilmm.py b/gpjax/models/oilmm.py index 28cf8f432..c0801efc5 100644 --- a/gpjax/models/oilmm.py +++ b/gpjax/models/oilmm.py @@ -298,11 +298,10 @@ def condition_on_observations(self, dataset: Dataset) -> OILMMPosterior: # Create likelihood with projected noise likelihood = Gaussian( - num_datapoints=dataset.n, obs_stddev=jnp.sqrt(projected_noise_vars[i]), ) - # Standard GPJax conditioning: Prior * Likelihood -> ConjugatePosterior + # Standard GPJax conditioning: Prior * Likelihood -> ConjugateModel latent_posteriors.append(self.latent_priors[i] * likelihood) return OILMMPosterior( @@ -315,7 +314,7 @@ def condition_on_observations(self, dataset: Dataset) -> OILMMPosterior: class OILMMPosterior: """Posterior distribution for OILMM. - Wraps M independent ConjugatePosterior objects and provides a unified + Wraps M independent ConjugateModel objects and provides a unified predict() interface that reconstructs predictions in output space. This is a plain class (not eqx.Module) because it holds Dataset objects @@ -323,7 +322,7 @@ class OILMMPosterior: are still eqx.Modules and participate in JAX transformations when accessed. Attributes: - latent_posteriors: Tuple of M independent ConjugatePosterior objects + latent_posteriors: Tuple of M independent ConjugateModel objects latent_datasets: Tuple of M projected training Datasets (one per latent GP) mixing_matrix: OrthogonalMixingMatrix for reconstruction num_latent_gps: Number of latent GPs (m) @@ -338,7 +337,7 @@ def __init__( """Initialize OILMM posterior. Args: - latent_posteriors: Tuple of M ConjugatePosterior objects + latent_posteriors: Tuple of M ConjugateModel objects latent_datasets: Tuple of M Dataset objects (projected training data) mixing_matrix: OrthogonalMixingMatrix containing H, T """ @@ -432,8 +431,6 @@ def oilmm_mll(model: OILMMModel, data: Dataset) -> ScalarFloat: Returns: Scalar log marginal likelihood. """ - from gpjax.linalg.utils import add_jitter - n = data.n p = model.num_outputs m = model.num_latent_gps @@ -461,23 +458,22 @@ def oilmm_mll(model: OILMMModel, data: Dataset) -> ScalarFloat: correction = term_log_S + term_noise + term_residual - # --- Latent GP log-likelihoods computed directly --- - # We compute each latent GP's MLL inline to avoid constructing Gaussian - # likelihood objects, which would trigger parameter validation checks - # that are incompatible with JAX's JIT tracing. + # --- Latent GP log-likelihoods via the conditioning module --- + from gpjax.conditioning import ExactPosterior + from gpjax.dataset import Dataset + from gpjax.likelihoods import Gaussian + X, y_projected = model._project_observations(data) # [N, D], [M, N] projected_noise_vars = mix.projected_noise_variance # [M] latent_lls = [] for i in range(m): - yi = y_projected[i] # [N] - prior_i = model.latent_priors[i] - mx = prior_i.mean_function(X).squeeze() # [N] - Kxx = prior_i.kernel.gram(X).as_matrix() # [N, N] - Kxx = add_jitter(Kxx, prior_i.jitter) - Sigma = Kxx + projected_noise_vars[i] * jnp.eye(n) - dist = GaussianDistribution(jnp.atleast_1d(mx), lx.MatrixLinearOperator(Sigma)) - latent_lls.append(dist.log_prob(jnp.atleast_1d(yi))) + latent_dataset = Dataset(X=X, y=y_projected[i][:, None]) + likelihood = Gaussian(obs_stddev=jnp.sqrt(projected_noise_vars[i])) + conditioned = ExactPosterior( + model.latent_priors[i], likelihood, latent_dataset + ) + latent_lls.append(conditioned.log_marginal_likelihood) return correction + jnp.sum(jnp.array(latent_lls)) diff --git a/gpjax/objectives.py b/gpjax/objectives.py index b52fa262d..7e524058a 100644 --- a/gpjax/objectives.py +++ b/gpjax/objectives.py @@ -5,15 +5,12 @@ import jax.numpy as jnp import jax.scipy as jsp from jaxtyping import Float -import lineax as lx -import numpyro.distributions as npd import typing_extensions as tpe from gpjax.dataset import Dataset -from gpjax.distributions import GaussianDistribution from gpjax.gps import ( - ConjugatePosterior, - NonConjugatePosterior, + ConjugateModel, + NonConjugateModel, ) from gpjax.likelihoods import ( AbstractHeteroscedasticLikelihood, @@ -36,7 +33,7 @@ Objective = tpe.Callable[[eqx.Module, Dataset], ScalarFloat] -def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat: +def conjugate_mll(model: ConjugateModel, data: Dataset) -> ScalarFloat: r"""Evaluate the marginal log-likelihood of the Gaussian process. Compute the marginal log-likelihood function of the Gaussian process. @@ -70,11 +67,11 @@ def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat: >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function = meanf, kernel=kernel) - >>> posterior = prior * likelihood + >>> model = prior * likelihood - >>> gpx.objectives.conjugate_mll(posterior, D) + >>> gpx.objectives.conjugate_mll(model, D) Our goal is to maximise the marginal log-likelihood. Therefore, when optimising the model's parameters with respect to the parameters, we use the negative @@ -83,46 +80,18 @@ def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat: >>> nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d) Args: - posterior (ConjugatePosterior): The posterior distribution for which - we want to compute the marginal log-likelihood. + model (ConjugateModel): The joint model for which we want to compute + the marginal log-likelihood. data: The training dataset used to compute the marginal log-likelihood. Returns: ScalarFloat: The marginal log-likelihood of the Gaussian process. """ + return model.condition(data).log_marginal_likelihood - from gpjax.kernels.multioutput.base import MultiOutputKernel - x, y = data.X, data.y - kernel = posterior.prior.kernel - mx = posterior.prior.mean_function(x) - - # Validation for multi-output models (user-facing error messages) - if isinstance(kernel, MultiOutputKernel): - if not data.multi_output: - raise ValueError("MultiOutputKernel requires multi-output data.") - if data.num_outputs != kernel.num_outputs: - raise ValueError( - f"Dataset has {data.num_outputs} outputs " - f"but kernel expects {kernel.num_outputs}." - ) - - # Unified path -- prepare_targets is identity for single-output, - # output-major reshape for multi-output - y_flat, mx_flat = posterior.likelihood.prepare_targets(y, mx) - noise = posterior.likelihood.noise_vector(data.n) - - Kxx = kernel.gram(x) - Kxx_dense = add_jitter(Kxx.as_matrix(), posterior.prior.jitter) - Sigma_dense = Kxx_dense + jnp.diag(noise) - Sigma = lx.MatrixLinearOperator(Sigma_dense) - - mll = GaussianDistribution(jnp.atleast_1d(mx_flat.squeeze()), Sigma) - return mll.log_prob(jnp.atleast_1d(y_flat.squeeze())).squeeze() - - -def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat: +def conjugate_loocv(model: ConjugateModel, data: Dataset) -> ScalarFloat: r"""Evaluate the leave-one-out log predictive probability of the Gaussian process following section 5.4.2 of Rasmussen et al. 2006 - Gaussian Processes for Machine Learning. This metric calculates the average performance of all models that can be obtained by training on all but one @@ -150,11 +119,11 @@ def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat ... >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function = meanf, kernel=kernel) - >>> posterior = prior * likelihood + >>> model = prior * likelihood ... - >>> gpx.objectives.conjugate_loocv(posterior, D) + >>> gpx.objectives.conjugate_loocv(model, D) Our goal is to maximise the leave-one-out log predictive probability. Therefore, when optimising the model's parameters with respect to the parameters, we use the negative @@ -163,48 +132,18 @@ def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat >>> nloocv = lambda p, d: -gpx.objectives.conjugate_loocv(p, d) Args: - posterior (ConjugatePosterior): The posterior distribution for which - we want to compute the marginal log-likelihood. + model (ConjugateModel): The joint model for which we want to compute + the leave-one-out predictive probability. data: The training dataset used to compute the - marginal log-likelihood. + leave-one-out predictive probability. Returns: - ScalarFloat: The marginal log-likelihood of the Gaussian process. + ScalarFloat: The leave-one-out log predictive probability. """ + return jnp.sum(model.condition(data).loo()) - x, y = data.X, data.y - - mx = posterior.prior.mean_function(x) - # Likelihood protocol: identity for single-output, output-major flatten + - # per-output noise for multi-output (mirrors conjugate_mll). - y_flat, mx_flat = posterior.likelihood.prepare_targets(y, mx) - noise = posterior.likelihood.noise_vector(data.n) - - # Sigma = Kxx + diag(noise) (+ jitter) - Kxx_dense = add_jitter( - posterior.prior.kernel.gram(x).as_matrix(), posterior.prior.jitter - ) - Sigma_dense = Kxx_dense + jnp.diag(noise) - L = jnp.linalg.cholesky(Sigma_dense) - - # diag(Sigma^-1) straight from L (R&W eq. 5.12) — no separate jnp.linalg.inv - # (folds in audit #662). - Linv = jsp.linalg.solve_triangular(L, jnp.eye(Sigma_dense.shape[0]), lower=True) - Sigma_inv_diag = jnp.sum(Linv**2, axis=0).reshape(-1, 1) # [NP, 1] - - resid = (y_flat - mx_flat).reshape(-1, 1) - Sigma_inv_y = jsp.linalg.cho_solve((L, True), resid) # [NP, 1] - - loocv_means = mx_flat.reshape(-1, 1) + resid - Sigma_inv_y / Sigma_inv_diag - loocv_stds = jnp.sqrt(1.0 / Sigma_inv_diag) - - loocv_posterior = npd.Normal(loc=loocv_means, scale=loocv_stds) - return jnp.sum(loocv_posterior.log_prob(y_flat.reshape(-1, 1))) - -def log_posterior_density( - posterior: NonConjugatePosterior, data: Dataset -) -> ScalarFloat: +def log_posterior_density(model: NonConjugateModel, data: Dataset) -> ScalarFloat: r"""The log-posterior density of a non-conjugate Gaussian process. This is sometimes referred to as the marginal log-likelihood. @@ -233,44 +172,27 @@ def log_posterior_density( >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Bernoulli() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) - >>> posterior = prior * likelihood + >>> model = (prior * likelihood).init_latent(D.n) - >>> gpx.objectives.log_posterior_density(posterior, D) + >>> gpx.objectives.log_posterior_density(model, D) Args: - posterior (NonConjugatePosterior): The posterior distribution for which - we want to compute the marginal log-likelihood. + model (NonConjugateModel): The joint model for which we want to + compute the log-posterior density. data: The training dataset used to compute the - marginal log-likelihood. + log-posterior density. Returns: ScalarFloat: The log-posterior density of the Gaussian process. """ - - x, y = data.X, data.y - - # Gram matrix - Kxx = posterior.prior.kernel.gram(x) - Kxx_dense = add_jitter(Kxx.as_matrix(), posterior.prior.jitter) - Lx = jnp.linalg.cholesky(Kxx_dense) - - # Compute the prior mean function - mx = posterior.prior.mean_function(x) - - # Whitened function values, wx, corresponding to the inputs, x - wx = _val(posterior.latent) - - # f(x) = mx + Lx wx - fx = mx + Lx @ wx - - # p(y | f(x), theta), where theta are the model hyperparameters - likelihood = posterior.likelihood.link_function(fx) - - # Whitened latent function values prior, p(wx | theta) = N(0, I) - latent_prior = npd.Normal(loc=0.0, scale=1.0) - return likelihood.log_prob(y).sum() + latent_prior.log_prob(wx).sum() + if model.latent is None: + raise ValueError( + "NonConjugateModel.latent is uninitialised: fit the model or call " + "model.init_latent(data.n) first." + ) + return model.condition(data).log_posterior_density non_conjugate_mll = log_posterior_density @@ -295,7 +217,7 @@ def elbo(variational_family: VF, data: Dataset) -> ScalarFloat: >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Bernoulli() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood @@ -346,7 +268,7 @@ def variational_expectation( >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Bernoulli() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood @@ -411,7 +333,7 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat: >>> meanf = gpx.mean_functions.Constant() >>> kernel = gpx.kernels.RBF() - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) + >>> likelihood = gpx.likelihoods.Gaussian() >>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) >>> posterior = prior * likelihood diff --git a/gpjax/state_space/fit.py b/gpjax/state_space/fit.py index 696387ccf..ac5b5912c 100644 --- a/gpjax/state_space/fit.py +++ b/gpjax/state_space/fit.py @@ -60,7 +60,7 @@ def fit_scipy( ... mean_function=gpx.mean_functions.Zero(), ... kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1) + >>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) >>> posterior = prior * likelihood >>> fitted, history = fit_scipy( ... model=posterior, @@ -105,7 +105,7 @@ def fit_lbfgs( ... mean_function=gpx.mean_functions.Zero(), ... kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1) + >>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) >>> posterior = prior * likelihood >>> fitted, history = fit_lbfgs( ... model=posterior, @@ -158,7 +158,7 @@ def fit( ... mean_function=gpx.mean_functions.Zero(), ... kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1) + >>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) >>> posterior = prior * likelihood >>> fitted, history = fit( ... model=posterior, diff --git a/gpjax/state_space/gps.py b/gpjax/state_space/gps.py index 9e3429dfe..e8b45b85c 100644 --- a/gpjax/state_space/gps.py +++ b/gpjax/state_space/gps.py @@ -10,7 +10,7 @@ import paramax from gpjax.distributions import GaussianDistribution -from gpjax.gps import ConjugatePosterior, Prior +from gpjax.gps import ConjugateModel, Prior from gpjax.likelihoods import Gaussian, MultiOutputGaussian @@ -25,7 +25,7 @@ class StateSpacePrior(Prior): covariance only; the marginals are exact. A dense joint predictive is not implemented in v1 and is tracked as a follow-up. This predictive is therefore not Liskov-substitutable for a dense - ``gpjax.gps.ConjugatePosterior`` predictive. + dense ``gpjax.gps.ConjugateModel`` predictive. Example: >>> import gpjax as gpx @@ -38,16 +38,16 @@ class StateSpacePrior(Prior): True """ - def __call__(self, test_inputs, *, return_covariance_type="diagonal"): - return self.predict(test_inputs, return_covariance_type=return_covariance_type) + def __call__(self, test_inputs, *, covariance="diagonal"): + return self.predict(test_inputs, covariance=covariance) - def predict(self, test_inputs, *, return_covariance_type="diagonal"): - if return_covariance_type != "diagonal": + def predict(self, test_inputs, *, covariance="diagonal"): + if covariance != "diagonal": raise NotImplementedError( "State-space prior prediction returns diagonal (marginal) covariance " "only; a dense joint predictive is not implemented in v1 and is " "tracked as a follow-up. The marginal variances returned are exact, " - "so for diagonal-only use pass return_covariance_type='diagonal'." + "so for diagonal-only use pass covariance='diagonal'." ) from gpjax.state_space.kernels import to_sde @@ -68,7 +68,7 @@ def __mul__(self, other): return StateSpaceConjugatePosterior(prior=self, likelihood=other) -class StateSpaceConjugatePosterior(ConjugatePosterior): +class StateSpaceConjugatePosterior(ConjugateModel): """Conjugate posterior for a state-space (Markovian) GP. v1 prediction surface: @@ -76,14 +76,14 @@ class StateSpaceConjugatePosterior(ConjugatePosterior): - ``predict_filter`` : causal filtered prediction (Phase 10) - ``__call__`` : delegates to ``predict`` - Both ``predict`` and ``predict_filter`` reject ``return_covariance_type="dense"`` + Both ``predict`` and ``predict_filter`` reject ``covariance="dense"`` in favour of v1's diagonal-only contract before any further dispatch. **Predictive contract (v1):** prediction returns diagonal (marginal) covariance only; the marginals are exact. A dense joint predictive is not implemented in v1 and is tracked as a follow-up. This predictive is therefore not Liskov-substitutable for a dense - ``gpjax.gps.ConjugatePosterior`` predictive. + dense ``gpjax.gps.ConjugateModel`` predictive. Example: >>> import gpjax as gpx @@ -92,7 +92,7 @@ class StateSpaceConjugatePosterior(ConjugatePosterior): ... mean_function=gpx.mean_functions.Zero(), ... kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1) + >>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) >>> posterior = prior * likelihood >>> posterior.__class__.__name__ 'StateSpaceConjugatePosterior' @@ -103,13 +103,13 @@ def __call__( test_inputs, train_data, *, - return_covariance_type="diagonal", + covariance="diagonal", observation_mask=None, ): return self.predict( test_inputs, train_data, - return_covariance_type=return_covariance_type, + covariance=covariance, observation_mask=observation_mask, ) @@ -118,16 +118,16 @@ def predict( test_inputs, train_data, *, - return_covariance_type="diagonal", + covariance="diagonal", observation_mask=None, ): - if return_covariance_type != "diagonal": + if covariance != "diagonal": raise NotImplementedError( "State-space posterior predict returns diagonal (marginal) " "covariance only; a dense joint predictive is not implemented in v1 " "and is tracked as a follow-up. The marginal variances returned are " "exact, so for diagonal-only use pass " - "return_covariance_type='diagonal'." + "covariance='diagonal'." ) from gpjax.state_space.prediction import predict_smoothed @@ -140,16 +140,16 @@ def predict_filter( test_inputs, train_data, *, - return_covariance_type="diagonal", + covariance="diagonal", observation_mask=None, ): - if return_covariance_type != "diagonal": + if covariance != "diagonal": raise NotImplementedError( "State-space posterior predict_filter returns diagonal (marginal) " "covariance only; a dense joint predictive is not implemented in v1 " "and is tracked as a follow-up. The marginal variances returned are " "exact, so for diagonal-only use pass " - "return_covariance_type='diagonal'." + "covariance='diagonal'." ) from gpjax.state_space.prediction import predict_filtered diff --git a/gpjax/state_space/objectives.py b/gpjax/state_space/objectives.py index 13ab2d271..e4f5546c7 100644 --- a/gpjax/state_space/objectives.py +++ b/gpjax/state_space/objectives.py @@ -45,7 +45,7 @@ def state_space_mll( ... mean_function=gpx.mean_functions.Zero(), ... kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ... ) - >>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1) + >>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) >>> posterior = prior * likelihood >>> train_data = gpx.Dataset(X=X, y=y) >>> mll = state_space_mll(posterior, train_data) diff --git a/gpjax/variational_families.py b/gpjax/variational_families.py index f1571e2bb..8ab0efe5c 100644 --- a/gpjax/variational_families.py +++ b/gpjax/variational_families.py @@ -29,10 +29,9 @@ from gpjax.dataset import Dataset from gpjax.distributions import GaussianDistribution from gpjax.gps import ( - AbstractPosterior, - AbstractPrior, - ChainedPosterior, - HeteroscedasticPosterior, + HeteroscedasticModel, + JointModel, + Prior, ) from gpjax.kernels.base import AbstractKernel from gpjax.likelihoods import ( @@ -60,9 +59,9 @@ NGL = tp.TypeVar("NGL", bound=NonGaussian) GL = tp.TypeVar("GL", bound=Gaussian) HL = tp.TypeVar("HL", bound=AbstractHeteroscedasticLikelihood) -P = tp.TypeVar("P", bound=AbstractPrior) -PP = tp.TypeVar("PP", bound=AbstractPosterior) -HP = tp.TypeVar("HP", HeteroscedasticPosterior, ChainedPosterior) +P = tp.TypeVar("P", bound=Prior) +PP = tp.TypeVar("PP", bound=JointModel) +HP = tp.TypeVar("HP", bound=HeteroscedasticModel) def _psd(matrix): @@ -81,9 +80,9 @@ class AbstractVariationalFamily(_SummaryMixin, eqx.Module, tp.Generic[L]): used within variational inference. """ - posterior: AbstractPosterior + posterior: JointModel - def __init__(self, posterior: AbstractPosterior[P, L]): + def __init__(self, posterior: JointModel): self.posterior = posterior def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> GaussianDistribution: @@ -126,7 +125,7 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]): def __init__( self, - posterior: AbstractPosterior[P, L], + posterior: JointModel, inducing_inputs: tp.Union[ Int[Array, "N D"], Float[Array, "N D"], @@ -163,7 +162,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]): def __init__( self, - posterior: AbstractPosterior[P, L], + posterior: JointModel, inducing_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]], variational_mean: tp.Union[Float[Array, "N 1"], None] = None, variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None, @@ -344,7 +343,7 @@ class GraphVariationalGaussian(VariationalGaussian[L]): def __init__( self, - posterior: AbstractPosterior[P, L], + posterior: JointModel, inducing_inputs: Int[Array, "N D"], variational_mean: tp.Union[Float[Array, "N 1"], None] = None, variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None, @@ -511,7 +510,7 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]): def __init__( self, - posterior: AbstractPosterior[P, L], + posterior: JointModel, inducing_inputs: Float[Array, "N D"], natural_vector: tp.Union[Float[Array, "M 1"], None] = None, natural_matrix: tp.Union[Float[Array, "M M"], None] = None, @@ -683,7 +682,7 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]): def __init__( self, - posterior: AbstractPosterior[P, L], + posterior: JointModel, inducing_inputs: Float[Array, "N D"], expectation_vector: tp.Union[Float[Array, "M 1"], None] = None, expectation_matrix: tp.Union[Float[Array, "M M"], None] = None, @@ -848,7 +847,7 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]): def __init__( self, - posterior: AbstractPosterior[P, GL], + posterior: JointModel, inducing_inputs: Float[Array, "N D"], jitter: ScalarFloat = 1e-6, ): @@ -1006,7 +1005,7 @@ def __init__( if noise_init is not None: self.noise_variational = VariationalGaussian( - posterior=posterior.noise_posterior, + posterior=posterior.noise_model, inducing_inputs=noise_init.inducing_inputs, variational_mean=noise_init.variational_mean, variational_root_covariance=noise_init.variational_root_covariance, @@ -1025,7 +1024,7 @@ def __init__( ) self.noise_variational = VariationalGaussian( - posterior=posterior.noise_posterior, + posterior=posterior.noise_model, inducing_inputs=noise_inducing, variational_mean=variational_mean_g, variational_root_covariance=variational_root_covariance_g, From b5c9884194b3890703a1fd25f80b89ea00adcbe4 Mon Sep 17 00:00:00 2001 From: Thomas Pinder Date: Thu, 6 Aug 2026 02:14:43 +0200 Subject: [PATCH 09/81] =?UTF-8?q?refactor!:=20rename=20sweep=20=E2=80=94?= =?UTF-8?q?=20tests,=20benchmarks,=20examples=20onto=20the=20v1.0=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mechanical: ConjugatePosterior->ConjugateModel and friends, construct_posterior->construct_model, return_covariance_type->covariance, num_datapoints/noise_prior constructor ceremony deleted (~240 sites). Semantic: heteroscedastic tests build HeteroscedasticModel directly; non-conjugate tests size the latent via init_latent; docs/index.md quickstart shows condition(); regression example narrates the condition API; StateSpaceConjugatePosterior renamed StateSpaceConjugateModel. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB --- benchmarks/compile.py | 2 +- benchmarks/objectives.py | 4 +- benchmarks/state_space.py | 6 +- docs/examples/backend.py | 2 +- docs/examples/barycentres.py | 4 +- docs/examples/classification.py | 4 +- docs/examples/collapsed_vi.py | 4 +- docs/examples/constructing_new_kernels.py | 2 +- docs/examples/deep_kernels.py | 2 +- docs/examples/graph_kernels.py | 2 +- docs/examples/heteroscedastic_inference.py | 21 ++--- docs/examples/intro_to_kernels.py | 4 +- docs/examples/likelihoods_guide.py | 9 +- docs/examples/multioutput.py | 4 +- docs/examples/numpyro_integration.py | 6 +- docs/examples/oak.py | 4 +- docs/examples/oceanmodelling.py | 2 +- docs/examples/poisson.py | 4 +- docs/examples/regression.py | 51 ++++++----- docs/examples/spatial_linear_gp.py | 2 +- docs/examples/state_space_gps.py | 4 +- docs/examples/uncollapsed_vi.py | 2 +- docs/examples/yacht.py | 4 +- docs/index.md | 18 +++- gpjax/state_space/__init__.py | 4 +- gpjax/state_space/gps.py | 8 +- tests/test_citations.py | 2 +- tests/test_conditioning_equivalence.py | 2 +- tests/test_conditioning_oracle.py | 4 +- tests/test_dtype.py | 4 +- tests/test_fit.py | 26 +++--- tests/test_gps.py | 96 +++++++++++--------- tests/test_heteroscedastic.py | 86 ++++++++++-------- tests/test_integration_equinox.py | 14 +-- tests/test_integrators.py | 4 +- tests/test_jit_compatibility.py | 2 +- tests/test_likelihoods.py | 18 ++-- tests/test_likelihoods_diagonal_fast_path.py | 14 +-- tests/test_mean_functions.py | 2 +- tests/test_multioutput_integration.py | 16 ++-- tests/test_numerical_stability.py | 12 +-- tests/test_numpyro_extras.py | 2 +- tests/test_objectives.py | 20 ++-- tests/test_state_space/test_filter.py | 2 +- tests/test_state_space/test_fit.py | 26 +++--- tests/test_state_space/test_gps.py | 30 +++--- tests/test_state_space/test_imports.py | 2 +- tests/test_state_space/test_objectives.py | 16 ++-- tests/test_state_space/test_prediction.py | 36 ++++---- tests/test_state_space/test_properties.py | 2 +- tests/test_state_space/test_slow.py | 10 +- tests/test_state_space/test_smoother.py | 4 +- tests/test_summary.py | 8 +- tests/test_variational_families.py | 14 +-- 54 files changed, 340 insertions(+), 313 deletions(-) diff --git a/benchmarks/compile.py b/benchmarks/compile.py index 0601c57ba..e03590543 100644 --- a/benchmarks/compile.py +++ b/benchmarks/compile.py @@ -39,7 +39,7 @@ def setup(self): kernel = gpx.kernels.RBF() mean = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=mean) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() self.posterior = prior * likelihood self.q = VariationalGaussian( posterior=self.posterior, inducing_inputs=X[:M_INDUCING] diff --git a/benchmarks/objectives.py b/benchmarks/objectives.py index 5689eab81..f8d7ebf64 100644 --- a/benchmarks/objectives.py +++ b/benchmarks/objectives.py @@ -50,7 +50,7 @@ def _conjugate_posterior(n: int): kernel = gpx.kernels.RBF() mean = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=mean) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() return prior * likelihood, data @@ -157,7 +157,7 @@ def setup(self, n): noise_prior = gpx.gps.Prior( kernel=noise_kernel, mean_function=gpx.mean_functions.Zero() ) - likelihood = HeteroscedasticGaussian(num_datapoints=n, noise_prior=noise_prior) + likelihood = HeteroscedasticGaussian(noise_prior=noise_prior) posterior = signal_prior * likelihood Z = data.X[:M_INDUCING] self.q = HeteroscedasticVariationalFamily( diff --git a/benchmarks/state_space.py b/benchmarks/state_space.py index f51a2ec77..cdffa7cb7 100644 --- a/benchmarks/state_space.py +++ b/benchmarks/state_space.py @@ -9,7 +9,7 @@ import gpjax as gpx from gpjax.state_space import StateSpacePrior, state_space_mll -from gpjax.state_space.gps import StateSpaceConjugatePosterior +from gpjax.state_space.gps import StateSpaceConjugateModel import jax.numpy as jnp import jax.random as jr @@ -35,8 +35,8 @@ def setup(self, n): mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n, obs_stddev=0.1) - self.posterior = StateSpaceConjugatePosterior( + likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) + self.posterior = StateSpaceConjugateModel( prior=prior, likelihood=likelihood ) self.data = _temporal_dataset(n) diff --git a/docs/examples/backend.py b/docs/examples/backend.py index 89b874eb5..8ddf0ea9a 100644 --- a/docs/examples/backend.py +++ b/docs/examples/backend.py @@ -157,7 +157,7 @@ def glue(*args, **kwargs): prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Gaussian(100) +likelihood = gpx.likelihoods.Gaussian() posterior = likelihood * prior print(posterior) diff --git a/docs/examples/barycentres.py b/docs/examples/barycentres.py index 591ec61c5..453e82546 100644 --- a/docs/examples/barycentres.py +++ b/docs/examples/barycentres.py @@ -132,7 +132,7 @@ def glue(*args, **kwargs): # ## Dataset # # We'll simulate five datasets and develop a Gaussian process -# [posterior](#gpjax.gps.ConjugatePosterior) before +# [posterior](#gpjax.gps.ConjugateModel) before # identifying the Gaussian process barycentre at a set of test points. Each dataset # will be a sine function with a different vertical shift, periodicity, and quantity # of noise. @@ -184,7 +184,7 @@ def fit_gp(x: jax.Array, y: jax.Array) -> gpx.distributions.GaussianDistribution y = y.reshape(-1, 1) D = gpx.Dataset(X=x, y=y) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() posterior = ( gpx.gps.Prior( mean_function=gpx.mean_functions.Constant(), kernel=gpx.kernels.RBF() diff --git a/docs/examples/classification.py b/docs/examples/classification.py index 74edeabdb..fe7d071b6 100644 --- a/docs/examples/classification.py +++ b/docs/examples/classification.py @@ -105,7 +105,7 @@ kernel = gpx.kernels.RBF() meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) +likelihood = gpx.likelihoods.Bernoulli() # %% [markdown] # We construct the posterior through the product of our prior and likelihood. @@ -144,7 +144,7 @@ ) # %% [markdown] -# From which we can [make predictions](#gpjax.gps.NonConjugatePosterior.predict) at +# From which we can [make predictions](#gpjax.gps.NonConjugateModel.predict) at # novel inputs, as illustrated in {numref}`fig-classification-map-predictive`. # %% mystnb={"figure": {"caption": "The MAP predictive mean and its one-sigma band over the binary observations.", "name": "fig-classification-map-predictive"}} diff --git a/docs/examples/collapsed_vi.py b/docs/examples/collapsed_vi.py index 5e8b92b05..a3761dc19 100644 --- a/docs/examples/collapsed_vi.py +++ b/docs/examples/collapsed_vi.py @@ -121,7 +121,7 @@ # %% meanf = gpx.mean_functions.Constant() kernel = gpx.kernels.RBF() # 1-dimensional inputs -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) posterior = prior * likelihood @@ -242,7 +242,7 @@ # %% full_rank_model = gpx.gps.Prior( mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF() -) * gpx.likelihoods.Gaussian(num_datapoints=D.n) +) * gpx.likelihoods.Gaussian() nmll = jit(lambda: -gpx.objectives.conjugate_mll(full_rank_model, D)) # %timeit nmll().block_until_ready() diff --git a/docs/examples/constructing_new_kernels.py b/docs/examples/constructing_new_kernels.py index 88da1e2f9..a2d3b3e2f 100644 --- a/docs/examples/constructing_new_kernels.py +++ b/docs/examples/constructing_new_kernels.py @@ -287,7 +287,7 @@ def __call__( # Define polar Gaussian process PKern = Polar() meanf = gpx.mean_functions.Zero() -likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) +likelihood = gpx.likelihoods.Gaussian() circular_posterior = gpx.gps.Prior(mean_function=meanf, kernel=PKern) * likelihood # Optimise GP's marginal log-likelihood using BFGS diff --git a/docs/examples/deep_kernels.py b/docs/examples/deep_kernels.py index 2a382e44c..6003285cc 100644 --- a/docs/examples/deep_kernels.py +++ b/docs/examples/deep_kernels.py @@ -196,7 +196,7 @@ def __call__(self, x: jax.Array) -> jax.Array: kernel = DeepKernelFunction(network=forward_linear, base_kernel=base_kernel) meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood # %% [markdown] # ### Optimisation diff --git a/docs/examples/graph_kernels.py b/docs/examples/graph_kernels.py index 50bdf60ee..6811745b3 100644 --- a/docs/examples/graph_kernels.py +++ b/docs/examples/graph_kernels.py @@ -174,7 +174,7 @@ def glue(*args, **kwargs): # [`fit_scipy`](#gpjax.fit.fit_scipy). # %% -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() kernel = gpx.kernels.GraphKernel(laplacian=L) prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=kernel) posterior = prior * likelihood diff --git a/docs/examples/heteroscedastic_inference.py b/docs/examples/heteroscedastic_inference.py index 3ae5079f7..3673cb64f 100644 --- a/docs/examples/heteroscedastic_inference.py +++ b/docs/examples/heteroscedastic_inference.py @@ -163,8 +163,9 @@ # \mathcal{N}\!\big(y_i \mid f(x_i), \exp(g(x_i))\big), # $$ (eq-heteroscedastic-likelihood) # -# to form the posterior target that we shall approximate variationally. The product -# syntax `signal_prior * likelihood` used below constructs this augmented GP model. +# to form the posterior target that we shall approximate variationally. Because the +# joint model holds *two* priors — one per latent process — it is constructed +# directly as a `HeteroscedasticModel` rather than via the two-operand product. # %% # Signal and noise priors. @@ -176,12 +177,10 @@ mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF(), ) -likelihood = HeteroscedasticGaussian( - num_datapoints=train.n, - noise_prior=noise_prior, - noise_transform=LogNormalTransform(), +likelihood = HeteroscedasticGaussian(noise_transform=LogNormalTransform()) +posterior = gpx.gps.HeteroscedasticModel( + prior=signal_prior, likelihood=likelihood, noise_prior=noise_prior ) -posterior = signal_prior * likelihood # Variational family over both processes. z = jnp.linspace(-3.2, 3.2, 25)[:, None] @@ -330,12 +329,10 @@ mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF(), ) -likelihood_adv = HeteroscedasticGaussian( - num_datapoints=data_adv.n, - noise_prior=noise_prior_adv, - noise_transform=SoftplusTransform(), +likelihood_adv = HeteroscedasticGaussian(noise_transform=SoftplusTransform()) +posterior_adv = gpx.gps.HeteroscedasticModel( + prior=mean_prior, likelihood=likelihood_adv, noise_prior=noise_prior_adv ) -posterior_adv = mean_prior * likelihood_adv # %% # Configure variational family diff --git a/docs/examples/intro_to_kernels.py b/docs/examples/intro_to_kernels.py index 945d0ed13..2bca91465 100644 --- a/docs/examples/intro_to_kernels.py +++ b/docs/examples/intro_to_kernels.py @@ -282,7 +282,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]: # noqa: F821 prior = gpx.gps.Prior(mean_function=mean, kernel=kernel) likelihood = gpx.likelihoods.Gaussian( - num_datapoints=D.n, obs_stddev=jnp.array(1e-3) + obs_stddev=jnp.array(1e-3) ) # Our function is noise-free, so we set the observation noise's standard deviation to a very small value no_opt_posterior = prior * likelihood @@ -561,7 +561,7 @@ def plot_ribbon(ax, x, dist, color): final_kernel = gpx.kernels.SumKernel(kernels=[rbf_kernel, sum_kernel]) prior = gpx.gps.Prior(mean_function=mean, kernel=final_kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood diff --git a/docs/examples/likelihoods_guide.py b/docs/examples/likelihoods_guide.py index c2f9cc9a7..8dc9ab9e8 100644 --- a/docs/examples/likelihoods_guide.py +++ b/docs/examples/likelihoods_guide.py @@ -119,7 +119,7 @@ # argument. # %% -gpx.likelihoods.Gaussian(num_datapoints=D.n) +gpx.likelihoods.Gaussian() # %% [markdown] # ### Likelihood parameters @@ -134,7 +134,7 @@ # this as follows: # %% -gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=0.5) +gpx.likelihoods.Gaussian(obs_stddev=0.5) # %% [markdown] # @@ -157,7 +157,7 @@ meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=0.1) +likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) posterior = prior * likelihood @@ -187,7 +187,7 @@ # Similarly, for a Bernoulli likelihood function, the samples of $y$ would be binary. # %% mystnb={"figure": {"caption": "The same latent draws passed through a Bernoulli likelihood, whose predictive samples are constrained to be binary.", "name": "fig-likelihoods-guide-bernoulli-samples"}} -likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) +likelihood = gpx.likelihoods.Bernoulli() fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(9, 2)) @@ -292,7 +292,6 @@ def q_moments(x): # %% lquad = gpx.likelihoods.Gaussian( - num_datapoints=D.n, obs_stddev=jnp.array([0.1]), integrator=gpx.integrators.GHQuadratureIntegrator(num_points=20), ) diff --git a/docs/examples/multioutput.py b/docs/examples/multioutput.py index 01260704f..2ece5d6ff 100644 --- a/docs/examples/multioutput.py +++ b/docs/examples/multioutput.py @@ -155,7 +155,7 @@ meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) likelihood = gpx.likelihoods.MultiOutputGaussian( - num_datapoints=N, num_outputs=P, obs_stddev=1.0 + num_outputs=P, obs_stddev=1.0 ) posterior = prior * likelihood @@ -442,7 +442,7 @@ meanf_lcm = gpx.mean_functions.Zero() prior_lcm = gpx.gps.Prior(mean_function=meanf_lcm, kernel=lcm_kernel) likelihood_lcm = gpx.likelihoods.MultiOutputGaussian( - num_datapoints=N_lcm, num_outputs=P_lcm, obs_stddev=1.0 + num_outputs=P_lcm, obs_stddev=1.0 ) posterior_lcm = prior_lcm * likelihood_lcm diff --git a/docs/examples/numpyro_integration.py b/docs/examples/numpyro_integration.py index a74e73014..6590bdb9c 100644 --- a/docs/examples/numpyro_integration.py +++ b/docs/examples/numpyro_integration.py @@ -131,7 +131,7 @@ # # We define a NumPyro model that samples all parameters directly using # ``numpyro.sample``, builds the GPJax -# [`ConjugatePosterior`](#gpjax.gps.ConjugatePosterior) from those samples, and +# [`ConjugateModel`](#gpjax.gps.ConjugateModel) from those samples, and # scores it with the conjugate marginal log-likelihood # ([`conjugate_mll`](#gpjax.objectives.conjugate_mll)) via ``numpyro.factor``. # No special registration step is needed -- GPJax constructors accept raw @@ -157,7 +157,7 @@ def model(X, Y, X_new=None): meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=N, obs_stddev=obs_noise) + likelihood = gpx.likelihoods.Gaussian(obs_stddev=obs_noise) posterior = prior * likelihood D_resid = gpx.Dataset(X=X, y=residuals) @@ -199,7 +199,7 @@ def model(X, Y, X_new=None): ) example_posterior = gpx.gps.Prior( mean_function=gpx.mean_functions.Constant(), kernel=example_kernel -) * gpx.likelihoods.Gaussian(num_datapoints=N, obs_stddev=1.0) +) * gpx.likelihoods.Gaussian(obs_stddev=1.0) parameter_priors = { "prior.kernel.kernels[0].lengthscale": dist.LogNormal(0.0, 1.0), diff --git a/docs/examples/oak.py b/docs/examples/oak.py index 1bbb3c6f0..2f610b191 100644 --- a/docs/examples/oak.py +++ b/docs/examples/oak.py @@ -273,7 +273,7 @@ def apply_flows(X_original: np.ndarray) -> jnp.ndarray: mean_function = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=mean_function, kernel=oak_kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=num_train) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood # %% @@ -286,7 +286,7 @@ def apply_flows(X_original: np.ndarray) -> jnp.ndarray: ) latent_dist = opt_posterior.predict( - X_test, train_data=train_data, return_covariance_type="diagonal" + X_test, train_data=train_data, covariance="diagonal" ) predictive_dist = opt_posterior.likelihood(latent_dist) predictive_mean = predictive_dist.mean diff --git a/docs/examples/oceanmodelling.py b/docs/examples/oceanmodelling.py index 14d5dbfa7..a4203fd49 100644 --- a/docs/examples/oceanmodelling.py +++ b/docs/examples/oceanmodelling.py @@ -334,7 +334,7 @@ def __call__( def initialise_gp(kernel, mean, dataset): prior = gpx.gps.Prior(mean_function=mean, kernel=kernel) likelihood = gpx.likelihoods.Gaussian( - num_datapoints=dataset.n, obs_stddev=jnp.array([1.0e-3], dtype=jnp.float64) + obs_stddev=jnp.array([1.0e-3], dtype=jnp.float64) ) posterior = prior * likelihood return posterior diff --git a/docs/examples/poisson.py b/docs/examples/poisson.py index 7c0700dc5..51a3ebab1 100644 --- a/docs/examples/poisson.py +++ b/docs/examples/poisson.py @@ -135,10 +135,10 @@ kernel = gpx.kernels.RBF() meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Poisson(num_datapoints=D.n) +likelihood = gpx.likelihoods.Poisson() # %% [markdown] -# We construct the [posterior](#gpjax.gps.NonConjugatePosterior) through the product of our +# We construct the [posterior](#gpjax.gps.NonConjugateModel) through the product of our # prior and likelihood. # %% diff --git a/docs/examples/regression.py b/docs/examples/regression.py index a827d4a03..6ca2c4ccb 100644 --- a/docs/examples/regression.py +++ b/docs/examples/regression.py @@ -136,11 +136,11 @@ # the evaluation of the GP's mean and covariance. # # Since we want to sample from the full posterior, we need to calculate the full covariance matrix. -# We can enforce this by including the `return_covariance_type = "dense"` attribute when predicting. +# We can enforce this by including the `covariance = "dense"` attribute when predicting. # Note this is what will be defaulted if left blank. # %% mystnb={"figure": {"caption": "Twenty function samples drawn from the zero-mean RBF prior, shown alongside the prior mean and variance band.", "name": "fig-regression-prior-samples"}} -prior_dist = prior.predict(xtest, return_covariance_type="dense") +prior_dist = prior.predict(xtest, covariance="dense") prior_mean = prior_dist.mean prior_std = prior_dist.variance @@ -181,19 +181,23 @@ # and what each one assumes about the observations. # %% -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() # %% [markdown] -# The posterior is proportional to the prior multiplied by the likelihood, written as +# The prior and likelihood together define the joint model over the latent +# function and the observations, # # $$ -# p(f(\cdot) | \mathcal{D}) \propto p(f(\cdot)) * p(\mathcal{D} | f(\cdot)). -# $$ (eq-regression-posterior) +# p(f(\cdot), \mathcal{D}) = p(\mathcal{D} | f(\cdot))\, p(f(\cdot)). +# $$ (eq-regression-joint) # -# Mimicking this construct, the posterior is established in GPJax through the `*` operator. +# Mimicking this equation, the joint model is established in GPJax through the +# `*` operator. Conditioning it on the data — `model.condition(D)`, or the +# operator form `model | D` that reads as $p(f \mid \mathcal{D})$ — yields the +# posterior process. # %% -posterior = prior * likelihood +model = prior * likelihood # %% [markdown] #