State-Space Models & Kalman

Layer 3 recipes for linear-Gaussian state-space models. Stationary 1-D GP kernels with rational spectral densities admit exact SDE representations

\[ \dot{x}(t) = F\,x(t) + L\,w(t), \qquad f(t) = H\,x(t), \]

turning \(O(N^3)\) GP inference into \(O(N d^3)\) Kalman filtering. This page covers the SDE kernel zoo, the filters and smoothers (sequential, parallel associative-scan, square-root, and steady-state), and the natural-parameter / site machinery for non-conjugate likelihoods.

SDE kernels

Structured linear algebra and Gaussian primitives for JAX.

SDEKernel

Bases: Module

Abstract base class for state-space kernel representations.

Subclasses implement sde_params to provide the continuous-time SDE matrices (F, L, H, Q_c, P_inf). The default discretise uses the matrix exponential for discretization; subclasses may override with closed-form solutions.

Source code in src/gaussx/_ssm/_sde_kernel.py

class SDEKernel(eqx.Module):
    """Abstract base class for state-space kernel representations.

    Subclasses implement `sde_params` to provide the continuous-time
    SDE matrices ``(F, L, H, Q_c, P_inf)``. The default `discretise`
    uses the matrix exponential for discretization; subclasses may override
    with closed-form solutions.
    """

    @property
    @abc.abstractmethod
    def state_dim(self) -> int:
        """Dimension of the latent state vector."""
        ...

    @property
    def stationary(self) -> bool:
        """Whether the process has a stationary distribution.

        ``True`` for every kernel in the zoo, which is why that is the
        default; a non-stationary kernel such as
        `gaussx.IntegratedWienerSDE` overrides it. Consumers should
        branch on this rather than on ``sde_params().P_inf is None``:
        the two are different questions, since a stationary kernel may
        report ``P_inf=None`` when it has no *closed form* for it (a
        learned drift, say).
        """
        return True

    @abc.abstractmethod
    def sde_params(self) -> SDEParams:
        """Return continuous-time SDE parameters."""
        ...

    def initial_covariance(self) -> Float[Array, "d d"]:
        r"""Return the covariance of the state at the first time point.

        For a stationary kernel the process is assumed started in its
        stationary distribution, so this is $P_\infty$ — the default
        implementation returns it and existing kernels need no change.
        A non-stationary kernel has no such limit to start from and must
        override this with an explicit choice.

        Returns:
            Initial state covariance, shape ``(d, d)``.

        Raises:
            ValueError: If ``sde_params().P_inf`` is ``None``, so there
                is no stationary covariance to fall back on.
        """
        P_inf = self.sde_params().P_inf
        if P_inf is None:
            msg = (
                f"{type(self).__name__} has no initial covariance: it "
                f"reports P_inf=None, and the default initial covariance "
                f"is the stationary one. Override initial_covariance() "
                f"with an explicit choice (a diffuse prior, typically), "
                f"or give the kernel a P_inf."
            )
            raise ValueError(msg)
        return P_inf

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        """Discretise the SDE at time step ``dt``.

        Default implementation computes:

            A = expm(F * dt)
            Q = P_inf - A @ P_inf @ A^T

        When ``sde_params()`` returns ``P_inf=None`` — a kernel with no
        closed-form stationary covariance, such as one whose drift is a
        learned parameter — this falls back to `gaussx.discretise_mfd`,
        which recovers both ``A`` and ``Q`` from one matrix exponential and
        is well defined for every ``F``. The fallback is chosen at trace
        time from a static ``None`` check, so kernels that do supply
        ``P_inf`` keep the stationary route and its precision exactly.

        Subclasses may override with closed-form expressions.

        Args:
            dt: Time step (scalar, non-negative).

        Returns:
            Tuple ``(A, Q)`` where A is the transition matrix and
            Q is the process noise covariance.
        """
        params = self.sde_params()
        if params.P_inf is None:
            diffusion = params.L @ params.Q_c @ params.L.T
            return discretise_mfd(params.F, diffusion, dt)
        A = jsl.expm(params.F * dt)
        Q = symmetrize(process_noise_covariance(A, params.P_inf))
        return A, Q

    def discretise_sequence(
        self,
        dt: Float[Array, " N"],
    ) -> tuple[Float[Array, "N d d"], Float[Array, "N d d"]]:
        """Discretise the SDE at multiple time steps.

        Args:
            dt: Time steps, shape ``(N,)``.

        Returns:
            Tuple ``(A_seq, Q_seq)`` with shapes ``(N, d, d)``.
        """
        return jax.vmap(self.discretise)(dt)

state_dim: int abstractmethod property

Dimension of the latent state vector.

stationary: bool property

Whether the process has a stationary distribution.

True for every kernel in the zoo, which is why that is the default; a non-stationary kernel such as gaussx.IntegratedWienerSDE overrides it. Consumers should branch on this rather than on sde_params().P_inf is None: the two are different questions, since a stationary kernel may report P_inf=None when it has no closed form for it (a learned drift, say).

sde_params() -> SDEParams abstractmethod

Return continuous-time SDE parameters.

Source code in src/gaussx/_ssm/_sde_kernel.py

@abc.abstractmethod
def sde_params(self) -> SDEParams:
    """Return continuous-time SDE parameters."""
    ...

initial_covariance() -> Float[Array, 'd d']

Return the covariance of the state at the first time point.

For a stationary kernel the process is assumed started in its stationary distribution, so this is \(P_\infty\) — the default implementation returns it and existing kernels need no change. A non-stationary kernel has no such limit to start from and must override this with an explicit choice.

Returns:

Type	Description
Float[Array, 'd d']	Initial state covariance, shape (d, d).

Raises:

Type	Description
ValueError	If sde_params().P_inf is None, so there is no stationary covariance to fall back on.

Source code in src/gaussx/_ssm/_sde_kernel.py

def initial_covariance(self) -> Float[Array, "d d"]:
    r"""Return the covariance of the state at the first time point.

    For a stationary kernel the process is assumed started in its
    stationary distribution, so this is $P_\infty$ — the default
    implementation returns it and existing kernels need no change.
    A non-stationary kernel has no such limit to start from and must
    override this with an explicit choice.

    Returns:
        Initial state covariance, shape ``(d, d)``.

    Raises:
        ValueError: If ``sde_params().P_inf`` is ``None``, so there
            is no stationary covariance to fall back on.
    """
    P_inf = self.sde_params().P_inf
    if P_inf is None:
        msg = (
            f"{type(self).__name__} has no initial covariance: it "
            f"reports P_inf=None, and the default initial covariance "
            f"is the stationary one. Override initial_covariance() "
            f"with an explicit choice (a diffuse prior, typically), "
            f"or give the kernel a P_inf."
        )
        raise ValueError(msg)
    return P_inf

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Discretise the SDE at time step dt.

Default implementation computes:

A = expm(F * dt)
Q = P_inf - A @ P_inf @ A^T

When sde_params() returns P_inf=None — a kernel with no closed-form stationary covariance, such as one whose drift is a learned parameter — this falls back to gaussx.discretise_mfd, which recovers both A and Q from one matrix exponential and is well defined for every F. The fallback is chosen at trace time from a static None check, so kernels that do supply P_inf keep the stationary route and its precision exactly.

Subclasses may override with closed-form expressions.

Parameters:

Name	Type	Description	Default
dt	Float[Array, '']	Time step (scalar, non-negative).	required

Returns:

Type	Description
Float[Array, 'd d']	Tuple (A, Q) where A is the transition matrix and
Float[Array, 'd d']	Q is the process noise covariance.

Source code in src/gaussx/_ssm/_sde_kernel.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    """Discretise the SDE at time step ``dt``.

    Default implementation computes:

        A = expm(F * dt)
        Q = P_inf - A @ P_inf @ A^T

    When ``sde_params()`` returns ``P_inf=None`` — a kernel with no
    closed-form stationary covariance, such as one whose drift is a
    learned parameter — this falls back to `gaussx.discretise_mfd`,
    which recovers both ``A`` and ``Q`` from one matrix exponential and
    is well defined for every ``F``. The fallback is chosen at trace
    time from a static ``None`` check, so kernels that do supply
    ``P_inf`` keep the stationary route and its precision exactly.

    Subclasses may override with closed-form expressions.

    Args:
        dt: Time step (scalar, non-negative).

    Returns:
        Tuple ``(A, Q)`` where A is the transition matrix and
        Q is the process noise covariance.
    """
    params = self.sde_params()
    if params.P_inf is None:
        diffusion = params.L @ params.Q_c @ params.L.T
        return discretise_mfd(params.F, diffusion, dt)
    A = jsl.expm(params.F * dt)
    Q = symmetrize(process_noise_covariance(A, params.P_inf))
    return A, Q

discretise_sequence(dt: Float[Array, ' N']) -> tuple[Float[Array, 'N d d'], Float[Array, 'N d d']]

Discretise the SDE at multiple time steps.

Parameters:

Name	Type	Description	Default
dt	Float[Array, ' N']	Time steps, shape (N,).	required

Returns:

Type	Description
tuple[Float[Array, 'N d d'], Float[Array, 'N d d']]	Tuple (A_seq, Q_seq) with shapes (N, d, d).

Source code in src/gaussx/_ssm/_sde_kernel.py

def discretise_sequence(
    self,
    dt: Float[Array, " N"],
) -> tuple[Float[Array, "N d d"], Float[Array, "N d d"]]:
    """Discretise the SDE at multiple time steps.

    Args:
        dt: Time steps, shape ``(N,)``.

    Returns:
        Tuple ``(A_seq, Q_seq)`` with shapes ``(N, d, d)``.
    """
    return jax.vmap(self.discretise)(dt)

SDEParams

Bases: NamedTuple

Continuous-time SDE parameters for a linear SDE.

Defines the linear time-invariant SDE:

dx = F x dt + L dW,   W ~ N(0, Q_c dt)

with observation model y = H x.

Attributes:

Name	Type	Description
F	Float[Array, 'd d']	Drift matrix, shape (d, d).
L	Float[Array, 'd s']	Diffusion matrix, shape (d, s).
H	Float[Array, '1 d']	Observation matrix, shape (1, d).
Q_c	Float[Array, 's s']	Spectral density, shape (s, s).
P_inf	Float[Array, 'd d'] | None	Stationary covariance, shape (d, d), or None when the kernel has no closed-form stationary covariance — as for a learned drift matrix, or a non-stationary kernel such as gaussx.IntegratedWienerSDE, which has no stationary covariance at all. SDEKernel.discretise then falls back to gaussx.discretise_mfd, which needs no P_inf, and the filter is started from SDEKernel.initial_covariance instead.

Source code in src/gaussx/_ssm/_sde_kernel.py

class SDEParams(NamedTuple):
    """Continuous-time SDE parameters for a linear SDE.

    Defines the linear time-invariant SDE:

        dx = F x dt + L dW,   W ~ N(0, Q_c dt)

    with observation model ``y = H x``.

    Attributes:
        F: Drift matrix, shape ``(d, d)``.
        L: Diffusion matrix, shape ``(d, s)``.
        H: Observation matrix, shape ``(1, d)``.
        Q_c: Spectral density, shape ``(s, s)``.
        P_inf: Stationary covariance, shape ``(d, d)``, or ``None`` when
            the kernel has no closed-form stationary covariance — as for a
            learned drift matrix, or a non-stationary kernel such as
            `gaussx.IntegratedWienerSDE`, which has no stationary
            covariance at all. `SDEKernel.discretise` then falls back to
            `gaussx.discretise_mfd`, which needs no ``P_inf``, and the
            filter is started from `SDEKernel.initial_covariance`
            instead.
    """

    F: Float[Array, "d d"]
    L: Float[Array, "d s"]
    H: Float[Array, "1 d"]
    Q_c: Float[Array, "s s"]
    P_inf: Float[Array, "d d"] | None = None

ConstantSDE

Bases: SDEKernel

State-space representation of a constant kernel.

Models \(k(\tau) = \sigma^2\) — a degenerate kernel with zero dynamics and zero diffusion. State dimension is 1.

Attributes:

Name	Type	Description
variance	Float[Array, '']	Signal variance \(\sigma^2\).

Source code in src/gaussx/_ssm/_constant.py

class ConstantSDE(SDEKernel):
    r"""State-space representation of a constant kernel.

    Models $k(\tau) = \sigma^2$ — a degenerate kernel with zero
    dynamics and zero diffusion. State dimension is 1.

    Attributes:
        variance: Signal variance $\sigma^2$.
    """

    variance: Float[Array, ""]

    @property
    def state_dim(self) -> int:
        return 1

    def sde_params(self) -> SDEParams:
        """Return SDE parameters for the constant kernel."""
        # Constant blocks follow the hyperparameter dtype; untyped
        # ``jnp.zeros``/``jnp.array`` are float64 under x64 (gh-224).
        dtype = jnp.result_type(self.variance)
        F = jnp.zeros((1, 1), dtype=dtype)
        L = jnp.zeros((1, 1), dtype=dtype)
        H = jnp.array([[1.0]], dtype=dtype)
        Q_c = jnp.zeros((1, 1), dtype=dtype)
        P_inf = jnp.array([[self.variance]])
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        """Closed-form: A = I, Q = 0 (no dynamics)."""
        dtype = jnp.result_type(self.variance)
        A = jnp.eye(1, dtype=dtype)
        Q = jnp.zeros((1, 1), dtype=dtype)
        return A, Q

sde_params() -> SDEParams

Return SDE parameters for the constant kernel.

Source code in src/gaussx/_ssm/_constant.py

def sde_params(self) -> SDEParams:
    """Return SDE parameters for the constant kernel."""
    # Constant blocks follow the hyperparameter dtype; untyped
    # ``jnp.zeros``/``jnp.array`` are float64 under x64 (gh-224).
    dtype = jnp.result_type(self.variance)
    F = jnp.zeros((1, 1), dtype=dtype)
    L = jnp.zeros((1, 1), dtype=dtype)
    H = jnp.array([[1.0]], dtype=dtype)
    Q_c = jnp.zeros((1, 1), dtype=dtype)
    P_inf = jnp.array([[self.variance]])
    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Closed-form: A = I, Q = 0 (no dynamics).

Source code in src/gaussx/_ssm/_constant.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    """Closed-form: A = I, Q = 0 (no dynamics)."""
    dtype = jnp.result_type(self.variance)
    A = jnp.eye(1, dtype=dtype)
    Q = jnp.zeros((1, 1), dtype=dtype)
    return A, Q

MaternSDE

Bases: SDEKernel

State-space representation of the Matern kernel.

Supports orders 0 (Matern-1/2), 1 (Matern-3/2), and 2 (Matern-5/2). The state dimension is order + 1.

Attributes:

Name	Type	Description
variance	Float[Array, '']	Signal variance \(\sigma^2\).
lengthscale	Float[Array, '']	Lengthscale \(\ell\).
order	int	Matern order (0, 1, or 2).

Source code in src/gaussx/_ssm/_matern.py

class MaternSDE(SDEKernel):
    r"""State-space representation of the Matern kernel.

    Supports orders 0 (Matern-1/2), 1 (Matern-3/2), and 2 (Matern-5/2).
    The state dimension is ``order + 1``.

    Attributes:
        variance: Signal variance $\sigma^2$.
        lengthscale: Lengthscale $\ell$.
        order: Matern order (0, 1, or 2).
    """

    variance: Float[Array, ""]
    lengthscale: Float[Array, ""]
    order: int = eqx.field(static=True)

    @property
    def state_dim(self) -> int:
        return self.order + 1

    def sde_params(self) -> SDEParams:
        """Compute SDE parameters for the Matern kernel."""
        if self.order == 0:
            return self._matern12()
        elif self.order == 1:
            return self._matern32()
        elif self.order == 2:
            return self._matern52()
        else:
            msg = f"Unsupported Matern order {self.order}; must be 0, 1, or 2"
            raise ValueError(msg)

    def _matern12(self) -> SDEParams:
        # Constant blocks follow the hyperparameter dtype; an untyped
        # ``jnp.array`` of Python floats is float64 under x64 (gh-224).
        dtype = jnp.result_type(self.variance, self.lengthscale)
        lam = 1.0 / self.lengthscale
        F = jnp.array([[-lam]])
        L = jnp.array([[1.0]], dtype=dtype)
        H = jnp.array([[1.0]], dtype=dtype)
        Q_c = jnp.array([[2.0 * lam * self.variance]])
        P_inf = jnp.array([[self.variance]])
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def _matern32(self) -> SDEParams:
        dtype = jnp.result_type(self.variance, self.lengthscale)
        lam = jnp.sqrt(3.0) / self.lengthscale
        F = jnp.array([[0.0, 1.0], [-(lam**2), -2.0 * lam]])
        L = jnp.array([[0.0], [1.0]], dtype=dtype)
        H = jnp.array([[1.0, 0.0]], dtype=dtype)
        q = 4.0 * lam**3 * self.variance
        Q_c = jnp.array([[q]])
        P_inf = jnp.array([[self.variance, 0.0], [0.0, lam**2 * self.variance]])
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def _matern52(self) -> SDEParams:
        dtype = jnp.result_type(self.variance, self.lengthscale)
        lam = jnp.sqrt(5.0) / self.lengthscale
        F = jnp.array(
            [
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 1.0],
                [-(lam**3), -3.0 * lam**2, -3.0 * lam],
            ]
        )
        L = jnp.array([[0.0], [0.0], [1.0]], dtype=dtype)
        H = jnp.array([[1.0, 0.0, 0.0]], dtype=dtype)
        kappa = 5.0 / 3.0 * self.variance / self.lengthscale**2
        q = 16.0 / 3.0 * lam**5 * self.variance
        Q_c = jnp.array([[q]])
        P_inf = jnp.array(
            [
                [self.variance, 0.0, -kappa],
                [0.0, kappa, 0.0],
                [-kappa, 0.0, lam**4 * self.variance],
            ]
        )
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

sde_params() -> SDEParams

Compute SDE parameters for the Matern kernel.

Source code in src/gaussx/_ssm/_matern.py

def sde_params(self) -> SDEParams:
    """Compute SDE parameters for the Matern kernel."""
    if self.order == 0:
        return self._matern12()
    elif self.order == 1:
        return self._matern32()
    elif self.order == 2:
        return self._matern52()
    else:
        msg = f"Unsupported Matern order {self.order}; must be 0, 1, or 2"
        raise ValueError(msg)

PeriodicSDE

Bases: SDEKernel

State-space representation of the periodic (MacKay) kernel.

Approximates the periodic kernel via Fourier series truncation to n_harmonics terms. State dimension is 2 * n_harmonics.

Attributes:

Name	Type	Description
variance	Float[Array, '']	Signal variance \(\sigma^2\).
lengthscale	Float[Array, '']	Lengthscale \(\ell\).
period	Float[Array, '']	Period \(T\).
n_harmonics	int	Number of Fourier harmonics (truncation order).

Source code in src/gaussx/_ssm/_periodic.py

class PeriodicSDE(SDEKernel):
    r"""State-space representation of the periodic (MacKay) kernel.

    Approximates the periodic kernel via Fourier series truncation
    to ``n_harmonics`` terms. State dimension is ``2 * n_harmonics``.

    Attributes:
        variance: Signal variance $\sigma^2$.
        lengthscale: Lengthscale $\ell$.
        period: Period $T$.
        n_harmonics: Number of Fourier harmonics (truncation order).
    """

    variance: Float[Array, ""]
    lengthscale: Float[Array, ""]
    period: Float[Array, ""]
    n_harmonics: int = eqx.field(static=True, default=6)

    @property
    def state_dim(self) -> int:
        return 2 * self.n_harmonics

    def sde_params(self) -> SDEParams:
        """Return SDE parameters for the periodic kernel."""
        dtype = jnp.result_type(self.variance, self.lengthscale, self.period)
        J = self.n_harmonics
        d = 2 * J
        w0 = 2.0 * jnp.pi / self.period

        inv_ell_sq = 1.0 / self.lengthscale**2
        # ``js`` feeds the Bessel series; an integer arange would promote it
        # to float64 under x64 (gh-224).
        js = jnp.arange(1, J + 1, dtype=dtype)
        log_ij = self._log_bessel_i(js, inv_ell_sq)
        log_q = jnp.log(2.0) + log_ij - inv_ell_sq
        q_j = self.variance * jnp.exp(log_q)

        F = jnp.zeros((d, d), dtype=dtype)
        P_inf = jnp.zeros((d, d), dtype=dtype)
        for j_idx in range(J):
            freq = (j_idx + 1) * w0
            block_start = 2 * j_idx
            F = F.at[block_start, block_start + 1].set(-freq)
            F = F.at[block_start + 1, block_start].set(freq)
            P_inf = P_inf.at[block_start, block_start].set(q_j[j_idx])
            P_inf = P_inf.at[block_start + 1, block_start + 1].set(q_j[j_idx])

        L = jnp.zeros((d, 1), dtype=dtype)
        H = jnp.zeros((1, d), dtype=dtype)
        for j_idx in range(J):
            H = H.at[0, 2 * j_idx].set(1.0)

        Q_c = jnp.zeros((1, 1), dtype=dtype)
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        """Closed-form: block-diagonal rotation matrices."""
        J = self.n_harmonics
        d = 2 * J
        w0 = 2.0 * jnp.pi / self.period

        A = jnp.zeros((d, d), dtype=jnp.result_type(w0, dt))
        for j_idx in range(J):
            freq = (j_idx + 1) * w0
            cos_val = jnp.cos(freq * dt)
            sin_val = jnp.sin(freq * dt)
            block_start = 2 * j_idx
            A = A.at[block_start, block_start].set(cos_val)
            A = A.at[block_start, block_start + 1].set(-sin_val)
            A = A.at[block_start + 1, block_start].set(sin_val)
            A = A.at[block_start + 1, block_start + 1].set(cos_val)

        Q = jnp.zeros((d, d), dtype=A.dtype)
        return A, Q

    @staticmethod
    def _log_bessel_i(
        order: Float[Array, " J"],
        x: Float[Array, ""],
    ) -> Float[Array, " J"]:
        """Log of modified Bessel function I_n(x) via series."""
        half_x = x / 2.0
        log_half_x = jnp.log(half_x)

        log_leading = order * log_half_x - jss.gammaln(order + 1.0)

        x2_over_4 = x**2 / 4.0
        K = 20
        log_sum = jnp.zeros_like(order)
        log_term = jnp.zeros_like(order)
        for k in range(1, K + 1):
            log_term = (
                log_term
                + jnp.log(x2_over_4)
                - jnp.log(jnp.array(k, dtype=order.dtype))
                - jnp.log(order + k)
            )
            log_sum = jnp.logaddexp(log_sum, log_term)

        return log_leading + log_sum

sde_params() -> SDEParams

Return SDE parameters for the periodic kernel.

Source code in src/gaussx/_ssm/_periodic.py

def sde_params(self) -> SDEParams:
    """Return SDE parameters for the periodic kernel."""
    dtype = jnp.result_type(self.variance, self.lengthscale, self.period)
    J = self.n_harmonics
    d = 2 * J
    w0 = 2.0 * jnp.pi / self.period

    inv_ell_sq = 1.0 / self.lengthscale**2
    # ``js`` feeds the Bessel series; an integer arange would promote it
    # to float64 under x64 (gh-224).
    js = jnp.arange(1, J + 1, dtype=dtype)
    log_ij = self._log_bessel_i(js, inv_ell_sq)
    log_q = jnp.log(2.0) + log_ij - inv_ell_sq
    q_j = self.variance * jnp.exp(log_q)

    F = jnp.zeros((d, d), dtype=dtype)
    P_inf = jnp.zeros((d, d), dtype=dtype)
    for j_idx in range(J):
        freq = (j_idx + 1) * w0
        block_start = 2 * j_idx
        F = F.at[block_start, block_start + 1].set(-freq)
        F = F.at[block_start + 1, block_start].set(freq)
        P_inf = P_inf.at[block_start, block_start].set(q_j[j_idx])
        P_inf = P_inf.at[block_start + 1, block_start + 1].set(q_j[j_idx])

    L = jnp.zeros((d, 1), dtype=dtype)
    H = jnp.zeros((1, d), dtype=dtype)
    for j_idx in range(J):
        H = H.at[0, 2 * j_idx].set(1.0)

    Q_c = jnp.zeros((1, 1), dtype=dtype)
    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Closed-form: block-diagonal rotation matrices.

Source code in src/gaussx/_ssm/_periodic.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    """Closed-form: block-diagonal rotation matrices."""
    J = self.n_harmonics
    d = 2 * J
    w0 = 2.0 * jnp.pi / self.period

    A = jnp.zeros((d, d), dtype=jnp.result_type(w0, dt))
    for j_idx in range(J):
        freq = (j_idx + 1) * w0
        cos_val = jnp.cos(freq * dt)
        sin_val = jnp.sin(freq * dt)
        block_start = 2 * j_idx
        A = A.at[block_start, block_start].set(cos_val)
        A = A.at[block_start, block_start + 1].set(-sin_val)
        A = A.at[block_start + 1, block_start].set(sin_val)
        A = A.at[block_start + 1, block_start + 1].set(cos_val)

    Q = jnp.zeros((d, d), dtype=A.dtype)
    return A, Q

QuasiPeriodicSDE

Bases: ProductSDE

Quasi-periodic kernel: product of Matern and Periodic SDE.

Attributes:

Name	Type	Description
kernel1	SDEKernel	Modulating kernel (typically Matern).
kernel2	SDEKernel	Periodic kernel.

Source code in src/gaussx/_ssm/_composition.py

class QuasiPeriodicSDE(ProductSDE):
    """Quasi-periodic kernel: product of Matern and Periodic SDE.

    Attributes:
        kernel1: Modulating kernel (typically Matern).
        kernel2: Periodic kernel.
    """

    pass

CosineSDE

Bases: SDEKernel

State-space representation of the cosine kernel.

Models \(k(\tau) = \sigma^2 \cos(\omega_0 \tau)\) via a 2-D rotation SDE. State dimension is 2.

Attributes:

Name	Type	Description
variance	Float[Array, '']	Signal variance \(\sigma^2\).
frequency	Float[Array, '']	Angular frequency \(\omega_0\).

Source code in src/gaussx/_ssm/_periodic.py

class CosineSDE(SDEKernel):
    r"""State-space representation of the cosine kernel.

    Models $k(\tau) = \sigma^2 \cos(\omega_0 \tau)$ via a 2-D
    rotation SDE. State dimension is 2.

    Attributes:
        variance: Signal variance $\sigma^2$.
        frequency: Angular frequency $\omega_0$.
    """

    variance: Float[Array, ""]
    frequency: Float[Array, ""]

    @property
    def state_dim(self) -> int:
        return 2

    def sde_params(self) -> SDEParams:
        """Return SDE parameters for the cosine kernel."""
        # Constant blocks follow the hyperparameter dtype; untyped
        # ``jnp.zeros``/``jnp.eye`` are float64 under x64 (gh-224).
        dtype = jnp.result_type(self.variance, self.frequency)
        w = self.frequency
        F = jnp.array([[0.0, -w], [w, 0.0]])
        L = jnp.zeros((2, 1), dtype=dtype)
        H = jnp.array([[1.0, 0.0]], dtype=dtype)
        Q_c = jnp.zeros((1, 1), dtype=dtype)
        P_inf = self.variance * jnp.eye(2, dtype=dtype)
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        """Closed-form rotation matrix discretization."""
        w = self.frequency
        cos_wdt = jnp.cos(w * dt)
        sin_wdt = jnp.sin(w * dt)
        A = jnp.array([[cos_wdt, -sin_wdt], [sin_wdt, cos_wdt]])
        Q = jnp.zeros((2, 2), dtype=A.dtype)
        return A, Q

sde_params() -> SDEParams

Return SDE parameters for the cosine kernel.

Source code in src/gaussx/_ssm/_periodic.py

def sde_params(self) -> SDEParams:
    """Return SDE parameters for the cosine kernel."""
    # Constant blocks follow the hyperparameter dtype; untyped
    # ``jnp.zeros``/``jnp.eye`` are float64 under x64 (gh-224).
    dtype = jnp.result_type(self.variance, self.frequency)
    w = self.frequency
    F = jnp.array([[0.0, -w], [w, 0.0]])
    L = jnp.zeros((2, 1), dtype=dtype)
    H = jnp.array([[1.0, 0.0]], dtype=dtype)
    Q_c = jnp.zeros((1, 1), dtype=dtype)
    P_inf = self.variance * jnp.eye(2, dtype=dtype)
    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Closed-form rotation matrix discretization.

Source code in src/gaussx/_ssm/_periodic.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    """Closed-form rotation matrix discretization."""
    w = self.frequency
    cos_wdt = jnp.cos(w * dt)
    sin_wdt = jnp.sin(w * dt)
    A = jnp.array([[cos_wdt, -sin_wdt], [sin_wdt, cos_wdt]])
    Q = jnp.zeros((2, 2), dtype=A.dtype)
    return A, Q

IntegratedWienerSDE

Bases: SDEKernel

State-space representation of an integrated Wiener process.

The \(p\)-times integrated Wiener process — the local linear trend prior at order=1 — with state \(x(t) = [f(t), f'(t), \dots, f^{(p)}(t)]\) and

\[ \mathrm{d} f^{(p)}(t) = \sqrt{q} \, \mathrm{d} W(t), \]

so the top derivative is white noise and every lower component is its integral. State dimension is order + 1. The SDE matrices are

\[ F = \begin{bmatrix} 0 & I_p \\ 0 & 0 \end{bmatrix}, \quad L = e_p, \quad H = e_0^\top, \quad Q_c = q , \]

with \(F\) nilpotent, which is what makes the process non-stationary: its marginal variance grows without bound, so no \(P_\infty\) exists and sde_params reports P_inf=None. discretise is overridden with the exact closed form (below) and never needs one, but the filter has to be started from an explicit initial_covariance instead.

That non-stationarity is the point: unlike a Matérn prior, the local linear trend commits to no lengthscale — it is smooth and linear unless the data push back — so a long record may drift without the model asserting a scale on which it must revert.

Attributes:

Name	Type	Description
diffusion	Float[Array, '']	Diffusion intensity \(q\), the spectral density of the white noise driving the top derivative. Must be non-negative — \(Q\) is linear in it, so a negative value returns something that is not a covariance. Like every hyperparameter in the kernel zoo this is assumed rather than checked; constrain it with a positive transform when it is learned.
order	int	Number of integrations \(p\). 0 is a Brownian random walk; 1 (the default) is the local linear trend, which is also the cubic-spline-equivalent prior — it is \(f''\) that is white noise there, and the smoothed posterior mean is the cubic smoothing spline. Each further order raises the spline by two degrees, so 2 is the quintic-spline prior.
P_0	Float[Array, 'd d'] | None	Initial state covariance, shape (order + 1, order + 1). A modelling choice rather than something the kernel can derive; None (the default) means a diffuse _default_diffuse_variance(dtype) * I.

Source code in src/gaussx/_ssm/_wiener.py

class IntegratedWienerSDE(SDEKernel):
    r"""State-space representation of an integrated Wiener process.

    The $p$-times integrated Wiener process — the local linear trend
    prior at ``order=1`` — with state
    $x(t) = [f(t), f'(t), \dots, f^{(p)}(t)]$ and

    $$
    \mathrm{d} f^{(p)}(t) = \sqrt{q} \, \mathrm{d} W(t),
    $$

    so the top derivative is white noise and every lower component is its
    integral. State dimension is ``order + 1``. The SDE matrices are

    $$
    F = \begin{bmatrix} 0 & I_p \\ 0 & 0 \end{bmatrix}, \quad
    L = e_p, \quad
    H = e_0^\top, \quad
    Q_c = q ,
    $$

    with $F$ nilpotent, which is what makes the process **non-stationary**:
    its marginal variance grows without bound, so no $P_\infty$ exists and
    `sde_params` reports ``P_inf=None``. `discretise` is overridden with
    the exact closed form (below) and never needs one, but the filter has
    to be started from an explicit `initial_covariance` instead.

    That non-stationarity is the point: unlike a Matérn prior, the local
    linear trend commits to no lengthscale — it is smooth and linear
    unless the data push back — so a long record may drift without the
    model asserting a scale on which it must revert.

    Attributes:
        diffusion: Diffusion intensity $q$, the spectral density of the
            white noise driving the top derivative. Must be
            non-negative — $Q$ is linear in it, so a negative value
            returns something that is not a covariance. Like every
            hyperparameter in the kernel zoo this is assumed rather
            than checked; constrain it with a positive transform when
            it is learned.
        order: Number of integrations $p$. ``0`` is a Brownian random
            walk; ``1`` (the default) is the local linear trend, which
            is also the cubic-spline-equivalent prior — it is $f''$ that
            is white noise there, and the smoothed posterior mean is the
            cubic smoothing spline. Each further order raises the spline
            by two degrees, so ``2`` is the quintic-spline prior.
        P_0: Initial state covariance, shape ``(order + 1, order + 1)``.
            A modelling choice rather than something the kernel can
            derive; ``None`` (the default) means a diffuse
            ``_default_diffuse_variance(dtype) * I``.
    """

    diffusion: Float[Array, ""]
    order: int = eqx.field(static=True, default=1)
    P_0: Float[Array, "d d"] | None = None

    def __check_init__(self) -> None:
        """Reject a negative order at construction.

        The state dimension is ``order + 1``, so a negative order gives
        an empty state and fails later, deep inside whichever method is
        called first, with an index error about an axis of size zero.
        """
        if self.order < 0:
            msg = (
                f"IntegratedWienerSDE order must be non-negative "
                f"(the state is [f, ..., f^(order)], of dimension "
                f"order + 1), got {self.order}."
            )
            raise ValueError(msg)

    @property
    def state_dim(self) -> int:
        return self.order + 1

    @property
    def stationary(self) -> bool:
        """``False`` — the marginal variance grows without bound."""
        return False

    def sde_params(self) -> SDEParams:
        """Return SDE parameters, with ``P_inf=None``.

        The drift is nilpotent, so no stationary covariance exists; see
        `initial_covariance` for what starts the filter instead.
        """
        # Constant blocks follow the hyperparameter dtype; untyped
        # ``jnp.zeros``/``jnp.eye`` are float64 under x64 (gh-224).
        dtype = _inexact_dtype(self.diffusion)
        d = self.state_dim
        F = jnp.eye(d, k=1, dtype=dtype)
        L = jnp.zeros((d, 1), dtype=dtype).at[d - 1, 0].set(1.0)
        H = jnp.zeros((1, d), dtype=dtype).at[0, 0].set(1.0)
        Q_c = jnp.reshape(self.diffusion, (1, 1)).astype(dtype)
        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=None)

    def initial_covariance(self) -> Float[Array, "d d"]:
        r"""Return the initial state covariance $P_0$.

        Defaults to a diffuse ``kappa * I`` when the ``P_0`` field is
        ``None``, with ``kappa`` set from the dtype's precision by
        `_default_diffuse_variance` — a vaguer prior than that is not
        merely wasteful but actively wrong, since the Kalman update
        cancels it to a zero-variance first estimate.

        Pass a ``P_0`` to encode what is actually known about the level
        and its derivatives at the first time point — e.g.
        ``diag(kappa, s^2)`` for a vague level and a slope of scale
        ``s``. Do so in particular when the observation noise is large:
        the default is diffuse relative to a noise variance of order
        one, not relative to every scale.
        """
        dtype = _inexact_dtype(self.diffusion)
        if self.P_0 is None:
            kappa = _default_diffuse_variance(dtype)
            return kappa * jnp.eye(self.state_dim, dtype=dtype)
        return jnp.asarray(self.P_0, dtype=dtype)

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        r"""Closed-form discretisation — no ``expm``, no $P_\infty$.

        The drift is nilpotent, so its exponential terminates:

        $$
        A(\Delta t)_{ij} = \frac{\Delta t^{\,j-i}}{(j-i)!} \;\; (j \ge i),
        \qquad
        Q(\Delta t)_{ij} = q \,
            \frac{\Delta t^{\,2p+1-i-j}}
                 {(p-i)!\,(p-j)!\,(2p+1-i-j)} ,
        $$

        which at ``order=1`` is the familiar

        $$
        A = \begin{bmatrix} 1 & \Delta t \\ 0 & 1 \end{bmatrix},
        \qquad
        Q = q \begin{bmatrix}
            \Delta t^3/3 & \Delta t^2/2 \\
            \Delta t^2/2 & \Delta t
        \end{bmatrix}.
        $$

        Both are exact, so this route is cheaper *and* more accurate than
        the `gaussx.discretise_mfd` fallback the ``P_inf=None`` base
        implementation would otherwise take.

        Note:
            The coefficients carry $1/(p-i)!$, which outruns float64
            from roughly ``order=90`` up and underflows to zero there
            (XLA flushes the subnormals in between). Nothing is
            silently wrong — the entries are zero rather than garbage —
            but a prior that integrates white noise ninety times is not
            a numerically meaningful object in any dtype.

        Args:
            dt: Time step. Must be non-negative — a negative step would
                return a ``Q`` that is not a covariance (at ``order=1``
                it is negative definite), rather than the harmless
                reverse-time transition the sign might suggest. Checked
                with `equinox.error_if`, matching
                `gaussx.discretise_mfd`, so under ``jit`` the error
                fires at evaluation rather than trace time.

        Returns:
            Tuple ``(A, Q)``, both shape ``(order + 1, order + 1)``.
        """
        dtype = _inexact_dtype(self.diffusion, dt)
        p = self.order
        d = self.state_dim

        # The closed form below is a polynomial in dt with no guard of
        # its own, unlike the ``expm`` routes; a negative step would run
        # it happily and hand back an indefinite Q.
        dt = eqx.error_if(
            dt, dt < 0, "IntegratedWienerSDE.discretise requires dt >= 0."
        )

        # Powers are taken with *static* Python exponents so JAX lowers
        # them to ``integer_pow``, whose derivative is exact at zero. A
        # traced exponent would go through the generic ``y * x**(y-1)``
        # rule instead, and dt = 0 -- which the natural
        # ``diff(times, prepend=times[0])`` produces at the first step --
        # would differentiate to 0 * inf = NaN.
        def power(exponent: int) -> Float[Array, ""]:
            if exponent == 0:
                return jnp.ones_like(dt)
            return dt**exponent

        # Only 2p+2 distinct powers appear across both matrices, so they
        # are formed once and gathered by a static index table. That is
        # O(d) traced operations rather than one per entry: at order 40
        # the per-entry version took over five seconds to trace.
        powers = jnp.stack([power(k) for k in range(2 * p + 2)]).astype(dtype)

        i = np.arange(d)[:, None]
        j = np.arange(d)[None, :]
        upper = j >= i
        # Zeroed through the *coefficient* rather than by masking a
        # negative power, so no entry raises dt to a negative exponent.
        a_index = np.maximum(j - i, 0)
        q_index = 2 * p + 1 - i - j

        # ``1 / n`` keeps both operands Python ints, which are unbounded
        # and divide to a correctly rounded float. ``1.0 / n`` would
        # convert first and raise OverflowError from order 98 up, where
        # the denominator exceeds the float range even though its
        # reciprocal is perfectly representable (down to a subnormal,
        # and to zero beyond that).
        factorial = [math.factorial(n) for n in range(d)]
        a_coeff = np.array(
            [
                [1 / factorial[j - i] if j >= i else 0.0 for j in range(d)]
                for i in range(d)
            ]
        )
        q_coeff = np.array(
            [
                [
                    1 / (factorial[p - i] * factorial[p - j] * (2 * p + 1 - i - j))
                    for j in range(d)
                ]
                for i in range(d)
            ]
        )

        A = jnp.asarray(a_coeff * upper, dtype=dtype) * powers[a_index]
        Q = self.diffusion * jnp.asarray(q_coeff, dtype=dtype) * powers[q_index]
        return A.astype(dtype), Q.astype(dtype)

stationary: bool property

False — the marginal variance grows without bound.

sde_params() -> SDEParams

Return SDE parameters, with P_inf=None.

The drift is nilpotent, so no stationary covariance exists; see initial_covariance for what starts the filter instead.

Source code in src/gaussx/_ssm/_wiener.py

def sde_params(self) -> SDEParams:
    """Return SDE parameters, with ``P_inf=None``.

    The drift is nilpotent, so no stationary covariance exists; see
    `initial_covariance` for what starts the filter instead.
    """
    # Constant blocks follow the hyperparameter dtype; untyped
    # ``jnp.zeros``/``jnp.eye`` are float64 under x64 (gh-224).
    dtype = _inexact_dtype(self.diffusion)
    d = self.state_dim
    F = jnp.eye(d, k=1, dtype=dtype)
    L = jnp.zeros((d, 1), dtype=dtype).at[d - 1, 0].set(1.0)
    H = jnp.zeros((1, d), dtype=dtype).at[0, 0].set(1.0)
    Q_c = jnp.reshape(self.diffusion, (1, 1)).astype(dtype)
    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=None)

