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

cuthbert/gaussian/moments/associative_filter.py

@@ -0,0 +1,180 @@
+"""Implements the associative linearized moments Kalman filter."""
+
+from jax import eval_shape, tree
+from jax import numpy as jnp
+
+from cuthbert.gaussian.kalman import GetInitParams
+from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
+from cuthbert.gaussian.types import (
+    LinearizedKalmanFilterState,
+)
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.kalman import filtering
+from cuthbertlib.linearize import linearize_moments
+from cuthbertlib.types import (
+    ArrayTreeLike,
+    KeyArray,
+)
+
+
+def init_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_params: GetInitParams,
+    get_observation_params: GetObservationMoments,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare the initial state for the linearized moments Kalman filter.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_params: Function to get m0, chol_P0 from model inputs.
+        get_observation_params: Function to get observation conditional mean,
+            (generalised) Cholesky covariance function, linearization point and
+            observation.
+        key: JAX random key - not used.
+
+    Returns:
+        State for the linearized moments Kalman filter.
+            Contains mean, chol_cov (generalised Cholesky factor of covariance)
+            and log_normalizing_constant.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    m0, chol_P0 = get_init_params(model_inputs)
+
+    prior_state = LinearizedKalmanFilterState(
+        elem=filtering.FilterScanElement(
+            A=jnp.zeros_like(chol_P0),
+            b=m0,
+            U=chol_P0,
+            eta=jnp.zeros_like(m0),
+            Z=jnp.zeros_like(chol_P0),
+            ell=jnp.array(0.0),
+        ),
+        model_inputs=model_inputs,
+        mean_prev=dummy_tree_like(m0),
+    )
+
+    mean_and_chol_cov_func, linearization_point, y = get_observation_params(
+        prior_state, model_inputs
+    )
+
+    H, d, chol_R = linearize_moments(mean_and_chol_cov_func, linearization_point)
+    (m, chol_P), ell = filtering.update(m0, chol_P0, H, d, chol_R, y)
+
+    elem = filtering.FilterScanElement(
+        A=jnp.zeros_like(chol_P),
+        b=m,
+        U=chol_P,
+        eta=jnp.zeros_like(m),
+        Z=jnp.zeros_like(chol_P),
+        ell=ell,
+    )
+
+    return LinearizedKalmanFilterState(
+        elem=elem,
+        model_inputs=model_inputs,
+        mean_prev=dummy_tree_like(m),
+    )
+
+
+def filter_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_params: GetInitParams,
+    get_dynamics_params: GetDynamicsMoments,
+    get_observation_params: GetObservationMoments,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare a state for a linearized moments Kalman filter step.
+
+    Just passes through model inputs.
+
+    `associative_scan` is supported but only accurate when `state` is ignored
+    in `get_dynamics_params` and `get_observation_params`.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_params: Function to get m0, chol_P0 from model inputs.
+            Only used to infer shape of mean and chol_cov.
+        get_dynamics_params: Function to get dynamics conditional mean and
+            (generalised) Cholesky covariance from linearization point and model inputs.
+            `associative_scan` only supported when `state` is ignored.
+        get_observation_params: Function to get observation conditional mean,
+            (generalised) Cholesky covariance and observation from linearization point
+            and model inputs.
+            `associative_scan` only supported when `state` is ignored.
+        key: JAX random key - not used.
+
+    Returns:
+        Prepared state for linearized moments Kalman filter.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    dummy_mean_struct = eval_shape(lambda mi: get_init_params(mi)[0], model_inputs)
+    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_chol_cov = dummy_tree_like(jnp.cov(dummy_mean[..., None]))
+
+    dummy_state = LinearizedKalmanFilterState(
+        elem=filtering.FilterScanElement(
+            A=dummy_chol_cov,
+            b=dummy_mean,
+            U=dummy_chol_cov,
+            eta=dummy_mean,
+            Z=dummy_chol_cov,
+            ell=jnp.array(0.0),
+        ),
+        model_inputs=model_inputs,
+        mean_prev=dummy_mean,
+    )
+
+    dynamics_mean_and_chol_cov_func, dynamics_linearization_point = get_dynamics_params(
+        dummy_state, model_inputs
+    )
+    F, c, chol_Q = linearize_moments(
+        dynamics_mean_and_chol_cov_func, dynamics_linearization_point
+    )
+
+    observation_mean_and_chol_cov_func, observation_linearization_point, y = (
+        get_observation_params(dummy_state, model_inputs)
+    )
+    H, d, chol_R = linearize_moments(
+        observation_mean_and_chol_cov_func, observation_linearization_point
+    )
+
+    elem = filtering.associative_params_single(F, c, chol_Q, H, d, chol_R, y)
+
+    return LinearizedKalmanFilterState(
+        elem=elem,
+        model_inputs=model_inputs,
+        mean_prev=dummy_mean,
+    )
+
+
+def filter_combine(
+    state_1: LinearizedKalmanFilterState,
+    state_2: LinearizedKalmanFilterState,
+) -> LinearizedKalmanFilterState:
+    """Combine previous filter state with state prepared with latest model inputs.
+
+    `associative_scan` is supported but only accurate when `state` is ignored
+    in `get_dynamics_params` and `get_observation_params`.
+
+    Applies standard associative Kalman filtering operator since dynamics and observation
+    parameters are extracted in filter_prepare.
+
+    Args:
+        state_1: State from previous time step.
+        state_2: State prepared (only access model_inputs attribute).
+
+    Returns:
+        Predicted and updated linearized moments Kalman filter state.
+            Contains mean, chol_cov (generalised Cholesky factor of covariance)
+            and log_normalizing_constant.
+    """
+    combined_elem = filtering.filtering_operator(
+        state_1.elem,
+        state_2.elem,
+    )
+    return LinearizedKalmanFilterState(
+        elem=combined_elem,
+        model_inputs=state_2.model_inputs,
+        mean_prev=state_1.mean,
+    )

