gp

gp#

Classes#

`ConditionedStates`	An object to hold the conditioned means and variances
`PredictedStates`	An object to hold the full predictive states
`GaussianProcess`	An interface for designing a Gaussian Process regression model.
`ConditionResult`	The result of conditioning a `GaussianProcess` on data

Functions#

`assign_unique_kernel_names`(...)	Return a new kernel where duplicated leaf kernel names are made unique by appending _1, _2, etc.
`assign_num_insts`(→ smolgp.kernels.StateSpaceModel)	Return a new kernel where every IntegratedStateSpaceModel component has
`_default_jitter`(→ tinygp.helpers.JAXArray)	Default to adding some amount of jitter to the diagonal, just in case,

Module Contents#

smolgp.gp.assign_unique_kernel_names(kernel: smolgp.kernels.StateSpaceModel) → smolgp.kernels.StateSpaceModel[source]#

Return a new kernel where duplicated leaf kernel names are made unique by appending _1, _2, etc.

For example, if the original kernel has three components named “SHO”, “Matern”, and “Matern”, they will be renamed to “SHO”, “Matern_1”, and “Matern_2”. This is useful for ensuring that the component kernels can be uniquely identified when making predictions at test points or when extracting component contributions.

smolgp.gp.assign_num_insts(kernel: smolgp.kernels.StateSpaceModel, num_insts: int) → smolgp.kernels.StateSpaceModel[source]#

Return a new kernel where every IntegratedStateSpaceModel component has num_insts set to the given value, reinitializing any that don’t already match.

Traverses Sum, Product, and Wrapper nodes to find IntegratedStateSpaceModel leaves anywhere in the kernel tree.

class smolgp.gp.ConditionedStates(X, y: tinygp.helpers.JAXArray, t_states: tinygp.helpers.JAXArray, instid: tinygp.helpers.JAXArray, obsid: tinygp.helpers.JAXArray, stateid: tinygp.helpers.JAXArray, m_pred: tinygp.helpers.JAXArray, P_pred: tinygp.helpers.JAXArray, m_filt: tinygp.helpers.JAXArray, P_filt: tinygp.helpers.JAXArray, m_smooth: tinygp.helpers.JAXArray, P_smooth: tinygp.helpers.JAXArray)[source]#

Bases: equinox.Module

An object to hold the conditioned means and variances

X: len(N) data coordinates y: len(N) observed data this was conditioned on t_states: len(K) time coordinates of all states instid : len(N) instrument ID for each measurement obsid : len(K) observation IDs corresponding to the measurement at each state stateid : len(K) state IDs corresponding to each state (0 for exposure-start, 1 for exposure-end) predicted_mean/var : len(K) Kalman predicted state filtered_mean/var : len(K) Kalman filtered state smoothed_mean/var : len(K) RTS smoothed state

X: tinygp.helpers.JAXArray#

y: tinygp.helpers.JAXArray#

t_states: tinygp.helpers.JAXArray#

instid: tinygp.helpers.JAXArray#

obsid: tinygp.helpers.JAXArray#

stateid: tinygp.helpers.JAXArray#

predicted_mean: tinygp.helpers.JAXArray#

filtered_mean: tinygp.helpers.JAXArray#

smoothed_mean: tinygp.helpers.JAXArray#

predicted_cov: tinygp.helpers.JAXArray#

filtered_cov: tinygp.helpers.JAXArray#

smoothed_cov: tinygp.helpers.JAXArray#

__call__()[source]#

project_at_data(observation_model)[source]#: Project the states with measurements (e.g. exposure-ends) and sort back into original order as the data

class smolgp.gp.PredictedStates(t_states: tinygp.helpers.JAXArray, m: tinygp.helpers.JAXArray, P: tinygp.helpers.JAXArray, kernel: smolgp.kernels.StateSpaceModel)[source]#

Bases: equinox.Module