initial_covariance() -> Float[Array, 'd d']

Return the initial state covariance \(P_0\).

Defaults to a diffuse kappa * I when the P_0 field is None, with kappa set from the dtype's precision by _default_diffuse_variance — a vaguer prior than that is not merely wasteful but actively wrong, since the Kalman update cancels it to a zero-variance first estimate.

Pass a P_0 to encode what is actually known about the level and its derivatives at the first time point — e.g. diag(kappa, s^2) for a vague level and a slope of scale s. Do so in particular when the observation noise is large: the default is diffuse relative to a noise variance of order one, not relative to every scale.

Source code in src/gaussx/_ssm/_wiener.py

def initial_covariance(self) -> Float[Array, "d d"]:
    r"""Return the initial state covariance $P_0$.

    Defaults to a diffuse ``kappa * I`` when the ``P_0`` field is
    ``None``, with ``kappa`` set from the dtype's precision by
    `_default_diffuse_variance` — a vaguer prior than that is not
    merely wasteful but actively wrong, since the Kalman update
    cancels it to a zero-variance first estimate.

    Pass a ``P_0`` to encode what is actually known about the level
    and its derivatives at the first time point — e.g.
    ``diag(kappa, s^2)`` for a vague level and a slope of scale
    ``s``. Do so in particular when the observation noise is large:
    the default is diffuse relative to a noise variance of order
    one, not relative to every scale.
    """
    dtype = _inexact_dtype(self.diffusion)
    if self.P_0 is None:
        kappa = _default_diffuse_variance(dtype)
        return kappa * jnp.eye(self.state_dim, dtype=dtype)
    return jnp.asarray(self.P_0, dtype=dtype)

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Closed-form discretisation — no expm, no \(P_\infty\).

The drift is nilpotent, so its exponential terminates:

\[ A(\Delta t)_{ij} = \frac{\Delta t^{\,j-i}}{(j-i)!} \;\; (j \ge i), \qquad Q(\Delta t)_{ij} = q \, \frac{\Delta t^{\,2p+1-i-j}} {(p-i)!\,(p-j)!\,(2p+1-i-j)} , \]

which at order=1 is the familiar

\[ A = \begin{bmatrix} 1 & \Delta t \\ 0 & 1 \end{bmatrix}, \qquad Q = q \begin{bmatrix} \Delta t^3/3 & \Delta t^2/2 \\ \Delta t^2/2 & \Delta t \end{bmatrix}. \]

Both are exact, so this route is cheaper and more accurate than the gaussx.discretise_mfd fallback the P_inf=None base implementation would otherwise take.

Note

The coefficients carry \(1/(p-i)!\), which outruns float64 from roughly order=90 up and underflows to zero there (XLA flushes the subnormals in between). Nothing is silently wrong — the entries are zero rather than garbage — but a prior that integrates white noise ninety times is not a numerically meaningful object in any dtype.

Parameters:

Name	Type	Description	Default
dt	Float[Array, '']	Time step. Must be non-negative — a negative step would return a Q that is not a covariance (at order=1 it is negative definite), rather than the harmless reverse-time transition the sign might suggest. Checked with equinox.error_if, matching gaussx.discretise_mfd, so under jit the error fires at evaluation rather than trace time.	required

Returns:

Type	Description
tuple[Float[Array, 'd d'], Float[Array, 'd d']]	Tuple (A, Q), both shape (order + 1, order + 1).

Source code in src/gaussx/_ssm/_wiener.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    r"""Closed-form discretisation — no ``expm``, no $P_\infty$.

    The drift is nilpotent, so its exponential terminates:

    $$
    A(\Delta t)_{ij} = \frac{\Delta t^{\,j-i}}{(j-i)!} \;\; (j \ge i),
    \qquad
    Q(\Delta t)_{ij} = q \,
        \frac{\Delta t^{\,2p+1-i-j}}
             {(p-i)!\,(p-j)!\,(2p+1-i-j)} ,
    $$

    which at ``order=1`` is the familiar

    $$
    A = \begin{bmatrix} 1 & \Delta t \\ 0 & 1 \end{bmatrix},
    \qquad
    Q = q \begin{bmatrix}
        \Delta t^3/3 & \Delta t^2/2 \\
        \Delta t^2/2 & \Delta t
    \end{bmatrix}.
    $$

    Both are exact, so this route is cheaper *and* more accurate than
    the `gaussx.discretise_mfd` fallback the ``P_inf=None`` base
    implementation would otherwise take.

    Note:
        The coefficients carry $1/(p-i)!$, which outruns float64
        from roughly ``order=90`` up and underflows to zero there
        (XLA flushes the subnormals in between). Nothing is
        silently wrong — the entries are zero rather than garbage —
        but a prior that integrates white noise ninety times is not
        a numerically meaningful object in any dtype.

    Args:
        dt: Time step. Must be non-negative — a negative step would
            return a ``Q`` that is not a covariance (at ``order=1``
            it is negative definite), rather than the harmless
            reverse-time transition the sign might suggest. Checked
            with `equinox.error_if`, matching
            `gaussx.discretise_mfd`, so under ``jit`` the error
            fires at evaluation rather than trace time.

    Returns:
        Tuple ``(A, Q)``, both shape ``(order + 1, order + 1)``.
    """
    dtype = _inexact_dtype(self.diffusion, dt)
    p = self.order
    d = self.state_dim

    # The closed form below is a polynomial in dt with no guard of
    # its own, unlike the ``expm`` routes; a negative step would run
    # it happily and hand back an indefinite Q.
    dt = eqx.error_if(
        dt, dt < 0, "IntegratedWienerSDE.discretise requires dt >= 0."
    )

    # Powers are taken with *static* Python exponents so JAX lowers
    # them to ``integer_pow``, whose derivative is exact at zero. A
    # traced exponent would go through the generic ``y * x**(y-1)``
    # rule instead, and dt = 0 -- which the natural
    # ``diff(times, prepend=times[0])`` produces at the first step --
    # would differentiate to 0 * inf = NaN.
    def power(exponent: int) -> Float[Array, ""]:
        if exponent == 0:
            return jnp.ones_like(dt)
        return dt**exponent

    # Only 2p+2 distinct powers appear across both matrices, so they
    # are formed once and gathered by a static index table. That is
    # O(d) traced operations rather than one per entry: at order 40
    # the per-entry version took over five seconds to trace.
    powers = jnp.stack([power(k) for k in range(2 * p + 2)]).astype(dtype)

    i = np.arange(d)[:, None]
    j = np.arange(d)[None, :]
    upper = j >= i
    # Zeroed through the *coefficient* rather than by masking a
    # negative power, so no entry raises dt to a negative exponent.
    a_index = np.maximum(j - i, 0)
    q_index = 2 * p + 1 - i - j

    # ``1 / n`` keeps both operands Python ints, which are unbounded
    # and divide to a correctly rounded float. ``1.0 / n`` would
    # convert first and raise OverflowError from order 98 up, where
    # the denominator exceeds the float range even though its
    # reciprocal is perfectly representable (down to a subnormal,
    # and to zero beyond that).
    factorial = [math.factorial(n) for n in range(d)]
    a_coeff = np.array(
        [
            [1 / factorial[j - i] if j >= i else 0.0 for j in range(d)]
            for i in range(d)
        ]
    )
    q_coeff = np.array(
        [
            [
                1 / (factorial[p - i] * factorial[p - j] * (2 * p + 1 - i - j))
                for j in range(d)
            ]
            for i in range(d)
        ]
    )

    A = jnp.asarray(a_coeff * upper, dtype=dtype) * powers[a_index]
    Q = self.diffusion * jnp.asarray(q_coeff, dtype=dtype) * powers[q_index]
    return A.astype(dtype), Q.astype(dtype)

ProductSDE

Bases: SDEKernel

Product of two SDE kernels via Kronecker composition.

Attributes:

Name	Type	Description
kernel1	SDEKernel	First component kernel.
kernel2	SDEKernel	Second component kernel.

Source code in src/gaussx/_ssm/_composition.py

class ProductSDE(SDEKernel):
    """Product of two SDE kernels via Kronecker composition.

    Attributes:
        kernel1: First component kernel.
        kernel2: Second component kernel.
    """

    kernel1: SDEKernel
    kernel2: SDEKernel

    @property
    def state_dim(self) -> int:
        return self.kernel1.state_dim * self.kernel2.state_dim

    @property
    def stationary(self) -> bool:
        """Stationary only if both factors are."""
        return self.kernel1.stationary and self.kernel2.stationary

    def sde_params(self) -> SDEParams:
        r"""Return Kronecker-structured SDE parameters.

        The drift is the Kronecker **sum** $F_1 \oplus F_2$ and the
        stationary covariance the Kronecker **product**
        $P_1 \otimes P_2$. Substituting those into the Lyapunov equation
        $F P + P F^\top + B = 0$ fixes the composite diffusion at

        $$
        B \;=\; B_1 \otimes P_2 \;+\; P_1 \otimes B_2,
        \qquad B_i = L_i Q_{c,i} L_i^\top,
        $$

        which is **not** $B_1 \otimes B_2$ — the value a naive
        $L_1 \otimes L_2$, $Q_{c,1} \otimes Q_{c,2}$ pair would imply.
        Reporting the latter used to hand out a tuple that failed its own
        Lyapunov equation (gh-219); for a Matérn ⊗ Cosine product it was
        identically zero, since `CosineSDE` has $Q_c = 0$.

        The sum of two Kronecker products is still expressible in the
        ``(L, Q_c)`` form, by widening the noise dimension and carrying
        each factor's stationary covariance in the spectral density:

        $$
        L = \bigl[\, L_1 \otimes I_{d_2} \;\;\big|\;\; I_{d_1} \otimes L_2 \,\bigr],
        \qquad
        Q_c = \operatorname{blockdiag}\!\bigl(
            Q_{c,1} \otimes P_2,\; P_1 \otimes Q_{c,2}
        \bigr),
        $$

        so that $L Q_c L^\top$ telescopes to exactly the $B$ above by the
        mixed-product property. Writing it this way rather than through
        square roots $S_i S_i^\top = P_i$ keeps the result exact for
        singular or zero $P_\infty$ (where a Cholesky would need jitter,
        and would then violate the very Lyapunov equation this enforces)
        and keeps ``sde_params`` reverse-mode differentiable.

        Note:
            ``SDEParams`` currently types its fields as dense
            ``jaxtyping.Float[Array, ...]``. The Kronecker products
            below are dense materializations of size
            ``(state_dim, state_dim)``, where ``state_dim`` is
            ``kernel1.state_dim * kernel2.state_dim`` — for typical SSM
            kernels (Matérn-3/2, periodic) this is ≤ 32, so the
            materialization is bounded and cheap. A future refactor
            could expose a parallel ``sde_operators()`` method that
            returns `gaussx.Kronecker` operators for downstream
            filters that can exploit the structure (issue #153).

        Raises:
            NotImplementedError: If either factor lacks a stationary
                covariance. The composite diffusion needs both, so the
                tuple cannot be built — and reporting the inconsistent
                Kronecker product instead is what gh-219 was about.
        """
        p1 = self.kernel1.sde_params()
        p2 = self.kernel2.sde_params()

        d1 = self.kernel1.state_dim
        d2 = self.kernel2.state_dim

        if p1.P_inf is None or p2.P_inf is None:
            msg = (
                f"ProductSDE cannot report SDE parameters when a factor has "
                f"no stationary covariance: "
                f"{type(self.kernel1).__name__}.P_inf is "
                f"{'None' if p1.P_inf is None else 'set'} and "
                f"{type(self.kernel2).__name__}.P_inf is "
                f"{'None' if p2.P_inf is None else 'set'}. The composite "
                f"diffusion of a product kernel is B1 (x) P2 + P1 (x) B2, "
                f"which needs both. Use the factors' own parameters, or "
                f"give the factor a P_inf."
            )
            raise NotImplementedError(msg)

        # Identities carry the factor dtypes: an untyped ``jnp.eye`` is
        # float64 under x64 and would promote float32 kernels.
        eye1 = jnp.eye(d1, dtype=p1.F.dtype)
        eye2 = jnp.eye(d2, dtype=p2.F.dtype)

        F = jnp.kron(p1.F, eye2) + jnp.kron(eye1, p2.F)
        H = jnp.kron(p1.H, p2.H)
        P_inf = jnp.kron(p1.P_inf, p2.P_inf)

        # B = B1 (x) P2 + P1 (x) B2, kept in the (L, Q_c) pair by putting
        # each factor's P_inf in the *spectral density* rather than taking
        # its square root:
        #
        #     B1 (x) P2 = (L1 (x) I) (Q_c1 (x) P2) (L1 (x) I)^T
        #     P1 (x) B2 = (I (x) L2) (P1 (x) Q_c2) (I (x) L2)^T
        #
        # by the mixed-product property. Exact for any PSD P_inf --
        # including singular or zero ones, where a Cholesky would need
        # jitter and stop satisfying the Lyapunov equation this is here to
        # enforce. The noise dimension widens from s1*s2 to s1*d2 + d1*s2.
        L = jnp.concatenate(
            [jnp.kron(p1.L, eye2), jnp.kron(eye1, p2.L)],
            axis=1,
        )
        Q_c = jsl.block_diag(
            jnp.kron(p1.Q_c, p2.P_inf),
            jnp.kron(p1.P_inf, p2.Q_c),
        )

        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

    def discretise(
        self,
        dt: Float[Array, ""],
    ) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
        r"""Discretise via the Kronecker matrix-exponential identity.

        For a product kernel ``F = F_1 \oplus F_2 = F_1 \otimes I + I \otimes F_2``,
        the factors ``F_1 \otimes I`` and ``I \otimes F_2`` commute, so

        $$
        \exp(F \, dt) = \exp(F_1 \, dt) \otimes \exp(F_2 \, dt).
        $$

        This computes two ``expm`` calls of size ``d_1`` and ``d_2``
        each, plus one Kronecker product, instead of one ``expm`` of
        size ``d_1 \cdot d_2``. Numerically equivalent to the dense
        ``expm`` on ``F`` but cheaper for moderate factor sizes.

        ``Q = P_\infty - A P_\infty A^T`` exploits the same factorisation.
        By the mixed-product property,

        $$
        (A_1 \otimes A_2)(P_1 \otimes P_2)(A_1 \otimes A_2)^\top
            = (A_1 P_1 A_1^\top) \otimes (A_2 P_2 A_2^\top),
        $$

        so the congruence is evaluated per factor via
        `gaussx.process_noise_covariance` on `gaussx.Kronecker`
        operands — ``O(d_1^3 + d_2^3)`` instead of the
        ``O((d_1 d_2)^3)`` triple product on the full matrix. Only the
        final ``Q`` is materialised, to keep the consumer-facing
        ``(A, Q)`` interface unchanged.

        Args:
            dt: Time step (scalar, positive).

        Returns:
            Tuple ``(A, Q)`` matching `SDEKernel.discretise`.
        """
        p1 = self.kernel1.sde_params()
        p2 = self.kernel2.sde_params()
        A1 = jsl.expm(p1.F * dt)
        A2 = jsl.expm(p2.F * dt)
        A = jnp.kron(A1, A2)

        # Use the per-factor stationary covariances directly; building
        # the full ``F`` via ``self.sde_params()`` would defeat the
        # whole point of this override. Keeping both operands as
        # ``Kronecker`` lets the shared helper contract each factor
        # separately instead of forming the (d1 d2)-square triple product.
        if p1.P_inf is None or p2.P_inf is None:
            # Deferring to the base MFD path here would be *wrong*, not
            # merely slower: the composite diffusion
            #
            #     B = B_1 (x) P_2  +  P_1 (x) B_2
            #
            # needs both factor stationary covariances -- exactly what is
            # missing -- so MFD has nothing correct to consume. Checked
            # against the factors directly rather than via
            # ``self.sde_params()`` (which now raises for the same reason)
            # so the message names the offending factor, and so the
            # happy path never builds the full-size drift it would
            # otherwise have to discard.
            msg = (
                f"ProductSDE cannot be discretised when a factor has no "
                f"stationary covariance: "
                f"{type(self.kernel1).__name__}.P_inf is "
                f"{'None' if p1.P_inf is None else 'set'} and "
                f"{type(self.kernel2).__name__}.P_inf is "
                f"{'None' if p2.P_inf is None else 'set'}. The composite "
                f"diffusion of a product kernel is B1 (x) P2 + P1 (x) B2, "
                f"which needs both. Discretise the factors separately, or "
                f"give the factor a P_inf."
            )
            raise NotImplementedError(msg)

        A_op = Kronecker(
            lx.MatrixLinearOperator(A1),
            lx.MatrixLinearOperator(A2),
        )
        P_op = Kronecker(
            lx.MatrixLinearOperator(p1.P_inf, lx.symmetric_tag),
            lx.MatrixLinearOperator(p2.P_inf, lx.symmetric_tag),
        )
        Q = symmetrize(process_noise_covariance(A_op, P_op).as_matrix())
        return A, Q

stationary: bool property

Stationary only if both factors are.

sde_params() -> SDEParams

Return Kronecker-structured SDE parameters.

The drift is the Kronecker sum \(F_1 \oplus F_2\) and the stationary covariance the Kronecker product \(P_1 \otimes P_2\). Substituting those into the Lyapunov equation \(F P + P F^\top + B = 0\) fixes the composite diffusion at

\[ B \;=\; B_1 \otimes P_2 \;+\; P_1 \otimes B_2, \qquad B_i = L_i Q_{c,i} L_i^\top, \]

which is not \(B_1 \otimes B_2\) — the value a naive \(L_1 \otimes L_2\), \(Q_{c,1} \otimes Q_{c,2}\) pair would imply. Reporting the latter used to hand out a tuple that failed its own Lyapunov equation (gh-219); for a Matérn ⊗ Cosine product it was identically zero, since CosineSDE has \(Q_c = 0\).

The sum of two Kronecker products is still expressible in the (L, Q_c) form, by widening the noise dimension and carrying each factor's stationary covariance in the spectral density:

\[ L = \bigl[\, L_1 \otimes I_{d_2} \;\;\big|\;\; I_{d_1} \otimes L_2 \,\bigr], \qquad Q_c = \operatorname{blockdiag}\!\bigl( Q_{c,1} \otimes P_2,\; P_1 \otimes Q_{c,2} \bigr), \]

so that \(L Q_c L^\top\) telescopes to exactly the \(B\) above by the mixed-product property. Writing it this way rather than through square roots \(S_i S_i^\top = P_i\) keeps the result exact for singular or zero \(P_\infty\) (where a Cholesky would need jitter, and would then violate the very Lyapunov equation this enforces) and keeps sde_params reverse-mode differentiable.

Note