cuthbert/gaussian/moments/filter.py

@@ -0,0 +1,95 @@
+r"""Linearized moments Kalman filter.
+
+Takes a user provided conditional `mean` and `chol_cov` functions to define a
+conditionally linear Gaussian state space model.
+
+I.e., we approximate conditional densities as
+
+$$
+p(y \mid x) \approx N(y \mid \mathrm{mean}(x), \mathrm{chol\_cov}(x) @ \mathrm{chol\_cov}(x)^\top).
+$$
+
+See `cuthbertlib.linearize` for more details.
+
+Parallelism via `associative_scan` is supported, but requires the `state` argument
+to be ignored in `get_dynamics_params` and `get_observation_params`.
+I.e. the linearization points are pre-defined or extracted from model inputs.
+"""
+
+from functools import partial
+
+from cuthbert.gaussian.moments import associative_filter, non_associative_filter
+from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
+from cuthbert.gaussian.types import GetInitParams
+from cuthbert.inference import Filter
+
+
+def build_filter(
+    get_init_params: GetInitParams,
+    get_dynamics_params: GetDynamicsMoments,
+    get_observation_params: GetObservationMoments,
+    associative: bool = False,
+) -> Filter:
+    """Build linearized moments Kalman inference filter.
+
+    If `associative` is True all filtering linearization points are pre-defined or
+    extracted from model inputs. The `state` argument should be ignored in
+    `get_dynamics_params` and `get_observation_params`.
+
+    If `associative` is False the linearization points can be extracted from the
+    previous filter state for dynamics parameters and the predict state for
+    observation parameters.
+
+    Args:
+        get_init_params: Function to get m0, chol_P0 from model inputs.
+        get_dynamics_params: Function to get dynamics conditional mean and
+            (generalised) Cholesky covariance from linearization point and model inputs.
+            and linearization points (for the previous and current time points)
+            If `associative` is True, the `state` argument should be ignored.
+        get_observation_params: Function to get observation conditional mean,
+            (generalised) Cholesky covariance and observation from linearization point
+            and model inputs.
+            If `associative` is True, the `state` argument should be ignored.
+        associative: If True, then the filter is suitable for associative scan, but
+            assumes that the `state` is ignored in `get_dynamics_params` and
+            `get_observation_params`.
+            If False, then the filter is suitable for non-associative scan, but
+            the user is free to use the `state` to extract the linearization points.
+
+    Returns:
+        Linearized moments Kalman filter object.
+    """
+    if associative:
+        return Filter(
+            init_prepare=partial(
+                associative_filter.init_prepare,
+                get_init_params=get_init_params,
+                get_observation_params=get_observation_params,
+            ),
+            filter_prepare=partial(
+                associative_filter.filter_prepare,
+                get_init_params=get_init_params,
+                get_dynamics_params=get_dynamics_params,
+                get_observation_params=get_observation_params,
+            ),
+            filter_combine=associative_filter.filter_combine,
+            associative=True,
+        )
+    else:
+        return Filter(
+            init_prepare=partial(
+                non_associative_filter.init_prepare,
+                get_init_params=get_init_params,
+                get_observation_params=get_observation_params,
+            ),
+            filter_prepare=partial(
+                non_associative_filter.filter_prepare,
+                get_init_params=get_init_params,
+            ),
+            filter_combine=partial(
+                non_associative_filter.filter_combine,
+                get_dynamics_params=get_dynamics_params,
+                get_observation_params=get_observation_params,
+            ),
+            associative=False,
+        )

