cuthbertlib 0.0.13__tar.gz → 0.0.15__tar.gz

{cuthbertlib-0.0.13 → cuthbertlib-0.0.15}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: cuthbertlib
-Version: 0.0.13
+Version: 0.0.15
 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.13 → cuthbertlib-0.0.15}/cuthbertlib/__init__.py

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

cuthbertlib-0.0.15/cuthbertlib/ensemble_kalman/README.md

@@ -0,0 +1,42 @@
+# Ensemble Kalman
+
+The high-level interfaces built on these functions are in
+[`cuthbert.ensemble_kalman.ensemble_kalman_filter`](../../cuthbert/ensemble_kalman/ensemble_kalman_filter.py)
+and
+[`cuthbert.ensemble_kalman.ensemble_rts_smoother`](../../cuthbert/ensemble_kalman/ensemble_rts_smoother.py).
+
+<!-- --8<-- [start:overview] -->
+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.
+- `construct_tapered_chol_innovation_covariance`: Construct a tapered innovation covariance factor without forming the covariance densely.
+- `gaspari_cohn`: Construct covariance tapers with the Gaspari-Cohn correlation function.
+- `gaussian`: Construct smooth, non-compact covariance tapers with the Gaussian correlation function.
+<!-- --8<-- [end:overview] -->
+
+## Ensemble Kalman filtering
+
+<!-- --8<-- [start:filtering] -->
+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.
+<!-- --8<-- [end:filtering] -->
+
+## Ensemble Rauch-Tung-Striebel smoothing
+
+<!-- --8<-- [start:smoothing] -->
+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.
+<!-- --8<-- [end:smoothing] -->
+
+## Covariance localization
+
+<!-- --8<-- [start:localization] -->
+`filter_update` accepts two independent low-level hooks. `cross_covariance_modifier` receives one empirical state-observation cross-covariance and returns its modified value; it defaults to `no_covariance_modifier`. For observation-space covariance localization, one may optionally use `construct_chol_innovation_covariance`, which receives normalized observation deviations `Y` and a Cholesky factor of the observation noise covariance, `chol_R`, and returns a generalized Cholesky factor of the localized innovation covariance $C_{yy} + R$. All localization happens before handling of missing data, so localization happens in the original coordinates.
+
+The most common form of localization is covariance tapering; cuthbertlib provides `gaspari_cohn` and `gaussian` correlation functions for this purpose. The Gaspari-Cohn correlation function is more classical, but has compact support which may cause difficulties in optimizing localization hyperparameters. The `gaussian` correlation function has infinite support, and may therefore have better gradient properties. For obsservation-space tapering, `cuthbertlib` provides a convenience `construct_tapered_chol_innovation_covariance`, which uses a Cholesky factor of the taper to construct a factor of the tapered innovation covariance. This comes at slightly higher cost, due to a larger QR solve.
+
+<!-- --8<-- [end:localization] -->

cuthbertlib-0.0.15/cuthbertlib/ensemble_kalman/__init__.py

@@ -0,0 +1,15 @@
+from cuthbertlib.ensemble_kalman.filtering import (
+    ConstructCholInnovationCovariance,
+    CrossCovarianceModifier,
+    no_covariance_modifier,
+    predict,
+)
+from cuthbertlib.ensemble_kalman.filtering import (
+    update as filter_update,
+)
+from cuthbertlib.ensemble_kalman.localization import (
+    construct_tapered_chol_innovation_covariance,
+    gaspari_cohn,
+    gaussian,
+)
+from cuthbertlib.ensemble_kalman.smoothing import update as smoother_update

{cuthbertlib-0.0.13/cuthbertlib/enkf → cuthbertlib-0.0.15/cuthbertlib/ensemble_kalman}/filtering.py

@@ -17,6 +17,23 @@ from cuthbertlib.types import Array, KeyArray, ScalarArray
 
 ObservationFn = Callable[[Array], Array]
 DynamicsFn = Callable[[Array, KeyArray], Array]