SDEParams currently types its fields as dense jaxtyping.Float[Array, ...]. The Kronecker products below are dense materializations of size (state_dim, state_dim), where state_dim is kernel1.state_dim * kernel2.state_dim — for typical SSM kernels (Matérn-3/2, periodic) this is ≤ 32, so the materialization is bounded and cheap. A future refactor could expose a parallel sde_operators() method that returns gaussx.Kronecker operators for downstream filters that can exploit the structure (issue #153).

Raises:

Type	Description
NotImplementedError	If either factor lacks a stationary covariance. The composite diffusion needs both, so the tuple cannot be built — and reporting the inconsistent Kronecker product instead is what gh-219 was about.

Source code in src/gaussx/_ssm/_composition.py

def sde_params(self) -> SDEParams:
    r"""Return Kronecker-structured SDE parameters.

    The drift is the Kronecker **sum** $F_1 \oplus F_2$ and the
    stationary covariance the Kronecker **product**
    $P_1 \otimes P_2$. Substituting those into the Lyapunov equation
    $F P + P F^\top + B = 0$ fixes the composite diffusion at

    $$
    B \;=\; B_1 \otimes P_2 \;+\; P_1 \otimes B_2,
    \qquad B_i = L_i Q_{c,i} L_i^\top,
    $$

    which is **not** $B_1 \otimes B_2$ — the value a naive
    $L_1 \otimes L_2$, $Q_{c,1} \otimes Q_{c,2}$ pair would imply.
    Reporting the latter used to hand out a tuple that failed its own
    Lyapunov equation (gh-219); for a Matérn ⊗ Cosine product it was
    identically zero, since `CosineSDE` has $Q_c = 0$.

    The sum of two Kronecker products is still expressible in the
    ``(L, Q_c)`` form, by widening the noise dimension and carrying
    each factor's stationary covariance in the spectral density:

    $$
    L = \bigl[\, L_1 \otimes I_{d_2} \;\;\big|\;\; I_{d_1} \otimes L_2 \,\bigr],
    \qquad
    Q_c = \operatorname{blockdiag}\!\bigl(
        Q_{c,1} \otimes P_2,\; P_1 \otimes Q_{c,2}
    \bigr),
    $$

    so that $L Q_c L^\top$ telescopes to exactly the $B$ above by the
    mixed-product property. Writing it this way rather than through
    square roots $S_i S_i^\top = P_i$ keeps the result exact for
    singular or zero $P_\infty$ (where a Cholesky would need jitter,
    and would then violate the very Lyapunov equation this enforces)
    and keeps ``sde_params`` reverse-mode differentiable.

    Note:
        ``SDEParams`` currently types its fields as dense
        ``jaxtyping.Float[Array, ...]``. The Kronecker products
        below are dense materializations of size
        ``(state_dim, state_dim)``, where ``state_dim`` is
        ``kernel1.state_dim * kernel2.state_dim`` — for typical SSM
        kernels (Matérn-3/2, periodic) this is ≤ 32, so the
        materialization is bounded and cheap. A future refactor
        could expose a parallel ``sde_operators()`` method that
        returns `gaussx.Kronecker` operators for downstream
        filters that can exploit the structure (issue #153).

    Raises:
        NotImplementedError: If either factor lacks a stationary
            covariance. The composite diffusion needs both, so the
            tuple cannot be built — and reporting the inconsistent
            Kronecker product instead is what gh-219 was about.
    """
    p1 = self.kernel1.sde_params()
    p2 = self.kernel2.sde_params()

    d1 = self.kernel1.state_dim
    d2 = self.kernel2.state_dim

    if p1.P_inf is None or p2.P_inf is None:
        msg = (
            f"ProductSDE cannot report SDE parameters when a factor has "
            f"no stationary covariance: "
            f"{type(self.kernel1).__name__}.P_inf is "
            f"{'None' if p1.P_inf is None else 'set'} and "
            f"{type(self.kernel2).__name__}.P_inf is "
            f"{'None' if p2.P_inf is None else 'set'}. The composite "
            f"diffusion of a product kernel is B1 (x) P2 + P1 (x) B2, "
            f"which needs both. Use the factors' own parameters, or "
            f"give the factor a P_inf."
        )
        raise NotImplementedError(msg)

    # Identities carry the factor dtypes: an untyped ``jnp.eye`` is
    # float64 under x64 and would promote float32 kernels.
    eye1 = jnp.eye(d1, dtype=p1.F.dtype)
    eye2 = jnp.eye(d2, dtype=p2.F.dtype)

    F = jnp.kron(p1.F, eye2) + jnp.kron(eye1, p2.F)
    H = jnp.kron(p1.H, p2.H)
    P_inf = jnp.kron(p1.P_inf, p2.P_inf)

    # B = B1 (x) P2 + P1 (x) B2, kept in the (L, Q_c) pair by putting
    # each factor's P_inf in the *spectral density* rather than taking
    # its square root:
    #
    #     B1 (x) P2 = (L1 (x) I) (Q_c1 (x) P2) (L1 (x) I)^T
    #     P1 (x) B2 = (I (x) L2) (P1 (x) Q_c2) (I (x) L2)^T
    #
    # by the mixed-product property. Exact for any PSD P_inf --
    # including singular or zero ones, where a Cholesky would need
    # jitter and stop satisfying the Lyapunov equation this is here to
    # enforce. The noise dimension widens from s1*s2 to s1*d2 + d1*s2.
    L = jnp.concatenate(
        [jnp.kron(p1.L, eye2), jnp.kron(eye1, p2.L)],
        axis=1,
    )
    Q_c = jsl.block_diag(
        jnp.kron(p1.Q_c, p2.P_inf),
        jnp.kron(p1.P_inf, p2.Q_c),
    )

    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

discretise(dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Discretise via the Kronecker matrix-exponential identity.

For a product kernel F = F_1 \oplus F_2 = F_1 \otimes I + I \otimes F_2, the factors F_1 \otimes I and I \otimes F_2 commute, so

\[ \exp(F \, dt) = \exp(F_1 \, dt) \otimes \exp(F_2 \, dt). \]

This computes two expm calls of size d_1 and d_2 each, plus one Kronecker product, instead of one expm of size d_1 \cdot d_2. Numerically equivalent to the dense expm on F but cheaper for moderate factor sizes.

Q = P_\infty - A P_\infty A^T exploits the same factorisation. By the mixed-product property,

\[ (A_1 \otimes A_2)(P_1 \otimes P_2)(A_1 \otimes A_2)^\top = (A_1 P_1 A_1^\top) \otimes (A_2 P_2 A_2^\top), \]

so the congruence is evaluated per factor via gaussx.process_noise_covariance on gaussx.Kronecker operands — O(d_1^3 + d_2^3) instead of the O((d_1 d_2)^3) triple product on the full matrix. Only the final Q is materialised, to keep the consumer-facing (A, Q) interface unchanged.

Parameters:

Name	Type	Description	Default
dt	Float[Array, '']	Time step (scalar, positive).	required

Returns:

Type	Description
tuple[Float[Array, 'd d'], Float[Array, 'd d']]	Tuple (A, Q) matching SDEKernel.discretise.

Source code in src/gaussx/_ssm/_composition.py

def discretise(
    self,
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    r"""Discretise via the Kronecker matrix-exponential identity.

    For a product kernel ``F = F_1 \oplus F_2 = F_1 \otimes I + I \otimes F_2``,
    the factors ``F_1 \otimes I`` and ``I \otimes F_2`` commute, so

    $$
    \exp(F \, dt) = \exp(F_1 \, dt) \otimes \exp(F_2 \, dt).
    $$

    This computes two ``expm`` calls of size ``d_1`` and ``d_2``
    each, plus one Kronecker product, instead of one ``expm`` of
    size ``d_1 \cdot d_2``. Numerically equivalent to the dense
    ``expm`` on ``F`` but cheaper for moderate factor sizes.

    ``Q = P_\infty - A P_\infty A^T`` exploits the same factorisation.
    By the mixed-product property,

    $$
    (A_1 \otimes A_2)(P_1 \otimes P_2)(A_1 \otimes A_2)^\top
        = (A_1 P_1 A_1^\top) \otimes (A_2 P_2 A_2^\top),
    $$

    so the congruence is evaluated per factor via
    `gaussx.process_noise_covariance` on `gaussx.Kronecker`
    operands — ``O(d_1^3 + d_2^3)`` instead of the
    ``O((d_1 d_2)^3)`` triple product on the full matrix. Only the
    final ``Q`` is materialised, to keep the consumer-facing
    ``(A, Q)`` interface unchanged.

    Args:
        dt: Time step (scalar, positive).

    Returns:
        Tuple ``(A, Q)`` matching `SDEKernel.discretise`.
    """
    p1 = self.kernel1.sde_params()
    p2 = self.kernel2.sde_params()
    A1 = jsl.expm(p1.F * dt)
    A2 = jsl.expm(p2.F * dt)
    A = jnp.kron(A1, A2)

    # Use the per-factor stationary covariances directly; building
    # the full ``F`` via ``self.sde_params()`` would defeat the
    # whole point of this override. Keeping both operands as
    # ``Kronecker`` lets the shared helper contract each factor
    # separately instead of forming the (d1 d2)-square triple product.
    if p1.P_inf is None or p2.P_inf is None:
        # Deferring to the base MFD path here would be *wrong*, not
        # merely slower: the composite diffusion
        #
        #     B = B_1 (x) P_2  +  P_1 (x) B_2
        #
        # needs both factor stationary covariances -- exactly what is
        # missing -- so MFD has nothing correct to consume. Checked
        # against the factors directly rather than via
        # ``self.sde_params()`` (which now raises for the same reason)
        # so the message names the offending factor, and so the
        # happy path never builds the full-size drift it would
        # otherwise have to discard.
        msg = (
            f"ProductSDE cannot be discretised when a factor has no "
            f"stationary covariance: "
            f"{type(self.kernel1).__name__}.P_inf is "
            f"{'None' if p1.P_inf is None else 'set'} and "
            f"{type(self.kernel2).__name__}.P_inf is "
            f"{'None' if p2.P_inf is None else 'set'}. The composite "
            f"diffusion of a product kernel is B1 (x) P2 + P1 (x) B2, "
            f"which needs both. Discretise the factors separately, or "
            f"give the factor a P_inf."
        )
        raise NotImplementedError(msg)

    A_op = Kronecker(
        lx.MatrixLinearOperator(A1),
        lx.MatrixLinearOperator(A2),
    )
    P_op = Kronecker(
        lx.MatrixLinearOperator(p1.P_inf, lx.symmetric_tag),
        lx.MatrixLinearOperator(p2.P_inf, lx.symmetric_tag),
    )
    Q = symmetrize(process_noise_covariance(A_op, P_op).as_matrix())
    return A, Q

SumSDE

Bases: SDEKernel

Sum of SDE kernels via block-diagonal composition.

Attributes:

Name	Type	Description
kernels	tuple[SDEKernel, ...]	Tuple of component SDE kernels.

Source code in src/gaussx/_ssm/_composition.py

class SumSDE(SDEKernel):
    """Sum of SDE kernels via block-diagonal composition.

    Attributes:
        kernels: Tuple of component SDE kernels.
    """

    kernels: tuple[SDEKernel, ...] = eqx.field()

    @property
    def state_dim(self) -> int:
        return sum(k.state_dim for k in self.kernels)

    @property
    def stationary(self) -> bool:
        """Stationary only if every component is."""
        return all(k.stationary for k in self.kernels)

    def initial_covariance(self) -> Float[Array, "d d"]:
        """Return the block-diagonal initial covariance.

        The components are independent, so their initial covariances
        stack block-diagonally — which lets a sum mix stationary and
        non-stationary components (a local linear trend plus a Matern
        seasonal, say), each started from its own.
        """
        return jsl.block_diag(*[k.initial_covariance() for k in self.kernels])

    def sde_params(self) -> SDEParams:
        """Return block-diagonal SDE parameters."""
        params_list = [k.sde_params() for k in self.kernels]

        F = jsl.block_diag(*[p.F for p in params_list])
        # A component with no closed-form stationary covariance leaves the
        # sum without one either; propagating ``None`` routes the composite
        # through ``discretise_mfd`` rather than fabricating a ``P_inf``.
        component_p_inf = [p.P_inf for p in params_list]
        P_inf = (
            None
            if any(block is None for block in component_p_inf)
            else jsl.block_diag(*component_p_inf)
        )

        L_blocks = [p.L for p in params_list]
        total_rows = sum(b.shape[0] for b in L_blocks)
        total_cols = sum(b.shape[1] for b in L_blocks)
        L = jnp.zeros((total_rows, total_cols))
        row_offset = 0
        col_offset = 0
        for block in L_blocks:
            r, c = block.shape
            L = L.at[row_offset : row_offset + r, col_offset : col_offset + c].set(
                block
            )
            row_offset += r
            col_offset += c

        Q_c = jsl.block_diag(*[p.Q_c for p in params_list])
        H = jnp.concatenate([p.H for p in params_list], axis=1)

        return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

stationary: bool property

Stationary only if every component is.

initial_covariance() -> Float[Array, 'd d']

Return the block-diagonal initial covariance.

The components are independent, so their initial covariances stack block-diagonally — which lets a sum mix stationary and non-stationary components (a local linear trend plus a Matern seasonal, say), each started from its own.

Source code in src/gaussx/_ssm/_composition.py

def initial_covariance(self) -> Float[Array, "d d"]:
    """Return the block-diagonal initial covariance.

    The components are independent, so their initial covariances
    stack block-diagonally — which lets a sum mix stationary and
    non-stationary components (a local linear trend plus a Matern
    seasonal, say), each started from its own.
    """
    return jsl.block_diag(*[k.initial_covariance() for k in self.kernels])

sde_params() -> SDEParams

Return block-diagonal SDE parameters.

Source code in src/gaussx/_ssm/_composition.py

def sde_params(self) -> SDEParams:
    """Return block-diagonal SDE parameters."""
    params_list = [k.sde_params() for k in self.kernels]

    F = jsl.block_diag(*[p.F for p in params_list])
    # A component with no closed-form stationary covariance leaves the
    # sum without one either; propagating ``None`` routes the composite
    # through ``discretise_mfd`` rather than fabricating a ``P_inf``.
    component_p_inf = [p.P_inf for p in params_list]
    P_inf = (
        None
        if any(block is None for block in component_p_inf)
        else jsl.block_diag(*component_p_inf)
    )

    L_blocks = [p.L for p in params_list]
    total_rows = sum(b.shape[0] for b in L_blocks)
    total_cols = sum(b.shape[1] for b in L_blocks)
    L = jnp.zeros((total_rows, total_cols))
    row_offset = 0
    col_offset = 0
    for block in L_blocks:
        r, c = block.shape
        L = L.at[row_offset : row_offset + r, col_offset : col_offset + c].set(
            block
        )
        row_offset += r
        col_offset += c

    Q_c = jsl.block_diag(*[p.Q_c for p in params_list])
    H = jnp.concatenate([p.H for p in params_list], axis=1)

    return SDEParams(F=F, L=L, H=H, Q_c=Q_c, P_inf=P_inf)

sde_autocovariance(kernel: SDEKernel, tau: Float[Array, ' *batch']) -> Float[Array, ' *batch']

Compute the stationary autocovariance of an SDE kernel.

Evaluates:

K(\tau) = H \, \exp(F |\tau|) \, P_\infty \, H^T

Parameters:

Name	Type	Description	Default
kernel	SDEKernel	An SDE kernel with sde_params() method.	required
tau	Float[Array, ' *batch']	Lag values, shape (*batch,).	required

Returns:

Type	Description
Float[Array, ' *batch']	Autocovariance values K(tau), shape (*batch,).

Source code in src/gaussx/_ssm/_autocovariance.py

def sde_autocovariance(
    kernel: SDEKernel,
    tau: Float[Array, " *batch"],
) -> Float[Array, " *batch"]:
    r"""Compute the stationary autocovariance of an SDE kernel.

    Evaluates:

        K(\tau) = H \, \exp(F |\tau|) \, P_\infty \, H^T

    Args:
        kernel: An SDE kernel with ``sde_params()`` method.
        tau: Lag values, shape ``(*batch,)``.

    Returns:
        Autocovariance values ``K(tau)``, shape ``(*batch,)``.
    """
    if not kernel.stationary:
        msg = (
            f"{type(kernel).__name__} is not stationary, so it has no "
            f"autocovariance K(tau): its covariance depends on both times "
            f"rather than on their difference. Evaluate the covariance "
            f"through the discretised recursion instead."
        )
        raise ValueError(msg)

    params = kernel.sde_params()
    if params.P_inf is None:
        msg = (
            f"{type(kernel).__name__} does not supply a stationary "
            f"covariance (sde_params().P_inf is None), which this "
            f"autocovariance is defined in terms of. Note P_inf=None means "
            f"'no closed form', not necessarily 'not stationary': a Hurwitz "
            f"drift does have a stationary covariance, recoverable as the "
            f"solution of F P + P F^T + L Q_c L^T = 0, but that Lyapunov "
            f"solve is exactly the fragile step gaussx.discretise_mfd "
            f"exists to avoid, so it is not done implicitly here. Supply "
            f"P_inf on the kernel if you have it; to discretise rather than "
            f"evaluate the autocovariance, use gaussx.discretise_mfd, which "
            f"needs no P_inf."
        )
        raise ValueError(msg)

    def _single_autocov(t: Float[Array, ""]) -> Float[Array, ""]:
        abs_t = jnp.abs(t)
        eF = jsl.expm(params.F * abs_t)
        cov_matrix = params.H @ eF @ params.P_inf @ params.H.T
        return cov_matrix.squeeze()

    orig_shape = tau.shape
    flat_tau = tau.ravel()
    flat_result = jax.vmap(_single_autocov)(flat_tau)
    return flat_result.reshape(orig_shape)

Stationary and non-stationary kernels

Most of the zoo is stationary: the process is assumed started in its stationary distribution, so SDEKernel.initial_covariance returns \(P_\infty\) and sde_autocovariance can report \(K(\tau)\).

IntegratedWienerSDE is not. Its drift is nilpotent, the marginal variance grows without bound, and the covariance depends on both times rather than on their difference — so there is no \(P_\infty\) and no \(K(\tau)\). It reports stationary = False and sde_params().P_inf = None, and the filter is started instead from an explicit initial_covariance (diffuse by default). That is the local linear trend prior at order=1: smooth, and linear unless the data push back, with no lengthscale to choose.

Consumers should branch on SDEKernel.stationary rather than on whether P_inf is None — the two are different questions, since a stationary kernel may have no closed form for \(P_\infty\) (a learned drift, say). A SumSDE may mix the two: it is stationary only if all its components are, and stacks their initial covariances block-diagonally. A ProductSDE needs both factors' \(P_\infty\) and rejects a non-stationary one.

Discretisation

Turning the continuous-time SDE into \(x_k = A x_{k-1} + q_k\) takes two routes, and which one applies depends on whether a stationary covariance exists.

The default, SDEKernel.discretise, uses the stationary route \(Q = P_\infty - A P_\infty A^\top\) (process_noise_covariance, documented under Process noise below). It is exact and cheap, and every stationary kernel in the zoo above supplies the \(P_\infty\) it needs. IntegratedWienerSDE overrides discretise with an exact closed form instead — its nilpotent drift makes the exponential terminate — so it needs neither route.

discretise_mfd is the fallback for when that covariance is not available — most importantly when \(F\) is a learned parameter rather than derived from a kernel. Recovering \(P_\infty\) then means solving the Lyapunov equation \(F P + P F^\top + Q_c = 0\), which has a unique solution only if \(\lambda_i(F) + \lambda_j(F) \neq 0\) for every pair of eigenvalues. That condition fails for any undamped oscillatory mode, where \(\lambda = \pm i\omega\) and so \(\lambda + \bar\lambda = 0\) identically — CosineSDE has exactly that drift, and sidesteps it with a closed form that a learned \(F\) cannot use.

Matrix-fraction decomposition needs no \(P_\infty\) at all: it recovers both \(A\) and \(Q\) from a single \(2d \times 2d\) matrix exponential (Van Loan 1978) and is well defined for every \(F\). Reach for it when the drift is fitted; keep the stationary route otherwise. Note the obstruction is degeneracy of that Sylvester system, not instability — constraining \(F\) to be Hurwitz would not fix it, and would forbid the oscillatory modes MFD exists to support.

Structured linear algebra and Gaussian primitives for JAX.

discretise_mfd(F: Float[Array, 'd d'], Q_c: Float[Array, 'd d'], dt: Float[Array, '']) -> tuple[Float[Array, 'd d'], Float[Array, 'd d']]

Discretise a linear SDE by matrix-fraction decomposition.

Returns \((A, Q)\) for the discrete-time model \(x_k = A x_{k-1} + q_k\) with \(q_k \sim N(0, Q)\), where

\[ A = e^{F\,\Delta t}, \qquad Q = \int_0^{\Delta t} e^{Fs}\, Q_c\, e^{F^\top s}\, ds . \]

Both come from one \(2d \times 2d\) matrix exponential (Van Loan 1978). For the augmented generator

\[ \Phi = \begin{bmatrix} F & Q_c \\ 0 & -F^\top \end{bmatrix}, \qquad e^{\Phi \Delta t} = \begin{bmatrix} A & C \\ 0 & D \end{bmatrix}, \]

the integral is \(Q = C D^{-1}\), and \(D = e^{-F^\top \Delta t}\) is always invertible with \(D^{-1} = A^\top\) — so no inverse or solve is formed here.

Unlike the stationary route in gaussx.SDEKernel.discretise, this needs no \(P_\infty\) and is well defined for every \(F\), including a learned \(F\) whose eigenvalues sum to zero in pairs. The stationary route recovers \(P_\infty\) from the Lyapunov equation \(F P + P F^\top + Q_c = 0\), which has a unique solution only when \(\lambda_i(F) + \lambda_j(F) \neq 0\) for all \(i, j\). An undamped oscillatory mode has \(\lambda = \pm i\omega\), so \(\lambda + \bar\lambda = 0\) always and the Lyapunov route degenerates.

Note

The obstruction is degeneracy of that Sylvester system, not instability. The identity \(Q = P_\infty - A P_\infty A^\top\) holds for unstable \(F\), and even when the Lyapunov solution is not PSD and so is not a valid covariance. Constraining \(F\) to be Hurwitz would not fix the degeneracy, and would forbid exactly the oscillatory modes this function exists to support.

Parameters:

Name	Type	Description	Default
F	Float[Array, 'd d']	Continuous-time drift matrix, shape (d, d).	required
Q_c	Float[Array, 'd d']	Continuous-time diffusion covariance \(L Q_c L^\top\), shape (d, d).	required
dt	Float[Array, '']	Time step. Must be non-negative; checked with equinox.error_if, so under jit the error fires at evaluation rather than trace time.	required

Returns:

Type	Description
Float[Array, 'd d']	Tuple (A, Q), both shape (d, d). Q is symmetrised —
Float[Array, 'd d']	C @ A.T is not symmetric to floating point.

Source code in src/gaussx/_ssm/_discretise.py

def discretise_mfd(
    F: Float[Array, "d d"],
    Q_c: Float[Array, "d d"],
    dt: Float[Array, ""],
) -> tuple[Float[Array, "d d"], Float[Array, "d d"]]:
    r"""Discretise a linear SDE by matrix-fraction decomposition.

    Returns $(A, Q)$ for the discrete-time model $x_k = A x_{k-1} + q_k$
    with $q_k \sim N(0, Q)$, where

    $$
    A = e^{F\,\Delta t}, \qquad
    Q = \int_0^{\Delta t} e^{Fs}\, Q_c\, e^{F^\top s}\, ds .
    $$

    Both come from one $2d \times 2d$ matrix exponential (Van Loan 1978).
    For the augmented generator

    $$
    \Phi = \begin{bmatrix} F & Q_c \\ 0 & -F^\top \end{bmatrix},
    \qquad
    e^{\Phi \Delta t}
        = \begin{bmatrix} A & C \\ 0 & D \end{bmatrix},
    $$

    the integral is $Q = C D^{-1}$, and $D = e^{-F^\top \Delta t}$ is
    always invertible with $D^{-1} = A^\top$ — so no inverse or solve is
    formed here.

    Unlike the stationary route in `gaussx.SDEKernel.discretise`, this
    needs no $P_\infty$ and is well defined for **every** $F$, including a
    learned $F$ whose eigenvalues sum to zero in pairs. The stationary
    route recovers $P_\infty$ from the Lyapunov equation
    $F P + P F^\top + Q_c = 0$, which has a unique solution only when
    $\lambda_i(F) + \lambda_j(F) \neq 0$ for all $i, j$. An undamped
    oscillatory mode has $\lambda = \pm i\omega$, so $\lambda + \bar\lambda
    = 0$ always and the Lyapunov route degenerates.

    Note:
        The obstruction is degeneracy of that Sylvester system, not
        instability. The identity $Q = P_\infty - A P_\infty A^\top$ holds
        for unstable $F$, and even when the Lyapunov solution is not PSD
        and so is not a valid covariance. Constraining $F$ to be Hurwitz
        would not fix the degeneracy, and would forbid exactly the
        oscillatory modes this function exists to support.

    Args:
        F: Continuous-time drift matrix, shape ``(d, d)``.
        Q_c: Continuous-time diffusion covariance $L Q_c L^\top$, shape
            ``(d, d)``.
        dt: Time step. Must be non-negative; checked with
            `equinox.error_if`, so under ``jit`` the error fires at
            evaluation rather than trace time.

    Returns:
        Tuple ``(A, Q)``, both shape ``(d, d)``. ``Q`` is symmetrised —
        ``C @ A.T`` is not symmetric to floating point.
    """
    # A negative step would silently run the exponential backwards and
    # return a Q that is not a covariance. error_if defers the check to
    # evaluation time so this stays traceable under jit.
    dt = eqx.error_if(dt, dt < 0, "discretise_mfd requires dt >= 0.")

    # How many doublings are needed to keep the augmented exponential in
    # range. The 1-norm bounds the spectral radius, so this is
    # conservative. Traced, but only ever used as a *predicate* below --
    # never as a trip count -- so the whole function stays reverse-mode
    # differentiable, unlike a lax.while_loop.
    drift_scale = jnp.abs(F).sum(axis=0).max() * dt
    required = jnp.ceil(jnp.log2(jnp.maximum(drift_scale / _MFD_EXPONENT_BUDGET, 1.0)))
    n_squarings = jnp.clip(required, 0.0, _MFD_MAX_SQUARINGS)

    # Beyond the cap the scaled step is still too stiff and the augmented
    # exponential would overflow anyway, so say so rather than returning a
    # silent NaN. In float32 this starts to bite around ||F|| dt ~ 1e6.
    dt = eqx.error_if(
        dt,
        required > _MFD_MAX_SQUARINGS,
        f"discretise_mfd: ||F|| * dt is too large to discretise in one step "
        f"(more than {_MFD_MAX_SQUARINGS} doublings would be needed). Split "
        f"the interval into shorter steps, or rescale time.",
    )

    # Q is *linear* in Q_c, so the diffusion's magnitude can be divided out
    # of the exponential and multiplied back afterwards. Without this the
    # augmented generator inherits ||Q_c|| in its off-diagonal block and
    # overflows for a large diffusion even when the drift is benign and the
    # answer trivial -- F = 0 with Q_c = 1e6 already returns NaN. Scaling
    # the *step* would not help there and would wrongly reject it: the
    # growth is linear in Q_c, not exponential.
    # Scale by a power of two, and only when the magnitude actually
    # threatens the exponential. A power of two is exact in binary floating
    # point, so it introduces no rounding of its own, and leaving small
    # diffusions alone matters: normalising unconditionally by the largest
    # entry pushes smaller modes down by the same factor, and one that
    # underflows is lost outright rather than merely rounded.
    #
    # A diffusion whose dynamic range exceeds the dtype's own cannot
    # survive any uniform scaling -- in float32, diag(1e30, 1e-20) has no
    # scaling that keeps both representable. That is a limit of the format,
    # not of the scheme.
    #
    # The budget applies to what the exponential actually sees, which is
    # Q_c scaled by the *substep*: `_van_loan` exponentiates Phi * dt, so a
    # modest diffusion over a long step is just as dangerous as a large one
    # over a short step. Q_c = 1e3 with dt = 1e5 presents the same 1e8
    # off-diagonal magnitude that NaNs at Q_c = 1e8, dt = 1.
    substep = dt / 2.0**n_squarings
    effective_diffusion = jnp.abs(Q_c).max() * substep
    diffusion_scale = jnp.where(
        effective_diffusion > _MFD_DIFFUSION_BUDGET,
        2.0 ** jnp.ceil(jnp.log2(effective_diffusion / _MFD_DIFFUSION_BUDGET)),
        1.0,
    )

    A, Q = _van_loan(F, Q_c / diffusion_scale, substep)

    # Compose back up. The unroll is static, but each doubling is placed
    # behind a lax.cond rather than a select so an inactive step is not
    # evaluated at all -- a well-scaled drift, which is the common case,
    # pays nothing here. lax.cond is used over a while_loop because it
    # stays reverse-mode differentiable.
    #
    # Under vmap (as in discretise_mfd_sequence) a cond with a batched
    # predicate lowers back to a select, so batched callers do pay for the
    # inactive branches; that is a JAX limitation, not an oversight.
    def _double(operands):
        A_i, Q_i = operands
        return A_i @ A_i, symmetrize(A_i @ Q_i @ A_i.T + Q_i)

    for i in range(_MFD_MAX_SQUARINGS):
        A, Q = jax.lax.cond(i < n_squarings, _double, lambda operands: operands, (A, Q))

    return A, Q * diffusion_scale

discretise_mfd_sequence(F: Float[Array, 'd d'], Q_c: Float[Array, 'd d'], dt: Float[Array, ' N']) -> tuple[Float[Array, 'N d d'], Float[Array, 'N d d']]

Vectorised gaussx.discretise_mfd over a vector of time steps.

Parameters:

Name	Type	Description	Default
F	Float[Array, 'd d']	Continuous-time drift matrix, shape (d, d).	required
Q_c	Float[Array, 'd d']	Continuous-time diffusion covariance, shape (d, d).	required
dt	Float[Array, ' N']	Time steps, shape (N,). All must be non-negative.	required

Returns:

Type	Description
tuple[Float[Array, 'N d d'], Float[Array, 'N d d']]	Tuple (A_seq, Q_seq), both shape (N, d, d).

Source code in src/gaussx/_ssm/_discretise.py

def discretise_mfd_sequence(
    F: Float[Array, "d d"],
    Q_c: Float[Array, "d d"],
    dt: Float[Array, " N"],
) -> tuple[Float[Array, "N d d"], Float[Array, "N d d"]]:
    """Vectorised `gaussx.discretise_mfd` over a vector of time steps.

    Args:
        F: Continuous-time drift matrix, shape ``(d, d)``.
        Q_c: Continuous-time diffusion covariance, shape ``(d, d)``.
        dt: Time steps, shape ``(N,)``. All must be non-negative.

    Returns:
        Tuple ``(A_seq, Q_seq)``, both shape ``(N, d, d)``.
    """
    return jax.vmap(lambda step: discretise_mfd(F, Q_c, step))(dt)

Nonlinear filters

nonlinear_kalman_filter and nonlinear_rts_smoother take callables rather than matrices, and moment-match them through an integrator. The choice of integrator is the choice of filter — one loop yields the whole textbook family:

integrator	filter
TaylorIntegrator	extended Kalman filter (EKF)
UnscentedIntegrator	unscented Kalman filter (UKF)
CubatureIntegrator	cubature Kalman filter (CKF), degree 3
FifthOrderCubatureIntegrator	degree-5 cubature filter
GaussHermiteIntegrator	Gauss-Hermite Kalman filter (GHKF)
MonteCarloIntegrator	Monte-Carlo Kalman filter

The gain is \(K = C S^{-1}\), built from the integrator's cross-covariance \(C\), so no Jacobian is formed anywhere — that is what makes the EKF and the UKF the same code. An integrator that does not supply a cross-covariance is rejected rather than silently producing a zero gain.

Two behaviours differ from kalman_filter, both deliberately:

• log_likelihood is a moment-matched surrogate, not the exact marginal likelihood, because \(S\) is the matched innovation covariance rather than the true one. It coincides with the exact value when the maps are affine. Maximising it to tune hyperparameters is standard practice for nonlinear Gaussian filters, but it is a surrogate.
• The covariance update defaults to Joseph form (joseph=True). \(K = C S^{-1}\) is only approximately the optimal gain, and \(P^- - K S K^\top\) is guaranteed PSD only for the optimal gain, whereas Joseph form is a sum of two PSD terms for any \(K\). The effective observation matrix it needs is the statistical-linearisation gain \(H_{\text{eff}} = C^\top (P^-)^{-1}\) — what statistical_linear_regression returns as A — and its noise is \(R + \Omega\), not \(R\), with \(\Omega\) the linearisation residual.

With that residual included the two forms are analytically identical for any consistent matched joint, not merely for affine maps. joseph is therefore a numerical choice, not a modelling one: it selects how the same covariance is computed, and should not change results beyond floating point. Dropping \(\Omega\) — the naive reading of Joseph here — would instead understate the posterior by \(K \Omega K^\top\).

With affine dynamics and obs_fn the filter reproduces kalman_filter — means, covariances and log-likelihood — and the smoother reproduces rts_smoother, for every deterministic rule, since each is exact for affine maps. MonteCarloIntegrator is the exception: it propagates finite-sample empirical moments, so it converges to the linear filter at the usual \(O(1/\sqrt{n})\) rate rather than matching it. A discrepancy there is sampling error, not a bug.

Driving your own loop

The wrappers above are a jax.lax.scan over three public per-step functions, which are usable on their own when the loop is the part you want to control — an irregular time grid, a custom gating rule, a filter interleaved with something else, or a bank of filters:

function	step
nonlinear_kalman_predict	\(m^-, P^- = \mathcal{T}[f](m, P)\), \(P^- \mathrel{+}= Q\)
nonlinear_kalman_update	match \(h\), then \(K = C S^{-1}\) and the Gaussian update
nonlinear_rts_step	one RTS backward correction

Each takes the same integrator, so the choice of filter carries through unchanged. nonlinear_kalman_update returns its log-likelihood increment rather than accumulating, leaving the accumulation to the caller.

Structured linear algebra and Gaussian primitives for JAX.

nonlinear_kalman_filter(dynamics: Callable[[Float[Array, ' N']], Float[Array, ' N']], obs_fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator, obs_noise: Float[Array, '*T M M'] | lx.AbstractLinearOperator, observations: Float[Array, 'T M'], init_mean: Float[Array, ' N'], init_cov: Float[Array, 'N N'], *, integrator: AbstractIntegrator | None = None, mask: Bool[Array, ' T'] | Bool[Array, 'T M'] | None = None, joseph: bool = True, solver: AbstractSolverStrategy | None = None) -> FilterState

Moment-matched nonlinear Kalman filter.

Propagates a Gaussian belief through nonlinear dynamics and obs_fn by moment matching, using integrator for both the predict and the update step. The choice of integrator is the choice of filter:

integrator	filter
gaussx.TaylorIntegrator	extended Kalman filter (EKF)
gaussx.UnscentedIntegrator	unscented Kalman filter (UKF)
gaussx.CubatureIntegrator	cubature Kalman filter (CKF), \(2N\) pts
gaussx.FifthOrderCubatureIntegrator	degree-5 cubature, \(2N^2+1\) pts
gaussx.GaussHermiteIntegrator	Gauss-Hermite Kalman filter (GHKF)
gaussx.MonteCarloIntegrator	Monte-Carlo Kalman filter

Each step is

\[ \begin{aligned} m^-, P^- &= \mathcal{T}[f](m, P), \quad P^- \mathrel{+}= Q, \\ \hat y, S_{yy}, C &= \mathcal{T}[h](m^-, P^-), \quad S = S_{yy} + R, \\ K &= C S^{-1}, \\ m^+ &= m^- + K(y - \hat y), \end{aligned} \]

where \(\mathcal{T}\) is the integrator's moment transform. The gain comes from the integrator's cross-covariance, so no Jacobian appears anywhere — that is what makes the EKF and the UKF the same code.

Note

log_likelihood is a moment-matched surrogate, not the exact marginal likelihood: \(S\) is the matched innovation covariance rather than the true one. It reduces to the exact value when dynamics and obs_fn are affine. Users maximising it to tune hyperparameters are maximising a surrogate, which is standard practice for nonlinear Gaussian filters but worth knowing.

Note

Unlike gaussx.kalman_filter, the covariance update defaults to Joseph form. \(K = C S^{-1}\) is only approximately the optimal gain, and \(P^- - K S K^\top\) is guaranteed PSD only for the optimal gain, whereas Joseph form is a sum of two PSD terms for any \(K\).

Joseph form needs an \(H\), which a moment-matched filter does not have; the stand-in is the statistical-linearisation gain \(H_{\text{eff}} = C^\top (P^-)^{-1}\) — what gaussx.statistical_linear_regression returns as A, and exactly \(H\) when obs_fn is linear. Its noise is \(R + \Omega\), with \(\Omega\) the linearisation residual, not \(R\): dropping \(\Omega\) would understate the posterior covariance by \(K \Omega K^\top\) on nonlinear maps. With it included the two forms agree analytically, so joseph selects how the same covariance is computed, not which covariance you get.

Parameters:

Name	Type	Description	Default
dynamics	Callable[[Float[Array, ' N']], Float[Array, ' N']]	State transition (N,) -> (N,). Deterministic; process noise is added separately via process_noise.	required
obs_fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Observation operator (N,) -> (M,).	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator	\(Q\), additive in state space. Shape (N, N), (T, N, N), or an operator (materialised once).	required
obs_noise	Float[Array, '*T M M'] | AbstractLinearOperator	\(R\), additive in observation space. Shape (M, M), (T, M, M), or an operator.	required
observations	Float[Array, 'T M']	Observed data, shape (T, M).	required
init_mean	Float[Array, ' N']	Initial state mean, shape (N,).	required
init_cov	Float[Array, 'N N']	Initial state covariance, shape (N, N).	required
integrator	AbstractIntegrator | None	Moment-matching rule. Defaults to UnscentedIntegrator(alpha=1.0), which is derivative-free and exact for affine maps. Must supply a cross-covariance. The alpha matters: gaussx.UnscentedIntegrator's own default of 1e-3 places the sigma points ~1e-3 from the mean and recovers the moments by cancellation, which costs roughly seven digits. That is invisible in float64 but ruinous in float32 — JAX's default — where it misplaces the log-likelihood of a linear problem by over one nat. Pass alpha=1e-3 explicitly only if you want the classic scaled transform and are running in x64.	None
mask	Bool[Array, ' T'] | Bool[Array, 'T M'] | None	Optional observation mask, with the same semantics as gaussx.kalman_filter — (T,) gates whole steps, (T, M) gates individual channels. Masked entries of observations are never read, so they may be NaN.	None
joseph	bool	Use the Joseph-form covariance update. Defaults to True; see Notes.	True
solver	AbstractSolverStrategy | None	Optional solver strategy for the innovation solve. When None, uses structural dispatch.	None

Returns:

Type	Description
FilterState	A gaussx.FilterState, identical in shape to
FilterState	gaussx.kalman_filter's output.

Raises:

Type	Description
TypeError	If integrator does not supply a cross-covariance (raised at trace time).
ValueError	If mask or the noise covariances are misshapen.

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def nonlinear_kalman_filter(
    dynamics: Callable[[Float[Array, " N"]], Float[Array, " N"]],
    obs_fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    obs_noise: Float[Array, "*T M M"] | lx.AbstractLinearOperator,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " N"],
    init_cov: Float[Array, "N N"],
    *,
    integrator: AbstractIntegrator | None = None,
    mask: Bool[Array, " T"] | Bool[Array, "T M"] | None = None,
    joseph: bool = True,
    solver: AbstractSolverStrategy | None = None,
) -> FilterState:
    r"""Moment-matched nonlinear Kalman filter.

    Propagates a Gaussian belief through nonlinear ``dynamics`` and
    ``obs_fn`` by moment matching, using ``integrator`` for both the
    predict and the update step. **The choice of integrator is the choice
    of filter:**

    | integrator | filter |
    |---|---|
    | `gaussx.TaylorIntegrator` | extended Kalman filter (EKF) |
    | `gaussx.UnscentedIntegrator` | unscented Kalman filter (UKF) |
    | `gaussx.CubatureIntegrator` | cubature Kalman filter (CKF), $2N$ pts |
    | `gaussx.FifthOrderCubatureIntegrator` | degree-5 cubature, $2N^2+1$ pts |
    | `gaussx.GaussHermiteIntegrator` | Gauss-Hermite Kalman filter (GHKF) |
    | `gaussx.MonteCarloIntegrator` | Monte-Carlo Kalman filter |

    Each step is

    $$
    \begin{aligned}
    m^-, P^- &= \mathcal{T}[f](m, P), \quad P^- \mathrel{+}= Q, \\
    \hat y, S_{yy}, C &= \mathcal{T}[h](m^-, P^-), \quad S = S_{yy} + R, \\
    K &= C S^{-1}, \\
    m^+ &= m^- + K(y - \hat y),
    \end{aligned}
    $$

    where $\mathcal{T}$ is the integrator's moment transform. The gain
    comes from the integrator's cross-covariance, so **no Jacobian appears
    anywhere** — that is what makes the EKF and the UKF the same code.

    Note:
        ``log_likelihood`` is a **moment-matched surrogate**, not the exact
        marginal likelihood: $S$ is the matched innovation covariance
        rather than the true one. It reduces to the exact value when
        ``dynamics`` and ``obs_fn`` are affine. Users maximising it to tune
        hyperparameters are maximising a surrogate, which is standard
        practice for nonlinear Gaussian filters but worth knowing.

    Note:
        Unlike `gaussx.kalman_filter`, the covariance update defaults to
        Joseph form. $K = C S^{-1}$ is only approximately the optimal gain,
        and $P^- - K S K^\top$ is guaranteed PSD only *for* the optimal
        gain, whereas Joseph form is a sum of two PSD terms for any $K$.

        Joseph form needs an $H$, which a moment-matched filter does not
        have; the stand-in is the statistical-linearisation gain
        $H_{\text{eff}} = C^\top (P^-)^{-1}$ — what
        `gaussx.statistical_linear_regression` returns as ``A``, and
        exactly $H$ when ``obs_fn`` is linear. Its noise is $R + \Omega$,
        with $\Omega$ the linearisation residual, **not** $R$: dropping
        $\Omega$ would understate the posterior covariance by
        $K \Omega K^\top$ on nonlinear maps. With it included the two
        forms agree analytically, so ``joseph`` selects how the same
        covariance is computed, not which covariance you get.

    Args:
        dynamics: State transition ``(N,) -> (N,)``. Deterministic; process
            noise is added separately via ``process_noise``.
        obs_fn: Observation operator ``(N,) -> (M,)``.
        process_noise: $Q$, additive in state space. Shape ``(N, N)``,
            ``(T, N, N)``, or an operator (materialised once).
        obs_noise: $R$, additive in observation space. Shape ``(M, M)``,
            ``(T, M, M)``, or an operator.
        observations: Observed data, shape ``(T, M)``.
        init_mean: Initial state mean, shape ``(N,)``.
        init_cov: Initial state covariance, shape ``(N, N)``.
        integrator: Moment-matching rule. Defaults to
            ``UnscentedIntegrator(alpha=1.0)``, which is derivative-free
            and exact for affine maps. Must supply a cross-covariance.

            The ``alpha`` matters: `gaussx.UnscentedIntegrator`'s own
            default of ``1e-3`` places the sigma points ~1e-3 from the mean
            and recovers the moments by cancellation, which costs roughly
            seven digits. That is invisible in float64 but ruinous in
            float32 — JAX's default — where it misplaces the
            log-likelihood of a *linear* problem by over one nat. Pass
            ``alpha=1e-3`` explicitly only if you want the classic scaled
            transform and are running in x64.
        mask: Optional observation mask, with the same semantics as
            `gaussx.kalman_filter` — ``(T,)`` gates whole steps, ``(T, M)``
            gates individual channels. Masked entries of ``observations``
            are never read, so they may be ``NaN``.
        joseph: Use the Joseph-form covariance update. Defaults to
            ``True``; see Notes.
        solver: Optional solver strategy for the innovation solve. When
            ``None``, uses structural dispatch.

    Returns:
        A `gaussx.FilterState`, identical in shape to
        `gaussx.kalman_filter`'s output.

    Raises:
        TypeError: If ``integrator`` does not supply a cross-covariance
            (raised at trace time).
        ValueError: If ``mask`` or the noise covariances are misshapen.
    """
    if integrator is None:
        integrator = UnscentedIntegrator(alpha=1.0)

    T, M = observations.shape

    Q_seq = _broadcast_noise(process_noise, T, init_mean.shape[-1], "process_noise")
    R_seq = _broadcast_noise(obs_noise, T, M, "obs_noise")
    mask_seq = _normalise_mask(mask, T, M)
    channel_mask = mask_seq.ndim == 2

    def step(carry, inputs):
        mean, cov, ll = carry
        Q_t, R_t, y_t, mask_t = inputs

        # The loop is exactly `predict` then `update`; both are public, so
        # a caller who wants a different loop can use them directly.
        mean_pred, cov_pred = nonlinear_kalman_predict(
            dynamics, mean, cov, Q_t, integrator=integrator
        )

        def _update(_):
            return nonlinear_kalman_update(
                obs_fn,
                mean_pred,
                cov_pred,
                y_t,
                R_t,
                integrator=integrator,
                mask=mask_t if channel_mask else None,
                joseph=joseph,
                solver=solver,
            )

        def _skip(_):
            # Gated-off step: keep the prediction and contribute no
            # likelihood. Filtered == predicted here, which is also what
            # makes the smoother's gain degenerate harmlessly at this step.
            return mean_pred, cov_pred, jnp.zeros((), dtype=cov_pred.dtype)

        if channel_mask:
            # No lax.cond needed: an all-False row already reduces the
            # update to the identity via the substitutions inside
            # `nonlinear_kalman_update`, so this path is branch-free.
            mean_new, cov_new, ll_inc = _update(None)
        else:
            # Gate the whole step so the predict-only branch evaluates
            # neither the update arithmetic nor its gradients.
            mean_new, cov_new, ll_inc = jax.lax.cond(
                mask_t, _update, _skip, operand=None
            )

        carry_new = (mean_new, cov_new, ll + ll_inc)
        return carry_new, (mean_new, cov_new, mean_pred, cov_pred)

    init_carry = (init_mean, init_cov, jnp.zeros((), dtype=init_cov.dtype))
    final_carry, (f_means, f_covs, p_means, p_covs) = jax.lax.scan(
        step, init_carry, (Q_seq, R_seq, observations, mask_seq)
    )

    return FilterState(
        filtered_means=f_means,
        filtered_covs=f_covs,
        predicted_means=p_means,
        predicted_covs=p_covs,
        log_likelihood=final_carry[2],
    )

nonlinear_rts_smoother(filter_state: FilterState, dynamics: Callable[[Float[Array, ' N']], Float[Array, ' N']], process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator | None = None, *, integrator: AbstractIntegrator | None = None, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, 'T N'], Float[Array, 'T N N']]

Moment-matched nonlinear Rauch-Tung-Striebel smoother.

The backward pass of gaussx.nonlinear_kalman_filter, sharing the same moment transform — the smoother gain uses the integrator's cross-covariance between \(x_t\) and \(x_{t+1}\):

\[ G_t = \mathrm{Cov}[x_t, f(x_t)]\, (P^-_{t+1})^{-1}, \qquad m^s_t = m_t + G_t (m^s_{t+1} - m^-_{t+1}), \]

with the matching covariance recursion. As in the filter, no Jacobian is formed; for linear dynamics the gain reduces to \(P_t A^\top (P^-_{t+1})^{-1}\) and the whole pass to gaussx.rts_smoother.

Parameters:

Name	Type	Description	Default
filter_state	FilterState	Output of gaussx.nonlinear_kalman_filter. Pass the same dynamics and integrator used to produce it.	required
dynamics	Callable[[Float[Array, ' N']], Float[Array, ' N']]	State transition (N,) -> (N,).	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator | None	Accepted for API symmetry with gaussx.rts_smoother and unused — the predicted covariances in filter_state already include it.	None
integrator	AbstractIntegrator | None	Moment-matching rule. Defaults to UnscentedIntegrator(alpha=1.0); use the one the filter used.	None
solver	AbstractSolverStrategy | None	Accepted for API symmetry with gaussx.rts_smoother and unused -- see gaussx.nonlinear_rts_step.	None

Returns:

Type	Description
tuple[Float[Array, 'T N'], Float[Array, 'T N N']]	Tuple (smoothed_means, smoothed_covs).

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def nonlinear_rts_smoother(
    filter_state: FilterState,
    dynamics: Callable[[Float[Array, " N"]], Float[Array, " N"]],
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator | None = None,
    *,
    integrator: AbstractIntegrator | None = None,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "T N"], Float[Array, "T N N"]]:
    r"""Moment-matched nonlinear Rauch-Tung-Striebel smoother.

    The backward pass of `gaussx.nonlinear_kalman_filter`, sharing the same
    moment transform — the smoother gain uses the integrator's
    cross-covariance between $x_t$ and $x_{t+1}$:

    $$
    G_t = \mathrm{Cov}[x_t, f(x_t)]\, (P^-_{t+1})^{-1},
    \qquad
    m^s_t = m_t + G_t (m^s_{t+1} - m^-_{t+1}),
    $$

    with the matching covariance recursion. As in the filter, no Jacobian
    is formed; for linear ``dynamics`` the gain reduces to
    $P_t A^\top (P^-_{t+1})^{-1}$ and the whole pass to
    `gaussx.rts_smoother`.

    Args:
        filter_state: Output of `gaussx.nonlinear_kalman_filter`. Pass the
            same ``dynamics`` and ``integrator`` used to produce it.
        dynamics: State transition ``(N,) -> (N,)``.
        process_noise: Accepted for API symmetry with
            `gaussx.rts_smoother` and unused — the predicted covariances in
            ``filter_state`` already include it.
        integrator: Moment-matching rule. Defaults to
            ``UnscentedIntegrator(alpha=1.0)``; use the one the filter
            used.
        solver: Accepted for API symmetry with `gaussx.rts_smoother` and
            unused -- see `gaussx.nonlinear_rts_step`.

    Returns:
        Tuple ``(smoothed_means, smoothed_covs)``.
    """
    del process_noise  # predicted covariances already include it

    if integrator is None:
        integrator = UnscentedIntegrator(alpha=1.0)

    T = filter_state.filtered_means.shape[0]

    def step(carry, inputs):
        mean_smooth, cov_smooth = carry
        mean_filt, cov_filt, mean_pred, cov_pred = inputs

        mean_new, cov_new = nonlinear_rts_step(
            dynamics,
            mean_filt,
            cov_filt,
            mean_pred,
            cov_pred,
            mean_smooth,
            cov_smooth,
            integrator=integrator,
            solver=solver,
        )

        return (mean_new, cov_new), (mean_new, cov_new)

    # The backward pass is seeded at the final step, where smoothed and
    # filtered coincide because there is no future left to condition on.
    init_carry = (
        filter_state.filtered_means[T - 1],
        filter_state.filtered_covs[T - 1],
    )

    # Step t consumes the filtered belief at t and the *predicted* belief at
    # t+1, hence the offset slices; both are reversed so the scan runs
    # backwards through time.
    inputs = (
        filter_state.filtered_means[:-1][::-1],
        filter_state.filtered_covs[:-1][::-1],
        filter_state.predicted_means[1:][::-1],
        filter_state.predicted_covs[1:][::-1],
    )

    _, (s_means_rev, s_covs_rev) = jax.lax.scan(step, init_carry, inputs)

    # Undo the reversal and re-attach the final step, which the scan never
    # produced because it was the seed.
    s_means = jnp.concatenate(
        [s_means_rev[::-1], filter_state.filtered_means[T - 1 :]], axis=0
    )
    s_covs = jnp.concatenate(
        [s_covs_rev[::-1], filter_state.filtered_covs[T - 1 :]], axis=0
    )
    return s_means, s_covs