An object to hold the full predictive states

t_states: time coordinates at each state mean : predictive mean vector for each state cov : predictive covariance for each state

t_states: tinygp.helpers.JAXArray#

mean: tinygp.helpers.JAXArray#

cov: tinygp.helpers.JAXArray#

kernel: smolgp.kernels.StateSpaceModel#

project_mean(observation_model) → tinygp.helpers.JAXArray[source]#: The projected mean at self.t_states given an observation model.

project_variance(observation_model) → tinygp.helpers.JAXArray[source]#: The projected variance at self.t_states given an observation model.

property loc: tinygp.helpers.JAXArray#: The overall mean at the predicted states

property variance: tinygp.helpers.JAXArray#: The overall variance at the predicted states

get_component(component: str | list[str], return_var: bool = False) → PredictedStates[source]#: Extract the predicted states corresponding to a component kernel

get_all_components(return_var: bool = False) → dict[str, Any][source]#: Extract the predicted mean/variance corresponding to each component kernel

class smolgp.gp.GaussianProcess(kernel: GaussianProcess.__init__.kernels, X: tinygp.helpers.JAXArray, *, noise: tinygp.helpers.JAXArray | float | None = None, mean: tinygp.means.MeanBase | collections.abc.Callable[[tinygp.helpers.JAXArray], tinygp.helpers.JAXArray] | tinygp.helpers.JAXArray | None = None, solver: Any | None = None, mean_value: tinygp.helpers.JAXArray | None = None, variance_value: tinygp.helpers.JAXArray | None = None, covariance_value: Any | None = None, states: tinygp.helpers.JAXArray | None = None, use_unique_names: bool = True, **solver_kwargs: Any)[source]#

Bases: equinox.Module

An interface for designing a Gaussian Process regression model.

Parameters:

kernel (Kernel) – The kernel function.
X (JAXArray) – The input coordinates — any PyTree compatible with kernel whose leading dimension has size N_data. For integrated kernels, pass (t, texp) where t is the array of exposure midpoints and texp is the array of exposure durations.
noise (JAXArray | float, optional) – Observation noise covariance matrices with shape (N, D, D), where N is the number of data points and D is the observation dimension (usually 1). Each slice noise[k] is the \(D \times D\) noise covariance for the k-th observation. Two shorthands are accepted and broadcast up to (N, 1, 1): a 1-D array of shape (N,) is interpreted as scalar per-observation variances, and a scalar as a single homoscedastic variance shared by every observation (i.e. noise=0.01 is equivalent to noise=jnp.full(N, 0.01)). Defaults to \(\sqrt{\varepsilon_{\mathrm{machine}}} \cdot I\) for all observations.
mean (Callable, optional) – A callable or constant mean function evaluated as mean(X).
solver – Solver class for filtering and smoothing. If None (default), selected automatically based on the kernel type.

num_data: int#

dtype: jax.numpy.dtype#

kernel: tinygp.kernels.Kernel#

X: tinygp.helpers.JAXArray#

mean_function: tinygp.means.MeanBase#

mean: tinygp.helpers.JAXArray#

var: tinygp.helpers.JAXArray | None#

noise: tinygp.helpers.JAXArray#

solver: smolgp.solvers.StateSpaceSolver#

states: ConditionedStates#

property loc: tinygp.helpers.JAXArray#: If conditioned, this will be the mean at the data points Otherwise, it is just the prior mean.

property variance: tinygp.helpers.JAXArray#

The marginal variance at each coordinate, i.e. \(\mathrm{diag}(\texttt{covariance})\).

If conditioned, this is the posterior variance at this GP’s coordinates. Otherwise it is the prior variance plus the observation noise.

Computed directly rather than constructing the full covariance matrix and taking its diagonal (more expensive).

_noise_diagonal() → tinygp.helpers.JAXArray[source]#: Per-observation noise variance, as a length-N vector.

