cuthbert 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

cuthbert/gaussian/taylor/smoother.py

@@ -0,0 +1,158 @@
+r"""Linearized Taylor Kalman smoother.
+
+Uses automatic differentiation to extract conditionally Gaussian parameters from log
+densities of the dynamics and observation distributions.
+
+This differs from `gaussian/moments`, which requires `mean` and `chol_cov`
+functions as input rather than log densities.
+
+I.e., we approximate conditional densities as
+
+$$
+p(y \mid x) \approx N(y \mid H x + d, L L^T),
+$$
+
+and potentials as
+
+$$
+G(x) \approx N(x \mid m, L L^T),
+$$
+
+where $L$ is the Cholesky factor of the covariance matrix.
+
+See `cuthbertlib.linearize` for more details.
+"""
+
+from functools import partial
+
+from jax import numpy as jnp
+from jax import tree
+
+from cuthbert.gaussian.kalman import (
+    KalmanSmootherState,
+    convert_filter_to_smoother_state,
+    smoother_combine,
+)
+from cuthbert.gaussian.taylor.types import (
+    GetDynamicsLogDensity,
+)
+from cuthbert.gaussian.types import (
+    LinearizedKalmanFilterState,
+)
+from cuthbert.inference import Smoother
+from cuthbertlib.kalman import smoothing
+from cuthbertlib.linearize import linearize_log_density
+from cuthbertlib.types import (
+    ArrayTreeLike,
+    KeyArray,
+)
+
+
+def build_smoother(
+    get_dynamics_log_density: GetDynamicsLogDensity,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+    store_gain: bool = False,
+    store_chol_cov_given_next: bool = False,
+) -> Smoother:
+    """Build linearized Taylor Kalman inference smoother.
+
+    Args:
+        get_dynamics_log_density: Function to get dynamics log density log p(x_t+1 | x_t)
+            and linearization points (for the previous and current time points)
+        rtol: The relative tolerance for the singular values of precision matrices
+            when passed to `symmetric_inv_sqrt` during linearization.
+            Cutoff for small singular values; singular values smaller than
+            `rtol * largest_singular_value` are treated as zero.
+            The default is determined based on the floating point precision of the dtype.
+            See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+        ignore_nan_dims: Whether to treat dimensions with NaN on the diagonal of the
+            precision matrices (found via linearization) as missing and ignore all rows
+            and columns associated with them.
+        store_gain: Whether to store the gain matrix in the smoother state.
+        store_chol_cov_given_next: Whether to store the chol_cov_given_next matrix
+            in the smoother state.
+
+    Returns:
+        Linearized Taylor Kalman smoother object, suitable for associative scan.
+    """
+    return Smoother(
+        smoother_prepare=partial(
+            smoother_prepare,
+            get_dynamics_log_density=get_dynamics_log_density,
+            rtol=rtol,
+            ignore_nan_dims=ignore_nan_dims,
+            store_gain=store_gain,
+            store_chol_cov_given_next=store_chol_cov_given_next,
+        ),
+        smoother_combine=smoother_combine,
+        convert_filter_to_smoother_state=partial(
+            convert_filter_to_smoother_state,
+            store_gain=store_gain,
+            store_chol_cov_given_next=store_chol_cov_given_next,
+        ),
+        associative=True,
+    )
+
+
+def smoother_prepare(
+    filter_state: LinearizedKalmanFilterState,
+    get_dynamics_log_density: GetDynamicsLogDensity,
+    model_inputs: ArrayTreeLike,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+    store_gain: bool = False,
+    store_chol_cov_given_next: bool = False,
+    key: KeyArray | None = None,
+) -> KalmanSmootherState:
+    """Prepare a state for a linearized Taylor Kalman smoother step.
+
+    Args:
+        filter_state: State generated by the linearized Taylor Kalman filter at the previous
+            time point.
+        get_dynamics_log_density: Function to get dynamics log density log p(x_t+1 | x_t)
+            and linearization points (for the previous and current time points)
+        model_inputs: Model inputs for the transition from t to t+1.
+        rtol: The relative tolerance for the singular values of precision matrices
+            when passed to `symmetric_inv_sqrt` during linearization.
+            Cutoff for small singular values; singular values smaller than
+            `rtol * largest_singular_value` are treated as zero.
+            The default is determined based on the floating point precision of the dtype.
+            See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.pinv.html.
+        ignore_nan_dims: Whether to treat dimensions with NaN on the diagonal of the
+            precision matrices (found via linearization) as missing and ignore all rows
+            and columns associated with them.
+        store_gain: Whether to store the gain matrix in the smoother state.
+        store_chol_cov_given_next: Whether to store the chol_cov_given_next matrix
+            in the smoother state.
+        key: JAX random key - not used.
+
+    Returns:
+        Prepared state for the Kalman smoother.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+
+    filter_mean = filter_state.mean
+    filter_chol_cov = filter_state.chol_cov
+
+    log_dynamics_density, linearization_point_prev, linearization_point_curr = (
+        get_dynamics_log_density(filter_state, model_inputs)
+    )
+
+    F, c, chol_Q = linearize_log_density(
+        log_dynamics_density,
+        linearization_point_prev,
+        linearization_point_curr,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    state = smoothing.associative_params_single(
+        filter_mean, filter_chol_cov, F, c, chol_Q
+    )
+    return KalmanSmootherState(
+        elem=state,
+        gain=state.E if store_gain else None,
+        chol_cov_given_next=state.D if store_chol_cov_given_next else None,
+        model_inputs=model_inputs,
+    )

cuthbert/gaussian/taylor/types.py

@@ -0,0 +1,86 @@
+"""Provides types for the Taylor-series linearization of Gaussian state-space models."""
+
+from typing import Protocol, TypeAlias
+
+from cuthbert.gaussian.types import (
+    LinearizedKalmanFilterState,
+)
+from cuthbertlib.types import (
+    Array,
+    ArrayTreeLike,
+    LogConditionalDensity,
+    LogDensity,
+)
+
+LogPotential: TypeAlias = LogDensity
+
+
+class GetInitLogDensity(Protocol):
+    """Protocol for extracting the initial specifications."""
+
+    def __call__(self, model_inputs: ArrayTreeLike) -> tuple[LogDensity, Array]:
+        """Get the initial log density and initial linearization point.
+
+        Args:
+            model_inputs: Model inputs.
+
+        Returns:
+            Tuple with initial log density and initial linearization point.
+        """
+        ...
+
+
+class GetDynamicsLogDensity(Protocol):
+    """Protocol for extracting the dynamics specifications."""
+
+    def __call__(
+        self,
+        state: LinearizedKalmanFilterState,
+        model_inputs: ArrayTreeLike,
+    ) -> tuple[LogConditionalDensity, Array, Array]:
+        """Get the dynamics log density and linearization points.
+
+        Linearization points required for both the previous and current time points
+
+        `associative_scan` only supported when `state` is ignored.
+
+        Args:
+            state: NamedTuple containing `mean` and `mean_prev` attributes.
+            model_inputs: Model inputs.
+
+        Returns:
+            Tuple with dynamics log density and linearization points.
+        """
+        ...
+
+
+class GetObservationFunc(Protocol):
+    """Protocol for extracting the required observation specifications."""
+
+    def __call__(
+        self,
+        state: LinearizedKalmanFilterState,
+        model_inputs: ArrayTreeLike,
+    ) -> tuple[LogConditionalDensity, Array, Array] | tuple[LogPotential, Array]:
+        """Extract observation function, linearization point and optional observation.
+
+        State is the predicted state after applying the Kalman dynamics propagation.
+
+        `associative_scan` only supported when `state` is ignored.
+
+        Two types of output are supported:
+        - Observation log density function log p(y | x) and points x and y
+            to linearize around.
+        - Log potential function log G(x) and a linearization point x.
+
+        Args:
+            state: NamedTuple containing `mean` and `mean_prev` attributes.
+                Predicted state after applying the Kalman dynamics propagation.
+            model_inputs: Model inputs.
+
+        Returns:
+            Either a tuple with observation function to linearize, linearization point
+                and observation, or a tuple with log potential function and linearization
+                point.
+        """
+        ...

cuthbert/gaussian/types.py

@@ -0,0 +1,57 @@
+"""Provides shared types for Gaussian representations in state-space models."""
+
+from typing import NamedTuple, Protocol
+
+from cuthbertlib.kalman import filtering
+from cuthbertlib.types import Array, ArrayTree, ArrayTreeLike
+
+
+### Kalman types
+class GetInitParams(Protocol):
+    """Protocol for defining the initial distribution of a linear Gaussian SSM."""
+
+    def __call__(self, model_inputs: ArrayTreeLike) -> tuple[Array, Array]:
+        """Get initial parameters (m0, chol_P0) from model inputs."""
+        ...
+
+
+class GetDynamicsParams(Protocol):
+    """Protocol for defining the dynamics model of a linear Gaussian SSM."""
+
+    def __call__(self, model_inputs: ArrayTreeLike) -> tuple[Array, Array, Array]:
+        """Get dynamics parameters (F, c, chol_Q) from model inputs."""
+        ...
+
+
+class GetObservationParams(Protocol):
+    """Protocol for defining the observation model of a linear Gaussian SSM."""
+
+    def __call__(
+        self, model_inputs: ArrayTreeLike
+    ) -> tuple[Array, Array, Array, Array]:
+        """Get observation parameters (H, d, chol_R, y) from model inputs."""
+        ...
+
+
+### Shared state type for linearized Kalman filters
+class LinearizedKalmanFilterState(NamedTuple):
+    """Linearized Kalman filter state."""
+
+    elem: filtering.FilterScanElement
+    model_inputs: ArrayTree
+    mean_prev: Array
+
+    @property
+    def mean(self) -> Array:
+        """Filtering mean."""
+        return self.elem.b
+
+    @property
+    def chol_cov(self) -> Array:
+        """Filtering generalised Cholesky covariance."""
+        return self.elem.U
+
+    @property
+    def log_normalizing_constant(self) -> Array:
+        """Log normalizing constant (cumulative)."""
+        return self.elem.ell

cuthbert/gaussian/utils.py

@@ -0,0 +1,41 @@
+"""Utility functions (dummy state generation) for the Gaussian inference."""
+
+from cuthbert.gaussian.types import LinearizedKalmanFilterState
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.kalman import filtering
+from cuthbertlib.types import Array, ArrayTree
+
+
+def linearized_kalman_filter_state_dummy_elem(
+    mean: Array,
+    chol_cov: Array,
+    log_normalizing_constant: Array,
+    model_inputs: ArrayTree,
+    mean_prev: Array,
+) -> LinearizedKalmanFilterState:
+    """Create a LinearizedKalmanFilterState with a dummy element.
+
+    I.e. when associated scan is not used.
+
+    Args:
+        mean: Mean of the state.
+        chol_cov: Cholesky covariance of the state.
+        log_normalizing_constant: Log normalizing constant of the state.
+        model_inputs: Model inputs.
+        mean_prev: Mean of the previous state.
+
+    Returns:
+        LinearizedKalmanFilterState with a dummy elem attribute.
+    """
+    return LinearizedKalmanFilterState(
+        elem=filtering.FilterScanElement(
+            A=dummy_tree_like(chol_cov),
+            b=mean,
+            U=chol_cov,
+            eta=dummy_tree_like(mean),
+            Z=dummy_tree_like(chol_cov),
+            ell=log_normalizing_constant,
+        ),
+        model_inputs=model_inputs,
+        mean_prev=mean_prev,
+    )

cuthbert/smc/backward_sampler.py

@@ -0,0 +1,193 @@
+"""Implements backward sampling for particle filters.
+
+Supports 3 different algorithms for backward sampling:
+
+- [`cuthbertlib.smc.smoothing.tracing.simulate`][].
+- [`cuthbertlib.smc.smoothing.exact_sampling.simulate`][].
+- [`cuthbertlib.smc.smoothing.mcmc.simulate`][].
+"""
+
+from functools import partial
+from typing import NamedTuple, cast
+
+import jax
+import jax.numpy as jnp
+from jax import Array, random
+
+from cuthbert.inference import Smoother
+from cuthbert.smc.particle_filter import LogPotential, ParticleFilterState
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.resampling import Resampling
+from cuthbertlib.smc.smoothing.protocols import BackwardSampling
+from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
+
+
+class ParticleSmootherState(NamedTuple):
+    """Particle smoother state."""
+
+    key: KeyArray
+    particles: ArrayTree
+    ancestor_indices: Array
+    model_inputs: ArrayTree
+    log_weights: Array
+
+    @property
+    def n_particles(self) -> int:
+        """Number of particles in the smoother state."""
+        return self.ancestor_indices.shape[-1]
+
+
+def build_smoother(
+    log_potential: LogPotential,
+    backward_sampling_fn: BackwardSampling,
+    resampling_fn: Resampling,
+    n_smoother_particles: int,
+) -> Smoother:
+    r"""Build a particle smoother object.
+
+    Args:
+        log_potential: Function to compute the JOINT log potential $\log G_t(x_{t-1}, x_t) + \log M_t(x_t \mid x_{t-1})$.
+        backward_sampling_fn: Backward sampling algorithm to use (e.g., genealogy tracing, exact backward sampling).
+            This choice specifies how to sample $x_{t-1} \sim p(x_{t-1} \mid x_t, y_{0:t-1})$ given
+            samples $x_{t} \sim p(x_t \mid y_{0:T})$. See `cuthbertlib/smc/smoothing/` for possible choices.
+        resampling_fn: Resampling algorithm to use (e.g., multinomial, systematic).
+        n_smoother_particles: Number of samples to draw from the backward sampling algorithm.
+
+    Returns:
+        Particle smoother object.
+    """
+    return Smoother(
+        convert_filter_to_smoother_state=partial(
+            convert_filter_to_smoother_state,
+            resampling=resampling_fn,
+            n_smoother_particles=n_smoother_particles,
+        ),
+        smoother_prepare=smoother_prepare,
+        smoother_combine=partial(
+            smoother_combine,
+            backward_sampling_fn=backward_sampling_fn,
+            log_potential=log_potential,
+        ),
+        associative=False,
+    )
+
+
+def convert_filter_to_smoother_state(
+    filter_state: ParticleFilterState,
+    resampling: Resampling,
+    n_smoother_particles: int,
+    model_inputs: ArrayTreeLike | None = None,
+    key: KeyArray | None = None,
+) -> ParticleSmootherState:
+    """Convert a particle filter state to a particle smoother state.
+
+    Args:
+        filter_state: Particle filter state.
+        resampling: Resampling algorithm to use (e.g., multinomial, systematic).
+        n_smoother_particles: Number of smoother samples to draw.
+        model_inputs: Only used to create an empty model_inputs tree
+            (the values are ignored).
+            Useful so that the final smoother state has the same structure as the rest.
+            By default, filter_state.model_inputs is used. So this
+            is only needed if the smoother model_inputs have a different tree
+            structure to filter_state.model_inputs.
+        key: JAX random key.
+
+    Returns:
+        Particle smoother state. Note that the model_inputs are set to dummy values.
+
+    Raises:
+        ValueError: If key is None.
+    """
+    if key is None:
+        raise ValueError("A JAX PRNG key must be provided.")
+
+    if model_inputs is None:
+        model_inputs = filter_state.model_inputs
+
+    dummy_model_inputs = dummy_tree_like(model_inputs)
+
+    key, resampling_key = random.split(key)
+    indices = resampling(resampling_key, filter_state.log_weights, n_smoother_particles)
+
+    return ParticleSmootherState(
+        key=cast(KeyArray, key),
+        particles=jax.tree.map(lambda z: z[indices], filter_state.particles),
+        ancestor_indices=filter_state.ancestor_indices[indices],
+        model_inputs=dummy_model_inputs,
+        log_weights=-jnp.log(n_smoother_particles) * jnp.ones(n_smoother_particles),
+    )
+
+
+def smoother_prepare(
+    filter_state: ParticleFilterState,
+    model_inputs: ArrayTreeLike,
+    key: KeyArray | None = None,
+) -> ParticleSmootherState:
+    """Prepare a state for a particle smoother step.
+
+    Note that the model_inputs here are different to filter_state.model_inputs.
+    The model_inputs required here are for the transition from t to t+1.
+    filter_state.model_inputs represents the transition from t-1 to t.
+
+    Args:
+        filter_state: Particle filter state from time t.
+        model_inputs: Model inputs for the transition from t to t+1.
+        key: JAX random key.
+
+    Returns:
+        Prepared state for the particle smoother.
+    """
+    if key is None:
+        raise ValueError("A JAX PRNG key must be provided.")
+
+    model_inputs = jax.tree.map(lambda x: jnp.asarray(x), model_inputs)
+
+    return ParticleSmootherState(
+        key,
+        filter_state.particles,
+        filter_state.ancestor_indices,
+        model_inputs,
+        filter_state.log_weights,
+    )
+
+
+def smoother_combine(
+    state_1: ParticleSmootherState,
+    state_2: ParticleSmootherState,
+    backward_sampling_fn: BackwardSampling,
+    log_potential: LogPotential,
+) -> ParticleSmootherState:
+    """Combine next smoother state with state prepared with latest model inputs.
+
+    Remember smoothing iterates backwards in time.
+
+    Args:
+        state_1: State prepared with model inputs at time t.
+        state_2: Smoother state at time t + 1.
+        backward_sampling_fn: Function to perform backward sampling from the joint distribution.
+        log_potential: Function to compute log potential.
+
+    Returns:
+        Combined particle smoother state.
+        Contains particles, the original ancestor indices of the particles, and model inputs.
+    """
+    new_particles_1, ancestors_1 = backward_sampling_fn(
+        state_1.key,
+        x0_all=state_1.particles,
+        x1_all=state_2.particles,
+        log_weight_x0_all=state_1.log_weights,
+        log_density=lambda s1, s2: log_potential(s1, s2, state_2.model_inputs),
+        x1_ancestor_indices=state_2.ancestor_indices,
+    )
+
+    n_particles = len(ancestors_1)
+    log_weights = -jnp.log(n_particles) * jnp.ones(n_particles)
+    new_state = ParticleSmootherState(
+        key=state_1.key,
+        particles=new_particles_1,
+        ancestor_indices=state_1.ancestor_indices[ancestors_1],
+        model_inputs=state_1.model_inputs,
+        log_weights=log_weights,
+    )
+    return new_state