nonlinear_kalman_predict(dynamics: Callable[[Float[Array, ' N']], Float[Array, ' N']], mean: Float[Array, ' N'], cov: Float[Array, 'N N'], process_noise: Float[Array, 'N N'], *, integrator: AbstractIntegrator | None = None) -> tuple[Float[Array, ' N'], Float[Array, 'N N']]

One moment-matched predict step.

\[ m^- = \mathbb{E}[f(x)], \qquad P^- = \mathrm{Cov}[f(x)] + Q . \]

Exposed alongside gaussx.nonlinear_kalman_update so a caller can drive their own loop -- an irregular time grid, a custom gating rule, a filter interleaved with something else -- without reimplementing the moment transform. gaussx.nonlinear_kalman_filter is exactly a jax.lax.scan over these two.

Parameters:

Name	Type	Description	Default
dynamics	Callable[[Float[Array, ' N']], Float[Array, ' N']]	State transition (N,) -> (N,). Deterministic.	required
mean	Float[Array, ' N']	Current mean, shape (N,).	required
cov	Float[Array, 'N N']	Current covariance, shape (N, N).	required
process_noise	Float[Array, 'N N']	\(Q\), shape (N, N).	required
integrator	AbstractIntegrator | None	Moment-matching rule. Defaults to UnscentedIntegrator(alpha=1.0) — see gaussx.nonlinear_kalman_filter on why not alpha=1e-3.	None

Returns:

Type	Description
tuple[Float[Array, ' N'], Float[Array, 'N N']]	Tuple (mean_pred, cov_pred).

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def nonlinear_kalman_predict(
    dynamics: Callable[[Float[Array, " N"]], Float[Array, " N"]],
    mean: Float[Array, " N"],
    cov: Float[Array, "N N"],
    process_noise: Float[Array, "N N"],
    *,
    integrator: AbstractIntegrator | None = None,
) -> tuple[Float[Array, " N"], Float[Array, "N N"]]:
    r"""One moment-matched predict step.

    $$
    m^- = \mathbb{E}[f(x)], \qquad P^- = \mathrm{Cov}[f(x)] + Q .
    $$

    Exposed alongside `gaussx.nonlinear_kalman_update` so a caller can
    drive their own loop -- an irregular time grid, a custom gating rule,
    a filter interleaved with something else -- without reimplementing the
    moment transform. `gaussx.nonlinear_kalman_filter` is exactly a
    `jax.lax.scan` over these two.

    Args:
        dynamics: State transition ``(N,) -> (N,)``. Deterministic.
        mean: Current mean, shape ``(N,)``.
        cov: Current covariance, shape ``(N, N)``.
        process_noise: $Q$, shape ``(N, N)``.
        integrator: Moment-matching rule. Defaults to
            ``UnscentedIntegrator(alpha=1.0)`` — see
            `gaussx.nonlinear_kalman_filter` on why not ``alpha=1e-3``.

    Returns:
        Tuple ``(mean_pred, cov_pred)``.
    """
    if integrator is None:
        integrator = UnscentedIntegrator(alpha=1.0)

    process_noise = _check_noise_shape(process_noise, mean.shape[-1], "process_noise")

    # The process noise is additive and independent of x, so it enters only
    # as an additive term on the covariance -- the moment transform sees
    # the *deterministic* dynamics alone. No cross-covariance is needed
    # here (nothing is being conditioned on yet); the smoother re-runs the
    # same transform precisely to recover it.
    mean_pred, cov_dyn, _ = moment_transform(dynamics, mean, cov, integrator=integrator)
    cov_pred = symmetrize(cov_dyn + process_noise)

    # Validate here as well as after the update. A negative-weight rule can
    # return an indefinite Cov[f(x)] that a small process noise does not
    # repair, and a step-level False mask would then expose it directly as
    # the filtered covariance. Nothing downstream would complain: the next
    # moment transform tags it PSD, and the dense square-root path clips
    # negative eigenvalues to zero, silently altering the belief rather
    # than reporting it.
    cov_pred = _reject_indefinite(
        cov_pred,
        "nonlinear_kalman_predict: the predicted covariance is not positive "
        "semi-definite. Cov[f(x)] came back indefinite, which a "
        "negative-weight quadrature rule can produce, and process_noise did "
        "not repair it. Use a positive-weight rule such as "
        "CubatureIntegrator or UnscentedIntegrator(alpha=1.0).",
    )
    return mean_pred, cov_pred

nonlinear_kalman_update(obs_fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], mean: Float[Array, ' N'], cov: Float[Array, 'N N'], observation: Float[Array, ' M'], obs_noise: Float[Array, 'M M'], *, integrator: AbstractIntegrator | None = None, mask: Bool[Array, ' M'] | None = None, joseph: bool = True, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' N'], Float[Array, 'N N'], Float[Array, '']]

One moment-matched update step.

Moment-matches obs_fn at the predicted belief and runs the ordinary linear-Gaussian update against that matched joint:

\[ \hat y, S_{yy}, C = \mathcal{T}[h](m^-, P^-), \quad S = S_{yy} + R, \quad K = C S^{-1}, \]

\[ m^+ = m^- + K(y - \hat y), \qquad \ell = -\tfrac{1}{2}\big(v^\top S^{-1} v + \log|S| + M \log 2\pi\big). \]

The gain is built from the cross-covariance directly, so no Jacobian appears. See gaussx.nonlinear_kalman_filter for the meaning of joseph and the caveat on the returned log-likelihood.

Parameters:

Name	Type	Description	Default
obs_fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Observation operator (N,) -> (M,).	required
mean	Float[Array, ' N']	Predicted mean \(m^-\), shape (N,).	required
cov	Float[Array, 'N N']	Predicted covariance \(P^-\), shape (N, N).	required
observation	Float[Array, ' M']	Observed vector \(y\), shape (M,).	required
obs_noise	Float[Array, 'M M']	\(R\), shape (M, M).	required
integrator	AbstractIntegrator | None	Moment-matching rule. Defaults to UnscentedIntegrator(alpha=1.0) — see gaussx.nonlinear_kalman_filter on why not alpha=1e-3.	None
mask	Bool[Array, ' M'] | None	Optional per-channel mask, shape (M,). False entries are marginalised out exactly and may be NaN in observation.	None
joseph	bool	Use the Joseph-form covariance update. Defaults to True.	True
solver	AbstractSolverStrategy | None	Optional solver strategy.	None

Returns:

Type	Description
Float[Array, ' N']	Tuple (mean_upd, cov_upd, log_likelihood_increment). The
Float[Array, 'N N']	increment is the exact marginal over the observed channels.

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def nonlinear_kalman_update(
    obs_fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    mean: Float[Array, " N"],
    cov: Float[Array, "N N"],
    observation: Float[Array, " M"],
    obs_noise: Float[Array, "M M"],
    *,
    integrator: AbstractIntegrator | None = None,
    mask: Bool[Array, " M"] | None = None,
    joseph: bool = True,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " N"], Float[Array, "N N"], Float[Array, ""]]:
    r"""One moment-matched update step.

    Moment-matches ``obs_fn`` at the predicted belief and runs the ordinary
    linear-Gaussian update against that matched joint:

    $$
    \hat y, S_{yy}, C = \mathcal{T}[h](m^-, P^-), \quad S = S_{yy} + R,
    \quad K = C S^{-1},
    $$

    $$
    m^+ = m^- + K(y - \hat y), \qquad
    \ell = -\tfrac{1}{2}\big(v^\top S^{-1} v + \log|S| + M \log 2\pi\big).
    $$

    The gain is built from the cross-covariance directly, so no Jacobian
    appears. See `gaussx.nonlinear_kalman_filter` for the meaning of
    ``joseph`` and the caveat on the returned log-likelihood.

    Args:
        obs_fn: Observation operator ``(N,) -> (M,)``.
        mean: Predicted mean $m^-$, shape ``(N,)``.
        cov: Predicted covariance $P^-$, shape ``(N, N)``.
        observation: Observed vector $y$, shape ``(M,)``.
        obs_noise: $R$, shape ``(M, M)``.
        integrator: Moment-matching rule. Defaults to
            ``UnscentedIntegrator(alpha=1.0)`` — see
            `gaussx.nonlinear_kalman_filter` on why not ``alpha=1e-3``.
        mask: Optional per-channel mask, shape ``(M,)``. ``False`` entries
            are marginalised out exactly and may be ``NaN`` in
            ``observation``.
        joseph: Use the Joseph-form covariance update. Defaults to ``True``.
        solver: Optional solver strategy.

    Returns:
        Tuple ``(mean_upd, cov_upd, log_likelihood_increment)``. The
        increment is the exact marginal over the *observed* channels.
    """
    if integrator is None:
        integrator = UnscentedIntegrator(alpha=1.0)

    M = observation.shape[-1]
    obs_noise = _check_noise_shape(obs_noise, M, "obs_noise")
    if mask is not None and jnp.shape(mask) != (M,):
        # A (1,) mask would broadcast across every channel: False would
        # suppress them all and return an identity update, True would enable
        # them all, in either case silently.
        msg = f"mask must have shape ({M},); got {jnp.shape(mask)}."
        raise ValueError(msg)

    y_hat, obs_cov, cross = moment_transform(obs_fn, mean, cov, integrator=integrator)

    if mask is None:
        obs_cov_e, cross_e, R_e = obs_cov, cross, obs_noise
        residual = observation - y_hat  # v = y - y_hat
        n_missing = jnp.zeros((), dtype=cov.dtype)
    else:
        # Rewrite the matched moments so masked channels carry no
        # information -- see `gaussx.masked_moment_inputs`.
        obs_cov_e, cross_e, R_e, residual, n_missing = masked_moment_inputs(
            obs_cov, cross, obs_noise, observation, y_hat, mask
        )

    # S = S_yy + R. Symmetrised because it is assembled from a weighted
    # outer-product sum, which drifts asymmetric.
    #
    # Tagged symmetric rather than positive-semidefinite: S_yy comes from a
    # quadrature rule, and rules with negative weights (the scaled
    # unscented transform, the degree-5 cubature rule above N=4) can return
    # an indefinite Cov[h(x)]. Claiming PSD would route the solve to a
    # Cholesky path that returns NaN on such a matrix instead of a solver
    # that copes.
    innovation = symmetrize(obs_cov_e + R_e)

    # A quadrature rule with negative weights can return an indefinite
    # Cov[h(x)], and R may be too small to repair it. Neither the update
    # nor the likelihood is defined then -- the quadratic form can go
    # negative and the log-determinant becomes log|det S| -- so this is
    # rejected rather than allowed to produce a plausible-looking but
    # meaningless number. The eigendecomposition is on the (M, M)
    # innovation and is negligible beside the moment transform that
    # produced it.
    smallest_eigenvalue = jnp.linalg.eigvalsh(innovation).min()
    innovation = eqx.error_if(
        innovation,
        smallest_eigenvalue <= 0.0,
        "nonlinear_kalman_update: the innovation covariance S = Cov[h(x)] + R "
        "is not positive definite. A negative-weight quadrature rule (the "
        "scaled unscented transform, or the degree-5 cubature rule above "
        "N = 4) can return an indefinite Cov[h(x)]. Use a positive-weight "
        "rule such as CubatureIntegrator or UnscentedIntegrator(alpha=1.0), "
        "or increase obs_noise.",
    )
    innovation_op = _symmetric(innovation)

    # K = C S^-1. solve_rows solves S x = c for each *row* of C, i.e. it
    # forms C S^-1 without inverting S.
    gain = solve_rows(innovation_op, cross_e, solver=solver)  # (N, M)

    # m+ = m- + K v
    mean_upd = mean + gain @ residual

    if joseph:
        # Joseph form: P+ = (I - K H)P-(I - K H)^T + K R K^T.
        #
        # That needs an H, which a moment-matched filter does not have. The
        # right stand-in is the statistical linearisation of h under the
        # predicted belief (gaussx#161 section 3.3.5): writing
        # h(x) ~ A x + b + eps, the regression gain is
        #
        #     A = C^T (P-)^-1
        #
        # which is exactly what statistical_linear_regression returns, and
        # exactly H when h is linear. So this reduces to the textbook
        # Joseph update in the linear case, while staying PSD for the
        # merely-approximate gain otherwise.
        #
        # Relative to the standard form the two differ by K Omega K^T, with
        # Omega = S_yy - A P- A^T the linearisation residual: PSD, and zero
        # for affine h. Hence switching this default cannot perturb the
        # linear reduction.
        # H_eff = C^T (P^-)^-1, via a least-squares solve rather than
        # `solve_rows`. P^- is legitimately singular for a deterministic
        # initial state, zero process noise, or dimension-reducing
        # dynamics, and a well-posed solver returns NaN on those even
        # though the update itself is perfectly well defined (R keeps S
        # invertible). The pseudo-inverse gives the minimum-norm H_eff,
        # which is the natural reading of the linearisation when the
        # belief is confined to a subspace.
        #
        # rcond=0.0 is deliberate: it discards only exactly-zero singular
        # values. lstsq's default cutoff scales with the dtype's epsilon,
        # which in float32 also discards small-but-real covariance modes --
        # for P = diag(1, 1e-8) it zeroes H_eff along the second axis, and
        # the update then *grows* that variance instead of shrinking it.
        obs_eff = jnp.linalg.lstsq(cov, cross_e, rcond=0.0)[0].T

        # The noise of that regression is R + Omega, *not* R: linearising
        # h leaves a residual eps ~ N(0, Omega) on top of the measurement
        # noise, and Joseph form must be given the noise of the model whose
        # H it is using. Omega = S_yy - H_eff P- H_eff^T, so
        #
        #     R + Omega = S - H_eff C
        #
        # which is free here -- both factors are already formed.
        #
        # Passing R alone would return the matched-joint posterior minus
        # K Omega K^T, i.e. systematically overconfident on nonlinear maps.
        # With the residual included the two covariance forms agree to
        # 2.8e-17, so Joseph is a numerically safer route to the *same*
        # answer rather than a different one.
        effective_noise = symmetrize(innovation - obs_eff @ cross_e)
        cov_upd = joseph_update(cov, gain, obs_eff, effective_noise)
    else:
        # P+ = P- - K S K^T. Correct and cheaper, but its PSD-ness relies
        # on the moment triple being mutually consistent (Omega >= 0),
        # which a negative-weight rule can violate.
        cov_upd = symmetrize(cov - gain @ innovation @ gain.T)

    # A positive-definite S is not on its own enough: the *joint* over
    # (x, h(x)) must be consistent. An inconsistent triple -- Omega =
    # S_yy - H_eff P^- H_eff^T indefinite, which a negative-weight rule can
    # produce even where S is fine -- leaves the posterior indefinite.
    #
    # The full spectrum is checked, not just the diagonal: an indefinite
    # covariance can have entirely positive variances, e.g.
    # [[0.36, -0.64], [-0.64, 0.36]], whose smallest eigenvalue is -0.28.
    # A diagonal test would pass that and hand the next predict step a
    # covariance it will go on to treat as PSD.
    #
    # The threshold is scaled to the size of the covariance rather than
    # being exactly zero. A legitimately rank-deficient posterior -- a
    # deterministic state, zero process noise -- is singular by
    # construction, and rounding puts its null directions a few ulps either
    # side of zero; a strict test would reject those. Genuine
    # inconsistency is not marginal: the example above sits at -0.28
    # against a trace of 0.72, many orders above this bound.
    cov_upd = _reject_indefinite(
        cov_upd,
        "nonlinear_kalman_update: the updated covariance is not positive "
        "semi-definite. The matched moments (Cov[h(x)], Cov[x, h(x)]) are "
        "not a consistent joint, which a negative-weight quadrature rule "
        "can produce. Use a positive-weight rule such as CubatureIntegrator "
        "or UnscentedIntegrator(alpha=1.0).",
    )

    # ll += -0.5 (v^T S^-1 v + log|S| + M log 2pi).
    #
    # An approximation, not the exact marginal: S is the *matched*
    # innovation covariance. Exact when the maps are affine.
    solved = dispatch_solve(innovation_op, residual, solver)
    logdet = dispatch_logdet(innovation_op, solver)
    # Each masked channel contributed a dummy unit block to S, worth
    # -0.5 log(2 pi) of the full-vector density. Adding it back makes the
    # result the exact marginal over the observed entries, and independent
    # of the dummy block's variance.
    ll_inc = (
        -0.5 * (residual @ solved + logdet + M * _LOG_2PI) + 0.5 * n_missing * _LOG_2PI
    )
    return mean_upd, cov_upd, ll_inc

nonlinear_rts_step(dynamics: Callable[[Float[Array, ' N']], Float[Array, ' N']], mean_filtered: Float[Array, ' N'], cov_filtered: Float[Array, 'N N'], mean_predicted: Float[Array, ' N'], cov_predicted: Float[Array, 'N N'], mean_smoothed: Float[Array, ' N'], cov_smoothed: Float[Array, 'N N'], *, integrator: AbstractIntegrator | None = None, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' N'], Float[Array, 'N N']]

One moment-matched RTS backward step.

\[ G_t = \Sigma_{xx_+} (P^-_{t+1})^{-1}, \qquad m^s_t = m_t + G_t (m^s_{t+1} - m^-_{t+1}), \]

with the matching covariance recursion. Exposed for the same reason as the filter steps: gaussx.nonlinear_rts_smoother is a jax.lax.scan over this.

Parameters:

Name	Type	Description	Default
dynamics	Callable[[Float[Array, ' N']], Float[Array, ' N']]	State transition (N,) -> (N,).	required
mean_filtered	Float[Array, ' N']	Filtered mean at \(t\).	required
cov_filtered	Float[Array, 'N N']	Filtered covariance at \(t\).	required
mean_predicted	Float[Array, ' N']	Predicted mean at \(t + 1\).	required
cov_predicted	Float[Array, 'N N']	Predicted covariance at \(t + 1\).	required
mean_smoothed	Float[Array, ' N']	Smoothed mean at \(t + 1\).	required
cov_smoothed	Float[Array, 'N N']	Smoothed covariance at \(t + 1\).	required
integrator	AbstractIntegrator | None	Moment-matching rule; use the one the filter used.	None
solver	AbstractSolverStrategy | None	Accepted for API symmetry with gaussx.rts_smoother and unused. The smoother gain is taken with a least-squares solve so that a singular predicted covariance -- a deterministic or rank-deficient process -- still yields the correction defined on its supported subspace, which supersedes the strategy.	None

Returns:

Type	Description
tuple[Float[Array, ' N'], Float[Array, 'N N']]	Tuple (mean, cov) smoothed at \(t\).

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def nonlinear_rts_step(
    dynamics: Callable[[Float[Array, " N"]], Float[Array, " N"]],
    mean_filtered: Float[Array, " N"],
    cov_filtered: Float[Array, "N N"],
    mean_predicted: Float[Array, " N"],
    cov_predicted: Float[Array, "N N"],
    mean_smoothed: Float[Array, " N"],
    cov_smoothed: Float[Array, "N N"],
    *,
    integrator: AbstractIntegrator | None = None,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " N"], Float[Array, "N N"]]:
    r"""One moment-matched RTS backward step.

    $$
    G_t = \Sigma_{xx_+} (P^-_{t+1})^{-1}, \qquad
    m^s_t = m_t + G_t (m^s_{t+1} - m^-_{t+1}),
    $$

    with the matching covariance recursion. Exposed for the same reason as
    the filter steps: `gaussx.nonlinear_rts_smoother` is a
    `jax.lax.scan` over this.

    Args:
        dynamics: State transition ``(N,) -> (N,)``.
        mean_filtered: Filtered mean at $t$.
        cov_filtered: Filtered covariance at $t$.
        mean_predicted: Predicted mean at $t + 1$.
        cov_predicted: Predicted covariance at $t + 1$.
        mean_smoothed: Smoothed mean at $t + 1$.
        cov_smoothed: Smoothed covariance at $t + 1$.
        integrator: Moment-matching rule; use the one the filter used.
        solver: Accepted for API symmetry with `gaussx.rts_smoother` and
            unused. The smoother gain is taken with a least-squares solve
            so that a singular predicted covariance -- a deterministic or
            rank-deficient process -- still yields the correction defined
            on its supported subspace, which supersedes the strategy.

    Returns:
        Tuple ``(mean, cov)`` smoothed at $t$.
    """
    if integrator is None:
        integrator = UnscentedIntegrator(alpha=1.0)

    # Re-run the *same* moment transform the filter used on the dynamics,
    # at the filtered belief for step t. Only the third element of the
    # triple is wanted here:
    #
    #     Sigma_xx+ = Cov[x_t, x_{t+1}^-] = Cov[x_t, f(x_t)]
    #
    # The two are equal because the additive process noise is independent
    # of x_t, so it contributes nothing to the cross term -- which is also
    # why Q never appears here, and why the predicted covariance passed in
    # already accounts for it.
    _, _, cross = moment_transform(
        dynamics, mean_filtered, cov_filtered, integrator=integrator
    )

    # G = Sigma_xx+ (P-_{t+1})^-1. For linear f this is
    # P_t A^T (P-_{t+1})^-1, the textbook RTS gain.
    #
    # Least-squares rather than a well-posed solve, for the same reason as
    # the filter's Joseph linearisation: a deterministic or rank-deficient
    # process leaves P^- singular while the RTS correction stays perfectly
    # well defined on the subspace the belief actually occupies. A
    # well-posed solver returns NaN there, so a filter run that handles a
    # singular covariance would still fail the moment its result reached
    # the smoother. rcond=0.0 keeps every representable mode.
    del solver  # the rank policy below supersedes the solver strategy
    gain = jnp.linalg.lstsq(cov_predicted, cross.T, rcond=0.0)[0].T  # (N, N)

    # The RTS corrections: push the filtered belief toward the smoothed
    # future, by however much that future disagreed with what was predicted
    # from here.
    #
    #     m^s_t = m_t + G (m^s_{t+1} - m^-_{t+1})
    #     P^s_t = P_t + G (P^s_{t+1} - P^-_{t+1}) G^T
    #
    # Note P^s_{t+1} - P^-_{t+1} is negative semi-definite in the exact
    # case, which is what makes smoothed variances no larger than filtered
    # ones.
    mean_new = mean_filtered + gain @ (mean_smoothed - mean_predicted)
    cov_new = symmetrize(cov_filtered + gain @ (cov_smoothed - cov_predicted) @ gain.T)

    # Validated for the same reason predict and update are. The filtered and
    # predicted covariances can each be PSD while an inconsistent
    # cross-covariance still drives the correction indefinite -- in one
    # dimension P_f = P_pred = 1 with cross = 2 gives G = 2 and a smoothed
    # variance of -2.6.
    cov_new = _reject_indefinite(
        cov_new,
        "nonlinear_rts_step: the smoothed covariance is not positive "
        "semi-definite. The dynamics moment triple is not a consistent "
        "joint, which a negative-weight quadrature rule can produce. Use a "
        "positive-weight rule such as CubatureIntegrator or "
        "UnscentedIntegrator(alpha=1.0).",
    )
    return mean_new, cov_new

masked_moment_inputs(obs_cov: Float[Array, 'M M'], cross_cov: Float[Array, 'N M'], obs_noise: Float[Array, 'M M'], y: Float[Array, ' M'], y_hat: Float[Array, ' M'], mask: Bool[Array, ' M']) -> tuple[Float[Array, 'M M'], Float[Array, 'N M'], Float[Array, 'M M'], Float[Array, ' M'], Float[Array, '']]

Make masked observation channels inert in a moment-matched update.

Marginalising channel \(i\) out of a Kalman update is equivalent to keeping it but making it carry no information. For a linear filter that means zeroing row \(i\) of \(H\) and substituting a unit block into \(R\). A moment-matched filter has no \(H\) to zero, so the same substitution is applied to the matched moments instead:

• zero the masked rows and columns of \(\mathrm{Cov}[h(x)]\) (jointly, what zeroing a row of \(H\) would do to \(H P H^\top\)),
• zero the masked columns of \(\mathrm{Cov}[x, h(x)]\) (likewise for \(P H^\top\)),
• substitute a unit block into \(R\),
• zero the residual entry.

The innovation then splits as

\[ S = \begin{bmatrix} S_{\mathrm{obs}} & 0 \\ 0 & I \end{bmatrix}, \]

so the gain \(K = C S^{-1}\) has zero columns on the masked channels and they cannot move the state. The posterior is exactly the channel-deleted filter, with no branching — which is what lets the per-channel path run inside a jax.lax.scan without a cond.

Exposed because the substitution is reusable: any moment-matched update — a custom filter loop, a smoother variant, an ensemble method — needs the same rewrite to handle partially observed vectors, and getting the joint row/column masking subtly wrong is easy.

Note

Residuals are formed from separately-masked \(y\) and \(\hat y\) rather than by masking their difference, so a masked \(y\) entry may be NaN — the usual "not measured" encoding — without poisoning the reverse-mode gradient.

Parameters:

Name	Type	Description	Default
obs_cov	Float[Array, 'M M']	Matched \(\mathrm{Cov}[h(x)]\), shape (M, M).	required
cross_cov	Float[Array, 'N M']	Matched \(\mathrm{Cov}[x, h(x)]\), shape (N, M).	required
obs_noise	Float[Array, 'M M']	Observation noise \(R\), shape (M, M).	required
y	Float[Array, ' M']	Observation vector, shape (M,). Masked entries are never read and may be NaN.	required
y_hat	Float[Array, ' M']	Matched \(\mathbb{E}[h(x)]\), shape (M,).	required
mask	Bool[Array, ' M']	Per-channel mask, shape (M,). True keeps the channel.	required

Returns:

Type	Description
Float[Array, 'M M']	Tuple ``(obs_cov_eff, cross_cov_eff, obs_noise_eff, residual,
Float[Array, 'N M']	n_missing).n_missing`` is the float count of masked channels;
Float[Array, 'M M']	each contributed \(-\tfrac{1}{2}\log 2\pi\) of dummy density to the
Float[Array, ' M']	full-vector log-likelihood, so adding
Float[Array, '']	0.5 * n_missing * log(2 pi) back recovers the exact marginal
tuple[Float[Array, 'M M'], Float[Array, 'N M'], Float[Array, 'M M'], Float[Array, ' M'], Float[Array, '']]	over the observed entries.

Source code in src/gaussx/_ssm/_nonlinear_kalman.py

def masked_moment_inputs(
    obs_cov: Float[Array, "M M"],
    cross_cov: Float[Array, "N M"],
    obs_noise: Float[Array, "M M"],
    y: Float[Array, " M"],
    y_hat: Float[Array, " M"],
    mask: Bool[Array, " M"],
) -> tuple[
    Float[Array, "M M"],
    Float[Array, "N M"],
    Float[Array, "M M"],
    Float[Array, " M"],
    Float[Array, ""],
]:
    r"""Make masked observation channels inert in a moment-matched update.

    Marginalising channel $i$ out of a Kalman update is equivalent to
    keeping it but making it carry no information. For a linear filter that
    means zeroing row $i$ of $H$ and substituting a unit block into $R$. A
    moment-matched filter has no $H$ to zero, so the same substitution is
    applied to the **matched moments** instead:

    - zero the masked rows and columns of $\mathrm{Cov}[h(x)]$ (jointly,
      what zeroing a row of $H$ would do to $H P H^\top$),
    - zero the masked columns of $\mathrm{Cov}[x, h(x)]$ (likewise for
      $P H^\top$),
    - substitute a unit block into $R$,
    - zero the residual entry.

    The innovation then splits as

    $$
    S = \begin{bmatrix} S_{\mathrm{obs}} & 0 \\ 0 & I \end{bmatrix},
    $$

    so the gain $K = C S^{-1}$ has zero columns on the masked channels and
    they cannot move the state. The posterior is *exactly* the
    channel-deleted filter, with no branching — which is what lets the
    per-channel path run inside a `jax.lax.scan` without a `cond`.

    Exposed because the substitution is reusable: any moment-matched
    update — a custom filter loop, a smoother variant, an ensemble
    method — needs the same rewrite to handle partially observed vectors,
    and getting the joint row/column masking subtly wrong is easy.

    Note:
        Residuals are formed from separately-masked $y$ and $\hat y$
        rather than by masking their difference, so a masked $y$ entry may
        be ``NaN`` — the usual "not measured" encoding — without poisoning
        the reverse-mode gradient.

    Args:
        obs_cov: Matched $\mathrm{Cov}[h(x)]$, shape ``(M, M)``.
        cross_cov: Matched $\mathrm{Cov}[x, h(x)]$, shape ``(N, M)``.
        obs_noise: Observation noise $R$, shape ``(M, M)``.
        y: Observation vector, shape ``(M,)``. Masked entries are never
            read and may be ``NaN``.
        y_hat: Matched $\mathbb{E}[h(x)]$, shape ``(M,)``.
        mask: Per-channel mask, shape ``(M,)``. ``True`` keeps the channel.

    Returns:
        Tuple ``(obs_cov_eff, cross_cov_eff, obs_noise_eff, residual,
        n_missing)``. ``n_missing`` is the float count of masked channels;
        each contributed $-\tfrac{1}{2}\log 2\pi$ of dummy density to the
        full-vector log-likelihood, so adding
        ``0.5 * n_missing * log(2 pi)`` back recovers the exact marginal
        over the observed entries.
    """
    M = y.shape[-1]
    # keep[i, j] is True only where *both* channels survive, so masked
    # rows and columns are cleared together.
    keep = mask[:, None] & mask[None, :]

    # Zeroing row i of H would zero row i and column i of H P H^T; do that
    # directly to the matched Cov[h(x)].
    obs_cov_eff = jnp.where(keep, obs_cov, jnp.zeros_like(obs_cov))

    # ... and substitute a unit block into R on the masked channels, so
    # S = blockdiag(S_obs, I) rather than becoming singular.
    obs_noise_eff = jnp.where(keep, obs_noise, jnp.eye(M, dtype=obs_noise.dtype))

    # Column j of C = Cov[x, h(x)] is what channel j uses to move the
    # state; zero it and the gain's column j vanishes with it.
    cross_cov_eff = jnp.where(mask[None, :], cross_cov, jnp.zeros_like(cross_cov))

    # Mask y and y_hat *separately* rather than masking their difference:
    # a masked y entry is commonly NaN, and NaN in the discarded branch of
    # a where still poisons the reverse-mode gradient.
    residual = jnp.where(mask, y, jnp.zeros_like(y)) - jnp.where(
        mask, y_hat, jnp.zeros_like(y_hat)
    )

    n_missing = M - jnp.sum(mask.astype(y.dtype))
    return obs_cov_eff, cross_cov_eff, obs_noise_eff, residual, n_missing

Kalman filtering & smoothing

The forward filter and RTS smoother, their \(O(\log N)\) parallel (associative-scan) counterparts, and the steady-state (infinite-horizon) variants built on the discrete algebraic Riccati equation.

Observation masks

kalman_filter and parallel_kalman_filter take an optional mask, dispatched on its rank:

shape	meaning
(T,)	Per-step gate. False runs the predict step only and contributes nothing to the log-likelihood — the usual way to predict on a merged train/test grid.
(T, M)	Per-channel gate, for partially observed multivariate series where different channels are measured at different times.

The per-channel path marginalises unobserved channels exactly: row \(i\) of \(H_t\) is zeroed, a unit block is substituted into \(R_t\), and the residual entry is set to zero. The innovation covariance is then block-diagonal in the observed/masked split, so column \(i\) of the gain vanishes and the masked channel cannot move the state — the posterior reproduces the row-deleted filter to machine precision, with no branching. A dummy block also contributes \(-\tfrac12 \log 2\pi\) per masked channel to the full-vector density, which is stripped per step, so log_likelihood is the exact marginal \(\log p(y_{\mathrm{obs}})\) and is invariant to both the mask pattern and the dummy variance.

An all-False row of a (T, M) mask is equivalent to a False entry in the (T,) form, and M == 1 is unambiguous either way. Masked entries of observations are never read, so they may be NaN. Operator-typed obs_model / obs_noise are materialised under a (T, M) mask, since zeroing rows is inherently dense; form="sqrt" supports the (T,) mask only and raises NotImplementedError otherwise. rts_smoother needs no mask of its own — it consumes filtered/predicted moments, which are already mask-aware.

