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

cuthbert/gaussian/taylor/associative_filter.py

@@ -0,0 +1,222 @@
+"""Implements the associative linearized Taylor Kalman filter."""
+
+from jax import eval_shape, tree
+from jax import numpy as jnp
+
+from cuthbert.gaussian.taylor.non_associative_filter import process_observation
+from cuthbert.gaussian.taylor.types import (
+    GetDynamicsLogDensity,
+    GetInitLogDensity,
+    GetObservationFunc,
+)
+from cuthbert.gaussian.types import (
+    LinearizedKalmanFilterState,
+)
+from cuthbert.utils import dummy_tree_like
+from cuthbertlib.kalman import filtering
+from cuthbertlib.linearize import linearize_log_density
+from cuthbertlib.types import (
+    ArrayTreeLike,
+    KeyArray,
+)
+
+
+def init_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_log_density: GetInitLogDensity,
+    get_observation_func: GetObservationFunc,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare the initial state for the linearized Taylor Kalman filter.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_log_density: Function that returns log density log p(x_0)
+            and linearization point.
+        get_observation_func: Function that returns either
+            - An observation log density
+                function log p(y_0 | x_0) as well as points x_0 and y_0
+                to linearize around.
+            - A log potential function log G(x_0) and a linearization point x_0.
+        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.
+        key: JAX random key - not used.
+
+    Returns:
+        State for the linearized Taylor 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)
+    init_log_density, linearization_point = get_init_log_density(model_inputs)
+
+    _, m0, chol_P0 = linearize_log_density(
+        lambda _, x: init_log_density(x),
+        linearization_point,
+        linearization_point,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    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),
+    )
+
+    observation_output = get_observation_func(prior_state, model_inputs)
+    H, d, chol_R, observation = process_observation(
+        observation_output,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    (m, chol_P), ell = filtering.update(m0, chol_P0, H, d, chol_R, observation)
+
+    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_log_density: GetInitLogDensity,
+    get_dynamics_log_density: GetDynamicsLogDensity,
+    get_observation_func: GetObservationFunc,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare a state for a linearized Taylor Kalman filter step.
+
+    `associative_scan` is supported but only accurate when `state` is ignored
+    in `get_dynamics_log_density` and `get_observation_func`.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_log_density: Function that returns log density log p(x_0)
+            and linearization point. Only used to infer shape of mean and chol_cov.
+        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)
+            `associative_scan` only supported when `state` is ignored.
+        get_observation_func: Function to get observation function (either conditional
+            log density or log potential), linearization point and optional observation
+            (not required for log potential functions).
+            `associative_scan` only supported when `state` is ignored.
+        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.
+        key: JAX random key - not used.
+
+    Returns:
+        Prepared state for linearized Taylor Kalman filter.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    dummy_mean_struct = eval_shape(lambda mi: get_init_log_density(mi)[1], 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,
+    )
+
+    log_dynamics_density, linearization_point_prev, linearization_point_curr = (
+        get_dynamics_log_density(dummy_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,
+    )
+
+    observation_output = get_observation_func(dummy_state, model_inputs)
+    H, d, chol_R, observation = process_observation(
+        observation_output,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    elem = filtering.associative_params_single(F, c, chol_Q, H, d, chol_R, observation)
+
+    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_log_density` and `get_observation_func`.
+
+    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 Taylor 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/taylor/filter.py

@@ -0,0 +1,129 @@
+r"""Linearized Taylor Kalman filter.
+
+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.
+
+Parallelism via `associative_scan` is supported, but requires the `state` argument
+to be ignored in `get_dynamics_log_density` and `get_observation_func`.
+I.e. the linearization points are pre-defined or extracted from model inputs.
+"""
+
+from functools import partial
+
+from cuthbert.gaussian.taylor import associative_filter, non_associative_filter
+from cuthbert.gaussian.taylor.types import (
+    GetDynamicsLogDensity,
+    GetInitLogDensity,
+    GetObservationFunc,
+)
+from cuthbert.inference import Filter
+
+
+def build_filter(
+    get_init_log_density: GetInitLogDensity,
+    get_dynamics_log_density: GetDynamicsLogDensity,
+    get_observation_func: GetObservationFunc,
+    associative: bool = False,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> Filter:
+    """Build linearized Taylor 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_log_density` and `get_observation_func`.
+
+    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_log_density: Function to get log density log p(x_0)
+            and linearization point.
+            Only takes `model_inputs` as input.
+        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)
+            If `associative` is True, the `state` argument should be ignored.
+        get_observation_func: Function to get observation function (either conditional
+            log density or log potential), linearization point and optional observation
+            (not required for log potential functions).
+            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_log_density` and
+            `get_observation_func`.
+            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.
+        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.
+
+    Returns:
+        Linearized Taylor Kalman filter object.
+    """
+    if associative:
+        return Filter(
+            init_prepare=partial(
+                associative_filter.init_prepare,
+                get_init_log_density=get_init_log_density,
+                get_observation_func=get_observation_func,
+                rtol=rtol,
+                ignore_nan_dims=ignore_nan_dims,
+            ),
+            filter_prepare=partial(
+                associative_filter.filter_prepare,
+                get_init_log_density=get_init_log_density,
+                get_dynamics_log_density=get_dynamics_log_density,
+                get_observation_func=get_observation_func,
+                rtol=rtol,
+                ignore_nan_dims=ignore_nan_dims,
+            ),
+            filter_combine=associative_filter.filter_combine,
+            associative=True,
+        )
+    else:
+        return Filter(
+            init_prepare=partial(
+                non_associative_filter.init_prepare,
+                get_init_log_density=get_init_log_density,
+                get_observation_func=get_observation_func,
+                rtol=rtol,
+                ignore_nan_dims=ignore_nan_dims,
+            ),
+            filter_prepare=partial(
+                non_associative_filter.filter_prepare,
+                get_init_log_density=get_init_log_density,
+            ),
+            filter_combine=partial(
+                non_associative_filter.filter_combine,
+                get_dynamics_log_density=get_dynamics_log_density,
+                get_observation_func=get_observation_func,
+                rtol=rtol,
+                ignore_nan_dims=ignore_nan_dims,
+            ),
+            associative=False,
+        )

cuthbert/gaussian/taylor/non_associative_filter.py

@@ -0,0 +1,246 @@
+"""Implements the non-associative linearized Taylor Kalman filter."""
+
+from jax import eval_shape, tree
+from jax import numpy as jnp
+
+from cuthbert.gaussian.taylor.types import (
+    GetDynamicsLogDensity,
+    GetInitLogDensity,
+    GetObservationFunc,
+    LogConditionalDensity,
+    LogPotential,
+)
+from cuthbert.gaussian.types import 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_log_density, linearize_taylor
+from cuthbertlib.types import (
+    Array,
+    ArrayTreeLike,
+    KeyArray,
+)
+
+
+def process_observation(
+    observation_output: tuple[LogConditionalDensity, Array, Array]
+    | tuple[LogPotential, Array],
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> tuple[Array, Array, Array, Array]:
+    """Process observation for linearized Taylor Kalman filter."""
+    if len(observation_output) == 3:
+        observation_cond_log_density, linearization_point, observation = (
+            observation_output
+        )
+        H, d, chol_R = linearize_log_density(
+            observation_cond_log_density,
+            linearization_point,
+            observation,
+            rtol=rtol,
+            ignore_nan_dims=ignore_nan_dims,
+        )
+    else:
+        observation_log_potential, linearization_point = observation_output
+        d, chol_R = linearize_taylor(
+            observation_log_potential,
+            linearization_point,
+            rtol=rtol,
+            ignore_nan_dims=ignore_nan_dims,
+        )
+        # dummy mat and observation as potential is unconditional
+        # Note the minus sign as linear potential is -0.5 (x - d)^T (R R^T)^{-1} (x - d)
+        # and kalman expects -0.5 (y - H @ x - d)^T (R R^T)^{-1} (y - H @ x - d)
+        H = -jnp.eye(d.shape[0])
+        observation = jnp.where(
+            jnp.isnan(jnp.diag(chol_R)) * ignore_nan_dims, jnp.nan, 0.0
+        )  # Tell the cuthbertlib.kalman to skip these dimensions
+    return H, d, chol_R, observation
+
+
+def init_prepare(
+    model_inputs: ArrayTreeLike,
+    get_init_log_density: GetInitLogDensity,
+    get_observation_func: GetObservationFunc,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare the initial state for the linearized Taylor Kalman filter.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_log_density: Function that returns log density log p(x_0)
+            and linearization point.
+        get_observation_func: Function that returns either
+            - An observation log density
+                function log p(y_0 | x_0) as well as points x_0 and y_0
+                to linearize around.
+            - A log potential function log G(x_0) and a linearization point x_0.
+        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.
+        key: JAX random key - not used.
+
+    Returns:
+        State for the linearized Taylor 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)
+    init_log_density, linearization_point = get_init_log_density(model_inputs)
+
+    _, m0, chol_P0 = linearize_log_density(
+        lambda _, x: init_log_density(x),
+        linearization_point,
+        linearization_point,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    prior_state = linearized_kalman_filter_state_dummy_elem(
+        mean=m0,
+        chol_cov=chol_P0,
+        log_normalizing_constant=jnp.array(0.0),
+        model_inputs=model_inputs,
+        mean_prev=dummy_tree_like(m0),
+    )
+
+    observation_output = get_observation_func(prior_state, model_inputs)
+
+    H, d, chol_R, observation = process_observation(
+        observation_output,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    (m, chol_P), ell = filtering.update(m0, chol_P0, H, d, chol_R, observation)
+
+    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_log_density: GetInitLogDensity,
+    key: KeyArray | None = None,
+) -> LinearizedKalmanFilterState:
+    """Prepare a state for a linearized Taylor Kalman filter step.
+
+    Just passes through model inputs.
+
+    Args:
+        model_inputs: Model inputs.
+        get_init_log_density: Function that returns log density log p(x_0)
+            and linearization point. Only used to infer shape of mean and chol_cov.
+        key: JAX random key - not used.
+
+    Returns:
+        Prepared state for linearized Taylor Kalman filter.
+    """
+    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
+    dummy_mean_struct = eval_shape(lambda mi: get_init_log_density(mi)[1], 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_log_density: GetDynamicsLogDensity,
+    get_observation_func: GetObservationFunc,
+    rtol: float | None = None,
+    ignore_nan_dims: bool = False,
+) -> LinearizedKalmanFilterState:
+    """Combine previous filter state with state prepared from latest model inputs.
+
+    Applies linearized Taylor 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_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)
+        get_observation_func: Function to get observation function (either conditional
+            log density or log potential), linearization point and optional observation
+            (not required for log potential functions).
+        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.
+
+    Returns:
+        Predicted and updated linearized Taylor Kalman filter state.
+            Contains mean, chol_cov (generalised Cholesky factor of covariance)
+            and log_normalizing_constant.
+    """
+    log_dynamics_density, linearization_point_prev, linearization_point_curr = (
+        get_dynamics_log_density(state_1, state_2.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,
+    )
+
+    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_output = get_observation_func(predict_state, state_2.model_inputs)
+
+    H, d, chol_R, observation = process_observation(
+        observation_output,
+        rtol=rtol,
+        ignore_nan_dims=ignore_nan_dims,
+    )
+
+    (update_mean, update_chol_cov), log_normalizing_constant = filtering.update(
+        predict_mean, predict_chol_cov, H, d, chol_R, observation
+    )
+
+    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,
+    )