smolgp 0.1.0

Gaussian Process State Space Models in Python/JAX

smolgp
State Space Models for O(Linear/Log) Gaussian Processes

smolgp is a Python/JAX standalone extension of the tinygp package that implements

1. A Kalman filter and RTS smoother as a StateSpaceSolver compatible with tinygp-like GP kernels (see smolgp.kernels)
2. An IntegratedStateSpaceSolver that can handle integrated (and possibly overlapping) measurements from multiple instruments (see Rubenzahl and Hattori et al. submitted)
3. Parallelized versions of 1 (ParallelStateSpaceSolver, see Särkkä and García-Fernández 2020) and 2 (ParallelIntegratedStateSpaceSolver, see Rubenzahl and Hattori et al. submitted) using jax.lax.associative_scan
   • see also Yaghoobi and Särkkä 2024 and its implementation
4. Approximations of popular GP kernels that lack quasiseparability (e.g., ExpSineSquared, Quasiperiodic) but can utilize the O(N) states space solvers.

This package (and its documentation) is still under heavy active development, with tutorials coming soon. Please raise issues here and/or reach out to Ryan Rubenzahl and/or So Hattori.

Installation

For the most up-to-date version of the code, clone this repository and install locally.

There is also a version on PyPI (TODO: will be auto-updated with most recent version of this repo):

uv add smolgp

Note that tinygp dependencies require the latest version of the tinygp GitHub repository, rather than the version on PyPI. uv should handle this automatically.