Mean-field (block-diagonal) filtering

When the state decomposes into \(L\) independent blocks of size \(d\) (e.g. a multi-output temporal GP with one SDE per output), meanfield_kalman_filter and meanfield_rts_smoother run \(L\) parallel \(d\)-state filters under jax.vmap — \(O(T\,L\,d^3)\) total instead of the full filter's \(O(T\,L^3 d^3)\). The trade-off is the mean-field approximation: inputs and posterior covariance are projected onto their diagonal blocks, so posterior cross-block covariance is dropped (the returned \((T, D, D)\) covariances are exactly zero off-block, and the log-likelihood is the sum of per-block log-likelihoods). The approximation is exact when the true cross-block dynamics are zero, which makes the decoupled case a useful consistency check against kalman_filter. A BlockDiag operator whose sub-operators match the blocking is split structurally, without materialising the full \((D, D)\) matrix; parallel=True routes each block through the associative-scan filter/smoother.

Structured linear algebra and Gaussian primitives for JAX.

EmissionModel

Bases: Module

Observation (emission) model wrapping a linear observation matrix.

Provides named methods for common Kalman filter projection operations with observation matrix H ∈ ℝᴹˣᴺ.

Attributes:

Name	Type	Description
H	Float[Array, 'M N']	Observation matrix, shape (M, N).

Source code in src/gaussx/_ssm/_emission.py

class EmissionModel(eqx.Module):
    """Observation (emission) model wrapping a linear observation matrix.

    Provides named methods for common Kalman filter projection
    operations with observation matrix H ∈ ℝᴹˣᴺ.

    Attributes:
        H: Observation matrix, shape ``(M, N)``.
    """

    H: Float[Array, "M N"]

    def project_mean(
        self,
        mean: Float[Array, " N"],
    ) -> Float[Array, " M"]:
        """Project state mean to observation space: ŷ = H x.

        Args:
            mean: State mean, shape ``(N,)``.

        Returns:
            Projected mean, shape ``(M,)``.
        """
        return self.H @ mean

    def project_covariance(
        self,
        cov: Float[Array, "N N"],
        noise: Float[Array, "M M"] | None = None,
    ) -> Float[Array, "M M"]:
        """Project state covariance: S = H P Hᵀ [+ R].

        Args:
            cov: State covariance P, shape ``(N, N)``.
            noise: Optional observation noise R, shape ``(M, M)``.

        Returns:
            Innovation covariance S, shape ``(M, M)``.
        """
        S = self.H @ cov @ self.H.T  # (M, M)
        if noise is not None:
            S = S + noise
        return S

    def innovation(
        self,
        y: Float[Array, " M"],
        x_pred: Float[Array, " N"],
    ) -> Float[Array, " M"]:
        """Compute innovation (measurement residual): v = y − H x.

        Args:
            y: Observation, shape ``(M,)``.
            x_pred: Predicted state mean, shape ``(N,)``.

        Returns:
            Innovation vector v, shape ``(M,)``.
        """
        return y - self.H @ x_pred

    def back_project_precision(
        self,
        noise_prec: Float[Array, "M M"],
    ) -> Float[Array, "N N"]:
        """Back-project observation precision: Hᵀ R⁻¹ H.

        Args:
            noise_prec: Observation noise precision R⁻¹, shape ``(M, M)``.

        Returns:
            Information matrix contribution, shape ``(N, N)``.
        """
        return self.H.T @ noise_prec @ self.H

    def back_project_info(
        self,
        y: Float[Array, " M"],
        noise_prec: Float[Array, "M M"],
    ) -> Float[Array, " N"]:
        """Back-project observation to information vector: Hᵀ R⁻¹ y.

        Args:
            y: Observation, shape ``(M,)``.
            noise_prec: Observation noise precision R⁻¹, shape ``(M, M)``.

        Returns:
            Information vector contribution, shape ``(N,)``.
        """
        return self.H.T @ noise_prec @ y

project_mean(mean: Float[Array, ' N']) -> Float[Array, ' M']

Project state mean to observation space: ŷ = H x.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	State mean, shape (N,).	required

Returns:

Type	Description
Float[Array, ' M']	Projected mean, shape (M,).

Source code in src/gaussx/_ssm/_emission.py

def project_mean(
    self,
    mean: Float[Array, " N"],
) -> Float[Array, " M"]:
    """Project state mean to observation space: ŷ = H x.

    Args:
        mean: State mean, shape ``(N,)``.

    Returns:
        Projected mean, shape ``(M,)``.
    """
    return self.H @ mean

project_covariance(cov: Float[Array, 'N N'], noise: Float[Array, 'M M'] | None = None) -> Float[Array, 'M M']

Project state covariance: S = H P Hᵀ [+ R].

Parameters:

Name	Type	Description	Default
cov	Float[Array, 'N N']	State covariance P, shape (N, N).	required
noise	Float[Array, 'M M'] | None	Optional observation noise R, shape (M, M).	None

Returns:

Type	Description
Float[Array, 'M M']	Innovation covariance S, shape (M, M).

Source code in src/gaussx/_ssm/_emission.py

def project_covariance(
    self,
    cov: Float[Array, "N N"],
    noise: Float[Array, "M M"] | None = None,
) -> Float[Array, "M M"]:
    """Project state covariance: S = H P Hᵀ [+ R].

    Args:
        cov: State covariance P, shape ``(N, N)``.
        noise: Optional observation noise R, shape ``(M, M)``.

    Returns:
        Innovation covariance S, shape ``(M, M)``.
    """
    S = self.H @ cov @ self.H.T  # (M, M)
    if noise is not None:
        S = S + noise
    return S

innovation(y: Float[Array, ' M'], x_pred: Float[Array, ' N']) -> Float[Array, ' M']

Compute innovation (measurement residual): v = y − H x.

Parameters:

Name	Type	Description	Default
y	Float[Array, ' M']	Observation, shape (M,).	required
x_pred	Float[Array, ' N']	Predicted state mean, shape (N,).	required

Returns:

Type	Description
Float[Array, ' M']	Innovation vector v, shape (M,).

Source code in src/gaussx/_ssm/_emission.py

def innovation(
    self,
    y: Float[Array, " M"],
    x_pred: Float[Array, " N"],
) -> Float[Array, " M"]:
    """Compute innovation (measurement residual): v = y − H x.

    Args:
        y: Observation, shape ``(M,)``.
        x_pred: Predicted state mean, shape ``(N,)``.

    Returns:
        Innovation vector v, shape ``(M,)``.
    """
    return y - self.H @ x_pred

back_project_precision(noise_prec: Float[Array, 'M M']) -> Float[Array, 'N N']

Back-project observation precision: Hᵀ R⁻¹ H.

Parameters:

Name	Type	Description	Default
noise_prec	Float[Array, 'M M']	Observation noise precision R⁻¹, shape (M, M).	required

Returns:

Type	Description
Float[Array, 'N N']	Information matrix contribution, shape (N, N).

Source code in src/gaussx/_ssm/_emission.py

def back_project_precision(
    self,
    noise_prec: Float[Array, "M M"],
) -> Float[Array, "N N"]:
    """Back-project observation precision: Hᵀ R⁻¹ H.

    Args:
        noise_prec: Observation noise precision R⁻¹, shape ``(M, M)``.

    Returns:
        Information matrix contribution, shape ``(N, N)``.
    """
    return self.H.T @ noise_prec @ self.H

back_project_info(y: Float[Array, ' M'], noise_prec: Float[Array, 'M M']) -> Float[Array, ' N']

Back-project observation to information vector: Hᵀ R⁻¹ y.

Parameters:

Name	Type	Description	Default
y	Float[Array, ' M']	Observation, shape (M,).	required
noise_prec	Float[Array, 'M M']	Observation noise precision R⁻¹, shape (M, M).	required

Returns:

Type	Description
Float[Array, ' N']	Information vector contribution, shape (N,).

Source code in src/gaussx/_ssm/_emission.py

def back_project_info(
    self,
    y: Float[Array, " M"],
    noise_prec: Float[Array, "M M"],
) -> Float[Array, " N"]:
    """Back-project observation to information vector: Hᵀ R⁻¹ y.

    Args:
        y: Observation, shape ``(M,)``.
        noise_prec: Observation noise precision R⁻¹, shape ``(M, M)``.

    Returns:
        Information vector contribution, shape ``(N,)``.
    """
    return self.H.T @ noise_prec @ y

FilterState

Bases: Module

Output of kalman_filter.

Attributes:

Name	Type	Description
filtered_means	Float[Array, 'T N']	Shape (T, N) — filtered state estimates.
filtered_covs	Float[Array, 'T N N']	Shape (T, N, N) — filtered covariances.
predicted_means	Float[Array, 'T N']	Shape (T, N) — predicted state estimates.
predicted_covs	Float[Array, 'T N N']	Shape (T, N, N) — predicted covariances.
log_likelihood	Float[Array, '']	Scalar — total log-likelihood.

Source code in src/gaussx/_ssm/_kalman.py

class FilterState(eqx.Module):
    """Output of ``kalman_filter``.

    Attributes:
        filtered_means: Shape ``(T, N)`` — filtered state estimates.
        filtered_covs: Shape ``(T, N, N)`` — filtered covariances.
        predicted_means: Shape ``(T, N)`` — predicted state estimates.
        predicted_covs: Shape ``(T, N, N)`` — predicted covariances.
        log_likelihood: Scalar — total log-likelihood.
    """

    filtered_means: Float[Array, "T N"]
    filtered_covs: Float[Array, "T N N"]
    predicted_means: Float[Array, "T N"]
    predicted_covs: Float[Array, "T N N"]
    log_likelihood: Float[Array, ""]

InfiniteHorizonState

Bases: Module

Output of infinite_horizon_filter.

Attributes:

Name	Type	Description
filtered_means	Float[Array, 'T N']	Filtered state estimates, shape (T, N).
filtered_covs	Float[Array, 'T N N']	Filtered covariances (constant), shape (T, N, N).
predicted_means	Float[Array, 'T N']	Predicted state estimates, shape (T, N).
predicted_covs	Float[Array, 'T N N']	Predicted covariances (constant), shape (T, N, N).
log_likelihood	Float[Array, '']	Total log-likelihood (scalar).

Source code in src/gaussx/_ssm/_infinite_horizon_kalman.py

class InfiniteHorizonState(eqx.Module):
    """Output of ``infinite_horizon_filter``.

    Attributes:
        filtered_means: Filtered state estimates, shape ``(T, N)``.
        filtered_covs: Filtered covariances (constant), shape ``(T, N, N)``.
        predicted_means: Predicted state estimates, shape ``(T, N)``.
        predicted_covs: Predicted covariances (constant), shape ``(T, N, N)``.
        log_likelihood: Total log-likelihood (scalar).
    """

    filtered_means: Float[Array, "T N"]
    filtered_covs: Float[Array, "T N N"]
    predicted_means: Float[Array, "T N"]
    predicted_covs: Float[Array, "T N N"]
    log_likelihood: Float[Array, ""]

DAREResult

Bases: Module

Result of DARE solver.

Attributes:

Name	Type	Description
P_inf	Float[Array, 'D D']	Steady-state covariance, shape (D, D).
K_inf	Float[Array, 'D M']	Steady-state Kalman gain, shape (D, M).
converged	Bool[Array, '']	Scalar boolean indicating convergence.

Source code in src/gaussx/_ssm/_dare.py

class DAREResult(eqx.Module):
    """Result of DARE solver.

    Attributes:
        P_inf: Steady-state covariance, shape ``(D, D)``.
        K_inf: Steady-state Kalman gain, shape ``(D, M)``.
        converged: Scalar boolean indicating convergence.
    """

    P_inf: Float[Array, "D D"]
    K_inf: Float[Array, "D M"]
    converged: Bool[Array, ""]

kalman_filter(transition: Float[Array, '*T N N'] | lx.AbstractLinearOperator, obs_model: Float[Array, '*T M N'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator, obs_noise: Float[Array, '*T M M'] | lx.AbstractLinearOperator, observations: Float[Array, 'T M'], init_mean: Float[Array, ' N'], init_cov: Float[Array, 'N N'], *, mask: Bool[Array, ' T'] | Bool[Array, 'T M'] | None = None, solver: AbstractSolverStrategy | None = None, woodbury_innovation: bool = False) -> FilterState

Kalman filter forward pass via jax.lax.scan.

Implements the predict-update cycle for a (possibly time-varying) linear-Gaussian state-space model:

x_t = A_t @ x_{t-1} + q_t,   q_t ~ N(0, Q_t)
y_t = H_t @ x_t + r_t,        r_t ~ N(0, R_t)

Time-invariant inputs (single (N, N) / (M, N) etc.) are automatically broadcast along the time axis. Time-varying inputs are passed as (T, …) stacks (e.g. from discretise_sequence).

Operator inputs (lineax BlockDiag / Kronecker / DiagonalLinearOperator / MaskedOperator / etc.) are accepted in the time-invariant signature only. The structural matvec (A @ x, H @ x) runs through the operator's mv; operator-typed Q / R are materialised to dense arrays once outside the scan (the per-step sandwiches A P A^T / H P H^T themselves run inside the scan because they depend on the evolving P_filt).

Parameters:

Name	Type	Description	Default
transition	Float[Array, '*T N N'] | AbstractLinearOperator	State transition matrix A. Shape (N, N), (T, N, N), or lineax.AbstractLinearOperator.	required
obs_model	Float[Array, '*T M N'] | AbstractLinearOperator	Observation matrix H. Shape (M, N), (T, M, N), or operator.	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator	Process noise covariance Q. Shape (N, N), (T, N, N), or operator.	required
obs_noise	Float[Array, '*T M M'] | AbstractLinearOperator	Observation noise covariance R. Shape (M, M), (T, M, M), or operator.	required
observations	Float[Array, 'T M']	Observed data, shape (T, M).	required
init_mean	Float[Array, ' N']	Initial state mean, shape (N,).	required
init_cov	Float[Array, 'N N']	Initial state covariance, shape (N, N).	required
mask	Bool[Array, ' T'] | Bool[Array, 'T M'] | None	Optional observation mask. Disambiguated by rank, so no extra keyword is needed (M == 1 is unambiguous either way, since a (T,) and a (T, 1) mask coincide). Shape (T,) — per-step gate. True (or 1) runs the full predict + update step; False (or 0) runs the predict step only and contributes nothing to the log-likelihood. Useful for prediction on merged train/test grids. Shape (T, M) — per-channel gate, for partially observed multivariate series. False entries are marginalised out exactly: the corresponding rows of H_t are zeroed, a unit block is substituted into R_t, and the residual entry is set to zero. The returned log_likelihood is the exact marginal \(\log p(y_{\mathrm{obs}})\), not the full-vector density. Masked entries of observations are never read, so they may be NaN. An all-False row is equivalent to a False entry in the (T,) form. Defaults to all-True. Operator-typed obs_model / obs_noise are materialised under a (T, M) mask, since zeroing rows is inherently a dense operation.	None
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None
woodbury_innovation	bool	When True, build the innovation covariance S = H P Hᵀ + R as a gaussx.LowRankUpdate so structured R can use Woodbury solves/log-determinants. Defaults to False to preserve the dense innovation path.	False

Raises:

Type	Description
TypeError	If operator-typed inputs are mixed with 3D (T, …) arrays. Operator inputs must come from the time-invariant signature (per-step structured stacks are not supported; pass dense (T, …) arrays for the time-varying path).

Returns:

Type	Description
FilterState	A FilterState with filtered/predicted means, covariances,
FilterState	and total log-likelihood.

Source code in src/gaussx/_ssm/_kalman.py

def kalman_filter(
    transition: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    obs_model: Float[Array, "*T M N"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    obs_noise: Float[Array, "*T M M"] | lx.AbstractLinearOperator,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " N"],
    init_cov: Float[Array, "N N"],
    *,
    mask: Bool[Array, " T"] | Bool[Array, "T M"] | None = None,
    solver: AbstractSolverStrategy | None = None,
    woodbury_innovation: bool = False,
) -> FilterState:
    r"""Kalman filter forward pass via ``jax.lax.scan``.

    Implements the predict-update cycle for a (possibly time-varying)
    linear-Gaussian state-space model:

        x_t = A_t @ x_{t-1} + q_t,   q_t ~ N(0, Q_t)
        y_t = H_t @ x_t + r_t,        r_t ~ N(0, R_t)

    **Time-invariant inputs** (single ``(N, N)`` / ``(M, N)`` etc.) are
    automatically broadcast along the time axis. **Time-varying inputs**
    are passed as ``(T, …)`` stacks (e.g. from
    `discretise_sequence`).

    **Operator inputs** (lineax ``BlockDiag`` / ``Kronecker`` /
    ``DiagonalLinearOperator`` / ``MaskedOperator`` / etc.) are accepted
    in the **time-invariant** signature only. The structural matvec
    (``A @ x``, ``H @ x``) runs through the operator's ``mv``;
    operator-typed ``Q`` / ``R`` are materialised to dense arrays once
    outside the scan (the per-step sandwiches ``A P A^T`` / ``H P H^T``
    themselves run inside the scan because they depend on the evolving
    ``P_filt``).

    Args:
        transition: State transition matrix ``A``. Shape ``(N, N)``,
            ``(T, N, N)``, or `lineax.AbstractLinearOperator`.
        obs_model: Observation matrix ``H``. Shape ``(M, N)``,
            ``(T, M, N)``, or operator.
        process_noise: Process noise covariance ``Q``. Shape ``(N, N)``,
            ``(T, N, N)``, or operator.
        obs_noise: Observation noise covariance ``R``. Shape ``(M, M)``,
            ``(T, M, M)``, or operator.
        observations: Observed data, shape ``(T, M)``.
        init_mean: Initial state mean, shape ``(N,)``.
        init_cov: Initial state covariance, shape ``(N, N)``.
        mask: Optional observation mask. Disambiguated by rank, so no
            extra keyword is needed (``M == 1`` is unambiguous either
            way, since a ``(T,)`` and a ``(T, 1)`` mask coincide).

            - Shape ``(T,)`` — per-step gate. ``True`` (or ``1``) runs
              the full predict + update step; ``False`` (or ``0``) runs
              the predict step only and contributes nothing to the
              log-likelihood. Useful for prediction on merged
              train/test grids.
            - Shape ``(T, M)`` — per-channel gate, for partially
              observed multivariate series. ``False`` entries are
              marginalised out exactly: the corresponding rows of
              ``H_t`` are zeroed, a unit block is substituted into
              ``R_t``, and the residual entry is set to zero. The
              returned ``log_likelihood`` is the exact marginal
              $\log p(y_{\mathrm{obs}})$, not the full-vector density.
              Masked entries of ``observations`` are never read, so
              they may be ``NaN``. An all-``False`` row is equivalent
              to a ``False`` entry in the ``(T,)`` form.

            Defaults to all-True. Operator-typed ``obs_model`` /
            ``obs_noise`` are materialised under a ``(T, M)`` mask,
            since zeroing rows is inherently a dense operation.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.
        woodbury_innovation: When ``True``, build the innovation
            covariance ``S = H P Hᵀ + R`` as a
            `gaussx.LowRankUpdate` so structured ``R`` can use
            Woodbury solves/log-determinants. Defaults to ``False`` to
            preserve the dense innovation path.

    Raises:
        TypeError: If operator-typed inputs are mixed with 3D ``(T, …)``
            arrays. Operator inputs must come from the time-invariant
            signature (per-step structured stacks are not supported;
            pass dense ``(T, …)`` arrays for the time-varying path).

    Returns:
        A ``FilterState`` with filtered/predicted means, covariances,
        and total log-likelihood.
    """
    M = observations.shape[-1]
    T = observations.shape[0]

    # Closure-friendly matvec: when an operator is supplied, prefer its
    # structural ``mv`` over the dense ``A @ x``. Otherwise the
    # broadcast 3D array contains ``A_seq[t]`` for each step.
    A_op = transition if isinstance(transition, lx.AbstractLinearOperator) else None
    H_op = obs_model if isinstance(obs_model, lx.AbstractLinearOperator) else None
    R_op = obs_noise if isinstance(obs_noise, lx.AbstractLinearOperator) else None

    # A per-channel mask rewrites the rows of H and the block structure
    # of R, neither of which survives as a structured operator — so the
    # (T, M) path drops to dense observation inputs.
    channel_mask = mask is not None and jnp.ndim(mask) == 2
    if channel_mask:
        H_op = None
        R_op = None

    A_seq, H_seq, Q_seq, R_seq, mask_seq, _ = _normalise_tv_inputs(
        transition,
        obs_model,
        process_noise,
        obs_noise,
        T=T,
        mask=mask,
        M=M,
        materialise_transition=A_op is None,
        materialise_obs=H_op is None,
        # Skip the O(T M²) dense broadcast of structured R when the
        # Woodbury path consumes the operator directly.
        materialise_obs_noise=not (woodbury_innovation and R_op is not None),
    )

    def step(carry, inputs):
        x_filt, P_filt, ll = carry
        A_t, H_t, Q_t, R_t, y_t, mask_t = inputs

        # --- Predict ---
        # Structural matvec when an operator was supplied; dense matmul otherwise.
        x_pred = A_op.mv(x_filt) if A_op is not None else A_t @ x_filt
        if A_op is not None:
            P_filt_op = lx.MatrixLinearOperator(P_filt, lx.positive_semidefinite_tag)
            P_pred = sandwich(A_op, P_filt_op).as_matrix() + Q_t
        else:
            P_pred = A_t @ P_filt @ A_t.T + Q_t

        # --- Update ---
        def _update(H_eff, R_eff, v, n_missing):
            """Shared update body for the gated and per-channel paths.

            ``H_eff`` is either a lineax operator (structural path) or a
            dense ``(M, N)`` array; ``n_missing`` is the count of masked
            channels, used to strip the dummy block's contribution from
            the log-likelihood.
            """
            H_is_op = isinstance(H_eff, lx.AbstractLinearOperator)
            S_op = _innovation_covariance(
                H_eff, P_pred, R_eff, woodbury=woodbury_innovation
            )

            PHt = (
                _right_matmul_transpose(P_pred, H_eff) if H_is_op else P_pred @ H_eff.T
            )  # (N, M)
            K = solve_rows(S_op, PHt, solver=solver)  # (N, M)

            x_upd = x_pred + K @ v
            if woodbury_innovation:
                # Avoid materialising S for the covariance update.
                HP_pred = _left_matmul(H_eff, P_pred) if H_is_op else H_eff @ P_pred
                P_upd = P_pred - K @ HP_pred
            else:
                P_upd = P_pred - K @ S_op.as_matrix() @ K.T

            Sinv_v = dispatch_solve(S_op, v, solver)
            ld = dispatch_logdet(S_op, solver)
            # The dummy unit block contributes -0.5 * log(2 pi) per
            # masked channel to the full-vector density; strip it here,
            # per step, so the result is the exact marginal over the
            # observed entries and is invariant to the dummy variance.
            ll_inc = (
                -0.5 * (v @ Sinv_v + ld + M * _LOG_2PI) + 0.5 * n_missing * _LOG_2PI
            )
            return x_upd, P_upd, ll_inc

        if channel_mask:
            # No lax.cond: an all-False row degenerates to the
            # predict-only step on its own, so the masked path is
            # branch-free.
            H_eff, R_eff, y_eff, n_missing = _masked_obs_inputs(H_t, R_t, y_t, mask_t)
            x_filt_new, P_filt_new, ll_inc = _update(
                H_eff, R_eff, y_eff - H_eff @ x_pred, n_missing
            )
        else:
            # Gate the whole step via lax.cond so the predict-only
            # branch evaluates neither the update arithmetic nor
            # produces gradients for the dropped path.
            def _do_update(_):
                v = y_t - (H_op.mv(x_pred) if H_op is not None else H_t @ x_pred)
                # Resolve ``R`` for innovation: operator path uses the
                # closed-over ``R_op`` (kept structural for Woodbury);
                # array path falls back to the per-step ``R_t``.
                R_innov = R_op if R_op is not None else R_t
                # Resolve ``H`` similarly so the operator preserves
                # structure in both the Woodbury and the
                # structural-sandwich paths.
                H_innov = H_op if H_op is not None else H_t
                return _update(H_innov, R_innov, v, 0.0)

            def _skip_update(_):
                # Match the update branch's dtype: a bare ``jnp.array(0.0)``
                # is float64 under x64 and makes ``lax.cond`` reject the
                # branches on float32 inputs.
                return x_pred, P_pred, jnp.zeros((), dtype=P_pred.dtype)

            x_filt_new, P_filt_new, ll_inc = jax.lax.cond(
                mask_t, _do_update, _skip_update, operand=None
            )
        ll_new = ll + ll_inc

        carry_new = (x_filt_new, P_filt_new, ll_new)
        outputs = (x_filt_new, P_filt_new, x_pred, P_pred)
        return carry_new, outputs

    init_carry = (init_mean, init_cov, jnp.zeros((), dtype=init_cov.dtype))
    final_carry, (f_means, f_covs, p_means, p_covs) = jax.lax.scan(
        step, init_carry, (A_seq, H_seq, Q_seq, R_seq, observations, mask_seq)
    )

    return FilterState(
        filtered_means=f_means,
        filtered_covs=f_covs,
        predicted_means=p_means,
        predicted_covs=p_covs,
        log_likelihood=final_carry[2],
    )

rts_smoother(filter_state: FilterState, transition: Float[Array, '*T N N'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, 'T N'], Float[Array, 'T N N']]

Rauch-Tung-Striebel backward smoother.

Accepts the same time-invariant / time-varying / operator forms for transition and process_noise as kalman_filter. When a step was masked off in the filter (mask[t] == 0), the smoother formula degenerates harmlessly because filtered == predicted at that step.

Parameters:

Name	Type	Description	Default
filter_state	FilterState	Output of kalman_filter.	required
transition	Float[Array, '*T N N'] | AbstractLinearOperator	State transition matrix or operator.	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator	Process noise covariance or operator. (Not currently used by the standard RTS recurrence — kept for API symmetry with kalman_filter.)	required
solver	AbstractSolverStrategy | None	Optional solver strategy.	None

Returns:

Type	Description
tuple[Float[Array, 'T N'], Float[Array, 'T N N']]	Tuple (smoothed_means, smoothed_covs).

Source code in src/gaussx/_ssm/_kalman.py

def rts_smoother(
    filter_state: FilterState,
    transition: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "T N"], Float[Array, "T N N"]]:
    """Rauch-Tung-Striebel backward smoother.

    Accepts the same time-invariant / time-varying / operator forms for
    ``transition`` and ``process_noise`` as `kalman_filter`. When
    a step was masked off in the filter (``mask[t] == 0``), the
    smoother formula degenerates harmlessly because filtered ==
    predicted at that step.

    Args:
        filter_state: Output of `kalman_filter`.
        transition: State transition matrix or operator.
        process_noise: Process noise covariance or operator. (Not
            currently used by the standard RTS recurrence — kept for
            API symmetry with `kalman_filter`.)
        solver: Optional solver strategy.

    Returns:
        Tuple ``(smoothed_means, smoothed_covs)``.
    """
    del process_noise  # not used in the standard RTS recurrence

    T = filter_state.filtered_means.shape[0]

    # Materialise once outside the scan for the sandwich; matvec stays
    # structural via the operator's mv.
    A_dense = _materialise(transition)
    A_op = transition if isinstance(transition, lx.AbstractLinearOperator) else None
    if A_dense.ndim == 2:
        A_seq = jnp.broadcast_to(A_dense, (T, *A_dense.shape))
    elif A_dense.ndim == 3:
        if A_op is not None:
            raise TypeError(
                "Operator-typed transition cannot have a leading time axis."
            )
        A_seq = A_dense
    else:
        raise ValueError(f"transition must have ndim 2 or 3, got {A_dense.ndim}.")

    def step(carry, inputs):
        x_smooth, P_smooth = carry
        x_filt, P_filt, x_pred, P_pred, A_next = inputs

        # Smoother gain: G = P_filt A_{t+1}^T P_pred_{t+1}^{-1}
        P_pred_op = lx.MatrixLinearOperator(P_pred, lx.positive_semidefinite_tag)
        G = P_filt @ A_next.T  # (N, N)
        G = solve_rows(P_pred_op, G, solver=solver)  # (N, N)

        x_smooth_new = x_filt + G @ (x_smooth - x_pred)
        P_smooth_new = P_filt + G @ (P_smooth - P_pred) @ G.T

        return (x_smooth_new, P_smooth_new), (x_smooth_new, P_smooth_new)

    init_carry = (
        filter_state.filtered_means[T - 1],
        filter_state.filtered_covs[T - 1],
    )

    # Reverse the sequences for backward pass (exclude last time step).
    # ``A_next[t]`` is the transition that maps step ``t`` to step ``t+1``,
    # i.e. ``A_seq[t+1]``.
    inputs = (
        filter_state.filtered_means[:-1][::-1],
        filter_state.filtered_covs[:-1][::-1],
        filter_state.predicted_means[1:][::-1],
        filter_state.predicted_covs[1:][::-1],
        A_seq[1:][::-1],
    )

    _, (s_means_rev, s_covs_rev) = jax.lax.scan(step, init_carry, inputs)

    # Reverse back and prepend last filtered state.
    s_means = jnp.concatenate(
        [s_means_rev[::-1], filter_state.filtered_means[T - 1 :]], axis=0
    )
    s_covs = jnp.concatenate(
        [s_covs_rev[::-1], filter_state.filtered_covs[T - 1 :]], axis=0
    )

    return s_means, s_covs

kalman_gain(P: lx.AbstractLinearOperator, H: lx.AbstractLinearOperator, R: lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None, woodbury_innovation: bool = False) -> Float[Array, 'N M']

Compute Kalman gain K = P @ H^T @ (H @ P @ H^T + R)^{-1}.

Parameters:

Name	Type	Description	Default
P	AbstractLinearOperator	Prior covariance operator, shape (N, N).	required
H	AbstractLinearOperator	Observation model operator, shape (M, N).	required
R	AbstractLinearOperator	Observation noise operator, shape (M, M).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None
woodbury_innovation	bool	When True, route the innovation covariance through gaussx.LowRankUpdate.	False

Returns:

Type	Description
Float[Array, 'N M']	Kalman gain matrix of shape (N, M).

Source code in src/gaussx/_ssm/_kalman.py

def kalman_gain(
    P: lx.AbstractLinearOperator,
    H: lx.AbstractLinearOperator,
    R: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
    woodbury_innovation: bool = False,
) -> Float[Array, "N M"]:
    """Compute Kalman gain ``K = P @ H^T @ (H @ P @ H^T + R)^{-1}``.

    Args:
        P: Prior covariance operator, shape ``(N, N)``.
        H: Observation model operator, shape ``(M, N)``.
        R: Observation noise operator, shape ``(M, M)``.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.
        woodbury_innovation: When ``True``, route the innovation
            covariance through `gaussx.LowRankUpdate`.

    Returns:
        Kalman gain matrix of shape ``(N, M)``.
    """
    P_mat = _materialise(P)
    H_mat = _materialise(H)

    S_op = _innovation_covariance(H, P, R, woodbury=woodbury_innovation)

    # K = P Hᵀ S⁻¹
    PHt = P_mat @ H_mat.T  # (N, M)
    return solve_rows(S_op, PHt, solver=solver)  # (N, M)