cuthbert/gaussian/moments/non_associative_filter.py

@@ -0,0 +1,161 @@
+"""Implements the non-associative linearized moments Kalman filter."""
+
+from jax import eval_shape, tree
+from jax import numpy as jnp
+
+from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
+from cuthbert.gaussian.types import GetInitParams, LinearizedKalmanFilterState
+from cuthbert.gaussian.utils import linearized_kalman_filter_state_dummy_elem
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.kalman import filtering
+from cuthbertlib.linearize import linearize_moments
+from cuthbertlib.types import ArrayTreeLike, KeyArray
+
+
+def init_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_params: GetInitParams,
+    get_observation_params: GetObservationMoments,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare the initial state for the linearized moments Kalman filter.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_params: Function to get m0, chol_P0 from model inputs.
+        get_observation_params: Function to get observation conditional mean,
+            (generalised) Cholesky covariance function, linearization point and
+            observation.
+        key: JAX random key - not used.
+
+    Returns:
+        State for the linearized moments Kalman filter.
+            Contains mean, chol_cov (generalised Cholesky factor of covariance)
+            and log_normalizing_constant.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    m0, chol_P0 = get_init_params(model_inputs)
+
+    prior_state = LinearizedKalmanFilterState(
+        elem=filtering.FilterScanElement(
+            A=jnp.zeros_like(chol_P0),
+            b=m0,
+            U=chol_P0,
+            eta=jnp.zeros_like(m0),
+            Z=jnp.zeros_like(chol_P0),
+            ell=jnp.array(0.0),
+        ),
+        model_inputs=model_inputs,
+        mean_prev=dummy_tree_like(m0),
+    )
+
+    mean_and_chol_cov_func, linearization_point, y = get_observation_params(
+        prior_state, model_inputs
+    )
+
+    H, d, chol_R = linearize_moments(mean_and_chol_cov_func, linearization_point)
+    (m, chol_P), ell = filtering.update(m0, chol_P0, H, d, chol_R, y)
+    return linearized_kalman_filter_state_dummy_elem(
+        mean=m,
+        chol_cov=chol_P,
+        log_normalizing_constant=ell,
+        model_inputs=model_inputs,
+        mean_prev=dummy_tree_like(m),
+    )
+
+
+def filter_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_params: GetInitParams,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare a state for a linearized moments Kalman filter step.
+
+    Just passes through model inputs.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_params: Function to get m0, chol_P0 from model inputs.
+            Only used to infer shape of mean and chol_cov.
+        key: JAX random key - not used.
+
+    Returns:
+        Prepared state for linearized moments Kalman filter.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    dummy_mean_struct = eval_shape(lambda mi: get_init_params(mi)[0], model_inputs)
+    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_chol_cov = dummy_tree_like(jnp.cov(dummy_mean[..., None]))
+
+    return linearized_kalman_filter_state_dummy_elem(
+        mean=dummy_mean,
+        chol_cov=dummy_chol_cov,
+        log_normalizing_constant=jnp.array(0.0),
+        model_inputs=model_inputs,
+        mean_prev=dummy_mean,
+    )
+
+
+def filter_combine(
+    state_1: LinearizedKalmanFilterState,
+    state_2: LinearizedKalmanFilterState,
+    get_dynamics_params: GetDynamicsMoments,
+    get_observation_params: GetObservationMoments,
+) -> LinearizedKalmanFilterState:
+    """Combine previous filter state with state prepared with latest model inputs.
+
+    Applies linearized moments Kalman predict + filter update in covariance square
+    root form.
+    Not suitable for associative scan.
+
+    Args:
+        state_1: State from previous time step.
+        state_2: State prepared (only access model_inputs attribute).
+        get_dynamics_params: Function to get dynamics conditional mean and
+            (generalised) Cholesky covariance from linearization point and model inputs.
+        get_observation_params: Function to get observation conditional mean,
+            (generalised) Cholesky covariance and observation from linearization point
+            and model inputs.
+
+    Returns:
+        Predicted and updated linearized moments Kalman filter state.
+            Contains mean, chol_cov (generalised Cholesky factor of covariance)
+            and log_normalizing_constant.
+    """
+    dynamics_mean_and_chol_cov_func, dynamics_linearization_point = get_dynamics_params(
+        state_1, state_2.model_inputs
+    )
+    F, c, chol_Q = linearize_moments(
+        dynamics_mean_and_chol_cov_func, dynamics_linearization_point
+    )
+
+    predict_mean, predict_chol_cov = filtering.predict(
+        state_1.mean, state_1.chol_cov, F, c, chol_Q
+    )
+    predict_state = linearized_kalman_filter_state_dummy_elem(
+        mean=predict_mean,
+        chol_cov=predict_chol_cov,
+        log_normalizing_constant=state_1.log_normalizing_constant,
+        model_inputs=state_2.model_inputs,
+        mean_prev=state_1.mean,
+    )
+
+    observation_mean_and_chol_cov_func, observation_linearization_point, y = (
+        get_observation_params(predict_state, state_2.model_inputs)
+    )
+    H, d, chol_R = linearize_moments(
+        observation_mean_and_chol_cov_func, observation_linearization_point
+    )
+
+    (update_mean, update_chol_cov), log_normalizing_constant = filtering.update(
+        predict_mean, predict_chol_cov, H, d, chol_R, y
+    )
+
+    return linearized_kalman_filter_state_dummy_elem(
+        mean=update_mean,
+        chol_cov=update_chol_cov,
+        log_normalizing_constant=state_1.log_normalizing_constant
+        + log_normalizing_constant,
+        model_inputs=state_2.model_inputs,
+        mean_prev=state_1.mean,
+    )

