smolgp

smolgp#

smolgp is designed to be a drop-in extension of the tinygp library for building Gaussian Process (GP) models in Python. As such, it is also built on top of jax. The driving design philosophy is to match the API of tinygp as closely as possible. With only a few exceptions, any existing code you have that uses tinygp should work with smolgp by simply by finding-and-replacing tiny with smol.

smolgp uses the state space representations of Gaussian Processes to implement linear-time (or up to logN with parallelization on GPU) solvers for GP regression and forecasting. It also implements “integrated” kernels that can model time-averaged measurements, such as those from long-exposure astronomical observations, which can also be solved in linear time and are also compatible with the parallel methods.

The primary way that you interact with smolgp is to construct “kernel” functions using the building blocks provided in the kernels subpackage (see smolgp.kernels), and then passing that to a GaussianProcess object to do all the computations. Check out the Tutorials for a more complete introduction.

Submodules#

Attributes#

__version__

Classes#

GaussianProcess

An interface for designing a Gaussian Process regression model.

Package Contents#

class smolgp.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.

smolgp.__version__#