parallel_kalman_filter(transition: Float[Array, '*T N N'] | lx.AbstractLinearOperator, obs_model: Float[Array, '*T M N'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator, obs_noise: Float[Array, '*T M M'] | lx.AbstractLinearOperator, observations: Float[Array, 'T M'], init_mean: Float[Array, ' N'], init_cov: Float[Array, 'N N'], *, mask: Bool[Array, ' T'] | Bool[Array, 'T M'] | None = None, solver: AbstractSolverStrategy | None = None, woodbury_innovation: bool = False, form: str = 'covariance') -> FilterState

Parallel Kalman filter via jax.lax.associative_scan.

Numerically equivalent to gaussx.kalman_filter but with O(log T) parallel depth on accelerators. Same generalised contract (TI / TV / operator-typed inputs, optional mask, scalar log-likelihood). Empty observation windows (T == 0) return a zero-length FilterState with log_likelihood == 0.

Parameters:

Name	Type	Description	Default
transition	Float[Array, '*T N N'] | AbstractLinearOperator	State transition matrix or operator.	required
obs_model	Float[Array, '*T M N'] | AbstractLinearOperator	Observation matrix or operator.	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator	Process noise covariance or operator.	required
obs_noise	Float[Array, '*T M M'] | AbstractLinearOperator	Observation noise covariance or operator.	required
observations	Float[Array, 'T M']	Observed data, shape (T, M).	required
init_mean	Float[Array, ' N']	Initial state mean, shape (N,).	required
init_cov	Float[Array, 'N N']	Initial state covariance, shape (N, N).	required
mask	Bool[Array, ' T'] | Bool[Array, 'T M'] | None	Optional observation mask, dispatched on rank exactly as in gaussx.kalman_filter. Shape (T,) gates whole steps (False runs predict-only and contributes 0 to the log-likelihood); shape (T, M) gates individual channels and yields the exact marginal log-likelihood over the observed entries. Defaults to all-True. Not supported by form="sqrt".	None
solver	AbstractSolverStrategy | None	Accepted for API symmetry with kalman_filter but not currently threaded through the per-element solves; the covariance-form combinator uses unstructured dense solves. The square-root form also uses dense solves for the affine terms.	None
woodbury_innovation	bool	When True, delegates to gaussx.kalman_filter with the same flag so structured R uses the Woodbury innovation path.	False
form	str	Either "covariance" (default) or "sqrt". The square-root form maintains lower-triangular covariance factors alongside the covariance updates and reconstructs PSD covariance matrices in the returned FilterState. Note: the associative-scan equations themselves still use the covariance form internally; the factor path is a PSD-safety net for ill-conditioned float32 chains rather than a fully factor-propagating combinator (see #165).	'covariance'

Raises:

Type	Description
ValueError	If form is not "covariance" or "sqrt".

Returns:

Type	Description
FilterState	FilterState with filtered / predicted means and covs
FilterState	and the total log-likelihood.

Source code in src/gaussx/_ssm/_parallel_kalman.py

def parallel_kalman_filter(
    transition: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    obs_model: Float[Array, "*T M N"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    obs_noise: Float[Array, "*T M M"] | lx.AbstractLinearOperator,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " N"],
    init_cov: Float[Array, "N N"],
    *,
    mask: Bool[Array, " T"] | Bool[Array, "T M"] | None = None,
    solver: AbstractSolverStrategy | None = None,
    woodbury_innovation: bool = False,
    form: str = "covariance",
) -> FilterState:
    """Parallel Kalman filter via `jax.lax.associative_scan`.

    Numerically equivalent to `gaussx.kalman_filter` but with
    ``O(log T)`` parallel depth on accelerators. Same generalised
    contract (TI / TV / operator-typed inputs, optional mask, scalar
    log-likelihood). Empty observation windows (``T == 0``) return a
    zero-length `FilterState` with ``log_likelihood == 0``.

    Args:
        transition: State transition matrix or operator.
        obs_model: Observation matrix or operator.
        process_noise: Process noise covariance or operator.
        obs_noise: Observation noise covariance or operator.
        observations: Observed data, shape ``(T, M)``.
        init_mean: Initial state mean, shape ``(N,)``.
        init_cov: Initial state covariance, shape ``(N, N)``.
        mask: Optional observation mask, dispatched on rank exactly as
            in `gaussx.kalman_filter`. Shape ``(T,)`` gates whole
            steps (``False`` runs predict-only and contributes 0 to the
            log-likelihood); shape ``(T, M)`` gates individual channels
            and yields the exact marginal log-likelihood over the
            observed entries. Defaults to all-True. Not supported by
            ``form="sqrt"``.
        solver: Accepted for API symmetry with `kalman_filter` but
            not currently threaded through the per-element solves; the
            covariance-form combinator uses unstructured dense solves.
            The square-root form also uses dense solves for the affine
            terms.
        woodbury_innovation: When ``True``, delegates to
            `gaussx.kalman_filter` with the same flag so structured
            ``R`` uses the Woodbury innovation path.
        form: Either ``"covariance"`` (default) or ``"sqrt"``. The
            square-root form maintains lower-triangular covariance
            factors alongside the covariance updates and reconstructs
            PSD covariance matrices in the returned `FilterState`.
            Note: the associative-scan equations themselves still use
            the covariance form internally; the factor path is a
            PSD-safety net for ill-conditioned float32 chains rather
            than a fully factor-propagating combinator (see #165).

    Raises:
        ValueError: If ``form`` is not ``"covariance"`` or ``"sqrt"``.

    Returns:
        `FilterState` with filtered / predicted means and covs
        and the total log-likelihood.
    """
    if form == "sqrt":
        from gaussx._ssm._parallel_kalman_sqrt import parallel_kalman_filter_sqrt

        return parallel_kalman_filter_sqrt(
            transition,
            obs_model,
            process_noise,
            obs_noise,
            observations,
            init_mean,
            init_cov,
            mask=mask,
            solver=solver,
        )
    if form != "covariance":
        raise ValueError("form must be 'covariance' or 'sqrt'.")

    if woodbury_innovation:
        return kalman_filter(
            transition,
            obs_model,
            process_noise,
            obs_noise,
            observations,
            init_mean,
            init_cov,
            mask=mask,
            solver=solver,
            woodbury_innovation=True,
        )

    del solver  # not currently threaded through; see docstring + #165

    M_obs = observations.shape[-1]
    T = observations.shape[0]
    N = init_mean.shape[0]

    # Empty observation window: match kalman_filter's empty-scan output.
    if T == 0:
        return FilterState(
            filtered_means=jnp.zeros((0, N), dtype=init_mean.dtype),
            filtered_covs=jnp.zeros((0, N, N), dtype=init_cov.dtype),
            predicted_means=jnp.zeros((0, N), dtype=init_mean.dtype),
            predicted_covs=jnp.zeros((0, N, N), dtype=init_cov.dtype),
            log_likelihood=jnp.zeros((), dtype=init_mean.dtype),
        )

    A_seq, H_seq, Q_seq, R_seq, mask_seq, _ = _normalise_tv_inputs(
        transition, obs_model, process_noise, obs_noise, T=T, mask=mask, M=M_obs
    )
    # Work per-channel throughout: a ``(T,)`` gate is the special case
    # where every channel of a step shares one flag, so broadcasting it
    # reproduces the whole-step path exactly.
    mask_ch = (
        mask_seq
        if mask_seq.ndim == 2
        else jnp.broadcast_to(mask_seq[:, None], (T, M_obs))
    )
    step_active = jnp.any(mask_ch, axis=-1)

    # Build per-step elements. ``vmap`` of ``lax.cond`` evaluates both
    # branches and selects, so we instead substitute mask-aware safe
    # inputs (zeroed H rows, unit R block, zeroed y) into a single
    # active path. For a fully-masked step those substitutions collapse
    # the active builder to (F, 0, Q, 0, 0) — exactly the predict-only
    # element — and the Cholesky operates on the well-conditioned
    # identity, so even garbage in masked H / R / y can't NaN the
    # gradient.
    def _build_step(F, H, Q, R, y, m):
        H_eff, R_eff, y_eff, _ = _masked_obs_inputs(H, R, y, m)
        return _generic_filter_element_active(F, H_eff, Q, R_eff, y_eff)

    elems = jax.vmap(_build_step)(A_seq, H_seq, Q_seq, R_seq, observations, mask_ch)

    # Patch element 0 to absorb the initial prior. Outer ``lax.cond``
    # genuinely skips the inactive branch (no ``vmap`` wrapping here);
    # a partially-observed step 0 takes the active branch on the
    # substituted inputs.
    H_first, R_first, y_first, _ = _masked_obs_inputs(
        H_seq[0], R_seq[0], observations[0], mask_ch[0]
    )
    first = jax.lax.cond(
        step_active[0],
        lambda: _first_filter_element_active(
            A_seq[0],
            H_first,
            Q_seq[0],
            R_first,
            y_first,
            init_mean,
            init_cov,
        ),
        lambda: _first_filter_element_masked(
            A_seq[0],
            Q_seq[0],
            init_mean,
            init_cov,
        ),
    )
    elems = tuple(arr.at[0].set(val) for arr, val in zip(elems, first, strict=True))

    # ----- Associative scan -----
    _A_out, b_out, C_out, _eta_out, _J_out = jax.lax.associative_scan(
        _filter_combine, elems
    )
    filtered_means = b_out
    filtered_covs = jax.vmap(_sym)(C_out)

    # Reconstruct predicted means / covs from filtered + transition.
    prev_means = jnp.concatenate([init_mean[None], filtered_means[:-1]], axis=0)
    prev_covs = jnp.concatenate([init_cov[None], filtered_covs[:-1]], axis=0)

    def _predict_step(F, m, P, Q):
        return F @ m, _sym(F @ P @ F.T + Q)

    predicted_means, predicted_covs = jax.vmap(_predict_step)(
        A_seq, prev_means, prev_covs, Q_seq
    )

    # Log-likelihood from innovations. Same safe substitution as the
    # element builder so masked steps don't drive the Cholesky through
    # ill-conditioned user-supplied R / NaN gradients.
    def _ll_contrib(y, m_pred, P_pred, H, R, m, active):
        H_eff, R_eff, y_eff, n_missing = _masked_obs_inputs(H, R, y, m)
        v = y_eff - H_eff @ m_pred
        S = _sym(H_eff @ P_pred @ H_eff.T + R_eff)
        L = jnp.linalg.cholesky(S)
        Sinv_v = jax.scipy.linalg.cho_solve((L, True), v)
        quad = v @ Sinv_v
        logdet = cholesky_logdet(L)
        # Strip the dummy unit block's -0.5 * log(2 pi) per masked
        # channel, so this is the exact marginal over observed entries.
        contrib = -0.5 * (quad + logdet + M_obs * _LOG_2PI) + 0.5 * n_missing * _LOG_2PI
        return jnp.where(active, contrib, jnp.zeros_like(contrib))

    ll_contribs = jax.vmap(_ll_contrib)(
        observations,
        predicted_means,
        predicted_covs,
        H_seq,
        R_seq,
        mask_ch,
        step_active,
    )
    log_likelihood = jnp.sum(ll_contribs)

    return FilterState(
        filtered_means=filtered_means,
        filtered_covs=filtered_covs,
        predicted_means=predicted_means,
        predicted_covs=predicted_covs,
        log_likelihood=log_likelihood,
    )

parallel_rts_smoother(filter_state: FilterState, transition: Float[Array, '*T N N'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T N N'] | lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None, form: str = 'covariance') -> tuple[Float[Array, 'T N'], Float[Array, 'T N N']]

Parallel RTS smoother via reverse jax.lax.associative_scan.

Pairs with parallel_kalman_filter. Numerically equivalent to gaussx.rts_smoother with O(log T) parallel depth.

Parameters:

Name	Type	Description	Default
filter_state	FilterState	Output of parallel_kalman_filter or gaussx.kalman_filter.	required
transition	Float[Array, '*T N N'] | AbstractLinearOperator	State transition matrix or operator.	required
process_noise	Float[Array, '*T N N'] | AbstractLinearOperator	Unused — kept for API symmetry with the sequential smoother.	required
solver	AbstractSolverStrategy | None	Accepted for API symmetry; not currently threaded through.	None
form	str	Either "covariance" (default) or "sqrt". The square-root form maintains lower-triangular factors alongside the smoother associative scan and returns PSD-reconstructed covariances (see parallel_kalman_filter for the same caveat about the internal combinator).	'covariance'

Raises:

Type	Description
ValueError	If form is not "covariance" or "sqrt".

Returns:

Type	Description
tuple[Float[Array, 'T N'], Float[Array, 'T N N']]	Tuple (smoothed_means, smoothed_covs).

Source code in src/gaussx/_ssm/_parallel_kalman.py

def parallel_rts_smoother(
    filter_state: FilterState,
    transition: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T N N"] | lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
    form: str = "covariance",
) -> tuple[Float[Array, "T N"], Float[Array, "T N N"]]:
    """Parallel RTS smoother via reverse `jax.lax.associative_scan`.

    Pairs with `parallel_kalman_filter`. Numerically equivalent to
    `gaussx.rts_smoother` with ``O(log T)`` parallel depth.

    Args:
        filter_state: Output of `parallel_kalman_filter` or
            `gaussx.kalman_filter`.
        transition: State transition matrix or operator.
        process_noise: Unused — kept for API symmetry with the sequential
            smoother.
        solver: Accepted for API symmetry; not currently threaded
            through.
        form: Either ``"covariance"`` (default) or ``"sqrt"``. The
            square-root form maintains lower-triangular factors
            alongside the smoother associative scan and returns
            PSD-reconstructed covariances (see `parallel_kalman_filter`
            for the same caveat about the internal combinator).

    Raises:
        ValueError: If ``form`` is not ``"covariance"`` or ``"sqrt"``.

    Returns:
        Tuple ``(smoothed_means, smoothed_covs)``.
    """
    if form == "sqrt":
        from gaussx._ssm._parallel_kalman_sqrt import parallel_rts_smoother_sqrt

        return parallel_rts_smoother_sqrt(
            filter_state,
            transition,
            process_noise,
            solver=solver,
        )
    if form != "covariance":
        raise ValueError("form must be 'covariance' or 'sqrt'.")

    del process_noise, solver

    f_means = filter_state.filtered_means
    f_covs = filter_state.filtered_covs
    p_means = filter_state.predicted_means
    p_covs = filter_state.predicted_covs
    T = f_means.shape[0]
    N = f_means.shape[-1]

    if T == 0:
        return (
            jnp.zeros((0, N), dtype=f_means.dtype),
            jnp.zeros((0, N, N), dtype=f_covs.dtype),
        )

    A_dense = _materialise(transition)
    A_op = transition if isinstance(transition, lx.AbstractLinearOperator) else None
    if A_dense.ndim == 2:
        A_seq = jnp.broadcast_to(A_dense, (T, *A_dense.shape))
    elif A_dense.ndim == 3:
        if A_op is not None:
            raise TypeError(
                "Operator-typed transition cannot have a leading time axis."
            )
        A_seq = A_dense
    else:
        raise ValueError(f"transition must have ndim 2 or 3, got {A_dense.ndim}.")

    def _build_inner(f_mean, f_cov, p_mean_next, p_cov_next, A_next):
        # G = f_cov @ A_next.T @ inv(p_cov_next); p_cov_next is symmetric.
        rhs = f_cov @ A_next.T  # (N, N)
        G = jnp.linalg.solve(p_cov_next, rhs.T).T
        E = G
        g = f_mean - G @ p_mean_next
        L = _sym(f_cov - G @ p_cov_next @ G.T)
        return E, g, L

    inner_E, inner_g, inner_L = jax.vmap(_build_inner)(
        f_means[:-1], f_covs[:-1], p_means[1:], p_covs[1:], A_seq[1:]
    )
    last_E = jnp.zeros((1, N, N), dtype=f_means.dtype)
    last_g = f_means[-1:]
    last_L = f_covs[-1:]

    E = jnp.concatenate([inner_E, last_E], axis=0)
    g = jnp.concatenate([inner_g, last_g], axis=0)
    L = jnp.concatenate([inner_L, last_L], axis=0)

    _E_out, smoothed_means, smoothed_covs = jax.lax.associative_scan(
        _smoother_combine, (E, g, L), reverse=True
    )
    smoothed_covs = jax.vmap(_sym)(smoothed_covs)
    return smoothed_means, smoothed_covs

meanfield_kalman_filter(transition: Float[Array, '*T D D'] | lx.AbstractLinearOperator, obs_model: Float[Array, '*T M D'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T D D'] | lx.AbstractLinearOperator, obs_noise: Float[Array, '*T M M'] | lx.AbstractLinearOperator, observations: Float[Array, 'T M'], init_mean: Float[Array, ' D'], init_cov: Float[Array, 'D D'], *, block_size: int, mask: Bool[Array, ' T'] | Bool[Array, 'T M'] | None = None, solver: AbstractSolverStrategy | None = None, parallel: bool = False) -> FilterState

Mean-field Kalman filter over L = D // block_size state blocks.

Runs L independent block_size-state Kalman filters under jax.vmap — O(T L d^3) total work instead of the full filter's O(T L^3 d^3) — at the cost of dropping posterior cross-block covariances. All inputs are projected onto their diagonal blocks; off-block entries are discarded, which is the mean-field projection (exact when the true cross-block couplings are zero). The returned log-likelihood is the sum of the per-block log-likelihoods,

\[\log p(y_{1:T}) \approx \sum_{\ell=1}^{L} \log p^{(\ell)}\bigl(y_{1:T}^{(\ell)}\bigr).\]

Accepts the same time-invariant / time-varying / operator input forms as gaussx.kalman_filter. A gaussx.BlockDiag whose sub-operators match the blocking is split structurally without materialising the full (D, D) matrix. The observation dimension must factorise the same way: channel block \ell (of size M // L) is observed through state block \ell.

Parameters:

Name	Type	Description	Default
transition	Float[Array, '*T D D'] | AbstractLinearOperator	State transition A. Shape (D, D), (T, D, D), or lineax.AbstractLinearOperator.	required
obs_model	Float[Array, '*T M D'] | AbstractLinearOperator	Observation matrix H. Shape (M, D), (T, M, D), or operator. M must be divisible by L.	required
process_noise	Float[Array, '*T D D'] | AbstractLinearOperator	Process noise covariance Q. Shape (D, D), (T, D, D), or operator.	required
obs_noise	Float[Array, '*T M M'] | AbstractLinearOperator	Observation noise covariance R. Shape (M, M), (T, M, M), or operator.	required
observations	Float[Array, 'T M']	Observed data, shape (T, M).	required
init_mean	Float[Array, ' D']	Initial state mean, shape (D,).	required
init_cov	Float[Array, 'D D']	Initial state covariance, shape (D, D).	required
block_size	int	State block size d; D must be divisible by it.	required
mask	Bool[Array, ' T'] | Bool[Array, 'T M'] | None	Optional observation mask, as in gaussx.kalman_filter. A (T,) mask gates whole steps for every block; a (T, M) mask gates individual channels and is split per block.	None
solver	AbstractSolverStrategy | None	Optional solver strategy for the per-block innovation solves. Sequential mode only (the parallel filter does not thread it through).	None
parallel	bool	When True, run each block through gaussx.parallel_kalman_filter (O(log T) depth on accelerators) instead of the sequential scan.	False

Raises:

Type	Description
ValueError	If D is not divisible by block_size, or the observation dimension is not divisible by the number of blocks.

Returns:

Type	Description
FilterState	A gaussx.FilterState over the full D-dimensional state.
FilterState	Covariances are block-diagonal (T, D, D) embeddings of the
FilterState	per-block covariances (exact zeros off-block).

Source code in src/gaussx/_ssm/_meanfield_kalman.py

def meanfield_kalman_filter(
    transition: Float[Array, "*T D D"] | lx.AbstractLinearOperator,
    obs_model: Float[Array, "*T M D"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T D D"] | lx.AbstractLinearOperator,
    obs_noise: Float[Array, "*T M M"] | lx.AbstractLinearOperator,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " D"],
    init_cov: Float[Array, "D D"],
    *,
    block_size: int,
    mask: Bool[Array, " T"] | Bool[Array, "T M"] | None = None,
    solver: AbstractSolverStrategy | None = None,
    parallel: bool = False,
) -> FilterState:
    r"""Mean-field Kalman filter over ``L = D // block_size`` state blocks.

    Runs ``L`` independent ``block_size``-state Kalman filters under
    `jax.vmap` — ``O(T L d^3)`` total work instead of the full filter's
    ``O(T L^3 d^3)`` — at the cost of dropping posterior cross-block
    covariances. All inputs are projected onto their diagonal blocks;
    off-block entries are discarded, which is the mean-field projection
    (exact when the true cross-block couplings are zero). The returned
    log-likelihood is the sum of the per-block log-likelihoods,

    $$\log p(y_{1:T}) \approx \sum_{\ell=1}^{L}
    \log p^{(\ell)}\bigl(y_{1:T}^{(\ell)}\bigr).$$

    Accepts the same time-invariant / time-varying / operator input
    forms as `gaussx.kalman_filter`. A `gaussx.BlockDiag` whose
    sub-operators match the blocking is split structurally without
    materialising the full ``(D, D)`` matrix. The observation dimension
    must factorise the same way: channel block ``\ell`` (of size
    ``M // L``) is observed through state block ``\ell``.

    Args:
        transition: State transition ``A``. Shape ``(D, D)``,
            ``(T, D, D)``, or `lineax.AbstractLinearOperator`.
        obs_model: Observation matrix ``H``. Shape ``(M, D)``,
            ``(T, M, D)``, or operator. ``M`` must be divisible by
            ``L``.
        process_noise: Process noise covariance ``Q``. Shape ``(D, D)``,
            ``(T, D, D)``, or operator.
        obs_noise: Observation noise covariance ``R``. Shape ``(M, M)``,
            ``(T, M, M)``, or operator.
        observations: Observed data, shape ``(T, M)``.
        init_mean: Initial state mean, shape ``(D,)``.
        init_cov: Initial state covariance, shape ``(D, D)``.
        block_size: State block size ``d``; ``D`` must be divisible by
            it.
        mask: Optional observation mask, as in `gaussx.kalman_filter`.
            A ``(T,)`` mask gates whole steps for every block; a
            ``(T, M)`` mask gates individual channels and is split
            per block.
        solver: Optional solver strategy for the per-block innovation
            solves. Sequential mode only (the parallel filter does not
            thread it through).
        parallel: When ``True``, run each block through
            `gaussx.parallel_kalman_filter` (``O(log T)`` depth on
            accelerators) instead of the sequential scan.

    Raises:
        ValueError: If ``D`` is not divisible by ``block_size``, or the
            observation dimension is not divisible by the number of
            blocks.

    Returns:
        A `gaussx.FilterState` over the full ``D``-dimensional state.
        Covariances are block-diagonal ``(T, D, D)`` embeddings of the
        per-block covariances (exact zeros off-block).
    """
    D = init_mean.shape[0]
    M = observations.shape[-1]
    d = block_size
    if D % d != 0:
        raise ValueError(f"State dimension {D} is not divisible by block_size {d}.")
    L = D // d
    if M % L != 0:
        raise ValueError(
            f"Observation dimension {M} is not divisible by the number of blocks {L}."
        )
    m = M // L

    A_blocks = _split_diag_blocks(transition, L, d, d)
    H_blocks = _split_diag_blocks(obs_model, L, m, d)
    Q_blocks = _split_diag_blocks(process_noise, L, d, d)
    R_blocks = _split_diag_blocks(obs_noise, L, m, m)
    y_blocks = rearrange(observations, "t (l m) -> l t m", l=L)
    m0_blocks = rearrange(init_mean, "(l d) -> l d", l=L)
    P0_blocks = _split_diag_blocks(init_cov, L, d, d)

    base_filter = parallel_kalman_filter if parallel else kalman_filter

    channel_mask = mask is not None and jnp.ndim(mask) == 2
    if channel_mask:
        mask_blocks = rearrange(mask, "t (l m) -> l t m", l=L)

        def _run(A_b, H_b, Q_b, R_b, y_b, m0_b, P0_b, mask_b):
            return base_filter(
                A_b, H_b, Q_b, R_b, y_b, m0_b, P0_b, mask=mask_b, solver=solver
            )

        states = jax.vmap(_run)(
            A_blocks,
            H_blocks,
            Q_blocks,
            R_blocks,
            y_blocks,
            m0_blocks,
            P0_blocks,
            mask_blocks,
        )
    else:

        def _run(A_b, H_b, Q_b, R_b, y_b, m0_b, P0_b):
            return base_filter(
                A_b, H_b, Q_b, R_b, y_b, m0_b, P0_b, mask=mask, solver=solver
            )

        states = jax.vmap(_run)(
            A_blocks, H_blocks, Q_blocks, R_blocks, y_blocks, m0_blocks, P0_blocks
        )

    return FilterState(
        filtered_means=rearrange(states.filtered_means, "l t d -> t (l d)"),
        filtered_covs=_embed_block_diag(states.filtered_covs),
        predicted_means=rearrange(states.predicted_means, "l t d -> t (l d)"),
        predicted_covs=_embed_block_diag(states.predicted_covs),
        log_likelihood=jnp.sum(states.log_likelihood),
    )

meanfield_rts_smoother(filter_state: FilterState, transition: Float[Array, '*T D D'] | lx.AbstractLinearOperator, process_noise: Float[Array, '*T D D'] | lx.AbstractLinearOperator, *, block_size: int, solver: AbstractSolverStrategy | None = None, parallel: bool = False) -> tuple[Float[Array, 'T D'], Float[Array, 'T D D']]

Mean-field RTS smoother over L = D // block_size state blocks.

Backward pass paired with meanfield_kalman_filter: runs L independent gaussx.rts_smoother passes under jax.vmap, one per diagonal block. transition / process_noise accept the same forms as the filter and are projected onto their diagonal blocks.

Parameters:

Name	Type	Description	Default
filter_state	FilterState	Output of meanfield_kalman_filter (or any gaussx.FilterState whose covariances are block-diagonal — off-block entries are discarded).	required
transition	Float[Array, '*T D D'] | AbstractLinearOperator	State transition matrix or operator.	required
process_noise	Float[Array, '*T D D'] | AbstractLinearOperator	Process noise covariance or operator. (Unused by the standard RTS recurrence — kept for API symmetry.)	required
block_size	int	State block size d; D must be divisible by it.	required
solver	AbstractSolverStrategy | None	Optional solver strategy for the per-block smoother gains. Sequential mode only.	None
parallel	bool	When True, run each block through gaussx.parallel_rts_smoother instead of the sequential scan.	False

Raises:

Type	Description
ValueError	If D is not divisible by block_size.

Returns:

Type	Description
Float[Array, 'T D']	Tuple (smoothed_means, smoothed_covs) with shapes
Float[Array, 'T D D']	(T, D) and (T, D, D); covariances are block-diagonal
tuple[Float[Array, 'T D'], Float[Array, 'T D D']]	embeddings of the per-block smoothed covariances.

Source code in src/gaussx/_ssm/_meanfield_kalman.py

def meanfield_rts_smoother(
    filter_state: FilterState,
    transition: Float[Array, "*T D D"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "*T D D"] | lx.AbstractLinearOperator,
    *,
    block_size: int,
    solver: AbstractSolverStrategy | None = None,
    parallel: bool = False,
) -> tuple[Float[Array, "T D"], Float[Array, "T D D"]]:
    """Mean-field RTS smoother over ``L = D // block_size`` state blocks.

    Backward pass paired with `meanfield_kalman_filter`: runs ``L``
    independent `gaussx.rts_smoother` passes under `jax.vmap`, one per
    diagonal block. ``transition`` / ``process_noise`` accept the same
    forms as the filter and are projected onto their diagonal blocks.

    Args:
        filter_state: Output of `meanfield_kalman_filter` (or any
            `gaussx.FilterState` whose covariances are block-diagonal —
            off-block entries are discarded).
        transition: State transition matrix or operator.
        process_noise: Process noise covariance or operator. (Unused by
            the standard RTS recurrence — kept for API symmetry.)
        block_size: State block size ``d``; ``D`` must be divisible by
            it.
        solver: Optional solver strategy for the per-block smoother
            gains. Sequential mode only.
        parallel: When ``True``, run each block through
            `gaussx.parallel_rts_smoother` instead of the sequential
            scan.

    Raises:
        ValueError: If ``D`` is not divisible by ``block_size``.

    Returns:
        Tuple ``(smoothed_means, smoothed_covs)`` with shapes
        ``(T, D)`` and ``(T, D, D)``; covariances are block-diagonal
        embeddings of the per-block smoothed covariances.
    """
    D = filter_state.filtered_means.shape[-1]
    d = block_size
    if D % d != 0:
        raise ValueError(f"State dimension {D} is not divisible by block_size {d}.")
    L = D // d

    A_blocks = _split_diag_blocks(transition, L, d, d)
    # The standard RTS recurrence never reads ``process_noise`` (both base
    # smoothers ``del`` it), so skip the block split — it would only force
    # dense materialisation of a structured Q — and pass a shape-compatible
    # placeholder instead.
    del process_noise
    Q_blocks = jnp.zeros_like(A_blocks)

    block_states = FilterState(
        filtered_means=rearrange(filter_state.filtered_means, "t (l d) -> l t d", l=L),
        filtered_covs=_split_diag_blocks(filter_state.filtered_covs, L, d, d),
        predicted_means=rearrange(
            filter_state.predicted_means, "t (l d) -> l t d", l=L
        ),
        predicted_covs=_split_diag_blocks(filter_state.predicted_covs, L, d, d),
        log_likelihood=jnp.broadcast_to(filter_state.log_likelihood, (L,)),
    )

    base_smoother = parallel_rts_smoother if parallel else rts_smoother

    def _run(state_b, A_b, Q_b):
        return base_smoother(state_b, A_b, Q_b, solver=solver)

    s_means, s_covs = jax.vmap(_run)(block_states, A_blocks, Q_blocks)

    return (
        rearrange(s_means, "l t d -> t (l d)"),
        _embed_block_diag(s_covs),
    )

infinite_horizon_filter(transition: Float[Array, 'N N'] | lx.AbstractLinearOperator, obs_model: Float[Array, 'M N'] | lx.AbstractLinearOperator, process_noise: Float[Array, 'N N'] | lx.AbstractLinearOperator, obs_noise: Float[Array, 'M M'] | lx.AbstractLinearOperator, observations: Float[Array, 'T M'], init_mean: Float[Array, ' N'] | None = None, *, dare_result: DAREResult | None = None, max_iter: int = 100, tol: float = 1e-08, solver: AbstractSolverStrategy | None = None, woodbury_innovation: bool = False) -> InfiniteHorizonState

Infinite-horizon Kalman filter with fixed steady-state gain.

Uses the DARE solution for a constant Kalman gain K∞, avoiding per-step Riccati updates. For dense matrices, the per-step cost is O(N² + MN + M²) instead of O(N³) for the standard Kalman filter:

Predict:  x⁻ₜ = A xₜ₋₁
Update:   vₜ  = yₜ − H x⁻ₜ
          xₜ  = x⁻ₜ + K∞ vₜ

All four operator/array arguments accept either a raw JAX array or a lineax.AbstractLinearOperator. Operator inputs preserve their structural matvec inside the per-step scan; the sandwiches materialise once outside the scan.

Parameters:

Name	Type	Description	Default
transition	Float[Array, 'N N'] | AbstractLinearOperator	State transition matrix or operator, shape (N, N).	required
obs_model	Float[Array, 'M N'] | AbstractLinearOperator	Observation matrix or operator, shape (M, N).	required
process_noise	Float[Array, 'N N'] | AbstractLinearOperator	Process noise covariance or operator, shape (N, N).	required
obs_noise	Float[Array, 'M M'] | AbstractLinearOperator	Observation noise covariance or operator, shape (M, M).	required
observations	Float[Array, 'T M']	Observed data y, shape (T, M).	required
init_mean	Float[Array, ' N'] | None	Initial state mean, shape (N,). Defaults to zeros.	None
dare_result	DAREResult | None	Precomputed DARE result. If None, calls dare() internally.	None
max_iter	int	Maximum DARE iterations (used only if dare_result is None).	100
tol	float	DARE convergence tolerance (used only if dare_result is None).	1e-08
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None
woodbury_innovation	bool	When True, build the steady-state innovation covariance as gaussx.LowRankUpdate so structured R can use Woodbury solves/log-determinants.	False

Returns:

Type	Description
InfiniteHorizonState	An InfiniteHorizonState with filtered/predicted means,
InfiniteHorizonState	covariances, and total log-likelihood.

Source code in src/gaussx/_ssm/_infinite_horizon_kalman.py

def infinite_horizon_filter(
    transition: Float[Array, "N N"] | lx.AbstractLinearOperator,
    obs_model: Float[Array, "M N"] | lx.AbstractLinearOperator,
    process_noise: Float[Array, "N N"] | lx.AbstractLinearOperator,
    obs_noise: Float[Array, "M M"] | lx.AbstractLinearOperator,
    observations: Float[Array, "T M"],
    init_mean: Float[Array, " N"] | None = None,
    *,
    dare_result: DAREResult | None = None,
    max_iter: int = 100,
    tol: float = 1e-8,
    solver: AbstractSolverStrategy | None = None,
    woodbury_innovation: bool = False,
) -> InfiniteHorizonState:
    """Infinite-horizon Kalman filter with fixed steady-state gain.

    Uses the DARE solution for a constant Kalman gain K∞, avoiding
    per-step Riccati updates.  For dense matrices, the per-step cost is
    O(N² + MN + M²) instead of O(N³) for the standard Kalman filter:

        Predict:  x⁻ₜ = A xₜ₋₁
        Update:   vₜ  = yₜ − H x⁻ₜ
                  xₜ  = x⁻ₜ + K∞ vₜ

    All four operator/array arguments accept either a raw JAX array or
    a `lineax.AbstractLinearOperator`. Operator inputs preserve
    their structural matvec inside the per-step scan; the sandwiches
    materialise once outside the scan.

    Args:
        transition: State transition matrix or operator, shape ``(N, N)``.
        obs_model: Observation matrix or operator, shape ``(M, N)``.
        process_noise: Process noise covariance or operator, shape ``(N, N)``.
        obs_noise: Observation noise covariance or operator, shape ``(M, M)``.
        observations: Observed data y, shape ``(T, M)``.
        init_mean: Initial state mean, shape ``(N,)``. Defaults to zeros.
        dare_result: Precomputed DARE result. If ``None``, calls
            ``dare()`` internally.
        max_iter: Maximum DARE iterations (used only if ``dare_result``
            is ``None``).
        tol: DARE convergence tolerance (used only if ``dare_result``
            is ``None``).
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.
        woodbury_innovation: When ``True``, build the steady-state
            innovation covariance as `gaussx.LowRankUpdate` so
            structured ``R`` can use Woodbury solves/log-determinants.

    Returns:
        An ``InfiniteHorizonState`` with filtered/predicted means,
        covariances, and total log-likelihood.
    """
    if dare_result is None:
        dare_result = dare(
            transition,
            obs_model,
            process_noise,
            obs_noise,
            max_iter=max_iter,
            tol=tol,
            solver=solver,
            woodbury_innovation=woodbury_innovation,
        )

    A_op = _as_operator(transition)
    H_op = _as_operator(obs_model)
    Q_dense = _materialise(process_noise)
    # Keep ``R`` lazy when the Woodbury innovation path will consume the
    # operator directly — avoids an O(M²) allocation for large structured
    # noise (e.g. ``DiagonalLinearOperator`` with large ``M``).
    R_for_innovation = (
        obs_noise
        if woodbury_innovation and isinstance(obs_noise, lx.AbstractLinearOperator)
        else _materialise(obs_noise)
    )

    P_inf = dare_result.P_inf  # (N, N)
    K_inf = dare_result.K_inf  # (N, M)
    T = observations.shape[0]
    M = observations.shape[-1]
    N = A_op.out_size()

    # Precompute steady-state quantities
    P_inf_op = lx.MatrixLinearOperator(P_inf, lx.positive_semidefinite_tag)
    P_pred_inf = sandwich(A_op, P_inf_op).as_matrix() + Q_dense  # (N, N)
    S_inf = _innovation_covariance(
        H_op, P_pred_inf, R_for_innovation, woodbury=woodbury_innovation
    )
    ld_inf = dispatch_logdet(S_inf, solver)  # scalar

    # Steady-state filtered covariance: P_filt = (I − K∞ H) P⁻pred
    HP_pred_inf = _left_matmul(H_op, P_pred_inf)
    P_filt_inf = P_pred_inf - K_inf @ HP_pred_inf  # (N, N)

    def step(carry, y_t):
        x_filt, ll = carry

        x_pred = _matvec(transition, x_filt)  # (N,)
        v = y_t - _matvec(obs_model, x_pred)  # (M,)  innovation
        x_filt_new = x_pred + K_inf @ v  # (N,)

        # Log-likelihood increment.
        Sinv_v = dispatch_solve(S_inf, v, solver)  # (M,)
        ll_inc = -0.5 * (v @ Sinv_v + ld_inf + M * _LOG_2PI)

        return (x_filt_new, ll + ll_inc), (x_filt_new, x_pred)

    if init_mean is None:
        init_mean = jnp.zeros(N)
    init_carry = (init_mean, jnp.array(0.0))
    (_, total_ll), (f_means, p_means) = jax.lax.scan(
        step,
        init_carry,
        observations,
    )

    # Broadcast constant covariances to (T, N, N)
    f_covs = repeat(P_filt_inf, "n1 n2 -> T n1 n2", T=T)
    p_covs = repeat(P_pred_inf, "n1 n2 -> T n1 n2", T=T)

    return InfiniteHorizonState(
        filtered_means=f_means,
        filtered_covs=f_covs,
        predicted_means=p_means,
        predicted_covs=p_covs,
        log_likelihood=total_ll,
    )

infinite_horizon_smoother(filter_state: InfiniteHorizonState, transition: Float[Array, 'N N'] | lx.AbstractLinearOperator, dare_result: DAREResult, process_noise: Float[Array, 'N N'] | lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, 'T N'], Float[Array, 'T N N']]

Infinite-horizon RTS smoother with fixed steady-state gain.

Precomputes the steady-state smoother gain G∞ = P∞ Aᵀ P⁻pred⁻¹, then runs a backward scan with fixed G∞. The steady-state smoothed covariance is the solution of the discrete Lyapunov equation:

P_smooth = P∞ + G∞ (P_smooth − P⁻pred) G∞ᵀ

Parameters:

Name	Type	Description	Default
filter_state	InfiniteHorizonState	Output of infinite_horizon_filter.	required
transition	Float[Array, 'N N'] | AbstractLinearOperator	State transition matrix or operator, shape (N, N).	required
dare_result	DAREResult	DARE result used in the filter.	required
process_noise	Float[Array, 'N N'] | AbstractLinearOperator	Process noise covariance or operator, shape (N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
Float[Array, 'T N']	Tuple (smoothed_means, smoothed_covs) with shapes
Float[Array, 'T N N']	(T, N) and (T, N, N).

Source code in src/gaussx/_ssm/_infinite_horizon_kalman.py

def infinite_horizon_smoother(
    filter_state: InfiniteHorizonState,
    transition: Float[Array, "N N"] | lx.AbstractLinearOperator,
    dare_result: DAREResult,
    process_noise: Float[Array, "N N"] | lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "T N"], Float[Array, "T N N"]]:
    """Infinite-horizon RTS smoother with fixed steady-state gain.

    Precomputes the steady-state smoother gain G∞ = P∞ Aᵀ P⁻pred⁻¹,
    then runs a backward scan with fixed G∞.  The steady-state smoothed
    covariance is the solution of the discrete Lyapunov equation:

        P_smooth = P∞ + G∞ (P_smooth − P⁻pred) G∞ᵀ

    Args:
        filter_state: Output of ``infinite_horizon_filter``.
        transition: State transition matrix or operator, shape ``(N, N)``.
        dare_result: DARE result used in the filter.
        process_noise: Process noise covariance or operator, shape ``(N, N)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(smoothed_means, smoothed_covs)`` with shapes
        ``(T, N)`` and ``(T, N, N)``.
    """
    A_op = _as_operator(transition)
    Q_dense = _materialise(process_noise)
    P_inf = dare_result.P_inf  # (N, N)
    P_inf_op = lx.MatrixLinearOperator(P_inf, lx.positive_semidefinite_tag)
    P_pred_inf = sandwich(A_op, P_inf_op).as_matrix() + Q_dense  # (N, N)

    # Steady-state smoother gain: G∞ = P∞ Aᵀ P⁻pred⁻¹
    P_pred_inf_op = lx.MatrixLinearOperator(P_pred_inf, lx.positive_semidefinite_tag)
    G_inf = solve_rows(
        P_pred_inf_op,
        _right_matmul_transpose(P_inf, A_op),
        solver=solver,
    )  # (N, N)

    # Solve discrete Lyapunov equation:
    # P_smooth = P∞ + G∞ (P_smooth − P⁻pred) G∞ᵀ
    # ⟺ P_smooth − G∞ P_smooth G∞ᵀ = P∞ − G∞ P⁻pred G∞ᵀ
    # Routed through `discrete_lyapunov_solve` which uses a
    # per-factor eigendecomposition of ``G∞`` instead of materializing
    # the ``(N², N²)`` Kronecker matrix ``I − G∞ ⊗ G∞``.
    rhs = P_inf - G_inf @ P_pred_inf @ G_inf.T  # (N, N)
    P_smooth_inf = discrete_lyapunov_solve(G_inf, rhs)
    P_smooth_inf = symmetrize(P_smooth_inf)

    T = filter_state.filtered_means.shape[0]

    def step(carry, inputs):
        x_smooth = carry
        x_filt, x_pred = inputs
        x_smooth_new = x_filt + G_inf @ (x_smooth - x_pred)  # (N,)
        return x_smooth_new, x_smooth_new

    init = filter_state.filtered_means[T - 1]
    inputs = (
        filter_state.filtered_means[:-1][::-1],
        filter_state.predicted_means[1:][::-1],
    )

    _, s_means_rev = jax.lax.scan(step, init, inputs)

    s_means = jnp.concatenate(
        [s_means_rev[::-1], filter_state.filtered_means[T - 1 :]],
        axis=0,
    )  # (T, N)
    s_covs = repeat(P_smooth_inf, "n1 n2 -> T n1 n2", T=T)  # (T, N, N)

    return s_means, s_covs

dare(A: Float[Array, 'D D'] | lx.AbstractLinearOperator, H: Float[Array, 'M D'] | lx.AbstractLinearOperator, Q: Float[Array, 'D D'] | lx.AbstractLinearOperator, R: Float[Array, 'M M'] | lx.AbstractLinearOperator, *, P_init: Float[Array, 'D D'] | None = None, max_iter: int = 100, tol: float = 1e-08, solver: AbstractSolverStrategy | None = None, woodbury_innovation: bool = False) -> DAREResult

Discrete Algebraic Riccati Equation solver.

Iterates the Kalman predict-update equations until convergence:

Predict:  P⁻ = A P Aᵀ + Q
Update:   S = H P⁻ Hᵀ + R
          K = P⁻ Hᵀ S⁻¹
          P = (I - KH) P⁻

Convergence is declared when max|P_new - P_old| < tol.

Parameters:

Name	Type	Description	Default
A	Float[Array, 'D D'] | AbstractLinearOperator	Transition matrix or operator, shape (D, D).	required
H	Float[Array, 'M D'] | AbstractLinearOperator	Observation matrix or operator, shape (M, D).	required
Q	Float[Array, 'D D'] | AbstractLinearOperator	Process noise covariance or operator, shape (D, D).	required
R	Float[Array, 'M M'] | AbstractLinearOperator	Observation noise covariance or operator, shape (M, M).	required
P_init	Float[Array, 'D D'] | None	Initial covariance guess, shape (D, D). Defaults to Q.	None
max_iter	int	Maximum number of iterations.	100
tol	float	Convergence tolerance on the element-wise max absolute change.	1e-08
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None
woodbury_innovation	bool	When True, build S = H P⁻ Hᵀ + R as a gaussx.LowRankUpdate so structured R uses Woodbury solves.	False

Returns:

Type	Description
DAREResult	A DAREResult containing the steady-state covariance,
DAREResult	Kalman gain, and convergence flag.

Source code in src/gaussx/_ssm/_dare.py

def dare(
    A: Float[Array, "D D"] | lx.AbstractLinearOperator,
    H: Float[Array, "M D"] | lx.AbstractLinearOperator,
    Q: Float[Array, "D D"] | lx.AbstractLinearOperator,
    R: Float[Array, "M M"] | lx.AbstractLinearOperator,
    *,
    P_init: Float[Array, "D D"] | None = None,
    max_iter: int = 100,
    tol: float = 1e-8,
    solver: AbstractSolverStrategy | None = None,
    woodbury_innovation: bool = False,
) -> DAREResult:
    """Discrete Algebraic Riccati Equation solver.

    Iterates the Kalman predict-update equations until convergence:

        Predict:  P⁻ = A P Aᵀ + Q
        Update:   S = H P⁻ Hᵀ + R
                  K = P⁻ Hᵀ S⁻¹
                  P = (I - KH) P⁻

    Convergence is declared when ``max|P_new - P_old| < tol``.

    Args:
        A: Transition matrix or operator, shape ``(D, D)``.
        H: Observation matrix or operator, shape ``(M, D)``.
        Q: Process noise covariance or operator, shape ``(D, D)``.
        R: Observation noise covariance or operator, shape ``(M, M)``.
        P_init: Initial covariance guess, shape ``(D, D)``. Defaults to ``Q``.
        max_iter: Maximum number of iterations.
        tol: Convergence tolerance on the element-wise max absolute change.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.
        woodbury_innovation: When ``True``, build ``S = H P⁻ Hᵀ + R``
            as a `gaussx.LowRankUpdate` so structured ``R`` uses
            Woodbury solves.

    Returns:
        A `DAREResult` containing the steady-state covariance,
        Kalman gain, and convergence flag.
    """
    A_op = _as_operator(A)
    H_op = _as_operator(H)
    Q_dense = _materialise(Q)
    # Keep ``R`` lazy when the Woodbury innovation path will consume the
    # operator directly — avoids an O(M²) allocation for large structured
    # noise (e.g. ``DiagonalLinearOperator`` with large ``M``).
    R_for_innovation = (
        R
        if woodbury_innovation and isinstance(R, lx.AbstractLinearOperator)
        else _materialise(R)
    )

    if P_init is None:
        P_init = Q_dense

    def _step(
        P: Float[Array, "D D"],
    ) -> tuple[Float[Array, "D D"], Float[Array, "D M"]]:
        """One predict-update step. Returns ``(P_new, K)``."""
        P_op = lx.MatrixLinearOperator(P, lx.positive_semidefinite_tag)
        P_pred = sandwich(A_op, P_op).as_matrix() + Q_dense
        # K = P_pred @ H.T @ S⁻¹, computed via a single factorization
        # on the matrix RHS for numerical stability and efficiency.
        S_op = _innovation_covariance(
            H_op, P_pred, R_for_innovation, woodbury=woodbury_innovation
        )
        HP_pred = _left_matmul(H_op, P_pred)
        K = solve_matrix(S_op, HP_pred, solver=solver).T
        P_new = P_pred - K @ HP_pred
        return P_new, K

    def _cond(
        state: tuple[Float[Array, "D D"], int, Bool[Array, ""]],
    ) -> Bool[Array, ""]:
        _, i, converged = state
        return (i < max_iter) & (~converged)

    def _body(
        state: tuple[Float[Array, "D D"], int, Bool[Array, ""]],
    ) -> tuple[Float[Array, "D D"], int, Bool[Array, ""]]:
        P_old, i, _ = state
        P_new, _ = _step(P_old)
        converged = jnp.max(jnp.abs(P_new - P_old)) < tol
        return P_new, i + 1, converged

    init_state = (P_init, 0, jnp.array(False))
    P_inf, _, converged = jax.lax.while_loop(_cond, _body, init_state)

    # Compute the final gain from the converged covariance.
    _, K_inf = _step(P_inf)

    return DAREResult(P_inf=P_inf, K_inf=K_inf, converged=converged)

pairwise_marginals(means: Float[Array, 'T d'], covariances: Float[Array, 'T d d'], cross_covariances: Float[Array, 'Tm1 d d']) -> tuple[Float[Array, 'Tm1 two_d'], Float[Array, 'Tm1 two_d two_d']]

Joint p(x_k, x_{k+1}) for each consecutive pair.

For each pair (k, k+1), the joint distribution is:

p(x_k, x_{k+1}) = N([mu_k; mu_{k+1}],
                     [[P_k,      C_k^T],
                      [C_k,      P_{k+1}]])

where C_k = Cov[x_{k+1}, x_k] is the pairwise cross-covariance.

Parameters:

Name	Type	Description	Default
means	Float[Array, 'T d']	Smoothed means, shape (T, d).	required
covariances	Float[Array, 'T d d']	Smoothed covariances, shape (T, d, d).	required
cross_covariances	Float[Array, 'Tm1 d d']	Pairwise cross-covariances Cov[x_{k+1}, x_k], shape (T-1, d, d).	required

Returns:

Type	Description
Float[Array, 'Tm1 two_d']	Tuple (joint_means, joint_covariances) where:
Float[Array, 'Tm1 two_d two_d']	joint_means: shape (T-1, 2*d)
tuple[Float[Array, 'Tm1 two_d'], Float[Array, 'Tm1 two_d two_d']]	joint_covariances: shape (T-1, 2*d, 2*d)

Source code in src/gaussx/_ssm/_pairwise_marginals.py

def pairwise_marginals(
    means: Float[Array, "T d"],
    covariances: Float[Array, "T d d"],
    cross_covariances: Float[Array, "Tm1 d d"],
) -> tuple[Float[Array, "Tm1 two_d"], Float[Array, "Tm1 two_d two_d"]]:
    r"""Joint p(x_k, x_{k+1}) for each consecutive pair.

    For each pair ``(k, k+1)``, the joint distribution is:

        p(x_k, x_{k+1}) = N([mu_k; mu_{k+1}],
                             [[P_k,      C_k^T],
                              [C_k,      P_{k+1}]])

    where ``C_k = Cov[x_{k+1}, x_k]`` is the pairwise cross-covariance.

    Args:
        means: Smoothed means, shape ``(T, d)``.
        covariances: Smoothed covariances, shape ``(T, d, d)``.
        cross_covariances: Pairwise cross-covariances
            ``Cov[x_{k+1}, x_k]``, shape ``(T-1, d, d)``.

    Returns:
        Tuple ``(joint_means, joint_covariances)`` where:

        - ``joint_means``: shape ``(T-1, 2*d)``
        - ``joint_covariances``: shape ``(T-1, 2*d, 2*d)``
    """

    def _single_pair(
        m_k: Float[Array, " d"],
        m_kp1: Float[Array, " d"],
        P_k: Float[Array, "d d"],
        P_kp1: Float[Array, "d d"],
        C_k: Float[Array, "d d"],
    ) -> tuple[Float[Array, " two_d"], Float[Array, "two_d two_d"]]:
        joint_mean = jnp.concatenate([m_k, m_kp1])
        joint_cov = jnp.block(
            [
                [P_k, C_k.T],
                [C_k, P_kp1],
            ]
        )
        return joint_mean, joint_cov

    joint_means, joint_covariances = jax.vmap(_single_pair)(
        means[:-1],
        means[1:],
        covariances[:-1],
        covariances[1:],
        cross_covariances,
    )

    return joint_means, joint_covariances

SpInGP

State-space (sparse-in-time) GP inference: marginal likelihood and posterior through the SSM representation.

Structured linear algebra and Gaussian primitives for JAX.

spingp_log_likelihood(prior_precision: BlockTriDiag, emission_model: Array, obs_noise: lx.AbstractLinearOperator, observations: Float[Array, 'N d_obs'], *, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']

Log marginal likelihood via sparse inverse GP formulation.

Computes the log marginal likelihood using the precision-form Kalman filter (SpInGP):

1. Likelihood precision sites: $\Lambda_{lik} = H^T R^{-1} H$
2. Posterior precision: $\Lambda_{post} = \Lambda_{prior} + \Lambda_{lik}$
3. log p(y) via banded Cholesky logdet and quadratic form

The full expression is:

log p(y) = -0.5 * (N_{obs} * log(2\pi) + log|R|_{total}
           + y^T R^{-1} y - \eta^T \Lambda_{post}^{-1} \eta
           + log|\Lambda_{post}| - log|\Lambda_{prior}|)

where \(\eta = H^T R^{-1} y\).

All operations exploit banded structure for O(Nd³) cost.

The solver parameter controls the algorithm used for the large-scale posterior precision operations (solve, logdet). Observation noise operations always use structural dispatch since obs_noise is typically a small dense matrix.

Parameters:

Name	Type	Description	Default
prior_precision	BlockTriDiag	Prior precision as BlockTriDiag, shape (N, d, d) diagonal and (N-1, d, d) sub-diagonal.	required
emission_model	Array	Emission matrix H. Shape (d_obs, d) for shared or (N, d_obs, d) per time step.	required
obs_noise	AbstractLinearOperator	Observation noise covariance R operator.	required
observations	Float[Array, 'N d_obs']	Observations y, shape (N, d_obs).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for posterior precision operations. When None, uses structural dispatch. Observation noise operations always use structural dispatch.	None

Returns:

Type	Description
Float[Array, '']	Scalar log marginal likelihood.

Source code in src/gaussx/_ssm/_spingp.py

def spingp_log_likelihood(
    prior_precision: BlockTriDiag,
    emission_model: Array,
    obs_noise: lx.AbstractLinearOperator,
    observations: Float[Array, "N d_obs"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> Float[Array, ""]:
    r"""Log marginal likelihood via sparse inverse GP formulation.

    Computes the log marginal likelihood using the precision-form
    Kalman filter (SpInGP):

        1. Likelihood precision sites: $\Lambda_{lik} = H^T R^{-1} H$
        2. Posterior precision: $\Lambda_{post} = \Lambda_{prior} + \Lambda_{lik}$
        3. log p(y) via banded Cholesky logdet and quadratic form

    The full expression is:

        log p(y) = -0.5 * (N_{obs} * log(2\pi) + log|R|_{total}
                   + y^T R^{-1} y - \eta^T \Lambda_{post}^{-1} \eta
                   + log|\Lambda_{post}| - log|\Lambda_{prior}|)

    where $\eta = H^T R^{-1} y$.

    All operations exploit banded structure for O(Nd³) cost.

    The ``solver`` parameter controls the algorithm used for the
    large-scale posterior precision operations (solve, logdet).
    Observation noise operations always use structural dispatch
    since ``obs_noise`` is typically a small dense matrix.

    Args:
        prior_precision: Prior precision as ``BlockTriDiag``,
            shape ``(N, d, d)`` diagonal and ``(N-1, d, d)`` sub-diagonal.
        emission_model: Emission matrix H. Shape ``(d_obs, d)`` for
            shared or ``(N, d_obs, d)`` per time step.
        obs_noise: Observation noise covariance R operator.
        observations: Observations y, shape ``(N, d_obs)``.
        solver: Optional solver strategy for posterior precision
            operations. When ``None``, uses structural dispatch.
            Observation noise operations always use structural dispatch.

    Returns:
        Scalar log marginal likelihood.
    """
    N = prior_precision._num_blocks
    d = prior_precision._block_size
    N_obs = observations.size
    log_2pi = jnp.log(2.0 * jnp.pi)

    # Build likelihood precision and posterior precision
    lik_prec = _build_likelihood_precision(emission_model, obs_noise, N, d)
    post_prec = prior_precision.add(lik_prec)

    # Data vector: eta = H^T R^{-1} y
    eta = _build_data_vector(emission_model, obs_noise, observations)

    # Quadratic term: eta^T Lambda_post^{-1} eta
    post_solve = dispatch_solve(post_prec, eta, solver)
    quad_term = jnp.dot(eta, post_solve)

    # Observation quadratic: y^T R^{-1} y (obs_noise is small, use inv)
    R_inv = inv(obs_noise).as_matrix()
    obs_quad = jnp.sum(jax.vmap(lambda y_k: y_k @ R_inv @ y_k)(observations))

    # Log determinants (posterior precision: may be large, use solver)
    ld_post = dispatch_logdet(post_prec, solver)
    ld_prior = dispatch_logdet(prior_precision, solver)

    # Total observation noise logdet: N * log|R| (small, structural dispatch)
    ld_R = logdet(obs_noise)
    ld_R_total = N * ld_R

    return -0.5 * (
        N_obs * log_2pi + ld_R_total + obs_quad - quad_term + ld_post - ld_prior
    )

spingp_posterior(prior_precision: BlockTriDiag, emission_model: Array, obs_noise: lx.AbstractLinearOperator, observations: Float[Array, 'N d_obs'], *, prior_mean: Float[Array, ' Nd'] | None = None, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' Nd'], BlockTriDiag]

Posterior mean and precision via SpInGP.

Computes the posterior by adding likelihood precision sites to the prior precision and solving for the posterior mean:

\Lambda_{post} = \Lambda_{prior} + H^T R^{-1} H
\mu_{post} = \Lambda_{post}^{-1} (H^T R^{-1} y + \Lambda_{prior} \mu_{prior})

With prior_mean=None the prior is taken to be zero-mean and the second term vanishes.

Parameters:

Name	Type	Description	Default
prior_precision	BlockTriDiag	Prior precision as BlockTriDiag.	required
emission_model	Array	Emission matrix H. Shape (d_obs, d) for shared or (N, d_obs, d) per time step.	required
obs_noise	AbstractLinearOperator	Observation noise covariance R operator.	required
observations	Float[Array, 'N d_obs']	Observations y, shape (N, d_obs).	required
prior_mean	Float[Array, ' Nd'] | None	Optional prior mean mu_prior, shape (N * d,) — e.g. the mean half of gaussx.MarkovGaussian.to_precision_form. Defaults to zero.	None
solver	AbstractSolverStrategy | None	Optional solver strategy for posterior precision operations. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, ' Nd']	Tuple (posterior_mean, posterior_precision) where
BlockTriDiag	posterior_mean has shape (N * d,) and
tuple[Float[Array, ' Nd'], BlockTriDiag]	posterior_precision is BlockTriDiag.

Source code in src/gaussx/_ssm/_spingp.py

def spingp_posterior(
    prior_precision: BlockTriDiag,
    emission_model: Array,
    obs_noise: lx.AbstractLinearOperator,
    observations: Float[Array, "N d_obs"],
    *,
    prior_mean: Float[Array, " Nd"] | None = None,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " Nd"], BlockTriDiag]:
    r"""Posterior mean and precision via SpInGP.

    Computes the posterior by adding likelihood precision sites to the
    prior precision and solving for the posterior mean:

        \Lambda_{post} = \Lambda_{prior} + H^T R^{-1} H
        \mu_{post} = \Lambda_{post}^{-1} (H^T R^{-1} y + \Lambda_{prior} \mu_{prior})

    With ``prior_mean=None`` the prior is taken to be zero-mean and the
    second term vanishes.

    Args:
        prior_precision: Prior precision as ``BlockTriDiag``.
        emission_model: Emission matrix H. Shape ``(d_obs, d)`` for
            shared or ``(N, d_obs, d)`` per time step.
        obs_noise: Observation noise covariance R operator.
        observations: Observations y, shape ``(N, d_obs)``.
        prior_mean: Optional prior mean ``mu_prior``, shape ``(N * d,)``
            — e.g. the ``mean`` half of
            `gaussx.MarkovGaussian.to_precision_form`. Defaults to
            zero.
        solver: Optional solver strategy for posterior precision
            operations. When ``None``, uses structural dispatch.

    Returns:
        Tuple ``(posterior_mean, posterior_precision)`` where
        ``posterior_mean`` has shape ``(N * d,)`` and
        ``posterior_precision`` is ``BlockTriDiag``.
    """
    N = prior_precision._num_blocks
    d = prior_precision._block_size

    # Build likelihood precision and posterior precision
    lik_prec = _build_likelihood_precision(emission_model, obs_noise, N, d)
    post_prec = prior_precision.add(lik_prec)

    # Data vector: eta = H^T R^{-1} y (+ Lambda_prior mu_prior)
    eta = _build_data_vector(emission_model, obs_noise, observations)
    if prior_mean is not None:
        eta = eta + prior_precision.mv(prior_mean)

    # Posterior mean: Lambda_post^{-1} eta
    post_mean = dispatch_solve(post_prec, eta, solver)

    return post_mean, post_prec

Precision-form chains

The joint prior or posterior of a Gauss-Markov chain has a block-tridiagonal precision \(\Lambda\). The UDL factorisation \(\Lambda = U \tilde{D} U^{\top}\) — unit upper block-bidiagonal \(U\), block-diagonal \(\tilde{D}\) — is computed last block first by Schur complementation,

\[ \tilde{D}_T = D_T, \qquad \tilde{D}_k = D_k - A_k^{\top} \tilde{D}_{k+1}^{-1} A_k, \qquad U_k^{\top} = \tilde{D}_{k+1}^{-1} A_k, \]

in \(O(T d^3)\). Unlike the banded Cholesky factor behind cholesky, its blocks have a physical reading: they are the state-space model with that precision, \(A^{\mathrm{ssm}}_k = -U_k^{\top}\) and \(Q_k^{-1} = \tilde{D}_k\). One pass therefore turns a precision-form posterior — the output of spingp_posterior, or a CVI / natural-gradient site update — into a chain that the Kalman machinery can sample, predict and interpolate from, without materialising the \((Td) \times (Td)\) covariance.

MarkovGaussian is the distribution that carries both views: its canonical parameterisation is the generative tuple \((A_k, b_k, Q_k, \mu_0, P_0)\), and to_precision_form / from_precision_form round-trip through the UDL factors exactly.

Structured linear algebra and Gaussian primitives for JAX.

UDLDecomposition

Bases: Module

Banded factorisation \(\Lambda = U \tilde{D} U^{\top}\) of a BlockTriDiag.

\(U\) is unit upper block-bidiagonal with super-diagonal blocks \(U_k\) and \(\tilde{D} = \mathrm{blockdiag}(\tilde{D}_1, \dots, \tilde{D}_T)\). The factors are produced by the backward Schur-complement recurrence

\[ \tilde{D}_T = D_T, \qquad \tilde{D}_k = D_k - A_k^{\top} \tilde{D}_{k+1}^{-1} A_k, \qquad U_k^{\top} = \tilde{D}_{k+1}^{-1} A_k, \]

where \(D_k\) / \(A_k\) are the diagonal / sub-diagonal blocks of \(\Lambda\). This is not the Cholesky factor of \(\Lambda\): for a Gauss-Markov chain \(x_{k+1} = A^{\mathrm{ssm}}_k x_k + \varepsilon_k\) the blocks are exactly \(U_k^{\top} = -A^{\mathrm{ssm}}_k\) and \(\tilde{D}_k = Q_k^{-1}\), so one pass converts a precision-form posterior into a sampleable state-space model (see udl_to_ssm_params).

Every method costs \(O(T d^3)\) (or \(O(T d^2)\) for solve) and never forms the dense \((Td) \times (Td)\) matrix.

Attributes:

Name	Type	Description
U_sub	Float[Array, 'Tm1 d d']	Sub-diagonal blocks \(U_k^{\top} = \tilde{D}_{k+1}^{-1} A_k\) of \(U^{\top}\), shape (T-1, d, d).
D_diag	Float[Array, 'T d d']	Block-diagonal factors \(\tilde{D}_k\), shape (T, d, d).
chol_D	Float[Array, 'T d d']	Lower Cholesky factor of each \(\tilde{D}_k\), shape (T, d, d), cached for solve and logdet.

Source code in src/gaussx/_ssm/_udl.py

class UDLDecomposition(eqx.Module):
    r"""Banded factorisation $\Lambda = U \tilde{D} U^{\top}$ of a `BlockTriDiag`.

    $U$ is unit upper block-bidiagonal with super-diagonal blocks $U_k$
    and $\tilde{D} = \mathrm{blockdiag}(\tilde{D}_1, \dots, \tilde{D}_T)$.
    The factors are produced by the backward Schur-complement recurrence

    $$
    \tilde{D}_T = D_T, \qquad
    \tilde{D}_k = D_k - A_k^{\top} \tilde{D}_{k+1}^{-1} A_k, \qquad
    U_k^{\top} = \tilde{D}_{k+1}^{-1} A_k,
    $$

    where $D_k$ / $A_k$ are the diagonal / sub-diagonal blocks of
    $\Lambda$. This is **not** the Cholesky factor of $\Lambda$: for a
    Gauss-Markov chain $x_{k+1} = A^{\mathrm{ssm}}_k x_k + \varepsilon_k$
    the blocks are exactly $U_k^{\top} = -A^{\mathrm{ssm}}_k$ and
    $\tilde{D}_k = Q_k^{-1}$, so one pass converts a precision-form
    posterior into a sampleable state-space model
    (see `udl_to_ssm_params`).

    Every method costs $O(T d^3)$ (or $O(T d^2)$ for `solve`) and never
    forms the dense $(Td) \times (Td)$ matrix.

    Attributes:
        U_sub: Sub-diagonal blocks $U_k^{\top} = \tilde{D}_{k+1}^{-1} A_k$
            of $U^{\top}$, shape ``(T-1, d, d)``.
        D_diag: Block-diagonal factors $\tilde{D}_k$, shape ``(T, d, d)``.
        chol_D: Lower Cholesky factor of each $\tilde{D}_k$, shape
            ``(T, d, d)``, cached for `solve` and `logdet`.
    """

    U_sub: Float[Array, "Tm1 d d"]
    D_diag: Float[Array, "T d d"]
    chol_D: Float[Array, "T d d"]

    @property
    def num_blocks(self) -> int:
        """Number of blocks ``T``."""
        return self.D_diag.shape[0]

    @property
    def block_size(self) -> int:
        """Block dimension ``d``."""
        return self.D_diag.shape[1]

    def solve(self, rhs: Float[Array, " Td"]) -> Float[Array, " Td"]:
        r"""Solve $\Lambda x = b$ through the factors in $O(T d^2)$.

        Three banded sweeps: $U z = b$ (backward), $\tilde{D} w = z$
        (block-wise, via the cached Cholesky factors) and
        $U^{\top} x = w$ (forward).

        Args:
            rhs: Right-hand side, shape ``(T * d,)``.

        Returns:
            Solution, shape ``(T * d,)``.
        """
        T, d = self.num_blocks, self.block_size
        b = rearrange(rhs, "(T d) -> T d", T=T, d=d)

        # Backward sweep: z_T = b_T, z_k = b_k - U_k z_{k+1}.
        def _backward(z_next, inputs):
            b_k, U_sub_k = inputs
            z_k = b_k - einsum(U_sub_k, z_next, "j i, j -> i")
            return z_k, z_k

        _, z_rest = jax.lax.scan(_backward, b[-1], (b[:-1], self.U_sub), reverse=True)
        z = jnp.concatenate([z_rest, b[-1:]], axis=0)

        # Block-diagonal solve with the cached Cholesky factors.
        w = jax.vmap(_cho_solve)(self.chol_D, z)

        # Forward sweep: x_1 = w_1, x_k = w_k - U_{k-1}^T x_{k-1}.
        def _forward(x_prev, inputs):
            w_k, U_sub_k = inputs
            x_k = w_k - U_sub_k @ x_prev
            return x_k, x_k

        _, x_rest = jax.lax.scan(_forward, w[0], (w[1:], self.U_sub))
        x = jnp.concatenate([w[:1], x_rest], axis=0)
        return rearrange(x, "T d -> (T d)")

    def logdet(self) -> Float[Array, ""]:
        r"""$\log|\Lambda| = \sum_k \log|\tilde{D}_k|$, from the cached factors."""
        log_diag = jnp.log(jnp.diagonal(self.chol_D, axis1=-2, axis2=-1))
        return 2.0 * jnp.sum(log_diag)

    def as_block_tridiag(self) -> BlockTriDiag:
        r"""Reassemble $\Lambda = U \tilde{D} U^{\top}$ as a `BlockTriDiag`.

        Inverts the recurrence: $A_k = \tilde{D}_{k+1} U_k^{\top}$ and
        $D_k = \tilde{D}_k + A_k^{\top} \tilde{D}_{k+1}^{-1} A_k
        = \tilde{D}_k + U_k \tilde{D}_{k+1} U_k^{\top}$.
        """
        sub = einsum(self.D_diag[1:], self.U_sub, "T i j, T j k -> T i k")
        future = einsum(self.U_sub, sub, "T j i, T j k -> T i k")
        diag = self.D_diag.at[:-1].add(future)
        return BlockTriDiag(diag, sub)

num_blocks: int property

Number of blocks T.

block_size: int property

Block dimension d.

solve(rhs: Float[Array, ' Td']) -> Float[Array, ' Td']

Solve \(\Lambda x = b\) through the factors in \(O(T d^2)\).

Three banded sweeps: \(U z = b\) (backward), \(\tilde{D} w = z\) (block-wise, via the cached Cholesky factors) and \(U^{\top} x = w\) (forward).

Parameters:

Name	Type	Description	Default
rhs	Float[Array, ' Td']	Right-hand side, shape (T * d,).	required

Returns:

Type	Description
Float[Array, ' Td']	Solution, shape (T * d,).

Source code in src/gaussx/_ssm/_udl.py

def solve(self, rhs: Float[Array, " Td"]) -> Float[Array, " Td"]:
    r"""Solve $\Lambda x = b$ through the factors in $O(T d^2)$.

    Three banded sweeps: $U z = b$ (backward), $\tilde{D} w = z$
    (block-wise, via the cached Cholesky factors) and
    $U^{\top} x = w$ (forward).

    Args:
        rhs: Right-hand side, shape ``(T * d,)``.

    Returns:
        Solution, shape ``(T * d,)``.
    """
    T, d = self.num_blocks, self.block_size
    b = rearrange(rhs, "(T d) -> T d", T=T, d=d)

    # Backward sweep: z_T = b_T, z_k = b_k - U_k z_{k+1}.
    def _backward(z_next, inputs):
        b_k, U_sub_k = inputs
        z_k = b_k - einsum(U_sub_k, z_next, "j i, j -> i")
        return z_k, z_k

    _, z_rest = jax.lax.scan(_backward, b[-1], (b[:-1], self.U_sub), reverse=True)
    z = jnp.concatenate([z_rest, b[-1:]], axis=0)

    # Block-diagonal solve with the cached Cholesky factors.
    w = jax.vmap(_cho_solve)(self.chol_D, z)

    # Forward sweep: x_1 = w_1, x_k = w_k - U_{k-1}^T x_{k-1}.
    def _forward(x_prev, inputs):
        w_k, U_sub_k = inputs
        x_k = w_k - U_sub_k @ x_prev
        return x_k, x_k

    _, x_rest = jax.lax.scan(_forward, w[0], (w[1:], self.U_sub))
    x = jnp.concatenate([w[:1], x_rest], axis=0)
    return rearrange(x, "T d -> (T d)")

logdet() -> Float[Array, '']

\(\log|\Lambda| = \sum_k \log|\tilde{D}_k|\), from the cached factors.

Source code in src/gaussx/_ssm/_udl.py

def logdet(self) -> Float[Array, ""]:
    r"""$\log|\Lambda| = \sum_k \log|\tilde{D}_k|$, from the cached factors."""
    log_diag = jnp.log(jnp.diagonal(self.chol_D, axis1=-2, axis2=-1))
    return 2.0 * jnp.sum(log_diag)

as_block_tridiag() -> BlockTriDiag

Reassemble \(\Lambda = U \tilde{D} U^{\top}\) as a BlockTriDiag.

Inverts the recurrence: \(A_k = \tilde{D}_{k+1} U_k^{\top}\) and \(D_k = \tilde{D}_k + A_k^{\top} \tilde{D}_{k+1}^{-1} A_k = \tilde{D}_k + U_k \tilde{D}_{k+1} U_k^{\top}\).

Source code in src/gaussx/_ssm/_udl.py

def as_block_tridiag(self) -> BlockTriDiag:
    r"""Reassemble $\Lambda = U \tilde{D} U^{\top}$ as a `BlockTriDiag`.

    Inverts the recurrence: $A_k = \tilde{D}_{k+1} U_k^{\top}$ and
    $D_k = \tilde{D}_k + A_k^{\top} \tilde{D}_{k+1}^{-1} A_k
    = \tilde{D}_k + U_k \tilde{D}_{k+1} U_k^{\top}$.
    """
    sub = einsum(self.D_diag[1:], self.U_sub, "T i j, T j k -> T i k")
    future = einsum(self.U_sub, sub, "T j i, T j k -> T i k")
    diag = self.D_diag.at[:-1].add(future)
    return BlockTriDiag(diag, sub)

udl_decomposition(precision: BlockTriDiag) -> UDLDecomposition

Factorise a symmetric block-tridiagonal precision as \(U \tilde{D} U^{\top}\).

Runs the backward Schur-complement recurrence as a jax.lax.scan from the last block to the first, at \(O(T d^3)\) time and \(O(T d^2)\) memory. Each \(\tilde{D}_k\) is Cholesky-factored on the fly, so the input must be positive definite.

Parameters:

Name	Type	Description	Default
precision	BlockTriDiag	Symmetric positive-definite BlockTriDiag with T diagonal blocks of size (d, d).	required

Returns:

Type	Description
UDLDecomposition	The UDLDecomposition of precision.

Examples:

>>> import jax.numpy as jnp, gaussx
>>> diag = jnp.broadcast_to(2.0 * jnp.eye(2), (4, 2, 2))
>>> sub = jnp.broadcast_to(-0.5 * jnp.eye(2), (3, 2, 2))
>>> udl = gaussx.udl_decomposition(gaussx.BlockTriDiag(diag, sub))
>>> udl.U_sub.shape, udl.D_diag.shape
((3, 2, 2), (4, 2, 2))

Source code in src/gaussx/_ssm/_udl.py

def udl_decomposition(precision: BlockTriDiag) -> UDLDecomposition:
    r"""Factorise a symmetric block-tridiagonal precision as $U \tilde{D} U^{\top}$.

    Runs the backward Schur-complement recurrence as a `jax.lax.scan`
    from the last block to the first, at $O(T d^3)$ time and $O(T d^2)$
    memory. Each $\tilde{D}_k$ is Cholesky-factored on the fly, so the
    input must be positive definite.

    Args:
        precision: Symmetric positive-definite `BlockTriDiag` with ``T``
            diagonal blocks of size ``(d, d)``.

    Returns:
        The `UDLDecomposition` of ``precision``.

    Examples:
        >>> import jax.numpy as jnp, gaussx
        >>> diag = jnp.broadcast_to(2.0 * jnp.eye(2), (4, 2, 2))
        >>> sub = jnp.broadcast_to(-0.5 * jnp.eye(2), (3, 2, 2))
        >>> udl = gaussx.udl_decomposition(gaussx.BlockTriDiag(diag, sub))
        >>> udl.U_sub.shape, udl.D_diag.shape
        ((3, 2, 2), (4, 2, 2))
    """
    D_last = precision.diagonal[-1]
    chol_last = jnp.linalg.cholesky(D_last)

    def _step(carry, inputs):
        _, chol_next = carry
        D_k, A_k = inputs
        # U_k^T = D~_{k+1}^{-1} A_k  and  D~_k = D_k - A_k^T D~_{k+1}^{-1} A_k.
        U_sub_k = _cho_solve(chol_next, A_k)
        D_tilde_k = D_k - einsum(A_k, U_sub_k, "j i, j k -> i k")
        chol_k = jnp.linalg.cholesky(D_tilde_k)
        return (D_tilde_k, chol_k), (U_sub_k, D_tilde_k, chol_k)

    _, (U_sub, D_rest, chol_rest) = jax.lax.scan(
        _step,
        (D_last, chol_last),
        (precision.diagonal[:-1], precision.sub_diagonal),
        reverse=True,
    )
    D_diag = jnp.concatenate([D_rest, D_last[None]], axis=0)
    chol_D = jnp.concatenate([chol_rest, chol_last[None]], axis=0)
    return UDLDecomposition(U_sub=U_sub, D_diag=D_diag, chol_D=chol_D)

udl_to_ssm_params(udl: UDLDecomposition) -> tuple[Float[Array, 'Tm1 d d'], Float[Array, 'T d d'], Float[Array, 'T d d']]

Read the equivalent Gauss-Markov chain off a UDLDecomposition.

A chain \(x_0 \sim \mathcal{N}(\mu_0, P_0)\), \(x_{k+1} = A_k x_k + \varepsilon_k\), \(\varepsilon_k \sim \mathcal{N}(0, Q_{k+1})\) has precision \(\Lambda = U \tilde{D} U^{\top}\) with

\[ A_k = -U_k^{\top}, \qquad Q_{k}^{-1} = \tilde{D}_{k}, \qquad P_0^{-1} = \tilde{D}_0, \]

so the factors are the SSM. The returned Q follows the ssm_to_naturals / naturals_to_ssm layout: Q[0] is \(P_0\) and Q[k] for \(k \ge 1\) is the process noise entering state \(k\).

Parameters:

Name	Type	Description	Default
udl	UDLDecomposition	Factorisation of the chain's precision.	required

Returns:

Type	Description
Float[Array, 'Tm1 d d']	Tuple (A, Q, chol_Q): transitions of shape (T-1, d, d),
Float[Array, 'T d d']	covariances of shape (T, d, d) and their lower Cholesky
Float[Array, 'T d d']	factors of shape (T, d, d).

Source code in src/gaussx/_ssm/_udl.py

def udl_to_ssm_params(
    udl: UDLDecomposition,
) -> tuple[
    Float[Array, "Tm1 d d"],
    Float[Array, "T d d"],
    Float[Array, "T d d"],
]:
    r"""Read the equivalent Gauss-Markov chain off a `UDLDecomposition`.

    A chain $x_0 \sim \mathcal{N}(\mu_0, P_0)$,
    $x_{k+1} = A_k x_k + \varepsilon_k$, $\varepsilon_k \sim \mathcal{N}(0, Q_{k+1})$
    has precision $\Lambda = U \tilde{D} U^{\top}$ with

    $$
    A_k = -U_k^{\top}, \qquad Q_{k}^{-1} = \tilde{D}_{k}, \qquad P_0^{-1} = \tilde{D}_0,
    $$

    so the factors *are* the SSM. The returned ``Q`` follows the
    `ssm_to_naturals` / `naturals_to_ssm` layout: ``Q[0]`` is $P_0$ and
    ``Q[k]`` for $k \ge 1$ is the process noise entering state $k$.

    Args:
        udl: Factorisation of the chain's precision.

    Returns:
        Tuple ``(A, Q, chol_Q)``: transitions of shape ``(T-1, d, d)``,
        covariances of shape ``(T, d, d)`` and their lower Cholesky
        factors of shape ``(T, d, d)``.
    """
    A = -udl.U_sub
    Q = jax.vmap(_cho_inv)(udl.chol_D)
    chol_Q = jnp.linalg.cholesky(Q)
    return A, Q, chol_Q

udl_from_ssm_params(A: Float[Array, 'Tm1 d d'], Q: Float[Array, 'T d d']) -> UDLDecomposition

Build the UDLDecomposition of a Gauss-Markov chain's precision.

Inverse of udl_to_ssm_params: sets \(U_k^{\top} = -A_k\) and \(\tilde{D}_k = Q_k^{-1}\) directly, so the chain's block-tridiagonal precision is available through UDLDecomposition.as_block_tridiag without ever factorising it.

Parameters:

Name	Type	Description	Default
A	Float[Array, 'Tm1 d d']	Transition matrices, shape (T-1, d, d).	required
Q	Float[Array, 'T d d']	Covariances, shape (T, d, d); Q[0] is the initial covariance \(P_0\) and Q[k] the process noise entering state \(k\).	required

Returns:

Type	Description
UDLDecomposition	The factorisation of the chain's precision.

Source code in src/gaussx/_ssm/_udl.py

def udl_from_ssm_params(
    A: Float[Array, "Tm1 d d"],
    Q: Float[Array, "T d d"],
) -> UDLDecomposition:
    r"""Build the `UDLDecomposition` of a Gauss-Markov chain's precision.

    Inverse of `udl_to_ssm_params`: sets $U_k^{\top} = -A_k$ and
    $\tilde{D}_k = Q_k^{-1}$ directly, so the chain's block-tridiagonal
    precision is available through `UDLDecomposition.as_block_tridiag`
    without ever factorising it.

    Args:
        A: Transition matrices, shape ``(T-1, d, d)``.
        Q: Covariances, shape ``(T, d, d)``; ``Q[0]`` is the initial
            covariance $P_0$ and ``Q[k]`` the process noise entering
            state $k$.

    Returns:
        The factorisation of the chain's precision.
    """
    chol_Q = jnp.linalg.cholesky(Q)
    D_diag = jax.vmap(_cho_inv)(chol_Q)
    chol_D = jnp.linalg.cholesky(D_diag)
    return UDLDecomposition(U_sub=-A, D_diag=D_diag, chol_D=chol_D)

Sites & natural parameters

Conjugate-computation VI (CVI) site updates and the conversions between SSM moment, expectation, and natural parameterizations used by non-conjugate temporal inference.

Structured linear algebra and Gaussian primitives for JAX.

GaussianSites

Bases: Module

Time-varying Gaussian likelihood sites in natural parameterization.

Stores per-timestep natural parameters for N Gaussian sites, following the \eta_2 = -\tfrac{1}{2}\Lambda convention (consistent with gaussx.mean_cov_to_natural).

Attributes:

Name	Type	Description
nat1	Float[Array, 'N d']	Natural location parameters, shape (N, d).
nat2	Float[Array, 'N d d']	Natural precision parameters, shape (N, d, d). Stores -\tfrac{1}{2}\Lambda_k at each time step.

Source code in src/gaussx/_ssm/_cvi.py

class GaussianSites(eqx.Module):
    r"""Time-varying Gaussian likelihood sites in natural parameterization.

    Stores per-timestep natural parameters for ``N`` Gaussian sites,
    following the ``\eta_2 = -\tfrac{1}{2}\Lambda`` convention
    (consistent with `gaussx.mean_cov_to_natural`).

    Attributes:
        nat1: Natural location parameters, shape ``(N, d)``.
        nat2: Natural precision parameters, shape ``(N, d, d)``.
            Stores ``-\tfrac{1}{2}\Lambda_k`` at each time step.
    """

    nat1: Float[Array, "N d"]
    nat2: Float[Array, "N d d"]

cvi_update_sites(sites: GaussianSites, grad_nat1: Float[Array, 'N d'], grad_nat2: Float[Array, 'N d d'], rho: float) -> GaussianSites

Natural gradient update for CVI sites.

Performs a damped update in natural parameter space:

\theta \leftarrow (1 - \rho) \theta + \rho \nabla

Parameters:

Name	Type	Description	Default
sites	GaussianSites	Current Gaussian sites.	required
grad_nat1	Float[Array, 'N d']	Natural gradient for location, shape (N, d).	required
grad_nat2	Float[Array, 'N d d']	Natural gradient for precision, shape (N, d, d).	required
rho	float	Step size / damping factor in [0, 1].	required

Returns:

Type	Description
GaussianSites	Updated GaussianSites.

Source code in src/gaussx/_ssm/_cvi.py

def cvi_update_sites(
    sites: GaussianSites,
    grad_nat1: Float[Array, "N d"],
    grad_nat2: Float[Array, "N d d"],
    rho: float,
) -> GaussianSites:
    r"""Natural gradient update for CVI sites.

    Performs a damped update in natural parameter space:

        \theta \leftarrow (1 - \rho) \theta + \rho \nabla

    Args:
        sites: Current Gaussian sites.
        grad_nat1: Natural gradient for location, shape ``(N, d)``.
        grad_nat2: Natural gradient for precision, shape ``(N, d, d)``.
        rho: Step size / damping factor in ``[0, 1]``.

    Returns:
        Updated `GaussianSites`.
    """
    new_nat1 = (1.0 - rho) * sites.nat1 + rho * grad_nat1
    new_nat2 = (1.0 - rho) * sites.nat2 + rho * grad_nat2
    return GaussianSites(nat1=new_nat1, nat2=new_nat2)

sites_to_precision(sites: GaussianSites) -> BlockTriDiag

Convert Gaussian sites to a block-tridiagonal precision.

Returns a block-diagonal BlockTriDiag (zero sub-diagonals) representing the precision contribution of the sites. This can be added to a prior precision via .add() or + to form the posterior precision:

\Lambda_{post} = \Lambda_{prior} + \Lambda_{sites}

Since nat2 stores -\tfrac{1}{2}\Lambda, the precision blocks are -2 \cdot nat2.

Parameters:

Name	Type	Description	Default
sites	GaussianSites	Gaussian sites with nat2 in eta2 convention.	required

Returns:

Type	Description
BlockTriDiag	Block-diagonal BlockTriDiag precision.

Source code in src/gaussx/_ssm/_cvi.py

def sites_to_precision(sites: GaussianSites) -> BlockTriDiag:
    r"""Convert Gaussian sites to a block-tridiagonal precision.

    Returns a block-diagonal `BlockTriDiag` (zero
    sub-diagonals) representing the precision contribution of the
    sites. This can be added to a prior precision via ``.add()``
    or ``+`` to form the posterior precision:

        \Lambda_{post} = \Lambda_{prior} + \Lambda_{sites}

    Since ``nat2`` stores ``-\tfrac{1}{2}\Lambda``, the precision
    blocks are ``-2 \cdot nat2``.

    Args:
        sites: Gaussian sites with ``nat2`` in eta2 convention.

    Returns:
        Block-diagonal `BlockTriDiag` precision.
    """
    N, d = sites.nat1.shape
    diag_blocks = -2.0 * sites.nat2  # (N, d, d)
    sub_diag_blocks = jnp.zeros((N - 1, d, d), dtype=diag_blocks.dtype)
    return BlockTriDiag(diag_blocks, sub_diag_blocks)

cavity_from_marginal(marg_mean: Float[Array, ' *batch'], marg_var: Float[Array, ' *batch'], site_nat1: Float[Array, ' *batch'], site_nat2: Float[Array, ' *batch']) -> tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]

Compute cavity distribution by removing a site from the marginal.

Parameters:

Name	Type	Description	Default
marg_mean	Float[Array, ' *batch']	Marginal distribution means.	required
marg_var	Float[Array, ' *batch']	Marginal distribution variances (positive).	required
site_nat1	Float[Array, ' *batch']	Site precision-weighted means to remove.	required
site_nat2	Float[Array, ' *batch']	Site precisions to remove.	required

Returns:

Type	Description
tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]	Tuple (cav_mean, cav_var) of the cavity distribution.

Source code in src/gaussx/_ssm/_site_natural.py

def cavity_from_marginal(
    marg_mean: Float[Array, " *batch"],
    marg_var: Float[Array, " *batch"],
    site_nat1: Float[Array, " *batch"],
    site_nat2: Float[Array, " *batch"],
) -> tuple[Float[Array, " *batch"], Float[Array, " *batch"]]:
    """Compute cavity distribution by removing a site from the marginal.

    Args:
        marg_mean: Marginal distribution means.
        marg_var: Marginal distribution variances (positive).
        site_nat1: Site precision-weighted means to remove.
        site_nat2: Site precisions to remove.

    Returns:
        Tuple ``(cav_mean, cav_var)`` of the cavity distribution.
    """
    cav_prec = jnp.reciprocal(marg_var) - site_nat2
    cav_var = jnp.reciprocal(cav_prec)
    cav_mean = (marg_mean * jnp.reciprocal(marg_var) - site_nat1) * cav_var
    return cav_mean, cav_var

site_natural_from_tilted(tilted_mean: Float[Array, ' *batch'], tilted_var: Float[Array, ' *batch'], cav_mean: Float[Array, ' *batch'], cav_var: Float[Array, ' *batch']) -> tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]

Compute site natural parameters from tilted and cavity moments.

Parameters:

Name	Type	Description	Default
tilted_mean	Float[Array, ' *batch']	Tilted distribution means.	required
tilted_var	Float[Array, ' *batch']	Tilted distribution variances (positive).	required
cav_mean	Float[Array, ' *batch']	Cavity distribution means.	required
cav_var	Float[Array, ' *batch']	Cavity distribution variances (positive).	required

Returns:

Type	Description
tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]	Tuple (site_nat1, site_nat2).

Source code in src/gaussx/_ssm/_site_natural.py

def site_natural_from_tilted(
    tilted_mean: Float[Array, " *batch"],
    tilted_var: Float[Array, " *batch"],
    cav_mean: Float[Array, " *batch"],
    cav_var: Float[Array, " *batch"],
) -> tuple[Float[Array, " *batch"], Float[Array, " *batch"]]:
    """Compute site natural parameters from tilted and cavity moments.

    Args:
        tilted_mean: Tilted distribution means.
        tilted_var: Tilted distribution variances (positive).
        cav_mean: Cavity distribution means.
        cav_var: Cavity distribution variances (positive).

    Returns:
        Tuple ``(site_nat1, site_nat2)``.
    """
    site_nat2 = jnp.reciprocal(tilted_var) - jnp.reciprocal(cav_var)
    site_nat1 = tilted_mean * jnp.reciprocal(tilted_var) - cav_mean * jnp.reciprocal(
        cav_var
    )
    return site_nat1, site_nat2

site_mean_var_from_natural(site_nat1: Float[Array, ' *batch'], site_nat2: Float[Array, ' *batch']) -> tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]

Convert per-site natural parameters to mean/variance.

Parameters:

Name	Type	Description	Default
site_nat1	Float[Array, ' *batch']	Site precision-weighted means.	required
site_nat2	Float[Array, ' *batch']	Site precisions (positive for valid Gaussians).	required

Returns:

Type	Description
tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]	Tuple (mean, var) of the equivalent Gaussian site.

