StateSpaceConjugateModel#

class gpjax.state_space.StateSpaceConjugateModel(prior, likelihood)[source]#

Bases: ConjugateModel

Joint model for a state-space (Markovian) GP: conditioning is Kalman.

Conditioning returns a StateSpacePosterior, whose queries run the square-root Kalman recursions in \(O(N d^3)\) time. The inherited ConjugateModel conditioning is deliberately overridden: its dense \(O(N^3)\) Cholesky is exactly the cost this model exists to avoid.

v1 prediction surface, all of it sugar over condition:

  • condition / __or__ : the conditioned process

  • predict / __call__ : condition(D)(t), the smoothed predictive

  • predict_filter : condition(D).filtered(t), the causal (filter-only) predictive

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 dense gpjax.gps.ConjugateModel predictive.

Example

>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> posterior.__class__.__name__
'StateSpaceConjugateModel'
Parameters:
condition(train_data, *, observation_mask=None)[source]#

Condition on data through the Kalman recursions.

Parameters:

  • train_data (Dataset) – The observations to condition on.

  • observation_mask (Bool[jaxlib._jax.Array, 'N'] | Bool[ndarray, 'N'] | None) – Optional boolean mask over the training points; False entries are not conditioned on. None conditions on every point.

Returns:

The conditioned process. Exposes the smoothed

predictive (via __call__), the causal predictive (via filtered), and log_marginal_likelihood.

Return type:

StateSpacePosterior

predict(test_inputs, train_data, *, covariance='diagonal', observation_mask=None)[source]#

Sugar for the smoothed predictive: condition(D)(t).

When making repeated predictions, condition once and reuse the returned posterior.

Parameters:

  • test_inputs (Num[jaxlib._jax.Array, 'M 1'] | Num[ndarray, 'M 1']) – Test timestamps of shape (M, 1).

  • train_data (Dataset) – The observations to condition on.

  • covariance (Literal['dense', 'diagonal']) – Must be "diagonal"; the v1 state-space predictive has no dense joint form.

  • observation_mask (Bool[jaxlib._jax.Array, 'N'] | Bool[ndarray, 'N'] | None) – Optional boolean mask over the training points.

Returns:

The smoothed predictive.

Return type:

GaussianDistribution

predict_filter(test_inputs, train_data, *, covariance='diagonal', observation_mask=None)[source]#

Sugar for the causal predictive: condition(D).filtered(t).

Each test point conditions only on training observations at timestamps less than or equal to its own, rather than on the whole training set.

Parameters:

  • test_inputs (Num[jaxlib._jax.Array, 'M 1'] | Num[ndarray, 'M 1']) – Test timestamps of shape (M, 1).

  • train_data (Dataset) – The observations to condition on.

  • covariance (Literal['dense', 'diagonal']) – Must be "diagonal"; the v1 state-space predictive has no dense joint form.

  • observation_mask (Bool[jaxlib._jax.Array, 'N'] | Bool[ndarray, 'N'] | None) – Optional boolean mask over the training points.

Returns:

The filtered predictive.

Return type:

GaussianDistribution

sample_approx(num_samples, train_data, key, num_features=100)[source]#

Not available for state-space models.

The inherited pathwise sampler is built on the dense conditioned process, which state-space models never form. Raising is deliberate: silently falling back would reintroduce the \(O(N^3)\) cost this model exists to avoid.

Raises:

NotImplementedError – Always.