property covariance: tinygp.helpers.JAXArray#

The full covariance matrix at this GP’s coordinates. For just the diagonal, use variance.

Warning

This materializes an \(N \times N\) matrix, which is exactly the cost smolgp exists to avoid: \(O(N^2)\) memory and, for the conditioned case, \(O(N^2 d^3)\) time. It is provided for small problems, for validation against dense references, and for cases genuinely needing the joint distribution. Prefer variance (\(O(N)\)) when the marginals suffice.

The unconditioned case is just the prior covariance plus measurement noise \(k(X, X) + \Sigma_n\)

The conditioned case returns the posterior covariance at the data using the smoother cross-covariance identity (Eq. 12.55 of Särkkä & Solin 2019). The diagonal is the observed smoothed variances,

\[\Sigma_{k,k} = H_k P_k^s H_k^T,\]

and the lower triangle of the symmetric covariance matrix is

\[\Sigma_{i,j} = H_i \left(\prod_{m=i}^{j-1} G_m\right) P^s_j H_j^T, \quad i < j,\]

with \(G_k\) the RTS smoothing gains. The gains are recomputed on demand from the cached filtered/predicted covariances in \(O(N d^3)\), negligible compared to forming the matrix itself.

Raises:: NotImplementedError – for a GP returned by condition(y, X_test=...), whose coordinates are the test points rather than the training states. The cross-covariance between two arbitrary test points is not produced by the current predict machinery.

_posterior_covariance() → tinygp.helpers.JAXArray[source]#

The joint posterior covariance at the data points, via the RTS smoother cross-covariance recursion (see covariance).

See Eq. 12.55 of Särkkä & Solin 2019. We expand this definition to include integrated SSMs by first building the (K, K, d, d) state-space cross-covariance, then selecting and projecting the N data-carrying states into observation space in data order.

log_probability(y: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Compute the log probability of this Gaussian Process, given the observed data y.

Parameters:: y (JAXArray) – The observed data. This should have the shape (N_data, D), where N_data is the number of data coordinates in X and D is the observation dimension.
Returns:: The marginal log probability of the GP, evaluated at y.

property state_coords: smolgp.solvers.state_coords.StateCoords#

The StateCoords for this GP’s states, shared by condition() and sample().

If already conditioned, rebuilds it from the cached self.states fields. Otherwise reuses self.solver.state_coords, which every solver builds once in its own __init__ (the instantaneous solvers via StateCoords.instantaneous()).

condition(y: tinygp.helpers.JAXArray, X_test: tinygp.helpers.JAXArray | None = None, *, include_mean: bool = True, kernel: tinygp.kernels.Kernel | None = None) → ConditionResult[source]#

Condition the model on observed data

Parameters:

y (JAXArray) – The observed data. This should have the shape (N_data,), where N_data was the zeroth axis of the X data provided when instantiating this object.
X_test (JAXArray, optional) – The coordinates where the prediction should be evaluated. This should have a data type compatible with the X data provided when instantiating this object. If it is not provided, X will be used by default, so the predictions will be made.
include_mean (bool, optional) – If True (default), the predicted values will include the mean function evaluated at X_test.
kernel (Kernel, optional) – A kernel to optionally specify the component kernel to be used for predicting after conditioning. See Multicomponent Kernels for an example.

Returns:

A named tuple where the first element log_probability is the log marginal probability of the model, and the second element gp is the GaussianProcess object describing the conditional distribution evaluated at X_test.

predict(X_test: tinygp.helpers.JAXArray | None = None, y: tinygp.helpers.JAXArray | None = None, *, return_full_state: bool = False, kernel: int | None = None, return_var: bool = False, observation_model: Any | None = None) → tinygp.helpers.JAXArray | tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]#

Predict the GP model at new test points conditioned on observed data

Parameters:

X_test (JAXArray, optional) – The coordinates where the prediction should be evaluated. This should have a data type compatible with the X data provided when instantiating this object. If it is not provided, X will be used by default, so the predictions will be made at the data coordinates.
y (JAXArray, optional) – The observed data. Only needs to be given if the GP has not yet been conditioned. Once conditioned, the data, if needed, is recalled automatically from self.states.y. This should have the shape (N_data,), where N_data was the zeroth axis of the X data provided when instantiating this object.
include_mean (bool, optional) – If True (default), the predicted values will include the mean function evaluated at X_test.
return_var (bool, optional) – If True (default), the variance of the predicted values at X_test will be returned.
return_cov (bool, optional) – If True, the covariance of the predicted values at X_test will be returned. If return_var is True, this flag will be ignored.
observation_model (Any, optional) – optionally provide a function of X_test to define the output observation model. Default will use that of the kernel.
return_full_state (bool, optional) – If True, return the full predicted state mean and covariance, rather than projecting to observation space. Default is False, i.e. the result is projected through kernel.observation_model.
kernel (int, optional) – If specified, the index of the kernel in a multi-component model (for example, a sum or product of kernels) to extract and project (if return_full_state is False) the prediction for.

Returns:

The mean of the predictive model evaluated at X_test, with shape (N_test,) where N_test is the zeroth dimension of X_test. If either return_var or return_cov is True, the variance or covariance of the predicted process will also be returned with shape (N_test,) or (N_test, N_test) respectively.

sample(key: jax.random.KeyArray, shape: collections.abc.Sequence[int] | None = None, X_test: tinygp.helpers.JAXArray | None = None, num_test_insts: int | None = None) → tinygp.helpers.JAXArray[source]#

Generate samples from the process.

If this GaussianProcess has not been conditioned, samples are drawn from the prior. If it was returned by condition(), samples are drawn from the posterior using Matheron’s rule (see _sample() and the references therein).

By default (X_test=None), samples are drawn at the training coordinates. Passing X_test draws samples at any, possibly out-of-sample, coordinates. If the GP is conditioned on exposure-integrated data (delta>0), then X_test needs to be either (t, delta) or (t, delta, instid).

For exposure-integrated sample points, instid says only which instrument project the result as (if the observation model depends on the instrument). As such, overlapping sample points are allowed to have the same instid. Internally, where overlap would not be allowed, we use a separate “probe group” ID to auto-assign each exposure to a non-overlapping integral accumulator. If samples are requested with no instid given, the default is to project the result as instrument 0.

Parameters:

key – A JAX random number key array.
shape (tuple, optional) – The number/shape of independent draws to generate. Each single draw is a complete joint sample across every coordinate in X_test (or the training coordinates), with the correct correlations between them. Defaults to a single draw.
X_test (JAXArray, optional) – New coordinates to sample at, instead of the training coordinates. For integrated samples with an integrated kernel, this should be either (t, delta) or (t, delta, instid) where t is the array of exposure midpoints, delta is the array of exposure durations, and instid is the array of instrument IDs for each exposure.
num_test_insts (int, optional) – Number of probe groups to allocate for exposure-integrated test points. Derived from the coordinates when omitted, which reads their values and so cannot run under jit. Pass it explicitly – it is 1 for instantaneous or non-overlapping test points – to keep this call jittable. Must be at least count_min_instids(), or overlapping windows will share a probe and the draw will be wrong.

Returns:

The sampled realizations from the process with shape (N_samples,) + shape (or plain (N_samples,) if shape is not given) where N_samples is the zeroth dimension of X_test or of self.X if X_test is not given). E.g. shape=(M,) returns M independent draws with shape (N_data, M).

numpyro_dist(**kwargs: Any) → tinygp.numpyro_support.TinyDistribution[source]#: Get the numpyro MultivariateNormal distribution for this process