Source code in src/gaussx/_ssm/_site_natural.py

def site_mean_var_from_natural(
    site_nat1: Float[Array, " *batch"],
    site_nat2: Float[Array, " *batch"],
) -> tuple[Float[Array, " *batch"], Float[Array, " *batch"]]:
    """Convert per-site natural parameters to mean/variance.

    Args:
        site_nat1: Site precision-weighted means.
        site_nat2: Site precisions (positive for valid Gaussians).

    Returns:
        Tuple ``(mean, var)`` of the equivalent Gaussian site.
    """
    var = jnp.reciprocal(site_nat2)
    mean = site_nat1 * var
    return mean, var

expectations_to_ssm(eta1: Float[Array, ' Nd'], eta2: BlockTriDiag) -> tuple[Float[Array, 'N d'], Float[Array, 'N d d'], Float[Array, 'Nm1 d d']]

Convert expectation parameters back to SSM marginals.

Recovers (means, covs, cross_covs) from the expectation parameters of the joint Gaussian.

Parameters:

Name	Type	Description	Default
eta1	Float[Array, ' Nd']	Concatenated means, shape (N*d,).	required
eta2	BlockTriDiag	Second-moment BlockTriDiag.	required

Returns:

Type	Description
Float[Array, 'N d']	Tuple (means, covs, cross_covs) where:
Float[Array, 'N d d']	means: shape (N, d)
Float[Array, 'Nm1 d d']	covs: shape (N, d, d)
tuple[Float[Array, 'N d'], Float[Array, 'N d d'], Float[Array, 'Nm1 d d']]	cross_covs: shape (N-1, d, d)