+CrossCovarianceModifier = Callable[[Array], Array]
+ConstructCholInnovationCovariance = Callable[[Array, Array], Array]
+
+
+def no_covariance_modifier(covariance: Array) -> Array:
+    """Return an empirical covariance unchanged.
+
+    The identity covariance modifier, used as the default when no modification
+    (e.g. localization) is requested.
+
+    Args:
+        covariance: Empirical covariance matrix.
+
+    Returns:
+        The covariance matrix, unchanged.
+    """
+    return covariance
 
 
 def predict(
@@ -56,6 +73,9 @@ def update(
     chol_R: Array,
     y: Array,
     perturbed_obs: bool = True,
+    cross_covariance_modifier: CrossCovarianceModifier = no_covariance_modifier,
+    construct_chol_innovation_covariance: ConstructCholInnovationCovariance
+    | None = None,
 ) -> tuple[Array, ScalarArray]:
     """Update ensemble members with an observation using the EnKF update.
 
@@ -71,6 +91,16 @@ def update(
         y: Observation vector, shape (y_dim,). NaNs indicate missing dimensions.
         perturbed_obs: If True, use perturbed observations (stochastic EnKF).
             If False, use deterministic update.
+        cross_covariance_modifier: Function that modifies the empirical
+            state-observation cross-covariance, shape (x_dim, y_dim), and returns
+            an array with the same shape. Defaults to the identity.
+        construct_chol_innovation_covariance: Optional function that
+            receives normalized observation deviations with shape (y_dim, N) and
+            ``chol_R`` with shape (y_dim, y_dim). It must return a generalized
+            Cholesky factor of the complete innovation covariance with shape
+            (y_dim, y_dim). The deviations have already been divided by
+            ``sqrt(N - 1)``. Both inputs use the original observation order.
+            ``None`` uses the standard, unlocalized square-root construction.
 
     Returns:
         Tuple of (updated_ensemble, log_likelihood).
@@ -79,27 +109,46 @@ def update(
 
     # Map ensemble to observation space
     y_pred = jax.vmap(observation_fn, (0,))(predicted_ensemble)
+    x_mean = jnp.mean(predicted_ensemble, axis=0)
+    x_dev = predicted_ensemble - x_mean
+
+    missing = jnp.isnan(y)
+
+    # Modify or construct covariances before reordering due to NaNs.
+    argsort = jnp.argsort(missing, stable=True)
+    original_y_dev = y_pred - jnp.mean(y_pred, axis=0)
+    normalized_original_y_dev = original_y_dev.T / jnp.sqrt(N - 1)
+
+    C_xy = x_dev.T @ original_y_dev / (N - 1)
+    C_xy = cross_covariance_modifier(C_xy)
+
+    if construct_chol_innovation_covariance is not None:
+        original_chol_S = construct_chol_innovation_covariance(
+            normalized_original_y_dev, chol_R
+        )
 
     # 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)
+    flag, chol_R, y, y_pred = collect_nans_chol(missing, 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
+    C_xy = C_xy[:, argsort]
+    C_xy = jnp.where(flag[None, :], 0.0, C_xy)
 
-    # 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)
+    if construct_chol_innovation_covariance is None:
+        chol_S = tria(jnp.concatenate([y_dev.T / jnp.sqrt(N - 1), chol_R], axis=1))
+    else:
+        # The constructor sees the original indexing. Only collect and refactor its
+        # result when dimensions are missing; otherwise preserve its returned factor.
+        chol_S = jax.lax.cond(
+            jnp.any(missing),
+            lambda chol: collect_nans_chol(missing, chol)[1],
+            lambda chol: chol,
+            original_chol_S,
+        )
 
     # Kalman gain: K = C_xy @ S^{-1} = C_xy @ cho_solve(chol_S, I)
     K = cho_solve((chol_S, True), C_xy.T).T

cuthbertlib-0.0.15/cuthbertlib/ensemble_kalman/localization.py

@@ -0,0 +1,89 @@
+"""Covariance localization utilities for ensemble Kalman methods."""
+
+import jax.numpy as jnp
+
+from cuthbertlib.linalg import tria
+from cuthbertlib.types import Array, ArrayLike, ScalarArrayLike
+
+
+def construct_tapered_chol_innovation_covariance(
+    Y: Array,
+    chol_taper: Array,
+    chol_R: Array,
+) -> Array:
+    """Construct a tapered innovation covariance factor without forming it densely.
+
+    If ``taper = chol_taper @ chol_taper.T`` and ``Y`` denotes the normalized
+    observation deviations, the returned generalized Cholesky factor ``chol_S``
+    satisfies ``chol_S @ chol_S.T = taper * (Y @ Y.T) + R``.
+
+    Args:
+        Y: Observation deviations transposed and divided by the square root of one
+            less than the ensemble size, shape (y_dim, n_particles).
+        chol_taper: Factor of a positive-semidefinite observation-space taper,
+            shape (y_dim, y_dim).
+        chol_R: Cholesky factor of the observation noise covariance, shape
+            (y_dim, y_dim).
+
+    Returns:
+        Generalized Cholesky factor of the complete tapered innovation covariance,
+        shape (y_dim, y_dim).
+    """
+    y_dim = Y.shape[0]
+    Y_tilde = (chol_taper[:, :, None] * Y[:, None, :]).reshape(y_dim, -1)
+    return tria(jnp.concatenate([Y_tilde, chol_R], axis=1))
+
+
+def gaussian(
+    distances: ArrayLike,
+    length_scale: ScalarArrayLike,
+) -> Array:
+    r"""Evaluates a Gaussian covariance taper.
+
+    The taper is the squared-exponential correlation function
+    $\rho(d; \ell) = \exp(-\frac{1}{2}(d / \ell)^2)$. This taper
+    has infinite support and is differentiable with respect
+    to the length scale.
+
+    Args:
+        distances: Distances at which to evaluate the taper.
+        length_scale: Positive characteristic distance of the taper.
+
+    Returns:
+        Taper values with the broadcast shape of the inputs.
+    """
+    scaled_distances = jnp.asarray(distances) / length_scale
+    return jnp.exp(-0.5 * jnp.square(scaled_distances))
+
+
+def gaspari_cohn(
+    distances: ArrayLike,
+    support_radius: ScalarArrayLike,
+) -> Array:
+    """Evaluates the compactly supported fifth-order Gaspari-Cohn taper.
+
+    This implements Eq. (4.10) of Gaspari and Cohn (1999),
+    https://doi.org/10.1002/qj.49712555417.
+
+    ``support_radius`` is the full support radius: taper values are exactly zero
+    where the absolute distance is greater than or equal to it.
+
+    Args:
+        distances: Distances at which to evaluate the taper.
+        support_radius: Positive distance at which the taper reaches zero.
+
+    Returns:
+        Taper values with the broadcast shape of the inputs.
+    """
+    distances = jnp.abs(distances)
+    q = 2 * distances / support_radius
+
+    inner = 1 - 5 / 3 * q**2 + 5 / 8 * q**3 + 1 / 2 * q**4 - 1 / 4 * q**5
+
+    # Avoid NaNs in computing the outer branch of the polynomial
+    safe_q = jnp.where(q > 0, q, 1)
+    # Factored version of Eq. (4.10) for numerical stability
+    outer = (2 - q) ** 4 * (2 * q**2 + 4 * q - 1) / (24 * safe_q)
+
+    within_support = jnp.where(q <= 1, inner, outer)
+    return jnp.where(distances < support_radius, within_support, 0)

cuthbertlib-0.0.15/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.13 → cuthbertlib-0.0.15}/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.13 → cuthbertlib-0.0.15}/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.13 → cuthbertlib-0.0.15}/pyproject.toml

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

cuthbertlib-0.0.13/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.13/cuthbertlib/enkf/__init__.py

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