cuthbertlib 0.0.12__tar.gz → 0.0.14__tar.gz

{cuthbertlib-0.0.12 → cuthbertlib-0.0.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuthbertlib
-Version: 0.0.12
+Version: 0.0.14
 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.12 → cuthbertlib-0.0.14}/cuthbertlib/__init__.py

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

cuthbertlib-0.0.14/cuthbertlib/ensemble_kalman/README.md

@@ -0,0 +1,15 @@
+# Ensemble Kalman
+
+This sub-repository provides modular functions for ensemble Kalman methods, including the ensemble Kalman filter (EnKF) and ensemble RTS smoother (EnRTS smoother).
+
+The core functions are:
+
+- `predict`: Propagate ensemble members through nonlinear dynamics with additive Gaussian noise.
+- `filter_update`: Update ensemble members with an observation using the EnKF update equation.
+- `smoother_update`: Apply one Ensemble Rauch-Tung-Striebel smoother update.
+
+Together, `predict` and `filter_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.
+
+The EnRTS algorithm provides a smoothing counterpart to the EnKF, based on the RTS smoother. It makes an empirical approximation of the RTS smoother gain, which is applied to the ensemble in a backwards pass to obtain a smoothing distribution. Note that, based on the outputs of the EnKF, the EnRTS step is entirely deterministic.

cuthbertlib-0.0.14/cuthbertlib/ensemble_kalman/__init__.py

@@ -0,0 +1,3 @@
+from cuthbertlib.ensemble_kalman.filtering import predict
+from cuthbertlib.ensemble_kalman.filtering import update as filter_update
+from cuthbertlib.ensemble_kalman.smoothing import update as smoother_update

cuthbertlib-0.0.14/cuthbertlib/ensemble_kalman/smoothing.py

@@ -0,0 +1,36 @@
+"""Implements the Ensemble Rauch-Tung-Striebel (EnRTS) smoother update.
+
+Cf. [Raanes (2016)[https://doi.org/10.1002/qj.2728].
+"""
+
+import jax.numpy as jnp
+
+from cuthbertlib.types import Array
+
+
+def update(
+    filtered_ensemble: Array,
+    predicted_ensemble: Array,
+    next_smoothed_ensemble: Array,
+) -> tuple[Array, Array]:
+    """Applies one EnRTS smoother update step.
+
+    Args:
+        filtered_ensemble: Filtered ensemble at time t, shape (N, x_dim).
+        predicted_ensemble: Paired forecast ensemble at time t + 1,
+            shape (N, next_x_dim).
+        next_smoothed_ensemble: Smoothed ensemble at time t + 1,
+            shape (N, next_x_dim).
+
+    Returns:
+        Tuple of the smoothed ensemble at time t and the EnRTS gain.
+    """
+    filtered_dev = filtered_ensemble - jnp.mean(filtered_ensemble, axis=0)
+    predicted_dev = predicted_ensemble - jnp.mean(predicted_ensemble, axis=0)
+
+    gain = filtered_dev.T @ jnp.linalg.pinv(predicted_dev.T)
+
+    smoothed_ensemble = (
+        filtered_ensemble + (next_smoothed_ensemble - predicted_ensemble) @ gain.T
+    )
+    return smoothed_ensemble, gain

{cuthbertlib-0.0.12 → cuthbertlib-0.0.14}/cuthbertlib/linearize/moments.py

@@ -63,14 +63,14 @@ def linearize_moments(
         - [sqrt-parallel-smoothers](https://github.com/EEA-sensors/sqrt-parallel-smoothers/blob/main/parsmooth/linearization/_extended.py)
     """
     if has_aux:
-        mean_and_chol_cov_function = cast(
+        mean_and_chol_cov_function_with_aux = cast(
             MeanAndCholCovFuncAux, mean_and_chol_cov_function
         )
 
         def mean_and_chol_cov_function_wrapper_aux(
             x: ArrayLike,
         ) -> tuple[Array, tuple[Array, Array, ArrayTree]]:
-            mean, chol_cov, aux = mean_and_chol_cov_function(x)
+            mean, chol_cov, aux = mean_and_chol_cov_function_with_aux(x)
             return mean, (mean, chol_cov, aux)
 
         F, (m, *extra) = jax.jacfwd(

{cuthbertlib-0.0.12 → cuthbertlib-0.0.14}/cuthbertlib/resampling/protocols.py

@@ -1,3 +1,6 @@
+# ty: ignore[invalid-return-type]
+# f-strings are not supported docstrings and ty is confused about this
+# See https://github.com/state-space-models/cuthbert/issues/261
 """Shared protocols for resampling algorithms."""
 
 from typing import Protocol, runtime_checkable

{cuthbertlib-0.0.12 → cuthbertlib-0.0.14}/pyproject.toml

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

cuthbertlib-0.0.12/cuthbertlib/enkf/README.md

@@ -1,14 +0,0 @@
-# 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.12/cuthbertlib/enkf/__init__.py

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