solver#

Classes#

ParallelIntegratedStateSpaceSolver

A solver that uses jax.lax.associative_scan to implement

Module Contents#

class smolgp.solvers.integrated.parallel.solver.ParallelIntegratedStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)[source]#

Bases: smolgp.solvers.integrated.solver.IntegratedStateSpaceSolver

A solver that uses jax.lax.associative_scan to implement parallel Kalman filtering and RTS smoothing for integrated measurements

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

The marginal log likelihood, reduced from this solver’s own filter.

Overrides IntegratedStateSpaceSolver.log_probability() deliberately, as that is an optimized sequential scan. The generic path to reuse the Kalman filtered v and S is better here, as those are determined via associative scan, hence the likelihood stays log-depth.

Kalman(y, return_v_S=True) Any[source]#

Wrapper for Kalman filter used with this solver

RTS(kalman_results) Any[source]#

Wrapper for RTS smoother used with this solver