cuthbert/gaussian/moments/smoother.py

@@ -0,0 +1,118 @@
+r"""Linearized moments Kalman smoother.
+
+Takes a user provided conditional `mean` and `chol_cov` functions to define a
+conditionally linear Gaussian state space model.
+
+I.e., we approximate conditional densities as
+
+$$
+p(y \mid x) \approx N(y \mid \mathrm{mean}(x), \mathrm{chol\_cov}(x) @ \mathrm{chol\_cov}(x)^\top).
+$$
+
+See `cuthbertlib.linearize` for more details.
+
+Parallelism via `associative_scan` is supported, but requires the `state` argument
+to be ignored in `get_dynamics_params`.
+I.e. the linearization points are pre-defined or extracted from model inputs.
+"""
+
+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.moments.types import GetDynamicsMoments
+from cuthbert.gaussian.types import LinearizedKalmanFilterState
+from cuthbert.inference import Smoother
+from cuthbertlib.kalman import smoothing
+from cuthbertlib.linearize import linearize_moments
+from cuthbertlib.types import ArrayTreeLike, KeyArray
+
+
+def build_smoother(
+    get_dynamics_params: GetDynamicsMoments,
+    store_gain: bool = False,
+    store_chol_cov_given_next: bool = False,
+) -> Smoother:
+    """Build linearized moments Kalman inference smoother for conditionally Gaussian SSMs.
+
+    Args:
+        get_dynamics_params: Function to get dynamics conditional mean and
+            (generalised) Cholesky covariance from linearization point and model inputs.
+        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 moments Kalman smoother object, suitable for associative scan.
+    """
+    return Smoother(
+        smoother_prepare=partial(
+            smoother_prepare,
+            get_dynamics_params=get_dynamics_params,
+            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_params: GetDynamicsMoments,
+    model_inputs: ArrayTreeLike,
+    store_gain: bool = False,
+    store_chol_cov_given_next: bool = False,
+    key: KeyArray | None = None,
+) -> KalmanSmootherState:
+    """Prepare a state for an extended Kalman 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: State generated by the extended Kalman filter at time t.
+        get_dynamics_params: Function to get dynamics conditional mean and
+            (generalised) Cholesky covariance from linearization point and model inputs.
+        model_inputs: Model inputs for the transition from t to t+1.
+        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
+
+    dynamics_mean_and_chol_cov_func, dynamics_linearization_point = get_dynamics_params(
+        filter_state, model_inputs
+    )
+
+    F, c, chol_Q = linearize_moments(
+        dynamics_mean_and_chol_cov_func, dynamics_linearization_point
+    )
+
+    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/moments/types.py

@@ -0,0 +1,51 @@
+"""Provides types for the moment-based linearization of Gaussian state-space models."""
+
+from typing import Protocol
+
+from cuthbert.gaussian.types import LinearizedKalmanFilterState
+from cuthbertlib.linearize.moments import MeanAndCholCovFunc
+from cuthbertlib.types import Array, ArrayTreeLike
+
+
+class GetDynamicsMoments(Protocol):
+    """Protocol for extracting the dynamics specifications."""
+
+    def __call__(
+        self,
+        state: LinearizedKalmanFilterState,
+        model_inputs: ArrayTreeLike,
+    ) -> tuple[MeanAndCholCovFunc, Array]:
+        """Get dynamics conditional mean and chol_cov function and linearization point.
+
+        `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 conditional mean and (generalised) Cholesky covariance
+                function and linearization point.
+        """
+        ...
+
+
+class GetObservationMoments(Protocol):
+    """Protocol for extracting the observation specifications."""
+
+    def __call__(
+        self, state: LinearizedKalmanFilterState, model_inputs: ArrayTreeLike
+    ) -> tuple[MeanAndCholCovFunc, Array, Array]:
+        """Get conditional mean and chol_cov function, linearization point and observation.
+
+        `associative_scan` only supported when `state` input is ignored.
+
+        Args:
+            state: NamedTuple containing `mean` and `mean_prev` attributes.
+            model_inputs: Model inputs.
+
+        Returns:
+            Tuple with conditional mean and chol_cov function, linearization point
+                and observation.
+        """
+        ...

cuthbert/gaussian/taylor/__init__.py

@@ -0,0 +1,14 @@
+from cuthbert.gaussian.taylor import (
+    associative_filter,
+    non_associative_filter,
+    smoother,
+)
+from cuthbert.gaussian.taylor.filter import build_filter
+from cuthbert.gaussian.taylor.smoother import build_smoother
+from cuthbert.gaussian.taylor.types import (
+    GetDynamicsLogDensity,
+    GetInitLogDensity,
+    GetObservationFunc,
+    LogPotential,
+)
+from cuthbert.gaussian.types import LinearizedKalmanFilterState