cuthbertlib 0.0.9__tar.gz → 0.0.11__tar.gz

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuthbertlib
-Version: 0.0.9
+Version: 0.0.11
 Summary: Atomic building blocks for state-space model inference with JAX
 Author-email: Sam Duffield <s@mduffield.com>, Sahel Iqbal <sahel13miqbal@proton.me>, Adrien Corenflos <adrien.corenflos.stats@gmail.com>
 License: Apache-2.0

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/__init__.py

@@ -5,6 +5,7 @@ del version
 
 from cuthbertlib import (
     discrete,
+    enkf,
     kalman,
     linalg,
     linearize,

cuthbertlib-0.0.11/cuthbertlib/enkf/README.md

@@ -0,0 +1,14 @@
+# Ensemble Kalman Filter (EnKF)
+
+This sub-repository provides modular functions for the Ensemble Kalman Filter.
+
+The core functions are:
+
+- `predict`: Propagate ensemble members through nonlinear dynamics with additive Gaussian noise.
+- `update`: Update ensemble members with an observation using the EnKF update equation.
+
+Together, `predict` and `update` can be used to perform an online EnKF filtering step.
+
+The EnKF uses an ensemble of particles with a Kalman-style measurement update based on
+empirical covariances. Unlike the EKF, it does not require Jacobians, while naturally
+handling nonlinear dynamics. 

cuthbertlib-0.0.11/cuthbertlib/enkf/__init__.py

@@ -0,0 +1 @@
+from cuthbertlib.enkf.filtering import predict, update

cuthbertlib-0.0.11/cuthbertlib/enkf/filtering.py

@@ -0,0 +1,119 @@
+"""Implements the Ensemble Kalman Filter (EnKF) predict and update steps.
+
+See Algorithm 10.2, [Sanz-Alonso et al., Inverse Problems and Data Assimilation](https://arxiv.org/abs/1810.06191).
+Based in part on the [CD-Dynamax implementation](https://github.com/hd-UQ/cd_dynamax/blob/public/cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_enkf.py).
+"""
+
+from typing import Callable
+
+import jax
+import jax.numpy as jnp
+from jax import random
+from jax.scipy.linalg import cho_solve
+
+from cuthbertlib.linalg import collect_nans_chol, tria
+from cuthbertlib.stats import multivariate_normal
+from cuthbertlib.types import Array, KeyArray, ScalarArray
+
+ObservationFn = Callable[[Array], Array]
+DynamicsFn = Callable[[Array, KeyArray], Array]
+
+
+def predict(
+    key: KeyArray,
+    ensemble: Array,
+    dynamics_fn: DynamicsFn,
+    inflation: float = 0.0,
+) -> Array:
+    """Propagate ensemble members through an arbitrary simulator p(x_{t+1} | x_t).
+
+    Args:
+        key: JAX PRNG key.
+        ensemble: Ensemble of state vectors, shape (N, x_dim).
+        dynamics_fn: Dynamics function mapping (state, key) -> state.
+        inflation: Multiplicative inflation factor applied to ensemble deviations.
+
+    Returns:
+        Predicted ensemble, shape (N, x_dim).
+    """
+    N, x_dim = ensemble.shape
+
+    # Propagate each member through the dynamics
+    keys = random.split(key, N)
+    propagated = jax.vmap(dynamics_fn, (0, 0))(ensemble, keys)
+
+    # Apply multiplicative inflation
+    mean = jnp.mean(propagated, axis=0)
+    propagated = mean + (1 + inflation) * (propagated - mean)
+
+    return propagated
+
+
+def update(
+    key: KeyArray,
+    predicted_ensemble: Array,
+    observation_fn: ObservationFn,
+    chol_R: Array,
+    y: Array,
+    perturbed_obs: bool = True,
+) -> tuple[Array, ScalarArray]:
+    """Update ensemble members with an observation using the EnKF update.
+
+    NaNs in ``y`` are treated as missing dimensions and are excluded from the
+    update. When ``y`` is entirely NaN, the update is a no-op: the predicted
+    ensemble is returned unchanged with zero log-likelihood contribution.
+
+    Args:
+        key: JAX PRNG key.
+        predicted_ensemble: Predicted ensemble, shape (N, x_dim).
+        observation_fn: Observation function mapping state -> obs.
+        chol_R: Cholesky factor of the observation noise covariance, shape (y_dim, y_dim).
+        y: Observation vector, shape (y_dim,). NaNs indicate missing dimensions.
+        perturbed_obs: If True, use perturbed observations (stochastic EnKF).
+            If False, use deterministic update.
+
+    Returns:
+        Tuple of (updated_ensemble, log_likelihood).
+    """
+    N, x_dim = predicted_ensemble.shape
+
+    # Map ensemble to observation space
+    y_pred = jax.vmap(observation_fn, (0,))(predicted_ensemble)
+
+    # Handle partially-missing observations by reordering and zeroing missing dims.
+    # Use y_pred.T because y_pred is (N, y_dim) and we want to reorder along axis 0.
+    flag = jnp.isnan(y)
+    flag, chol_R, y, y_pred = collect_nans_chol(flag, chol_R, y, y_pred.T)
+    y_pred = y_pred.T
+    y_dim = y.shape[0]
+
+    # Ensemble means
+    x_mean = jnp.mean(predicted_ensemble, axis=0)
+    y_mean = jnp.mean(y_pred, axis=0)
+
+    # Deviations from ensemble mean
+    x_dev = predicted_ensemble - x_mean
+    y_dev = y_pred - y_mean
+
+    # Square-root innovation covariance via tria
+    chol_S = tria(jnp.concatenate([y_dev.T / jnp.sqrt(N - 1), chol_R], axis=1))
+
+    # Cross-covariance
+    C_xy = x_dev.T @ y_dev / (N - 1)
+
+    # Kalman gain: K = C_xy @ S^{-1} = C_xy @ cho_solve(chol_S, I)
+    K = cho_solve((chol_S, True), C_xy.T).T
+
+    # Innovation per member
+    if perturbed_obs:
+        y_n = y[None, :] + (chol_R @ random.normal(key, (y_dim, N))).T
+    else:
+        y_n = jnp.broadcast_to(y[None, :], (N, y_dim))
+
+    # Update ensemble
+    updated = predicted_ensemble + (y_n - y_pred) @ K.T
+
+    # Log-likelihood
+    ll = multivariate_normal.logpdf(y, y_mean, chol_S, nan_support=False)
+
+    return updated, jnp.asarray(ll)

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/linalg/tria.py

@@ -85,8 +85,8 @@ def _tria_jvp(primals, tangents):
     K_T = jnp.swapaxes(K, -1, -2)
 
     # Solve for lower triangular perturbation dM + dM^T = K + K^T
-    I = jnp.eye(K.shape[-1], dtype=K.dtype)
-    dM = jnp.tril(K + K_T) - K * I
+    Id = jnp.eye(K.shape[-1], dtype=K.dtype)
+    dM = jnp.tril(K + K_T) - K * Id
 
     # Compute the null-space part
     dR_null = (jnp.eye(R.shape[-2], dtype=R.dtype) - R @ R_pinv) @ dA @ Q

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/linearize/log_density.py

@@ -63,6 +63,7 @@ def linearize_log_density(
 
     Args:
         log_density: A conditional log density of y given x. Returns a scalar.
+            x must be the first argument and y the second.
         x: The input points.
         y: The output points.
         has_aux: Whether `log_density` returns an auxiliary value.
@@ -137,6 +138,7 @@ def linearize_log_density_given_chol_cov(
 
     Args:
         log_density: A conditional log density of y given x. Returns a scalar.
+            x must be the first argument and y the second.
         x: The input points.
         y: The output points.
         chol_cov: The Cholesky factor of the covariance matrix of the Gaussian.

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/resampling/__init__.py

@@ -3,6 +3,7 @@ from cuthbertlib.resampling import (
     autodiff,
     killing,
     multinomial,
+    no_resampling,
     systematic,
 )
 from cuthbertlib.resampling.adaptive import ess_decorator

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/resampling/autodiff.py

@@ -15,7 +15,6 @@ import jax.numpy as jnp
 
 from cuthbertlib.resampling.protocols import Resampling
 from cuthbertlib.resampling.utils import apply_resampling_indices
-from cuthbertlib.smc.ess import log_ess
 from cuthbertlib.types import Array, ArrayLike, ArrayTree, ArrayTreeLike
 
 

cuthbertlib-0.0.11/cuthbertlib/resampling/no_resampling.py

@@ -0,0 +1,57 @@
+"""No resampling dummy implementation."""
+
+from functools import partial
+
+from jax import numpy as jnp
+
+from cuthbertlib.resampling.protocols import (
+    conditional_resampling_decorator,
+    resampling_decorator,
+)
+from cuthbertlib.resampling.utils import apply_resampling_indices
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    ScalarArrayLike,
+)
+
+_DESCRIPTION = """
+No resampling is performed.
+Useful for factorial SMC where resampling is applied during `join` rather than
+`filter_combine`."""
+
+
+@partial(resampling_decorator, name="No Resampling", desc=_DESCRIPTION)
+def resampling(
+    key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+) -> tuple[Array, Array, ArrayTree]:
+    logits = jnp.asarray(logits)
+    if n != logits.shape[0]:
+        raise AssertionError(
+            "The number of sampled indices must be equal to the number of "
+            "output particles for `No Resampling` resampling."
+        )
+    return jnp.arange(n), logits, positions
+
+
+@partial(conditional_resampling_decorator, name="No Resampling", desc=_DESCRIPTION)
+def conditional_resampling(
+    key: Array,
+    logits: ArrayLike,
+    positions: ArrayTreeLike,
+    n: int,
+    pivot_in: ScalarArrayLike,
+    pivot_out: ScalarArrayLike,
+) -> tuple[Array, Array, ArrayTree]:
+    logits = jnp.asarray(logits)
+    if n != logits.shape[0]:
+        raise AssertionError(
+            "The number of sampled indices must be equal to the number of "
+            "output particles for `No Resampling` resampling."
+        )
+    pivot_in = jnp.asarray(pivot_in, dtype=jnp.int32)
+    pivot_out = jnp.asarray(pivot_out, dtype=jnp.int32)
+    idx = jnp.arange(n).at[pivot_in].set(pivot_out)
+    return idx, logits, apply_resampling_indices(positions, idx)

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/smc/smoothing/exact_sampling.py

@@ -28,6 +28,7 @@ def log_weights_single(
         x1: The current state.
         log_weight_x0: The log weights of the previous state.
         log_density: The log density function of x1 given x0.
+            x0 must be the first argument and x1 the second.
 
     Returns:
         The smoothing weight for sample x0 given a single sample x1.
@@ -97,6 +98,7 @@ def simulate(
         x1_all: A collection of current states $x_1$.
         log_weight_x0_all: The log weights of $x_0$.
         log_density: The log density function of $x_1$ given $x_0$.
+            $x_0$ must be the first argument and $x_1$ the second.
         x1_ancestor_indices: The ancestor indices of $x_1$. Not used.
 
     Returns:

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/smc/smoothing/mcmc.py

@@ -33,6 +33,7 @@ def simulate(
         x1_all: A collection of current states $x_1$.
         log_weight_x0_all: The log weights of $x_0$.
         log_density: The log density function of $x_1$ given $x_0$.
+            $x_0$ must be the first argument and $x_1$ the second.
         x1_ancestor_indices: The ancestor indices of $x_1$.
         n_steps: Number of MCMC steps to perform.
 
@@ -71,7 +72,7 @@ def simulate(
         return (idx, x0_res, idx_log_p), None
 
     x0_init = jax.tree.map(lambda z: z[x1_ancestor_indices], x0_all)
-    init_log_p = jax.vmap(log_density)(x1_all, x0_init)
+    init_log_p = jax.vmap(log_density)(x0_init, x1_all)
     init = (x1_ancestor_indices, x0_init, init_log_p)
     (out_index, out_samples, _), _ = jax.lax.scan(body, init, keys)
     return out_samples, out_index

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/smc/smoothing/protocols.py

@@ -36,6 +36,7 @@ class BackwardSampling(Protocol):
             x1_all: A collection of current states $x_1$.
             log_weight_x0_all: The log weights of $x_0$.
             log_density: The log density function of $x_1$ given $x_0$.
+                $x_0$ must be the first argument and $x_1$ the second.
             x1_ancestor_indices: The ancestor indices of $x_1$.
 
         Returns:

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/cuthbertlib/types.py

@@ -14,7 +14,9 @@ ScalarArray: TypeAlias = (
 ScalarArrayLike: TypeAlias = ArrayLike  # Object that will be cast to a ScalarArray
 
 LogDensity: TypeAlias = Callable[[ArrayTreeLike], ScalarArray]
-LogConditionalDensity: TypeAlias = Callable[[ArrayTreeLike, ArrayTreeLike], ScalarArray]
+LogConditionalDensity: TypeAlias = Callable[
+    [ArrayTreeLike, ArrayTreeLike], ScalarArray
+]  # p(x_1 | x_0), where x_0 is the first argument and x_1 the second.
 LogConditionalDensityAux: TypeAlias = Callable[
     [ArrayTreeLike, ArrayTreeLike], tuple[ScalarArray, ArrayTree]
 ]

{cuthbertlib-0.0.9 → cuthbertlib-0.0.11}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cuthbertlib"
-version = "0.0.9"
+version = "0.0.11"
 description = "Atomic building blocks for state-space model inference with JAX"
 requires-python = ">=3.10"
 readme = "README.md"