Source code in src/gaussx/_ssm/_ssm_natural.py

def expectations_to_ssm(
    eta1: Float[Array, " Nd"],
    eta2: BlockTriDiag,
) -> tuple[
    Float[Array, "N d"],
    Float[Array, "N d d"],
    Float[Array, "Nm1 d d"],
]:
    r"""Convert expectation parameters back to SSM marginals.

    Recovers ``(means, covs, cross_covs)`` from the expectation
    parameters of the joint Gaussian.

    Args:
        eta1: Concatenated means, shape ``(N*d,)``.
        eta2: Second-moment `BlockTriDiag`.

    Returns:
        Tuple ``(means, covs, cross_covs)`` where:

        - ``means``: shape ``(N, d)``
        - ``covs``: shape ``(N, d, d)``
        - ``cross_covs``: shape ``(N-1, d, d)``
    """
    d = eta2._block_size
    N = eta2._num_blocks

    means = rearrange(eta1, "(N d) -> N d", N=N, d=d)

    # covs = E[xₖ xₖᵀ] − mₖ mₖᵀ
    covs = eta2.diagonal - einsum(means, means, "N i, N j -> N i j")

    # cross_covs = E[xₖ₊₁ xₖᵀ] − mₖ₊₁ mₖᵀ
    cross_covs = eta2.sub_diagonal - einsum(means[1:], means[:-1], "N i, N j -> N i j")

    return means, covs, cross_covs

naturals_to_ssm(theta_linear: Float[Array, ' Nd'], theta_precision: BlockTriDiag, *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, 'Nm1 d d'], Float[Array, 'N d d'], Float[Array, ' d'], Float[Array, 'd d']]

Convert natural parameters back to SSM parameters.

Recovers (A, Q, \mu_0, P_0) from the block-tridiagonal natural parameters via a backward recurrence on the precision blocks.

Parameters:

Name	Type	Description	Default
theta_linear	Float[Array, ' Nd']	Natural location parameter, shape (N*d,).	required
theta_precision	BlockTriDiag	Natural precision parameter as BlockTriDiag (eta2 convention).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch. This parameter is accepted for API consistency but is not currently used by the matrix inverse operations in this function.	None

Returns:

Type	Description
Float[Array, 'Nm1 d d']	Tuple (A, Q, mu_0, P_0) where:
Float[Array, 'N d d']	A: Transition matrices, shape (N-1, d, d).
Float[Array, ' d']	Q: Process noise covariances, shape (N, d, d).
Float[Array, 'd d']	mu_0: Initial mean, shape (d,).
tuple[Float[Array, 'Nm1 d d'], Float[Array, 'N d d'], Float[Array, ' d'], Float[Array, 'd d']]	P_0: Initial covariance, shape (d, d).

Source code in src/gaussx/_ssm/_ssm_natural.py

def naturals_to_ssm(
    theta_linear: Float[Array, " Nd"],
    theta_precision: BlockTriDiag,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[
    Float[Array, "Nm1 d d"],
    Float[Array, "N d d"],
    Float[Array, " d"],
    Float[Array, "d d"],
]:
    r"""Convert natural parameters back to SSM parameters.

    Recovers ``(A, Q, \mu_0, P_0)`` from the block-tridiagonal natural
    parameters via a backward recurrence on the precision blocks.

    Args:
        theta_linear: Natural location parameter, shape ``(N*d,)``.
        theta_precision: Natural precision parameter as
            `BlockTriDiag` (eta2 convention).
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch. This parameter
            is accepted for API consistency but is not currently used by the
            matrix inverse operations in this function.

    Returns:
        Tuple ``(A, Q, mu_0, P_0)`` where:
        - ``A``: Transition matrices, shape ``(N-1, d, d)``.
        - ``Q``: Process noise covariances, shape ``(N, d, d)``.
        - ``mu_0``: Initial mean, shape ``(d,)``.
        - ``P_0``: Initial covariance, shape ``(d, d)``.
    """
    del solver  # inv does not accept a solver; parameter reserved for future use
    d = theta_precision._block_size

    # Convert from eta2 to raw precision
    prec_diag = -2.0 * theta_precision.diagonal  # (N, d, d)
    prec_sub = -2.0 * theta_precision.sub_diagonal  # (N-1, d, d)

    # Backward recurrence to recover Q and A
    # Start from last block: Q[N-1] = inv(prec_diag[N-1])
    # Then for k = N-2 down to 0:
    #   A[k] = Q[k+1] @ (-prec_sub[k])  (sub-diag was -Q_{k+1}^{-1} A_k)
    #   Q[k] = inv(prec_diag[k] - A[k]^T @ Q[k+1]^{-1} @ A[k])

    def _backward_step(Q_next_inv, inputs):
        diag_k, sub_k = inputs
        Q_next = inv(
            lx.MatrixLinearOperator(Q_next_inv, lx.positive_semidefinite_tag)
        ).as_matrix()
        # sub_k = -Q_{k+1}^{-1} A_k, so A_k = -Q_{k+1} @ sub_k
        A_k = -Q_next @ sub_k
        # Q_k^{-1} = diag_k - A_k^T @ Q_next_inv @ A_k
        Q_k_inv = diag_k - A_k.T @ Q_next_inv @ A_k
        return Q_k_inv, (A_k, Q_k_inv)

    Q_last_inv = prec_diag[-1]

    # Reverse scan: iterate from k=N-2 down to 0
    _, (A_rev, Q_inv_rev) = jax.lax.scan(
        _backward_step,
        Q_last_inv,
        (prec_diag[:-1], prec_sub),
        reverse=True,
    )

    # A_rev is (N-1, d, d), Q_inv_rev is (N-1, d, d) for k=0..N-2
    A = A_rev

    # Q: invert all Q_inv values (batch over N)
    def _inv_single(q_inv):
        return inv(
            lx.MatrixLinearOperator(q_inv, lx.positive_semidefinite_tag)
        ).as_matrix()

    Q_inv_all = jnp.concatenate([Q_inv_rev, Q_last_inv[None]], axis=0)
    Q = jax.vmap(_inv_single)(Q_inv_all)

    # Recover initial conditions
    P_0 = Q[0]
    mu_0 = P_0 @ theta_linear[:d]

    return A, Q, mu_0, P_0

ssm_to_expectations(means: Float[Array, 'N d'], covs: Float[Array, 'N d d'], cross_covs: Float[Array, 'Nm1 d d']) -> tuple[Float[Array, ' Nd'], BlockTriDiag]

Convert SSM marginals to expectation parameters.

Given filtered or smoothed marginals, computes the expectation parameters (eta1, eta2) of the joint Gaussian where:

• eta1 = E[x] (concatenated means)
• eta2 is a BlockTriDiag storing the block-tridiagonal subset of E[xx^T] (second moments matching the Gauss-Markov sparsity pattern, not the full dense matrix)

The diagonal blocks of eta2 are E[x_k x_k^T] = P_k + m_k m_k^T and the sub-diagonal blocks are E[x_{k+1} x_k^T] = C_k + m_{k+1} m_k^T where C_k is the cross-covariance Cov(x_{k+1}, x_k).

Parameters:

Name	Type	Description	Default
means	Float[Array, 'N d']	Marginal means, shape (N, d).	required
covs	Float[Array, 'N d d']	Marginal covariances, shape (N, d, d).	required
cross_covs	Float[Array, 'Nm1 d d']	Cross-covariances Cov(x_{k+1}, x_k), shape (N-1, d, d).	required

Returns:

Type	Description
Float[Array, ' Nd']	Tuple (eta1, eta2) where eta1 has shape (N*d,)
BlockTriDiag	and eta2 is a BlockTriDiag.

Source code in src/gaussx/_ssm/_ssm_natural.py

def ssm_to_expectations(
    means: Float[Array, "N d"],
    covs: Float[Array, "N d d"],
    cross_covs: Float[Array, "Nm1 d d"],
) -> tuple[Float[Array, " Nd"], BlockTriDiag]:
    r"""Convert SSM marginals to expectation parameters.

    Given filtered or smoothed marginals, computes the expectation
    parameters ``(eta1, eta2)`` of the joint Gaussian where:

    - ``eta1 = E[x]`` (concatenated means)
    - ``eta2`` is a `BlockTriDiag` storing the
      block-tridiagonal subset of ``E[xx^T]`` (second moments matching
      the Gauss-Markov sparsity pattern, not the full dense matrix)

    The diagonal blocks of ``eta2`` are ``E[x_k x_k^T] = P_k + m_k m_k^T``
    and the sub-diagonal blocks are
    ``E[x_{k+1} x_k^T] = C_k + m_{k+1} m_k^T`` where ``C_k`` is the
    cross-covariance ``Cov(x_{k+1}, x_k)``.

    Args:
        means: Marginal means, shape ``(N, d)``.
        covs: Marginal covariances, shape ``(N, d, d)``.
        cross_covs: Cross-covariances ``Cov(x_{k+1}, x_k)``,
            shape ``(N-1, d, d)``.

    Returns:
        Tuple ``(eta1, eta2)`` where ``eta1`` has shape ``(N*d,)``
        and ``eta2`` is a `BlockTriDiag`.
    """
    _N, _d = means.shape

    # eta1 = concatenated means
    eta1 = rearrange(means, "N d -> (N d)")

    # Diagonal blocks: E[xₖ xₖᵀ] = Pₖ + mₖ mₖᵀ
    diag = covs + einsum(means, means, "N i, N j -> N i j")  # (N, d, d)

    # Sub-diagonal blocks: E[xₖ₊₁ xₖᵀ] = Cₖ + mₖ₊₁ mₖᵀ
    sub_diag = cross_covs + einsum(
        means[1:], means[:-1], "N i, N j -> N i j"
    )  # (N-1, d, d)

    eta2 = BlockTriDiag(diag, sub_diag)
    return eta1, eta2

ssm_to_naturals(A: Float[Array, 'Nm1 d d'], Q: Float[Array, 'N d d'], mu_0: Float[Array, ' d'], P_0: Float[Array, 'd d'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' Nd'], BlockTriDiag]

Convert SSM parameters to natural parameters.

For a linear-Gaussian state-space model:

x_0 \sim N(\mu_0, P_0)
x_{k+1} = A_k x_k + \epsilon_k,\quad \epsilon_k \sim N(0, Q_{k+1})

the joint prior p(x_0, \ldots, x_{N-1}) has a block-tridiagonal precision matrix. This function returns its natural parameters (\theta_1, \theta_2) where \theta_2 = -\tfrac{1}{2}\Lambda (matching the convention in gaussx.mean_cov_to_natural).

Parameters:

Name	Type	Description	Default
A	Float[Array, 'Nm1 d d']	Transition matrices, shape (N-1, d, d).	required
Q	Float[Array, 'N d d']	Process noise covariances, shape (N, d, d). Q[0] must equal P_0 and Q[k] for k >= 1 is the process noise at step k.	required
mu_0	Float[Array, ' d']	Initial mean, shape (d,).	required
P_0	Float[Array, 'd d']	Initial covariance, shape (d, d).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
Float[Array, ' Nd']	Tuple (theta_linear, theta_precision) where
BlockTriDiag	theta_linear has shape (N*d,) and
tuple[Float[Array, ' Nd'], BlockTriDiag]	theta_precision is a BlockTriDiag
tuple[Float[Array, ' Nd'], BlockTriDiag]	in the eta_2 = -0.5 * Lambda convention.

Source code in src/gaussx/_ssm/_ssm_natural.py

def ssm_to_naturals(
    A: Float[Array, "Nm1 d d"],
    Q: Float[Array, "N d d"],
    mu_0: Float[Array, " d"],
    P_0: Float[Array, "d d"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " Nd"], BlockTriDiag]:
    r"""Convert SSM parameters to natural parameters.

    For a linear-Gaussian state-space model:

        x_0 \sim N(\mu_0, P_0)
        x_{k+1} = A_k x_k + \epsilon_k,\quad \epsilon_k \sim N(0, Q_{k+1})

    the joint prior ``p(x_0, \ldots, x_{N-1})`` has a block-tridiagonal
    precision matrix. This function returns its natural parameters
    ``(\theta_1, \theta_2)`` where ``\theta_2 = -\tfrac{1}{2}\Lambda``
    (matching the convention in `gaussx.mean_cov_to_natural`).

    Args:
        A: Transition matrices, shape ``(N-1, d, d)``.
        Q: Process noise covariances, shape ``(N, d, d)``.
            ``Q[0]`` must equal ``P_0`` and ``Q[k]`` for ``k >= 1`` is the
            process noise at step ``k``.
        mu_0: Initial mean, shape ``(d,)``.
        P_0: Initial covariance, shape ``(d, d)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(theta_linear, theta_precision)`` where
        ``theta_linear`` has shape ``(N*d,)`` and
        ``theta_precision`` is a `BlockTriDiag`
        in the ``eta_2 = -0.5 * Lambda`` convention.
    """
    N = Q.shape[0]
    d = Q.shape[1]

    try:
        q0_matches_p0 = bool(jnp.allclose(Q[0], P_0))
    except jax.errors.TracerBoolConversionError:
        q0_matches_p0 = True  # skip validation under jax.jit

    if not q0_matches_p0:
        msg = "Q[0] must match P_0 so the returned natural parameters are consistent"
        raise ValueError(msg)

    # Invert all process noise covariances (batch over N)
    def _inv_single(q):
        return inv(lx.MatrixLinearOperator(q, lx.positive_semidefinite_tag)).as_matrix()

    Q_inv = jax.vmap(_inv_single)(Q)  # (N, d, d)
    P_0_op = lx.MatrixLinearOperator(P_0, lx.positive_semidefinite_tag)
    P_0_inv = inv(P_0_op).as_matrix()

    # Future contributions: A_k^T Q_{k+1}^{-1} A_k for k = 0..N-2
    future = jax.vmap(lambda Ak, Qinv_kp1: Ak.T @ Qinv_kp1 @ Ak)(
        A, Q_inv[1:]
    )  # (N-1, d, d)

    # Precision diagonal blocks (raw Lambda, not eta2)
    # D[0] = P_0^{-1} + A[0]^T Q[1]^{-1} A[0]
    # D[k] = Q[k]^{-1} + A[k]^T Q[k+1]^{-1} A[k]  for k=1..N-2
    # D[N-1] = Q[N-1]^{-1}
    diag = jnp.zeros((N, d, d), dtype=Q.dtype)
    diag = diag.at[0].set(P_0_inv + future[0] if N > 1 else P_0_inv)
    if N > 2:
        diag = diag.at[1:-1].set(Q_inv[1:-1] + future[1:])
    diag = diag.at[-1].set(Q_inv[-1])

    # Sub-diagonal blocks (raw precision off-diagonal)
    # S[k] = -Q[k+1]^{-1} A[k]  for k=0..N-2
    # (negative because precision cross-terms are negative for transitions)
    sub_diag = jax.vmap(lambda Qinv_kp1, Ak: -Qinv_kp1 @ Ak)(
        Q_inv[1:], A
    )  # (N-1, d, d)

    # Convert to eta2 convention: theta_precision = -0.5 * Lambda
    theta_precision = BlockTriDiag(-0.5 * diag, -0.5 * sub_diag)

    # Linear natural parameter: eta1 = Lambda @ mu
    # For zero-mean transitions, only the initial condition contributes
    theta_linear = jnp.zeros(N * d, dtype=Q.dtype)
    eta1_0 = dispatch_solve(P_0_op, mu_0, solver)
    theta_linear = theta_linear.at[:d].set(eta1_0)

    return theta_linear, theta_precision

Process noise

The exact discretisation \(Q = P_\infty - A P_\infty A^\top\), which is the forward direction of the discrete Lyapunov equation that discrete_lyapunov_solve inverts. The congruence is delegated to cov_transform, so passing operators returns a lazy operator with structure intact — matched Kronecker factors stay factorised, a diagonal \(P_\infty\) skips its \((N, N)\) materialization. Passing arrays returns an array.

Structured linear algebra and Gaussian primitives for JAX.

process_noise_covariance(A: Float[Array, 'N N'] | lx.AbstractLinearOperator, Pinf: Float[Array, 'N N'] | lx.AbstractLinearOperator) -> Float[Array, 'N N'] | lx.AbstractLinearOperator

process_noise_covariance(A: Float[Array, 'N N'], Pinf: Float[Array, 'N N']) -> Float[Array, 'N N']

process_noise_covariance(A: lx.AbstractLinearOperator, Pinf: Float[Array, 'N N'] | lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

process_noise_covariance(A: Float[Array, 'N N'] | lx.AbstractLinearOperator, Pinf: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

Compute process noise from stationary covariance.

Computes:

Q = Pinf - A @ Pinf @ A^T

For a discrete-time state-space model with stationary covariance Pinf and transition matrix A. This is the standard "steady-state trick": it follows from stationarity of the discretised recursion Pinf = A Pinf A^T + Q, and avoids matrix-fraction integration (Särkkä & Solin 2019, §6.4).

It is the forward direction of the discrete Lyapunov equation that gaussx.discrete_lyapunov_solve inverts: this maps Pinf -> Q, that one maps Q -> Pinf for the same A.

The A Pinf A^T term is delegated to gaussx.cov_transform, so structure is exploited rather than re-derived: a diagonal Pinf skips its (N, N) materialization, and matched gaussx.Kronecker / gaussx.BlockDiag operands stay factorised through gaussx.sandwich.

Passing operators returns a lazy operator — nothing is materialized, and the structural class of the result follows gaussx.sandwich. Passing arrays returns an array, as before.

The result is not symmetrised. Callers that need a covariance clean of floating-point asymmetry should wrap it in gaussx.symmetrize, as gaussx.SDEKernel.discretise does.

Parameters:

Name	Type	Description	Default
A	Float[Array, 'N N'] | AbstractLinearOperator	State transition matrix, shape (N, N) — array or operator.	required
Pinf	Float[Array, 'N N'] | AbstractLinearOperator	Stationary covariance, shape (N, N) — array or operator.	required

Returns:

Type	Description
Float[Array, 'N N'] | AbstractLinearOperator	Process noise covariance Q, shape (N, N). An operator when
Float[Array, 'N N'] | AbstractLinearOperator	either argument is an operator, otherwise an array.

Source code in src/gaussx/_ssm/_discretise.py

def process_noise_covariance(
    A: Float[Array, "N N"] | lx.AbstractLinearOperator,
    Pinf: Float[Array, "N N"] | lx.AbstractLinearOperator,
) -> Float[Array, "N N"] | lx.AbstractLinearOperator:
    r"""Compute process noise from stationary covariance.

    Computes:

        Q = Pinf - A @ Pinf @ A^T

    For a discrete-time state-space model with stationary covariance
    ``Pinf`` and transition matrix ``A``. This is the standard
    "steady-state trick": it follows from stationarity of the discretised
    recursion ``Pinf = A Pinf A^T + Q``, and avoids matrix-fraction
    integration (Särkkä & Solin 2019, §6.4).

    It is the forward direction of the discrete Lyapunov equation that
    `gaussx.discrete_lyapunov_solve` inverts: this maps
    ``Pinf -> Q``, that one maps ``Q -> Pinf`` for the same ``A``.

    The ``A Pinf A^T`` term is delegated to
    `gaussx.cov_transform`, so structure is exploited rather than
    re-derived: a diagonal ``Pinf`` skips its ``(N, N)`` materialization, and
    matched `gaussx.Kronecker` / `gaussx.BlockDiag`
    operands stay factorised through `gaussx.sandwich`.

    Passing **operators** returns a lazy operator — nothing is materialized,
    and the structural class of the result follows `gaussx.sandwich`.
    Passing **arrays** returns an array, as before.

    The result is not symmetrised. Callers that need a covariance clean of
    floating-point asymmetry should wrap it in `gaussx.symmetrize`,
    as `gaussx.SDEKernel.discretise` does.

    Args:
        A: State transition matrix, shape ``(N, N)`` — array or operator.
        Pinf: Stationary covariance, shape ``(N, N)`` — array or operator.

    Returns:
        Process noise covariance Q, shape ``(N, N)``. An operator when
        either argument is an operator, otherwise an array.
    """
    Pinf_op = (
        Pinf
        if isinstance(Pinf, lx.AbstractLinearOperator)
        else lx.MatrixLinearOperator(Pinf, lx.symmetric_tag)
    )
    # One implementation of the congruence, for both paths. Wrapping a dense
    # ``Pinf`` costs nothing at runtime — it is pytree bookkeeping at trace
    # time, and the dense branch of ``cov_transform`` evaluates the identical
    # ``A @ Pinf @ A.T`` expression.
    congruence = cov_transform(A, Pinf_op)

    if isinstance(A, lx.AbstractLinearOperator) or isinstance(
        Pinf, lx.AbstractLinearOperator
    ):
        return Pinf_op - congruence
    return Pinf - congruence.as_matrix()