_sample(key: jax.random.KeyArray, shape: collections.abc.Sequence[int] | None, X_test: tinygp.helpers.JAXArray | None, num_test_insts: int, instid_proj: tinygp.helpers.JAXArray | None = None) → tinygp.helpers.JAXArray[source]#

Draw samples via the residual/Matheron’s-rule/Durbin-Koopman method:

See the following references for details: - Doucet, “A Note on Efficient Conditional Simulation of Gaussian

Distributions” (https://www.stats.ox.ac.uk/~doucet/doucet_simulationconditionalgaussian.pdf).

Durbin & Koopman (2002), “A simple and efficient simulation smoother for state space time series analysis,” Biometrika 89(3):603-616.
Wilson et al. (2020/2021), “Efficiently sampling functions from Gaussian process posteriors” / “Pathwise Conditioning of Gaussian Processes.”

A prior sample is a plain forward SDE simulation projected to observation space. A posterior sample corrects that draw toward the data via Matheron’s rule:

x_prior_traj ~ forward SDE simulation of the full latent trajectory prior_obs = x_prior_traj projected to observation space noise_sample ~ N(0, self.noise) residual = y_obs - (prior_obs + noise_sample) posterior_sample = prior_obs + project(condition(residual).smoothed_mean)

That is, an ordinary conditioning pass (the same Kalman filter + RTS smoother condition() uses) run on the residual instead of on the data.

Implementation notes: - The covariance/gain recursion does not depend on the residual’s

values, so StateSpaceSolver/IntegratedStateSpaceSolver compute it once and share it across all shape samples (solver.condition_batched_mean). The parallel solvers instead call solver.condition once per sample: their associative-scan operator fuses the covariance and mean paths at every node, so the same split is possible but not yet implemented (TODO: if needed).

With X_test, the prior trajectory covers the training states and the test points in one joint pass (merge_test_coords()), and the residual correction is propagated out to those points by the ordinary solver.predict. Not yet batched as above (TODO: if needed).

get_component_mean(component: list | str, return_var: bool = False, **kwargs) → Any[source]#

Get the predictive mean (and variance) of a particular (or sum of) component kernel in a multi-component model evaluated at self.X

Parameters:

X (JAXArray, optional) – The coordinates where the prediction should be evaluated. This should have a data type compatible with the X data provided when instantiating this object.
component (list | str) – The name(s) of the component kernel(s) to extract the mean for. If a list of names is provided, the joint mean and variance for that collection of kernels will be returned.
return_var (bool, optional) – If True, also return the variances of each component. Default is False.

Returns:

component_mean (JAXArray) If return_var is True:

component_mean (JAXArray) component_var (JAXArray)

Return type:

If return_var is False

get_all_component_means(return_var: bool = False, **kwargs) → Any[source]#

Get the predictive mean (and optionally variance) of each component kernels individually, evaluated at self.X

Parameters:: return_var (bool, optional) – If True, also return the variances of each component. Default is False.
Returns:: If return_var is False, a list of JAX arrays containing the means of each component kernel evaluated at the data points. If return_var is True, a tuple where the first element is the list of means as before, and the second element is a list of JAX arrays containing the variances of each component kernel evaluated at the data points.

class smolgp.gp.ConditionResult[source]#

Bases: NamedTuple

The result of conditioning a GaussianProcess on data

This has two entries, log_probability and gp, that are described below.

log_probability: tinygp.helpers.JAXArray#

The log probability of the conditioned model

In other words, this is the marginal likelihood for the kernel parameters, given the observed data, or the multivariate normal log probability evaluated at the given data.

gp: GaussianProcess#

A GaussianProcess describing the conditional distribution

This will have a mean and covariance conditioned on the observed data, but it is otherwise a fully functional GP that can sample from or condition further (although that’s probably not going to be very efficient).

smolgp.gp._default_jitter(reference: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#: Default to adding some amount of jitter to the diagonal, just in case, we use sqrt(eps) for the dtype of the mean function because that seems to give sensible results in general.