cuthbertlib 0.0.8__tar.gz → 0.0.10__tar.gz

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/PKG-INFO

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

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

cuthbertlib-0.0.10/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.10/cuthbertlib/enkf/__init__.py

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

cuthbertlib-0.0.10/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, ArrayTreeLike, 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.8 → cuthbertlib-0.0.10}/cuthbertlib/kalman/filtering.py

@@ -208,6 +208,42 @@ def filtering_operator(
 
     mu = cho_solve((U1, True), b1)
     t1 = b1 @ mu - (eta2 + mu) @ tmp_2
-    ell = ell1 + ell2 - 0.5 * t1 + 0.5 * jnp.linalg.slogdet(D_inv)[1]
+
+    # Derivation for O(nx) log-determinant computation of D_inv:
+    # This is a long comment but I wanted to include the full derivation for clarity and future reference.
+    # The key idea is to express D_inv in terms of the blocks of Xi and then apply Sylvester's determinant theorem
+    # to compute its determinant efficiently.
+    #
+    # 1. Expand the blocks of Xi @ Xi.T:
+    #    (Xi @ Xi.T)[1,1] = I + U1.T @ Z2 @ Z2.T @ U1
+    #    (Xi @ Xi.T)[2,1] = Z2 @ Z2.T @ U1
+    #
+    # 2. Equate to the corresponding blocks of L @ L.T:
+    #    Xi11 @ Xi11.T = I + U1.T @ Z2 @ Z2.T @ U1
+    #    Xi21 @ Xi11.T = Z2 @ Z2.T @ U1
+    #
+    # 3. Expand D_inv using tmp_1 = Xi11^{-1} @ U1.T:
+    #    D_inv = I - tmp_1.T @ Xi21.T
+    #          = I - U1 @ Xi11^{-T} @ Xi21.T
+    #
+    # 4. Apply Sylvester's determinant theorem:
+    #    det(D_inv) = det(I - Xi11^{-T} @ Xi21.T @ U1)
+    #
+    # 5. Multiply interior by Xi11^{-1} @ Xi11 and substitute block identities:
+    #    Let P = U1.T @ Z2 @ Z2.T @ U1
+    #    det(D_inv) = det(I - (Xi11 @ Xi11.T)^{-1} @ (Xi21 @ Xi11.T).T @ U1)
+    #               = det(I - (I + P)^{-1} @ P)
+    #               = det((I + P)^{-1})
+    #               = 1 / det(Xi11 @ Xi11.T)
+    #               = det(Xi11)^{-2}
+    #
+    # 6. Simplify the log-determinant term in the log-likelihood:
+    #    0.5 * log(det(D_inv)) = -log(|det(Xi11)|)
+    #
+    # Since Xi11 is lower triangular, the log-determinant is the sum of the logs of its diagonal.
+    # Replace `0.5 * jnp.linalg.slogdet(D_inv)[1]` with:
+    # -jnp.sum(jnp.log(jnp.abs(jnp.diag(Xi11))))
+    # ell = ell1 + ell2 - 0.5 * t1 + 0.5 * jnp.linalg.slogdet(D_inv)[1]
+    ell = ell1 + ell2 - 0.5 * t1 - jnp.sum(jnp.log(jnp.abs(jnp.diag(Xi11))))
 
     return FilterScanElement(A, b, U, eta, Z, ell)

cuthbertlib-0.0.10/cuthbertlib/linalg/tria.py

@@ -0,0 +1,97 @@
+"""Implements triangularization operator a matrix via QR decomposition."""
+
+import jax
+import jax.numpy as jnp
+
+from cuthbertlib.types import Array
+
+
+def _adj(x: Array) -> Array:
+    """Conjugate transpose for batched matrices."""
+    return jnp.swapaxes(x.conj(), -1, -2)
+
+
+@jax.custom_jvp
+def tria(A: Array) -> Array:
+    """A triangularization operator using QR decomposition.
+
+    Args:
+        A: The matrix to triangularize.
+
+    Returns:
+        A lower triangular matrix R such that R @ R.T = A @ A.T.
+
+    References:
+        Paper: Arasaratnam and Haykin (2008): Square-Root Quadrature Kalman Filtering
+            https://ieeexplore.ieee.org/document/4524036
+    """
+    _, R_qr = jnp.linalg.qr(_adj(A), mode="reduced")
+    return _adj(R_qr)
+
+
+@tria.defjvp
+def _tria_jvp(primals, tangents):
+    # Derivation of the exact analytical JVP for the lower-triangularization operator:
+    #
+    # 1. The operation computes a lower triangular R such that:
+    #    R @ R.T = A @ A.T
+    #
+    # 2. Taking the differential of both sides yields the Gramian differential identity:
+    #    dR @ R.T + R @ dR.T = dA @ A.T + A @ dA.T
+    #
+    # 3. For rank-deficient A, the exact lower-triangular tangent dR decomposes into
+    #    column-space and null-space components: dR = dR_col + dR_null.
+    #
+    # 4. To find dR_col, multiply the identity from the left by R^\dagger (pseudoinverse)
+    #    and from the right by R^{\dagger T}:
+    #    R^\dagger @ dR_col + dR_col.T @ R^{\dagger T} = R^\dagger @ dA @ A.T @ R^{\dagger T} + R^\dagger @ A @ dA.T @ R^{\dagger T}
+    #
+    # 5. Let Q be the active subspace orthogonal factor such that A = R @ Q.T.
+    #    Substituting A.T @ R^{\dagger T} = Q and R^\dagger @ A = Q.T:
+    #    R^\dagger @ dR_col + dR_col.T @ R^{\dagger T} = R^\dagger @ dA @ Q + Q.T @ dA.T @ R^{\dagger T}
+    #
+    # 6. Define K = R^\dagger @ dA @ Q. The right hand side becomes K + K.T:
+    #    R^\dagger @ dR_col + (R^\dagger @ dR_col).T = K + K.T
+    #
+    # 7. Define dM_col = R^\dagger @ dR_col. Solving for the lower-triangular dM_col:
+    #    dM_col = tril(K + K.T) - diag(K)
+    #
+    # 8. Recover the column-space differential dR_col by left-multiplying by R:
+    #    dR_col = R @ dM_col
+    #
+    # 9. To satisfy the orthogonal cross-terms for arbitrary perturbations outside
+    #    the column space of A, add the null-space component projected via (I - R @ R^\dagger):
+    #    dR_null = (I - R @ R^\dagger) @ dA @ Q
+    #
+    # 10. The complete, exact JVP is the sum of both components:
+    #     dR = dR_col + dR_null
+
+    (A,) = primals
+    (dA,) = tangents
+
+    A_T = jnp.swapaxes(A, -1, -2)
+    Q, R_qr = jnp.linalg.qr(A_T, mode="reduced")
+
+    R = jnp.swapaxes(R_qr, -1, -2)
+
+    # Q has shape (..., M, N). A is (..., N, M).
+    # A^T = Q R_qr => A = R_qr^T Q^T = R Q^T
+
+    R_pinv = jnp.linalg.pinv(R)
+
+    # K = R^{-1} dA Q
+    # R can be degenerate so we use the pseudoinverse to ensure the JVP is well-defined everywhere,
+    K = R_pinv @ dA @ Q
+    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
+
+    # Compute the null-space part
+    dR_null = (jnp.eye(R.shape[-2], dtype=R.dtype) - R @ R_pinv) @ dA @ Q
+
+    # Apply to get the tangent at R
+    dR = R @ dM + dR_null
+
+    return R, dR

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/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.10/cuthbertlib/resampling/README.md

@@ -0,0 +1,53 @@
+# Resampling
+
+This sub-repository provides a unified interface for a variety of resampling
+methods, which convert a set of weighted samples into an unweighted one which
+likely contains duplicates.
+
+A typical call to the library would be:
+
+```python
+sampling_key, resampling_key = jax.random.split(jax.random.key(0))
+particles = jax.random.normal(sampling_key, (100, 2))
+logits = jax.vmap(lambda x: jnp.where(jnp.all(x > 0), 0, -jnp.inf))(particles)
+
+resampled_indices, _, resampled_particles = resampling.multinomial.resampling(resampling_key, logits, particles, 100)
+```
+
+Or for conditional resampling:
+
+```python
+# Here we resample but keep particle at index 0 fixed
+conditional_resampled_indices, _, conditional_resampled_particles = resampling.multinomial.conditional_resampling(
+    resampling_key, logits, particles, 100, pivot_in=0, pivot_out=0
+)
+```
+
+Adaptive resampling (i.e. resampling only when the effective sample size is below a
+threshold) is also supported via a decorator:
+
+```python
+adaptive_resampling = resampling.adaptive.ess_decorator(
+    resampling.multinomial.resampling,
+    threshold=0.5,
+)
+adaptive_resampled_indices, _, adaptive_resampled_particles = adaptive_resampling(
+    resampling_key, logits, particles, 100
+)
+```
+
+For consistent gradient estimates with respect to model parameters, the [stop-gradient particle filter](https://arxiv.org/abs/2106.10314) is also implemented as a decorator.
+
+```python
+differentiable_resampling = resampling.stop_gradient.stop_gradient_decorator(
+    resampling.multinomial.resampling
+)
+# can be combined with adaptive resampling
+adaptive_and_differentiable_resampling = resampling.adaptive.ess_decorator(
+    differentiable_resampling,
+    threshold=0.5,
+)
+resampled_indices, _, resampled_particles = adaptive_and_differentiable_resampling(
+    resampling_key, logits, particles, 100
+)
+```

cuthbertlib-0.0.10/cuthbertlib/resampling/__init__.py

@@ -0,0 +1,11 @@
+from cuthbertlib.resampling import (
+    adaptive,
+    autodiff,
+    killing,
+    multinomial,
+    systematic,
+)
+from cuthbertlib.resampling.adaptive import ess_decorator
+from cuthbertlib.resampling.autodiff import stop_gradient_decorator
+from cuthbertlib.resampling.protocols import ConditionalResampling, Resampling
+from cuthbertlib.resampling.utils import inverse_cdf

cuthbertlib-0.0.10/cuthbertlib/resampling/adaptive.py

@@ -0,0 +1,72 @@
+"""Adaptive resampling decorator.
+
+Provides a decorator to turn any Resampling function into an adaptive resampling
+function which performs resampling only when the effective sample size (ESS)
+falls below a threshold.
+"""
+
+from functools import wraps
+
+import jax
+import jax.numpy as jnp
+
+from cuthbertlib.resampling.protocols import Resampling
+from cuthbertlib.smc.ess import log_ess
+from cuthbertlib.types import Array, ArrayLike, ArrayTree, ArrayTreeLike
+
+
+def ess_decorator(func: Resampling, threshold: float) -> Resampling:
+    """Wrap a Resampling function so that it only resamples when ESS < threshold.
+
+    The returned function is jitted and has `n` as a static argument. The
+    original resampler's docstring is appended to this wrapper's docstring so
+    IDEs and users can see the underlying algorithm documentation.
+
+    Args:
+        func: A resampling function with signature
+              (key, logits, positions, n) -> (indices, logits_out, positions_out).
+        threshold: Fraction of particle count specifying when to resample.
+            Resampling is triggered when ESS < ess_threshold * n.
+
+    Returns:
+        A Resampling function implementing adaptive resampling.
+    """
+    # Build a descriptive docstring that includes the wrapped function doc
+    wrapped_doc = func.__doc__ or ""
+    doc = f"""
+    Adaptive resampling decorator (threshold={threshold}).
+
+    This wrapper will call the provided resampling function only when the
+    effective sample size (ESS) is below `ess_threshold * n`.
+
+    Wrapped resampler documentation:
+    {wrapped_doc}
+    """
+
+    @wraps(func)
+    def _wrapped(
+        key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+    ) -> tuple[Array, Array, ArrayTree]:
+        logits_arr = jnp.asarray(logits)
+        N = logits_arr.shape[0]
+        if n != N:
+            raise AssertionError(
+                "The number of sampled indices must be equal to the number of "
+                f"particles for `adaptive` resampling. Got {n} instead of {N}."
+            )
+
+        def _do_resample():
+            return func(key, logits_arr, positions, n)
+
+        def _no_resample():
+            return jnp.arange(n), logits_arr, positions
+
+        return jax.lax.cond(
+            log_ess(logits_arr) < jnp.log(threshold * n),
+            _do_resample,
+            _no_resample,
+        )
+
+    # Attach the composed docstring and return a jitted version
+    _wrapped.__doc__ = doc
+    return jax.jit(_wrapped, static_argnames=("n",))

cuthbertlib-0.0.10/cuthbertlib/resampling/autodiff.py

@@ -0,0 +1,64 @@
+"""Implements decorators for automatic differentiation of resampling schemes.
+
+Current supported is the stop_gradient resampling scheme, which provides the
+classical Fisher estimates for the score function via automatic differentiation.
+This can be wrapped around a resampling scheme such as multinomial or systematic
+resampling.
+
+See [Scibior and Wood (2021)](https://arxiv.org/abs/2106.10314) for more details.
+"""
+
+from functools import wraps
+
+import jax
+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
+
+
+def stop_gradient_decorator(func: Resampling) -> Resampling:
+    """Wrap a Resampling function to use stop gradient resampling.
+
+    Args:
+        func: A resampling function with signature
+              (key, logits, positions, n) -> (indices, logits_out, positions_out).
+
+    Returns:
+        A Resampling function implementing stop gradient resampling.
+    """
+    # Build a descriptive docstring that includes the wrapped function doc
+    wrapped_doc = func.__doc__ or ""
+    doc = f"""
+    Stop gradient resampling decorator.
+
+    This wrapper will call the provided resampling function, and then apply 
+    the stop gradient trick of [Scibior and Wood (2021)](https://arxiv.org/abs/2106.10314). 
+    Resulting estimates of the score function (i.e., the gradient of the 
+    log-likelihood with respect to model parameters) are unbiased, 
+    corresponding to the classical Fisher estimate.
+
+    Wrapped resampler documentation:
+    {wrapped_doc}
+    """
+
+    @wraps(func)
+    def _wrapped(
+        key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+    ) -> tuple[Array, Array, ArrayTree]:
+        idx_base, logits_base, positions_base = func(
+            key, jax.lax.stop_gradient(logits), positions, n
+        )
+
+        logits = jnp.asarray(
+            logits_base
+            + apply_resampling_indices(logits, idx_base)
+            - jax.lax.stop_gradient(apply_resampling_indices(logits, idx_base))
+        )
+        return idx_base, logits, positions_base
+
+    # Attach the composed docstring and return a jitted version
+    _wrapped.__doc__ = doc
+    return jax.jit(_wrapped, static_argnames=("n",))

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/cuthbertlib/resampling/killing.py

@@ -11,7 +11,14 @@ from cuthbertlib.resampling.protocols import (
     conditional_resampling_decorator,
     resampling_decorator,
 )
-from cuthbertlib.types import Array, ArrayLike, ScalarArrayLike
+from cuthbertlib.resampling.utils import apply_resampling_indices
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    ScalarArrayLike,
+)
 
 _DESCRIPTION = """
 The Killing resampling is a simple resampling mechanism that checks if 
@@ -28,7 +35,9 @@ number of particles `logits.shape[0]`.
 
 
 @partial(resampling_decorator, name="Killing", desc=_DESCRIPTION)
-def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+def resampling(
+    key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+) -> tuple[Array, Array, ArrayTree]:
     logits = jnp.asarray(logits)
     key_1, key_2 = random.split(key)
     N = logits.shape[0]
@@ -43,27 +52,30 @@ def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
 
     survived = log_uniforms <= logits - max_logit
     if_survived = jnp.arange(N)  # If the particle survives, it keeps its index
-    otherwise = multinomial.resampling(
-        key_2, logits, N
+    otherwise_idx, _, _ = multinomial.resampling(
+        key_2, logits, positions, N
     )  # otherwise, it is replaced by another particle
-    idx = jnp.where(survived, if_survived, otherwise)
-    return idx
+    idx = jnp.where(survived, if_survived, otherwise_idx)
+    # After resampling, all particles have equal weight
+    logits_out = jnp.zeros_like(logits)
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 @partial(conditional_resampling_decorator, name="Killing", desc=_DESCRIPTION)
 def conditional_resampling(
     key: Array,
     logits: ArrayLike,
+    positions: ArrayTreeLike,
     n: int,
     pivot_in: ScalarArrayLike,
     pivot_out: ScalarArrayLike,
-) -> Array:
+) -> tuple[Array, Array, ArrayTree]:
     pivot_in = jnp.asarray(pivot_in)
     pivot_out = jnp.asarray(pivot_out)
 
     # Unconditional resampling
     key_resample, key_shuffle = random.split(key)
-    idx = resampling(key_resample, logits, n)
+    idx_uncond, _, _ = resampling(key_resample, logits, positions, n)
 
     # Conditional rolling pivot
     max_logit = jnp.max(logits)
@@ -76,9 +88,11 @@ def conditional_resampling(
 
     pivot_weights = jnp.exp(pivot_logits - logsumexp(pivot_logits))
     pivot = random.choice(key_shuffle, n, p=pivot_weights)
-    idx = jnp.roll(idx, pivot_in - pivot)
+    idx = jnp.roll(idx_uncond, pivot_in - pivot)
     idx = idx.at[pivot_in].set(pivot_out)
-    return idx
+    # After resampling, all particles have equal weight
+    logits_out = jnp.zeros_like(logits)
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 def _log1mexp(x: ArrayLike) -> Array:

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/cuthbertlib/resampling/multinomial.py

@@ -10,8 +10,14 @@ from cuthbertlib.resampling.protocols import (
     conditional_resampling_decorator,
     resampling_decorator,
 )
-from cuthbertlib.resampling.utils import inverse_cdf
-from cuthbertlib.types import Array, ArrayLike, ScalarArrayLike
+from cuthbertlib.resampling.utils import apply_resampling_indices, inverse_cdf
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    ScalarArrayLike,
+)
 
 _DESCRIPTION = """
 This has higher variance than other resampling schemes as it samples from
@@ -21,7 +27,9 @@ As a rule of thumb, you often don't."""
 
 
 @partial(resampling_decorator, name="Multinomial", desc=_DESCRIPTION)
-def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+def resampling(
+    key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+) -> tuple[Array, Array, ArrayTree]:
     # In practice we don't have to sort the generated uniforms, but searchsorted
     # works faster and is more stable if both inputs are sorted, so we use the
     # _sorted_uniforms from N. Chopin, but still use searchsorted instead of his
@@ -32,23 +40,26 @@ def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
     key_uniforms, key_shuffle = random.split(key)
     sorted_uniforms = _sorted_uniforms(key_uniforms, n)
     idx = inverse_cdf(sorted_uniforms, logits)
-    return random.permutation(key_shuffle, idx)
+    idx = random.permutation(key_shuffle, idx)
+    logits_out = jnp.zeros_like(sorted_uniforms)
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 @partial(conditional_resampling_decorator, name="Multinomial", desc=_DESCRIPTION)
 def conditional_resampling(
     key: Array,
     logits: ArrayLike,
+    positions: ArrayTreeLike,
     n: int,
     pivot_in: ScalarArrayLike,
     pivot_out: ScalarArrayLike,
-) -> Array:
+) -> tuple[Array, Array, ArrayTree]:
     pivot_in = jnp.asarray(pivot_in)
     pivot_out = jnp.asarray(pivot_out)
 
-    idx = resampling(key, logits, n)
+    idx, logits_out, _ = resampling(key, logits, positions, n)
     idx = idx.at[pivot_in].set(pivot_out)
-    return idx
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 @partial(jax.jit, static_argnames=("n",))

cuthbertlib-0.0.10/cuthbertlib/resampling/protocols.py

@@ -0,0 +1,98 @@
+"""Shared protocols for resampling algorithms."""
+
+from typing import Protocol, runtime_checkable
+
+import jax
+
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    KeyArray,
+    ScalarArrayLike,
+)
+
+_RESAMPLING_DOC = """
+Args:
+    key: JAX PRNG key.
+    logits: Logits.
+    positions: ArrayTreeLike
+    n: Number of indices to sample.
+
+Returns:
+    ancestors: Array of resampling indices.
+    logits: Array of log-weights after resampling.
+    positions: ArrayTreeLike of resampled positions.
+"""
+
+_CONDITIONAL_RESAMPLING_DOC = """
+Args:
+    key: JAX PRNG key.
+    logits: Log-weights, possibly unnormalized.
+    positions: ArrayTreeLike
+    n: Number of indices to sample.
+    pivot_in: Index of the particle to keep.
+    pivot_out: Value of the output at index `pivot_in`.
+
+Returns:
+    ancestors: Array of size n with indices to use for resampling.
+    logits: Array of log-weights after resampling.
+    positions: ArrayTreeLike of resampled positions.
+"""
+
+
+@runtime_checkable
+class Resampling(Protocol):
+    """Protocol for resampling operations."""
+
+    def __call__(
+        self, key: KeyArray, logits: ArrayLike, positions: ArrayTreeLike, n: int
+    ) -> tuple[Array, Array, ArrayTree]:
+        f"""Computes resampling indices according to given logits.
+        {_RESAMPLING_DOC}
+        """
+        ...
+
+
+@runtime_checkable
+class ConditionalResampling(Protocol):
+    """Protocol for conditional resampling operations."""
+
+    def __call__(
+        self,
+        key: KeyArray,
+        logits: ArrayLike,
+        positions: ArrayTreeLike,
+        n: int,
+        pivot_in: ScalarArrayLike,
+        pivot_out: ScalarArrayLike,
+    ) -> tuple[Array, Array, ArrayTree]:
+        f"""Conditional resampling.
+        {_CONDITIONAL_RESAMPLING_DOC}
+        """
+        ...
+
+
+def resampling_decorator(func: Resampling, name: str, desc: str = "") -> Resampling:
+    """Decorate Resampling function with unified docstring."""
+    doc = f"""
+    {name} resampling. {desc}
+    {_RESAMPLING_DOC}
+    """
+
+    func.__doc__ = doc
+    return jax.jit(func, static_argnames=("n",))
+
+
+def conditional_resampling_decorator(
+    func: ConditionalResampling, name: str, desc: str = ""
+) -> ConditionalResampling:
+    """Decorate ConditionalResampling function with unified docstring."""
+    doc = f"""
+    {name} conditional resampling. {desc}
+    {_CONDITIONAL_RESAMPLING_DOC}
+    """
+
+    func.__doc__ = doc
+    return jax.jit(func, static_argnames=("n",))

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/cuthbertlib/resampling/systematic.py

@@ -11,8 +11,14 @@ from cuthbertlib.resampling.protocols import (
     conditional_resampling_decorator,
     resampling_decorator,
 )
-from cuthbertlib.resampling.utils import inverse_cdf
-from cuthbertlib.types import Array, ArrayLike, ScalarArrayLike
+from cuthbertlib.resampling.utils import apply_resampling_indices, inverse_cdf
+from cuthbertlib.types import (
+    Array,
+    ArrayLike,
+    ArrayTree,
+    ArrayTreeLike,
+    ScalarArrayLike,
+)
 
 _DESCRIPTION = """
 The Systematic resampling is a variance reduction which places marginally
@@ -21,19 +27,24 @@ uniform samples into the [0, 1] interval but only requires one uniform random.
 
 
 @partial(resampling_decorator, name="Systematic", desc=_DESCRIPTION)
-def resampling(key: Array, logits: ArrayLike, n: int) -> Array:
+def resampling(
+    key: Array, logits: ArrayLike, positions: ArrayTreeLike, n: int
+) -> tuple[Array, Array, ArrayTree]:
     us = (random.uniform(key, ()) + jnp.arange(n)) / n
-    return inverse_cdf(us, logits)
+    idx = inverse_cdf(us, logits)
+    logits_out = jnp.zeros_like(us)
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 @partial(conditional_resampling_decorator, name="Systematic", desc=_DESCRIPTION)
 def conditional_resampling(
     key: Array,
     logits: ArrayLike,
+    positions: ArrayTreeLike,
     n: int,
     pivot_in: ScalarArrayLike,
     pivot_out: ScalarArrayLike,
-) -> Array:
+) -> tuple[Array, Array, ArrayTree]:
     logits = jnp.asarray(logits)
     pivot_in = jnp.asarray(pivot_in)
     pivot_out = jnp.asarray(pivot_out)
@@ -46,17 +57,17 @@ def conditional_resampling(
     logits = jnp.roll(logits, -pivot_out)
     arange = jnp.roll(arange, -pivot_out)
 
-    idx = conditional_resampling_0_to_0(key, logits, n)
+    idx, logits_out = conditional_resampling_0_to_0(key, logits, n)
     idx = arange[idx]
     idx = jnp.roll(idx, pivot_in)
-    return idx
+    return idx, logits_out, apply_resampling_indices(positions, idx)
 
 
 def conditional_resampling_0_to_0(
     key: Array,
     logits: ArrayLike,
     n: int,
-) -> Array:
+) -> tuple[Array, Array]:
     logits = jnp.asarray(logits)
 
     N = logits.shape[0]
@@ -81,4 +92,4 @@ def conditional_resampling_0_to_0(
     roll_idx = jnp.floor(n_zero * W).astype(int)
 
     idx = select(n_zero == 1, idx, jnp.roll(idx, -zero_loc[roll_idx]))
-    return jnp.clip(idx, 0, N - 1)
+    return jnp.clip(idx, 0, N - 1), jnp.zeros_like(linspace)

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/cuthbertlib/resampling/utils.py

@@ -4,10 +4,11 @@ import jax
 import jax.numpy as jnp
 import numba as nb
 import numpy as np
-from jax.lax import platform_dependent
+from jax.lax import platform_dependent, stop_gradient
 from jax.scipy.special import logsumexp
+from jax.tree_util import tree_map
 
-from cuthbertlib.types import Array, ArrayLike
+from cuthbertlib.types import Array, ArrayLike, ArrayTree, ArrayTreeLike
 
 
 @jax.jit
@@ -31,7 +32,10 @@ def inverse_cdf(sorted_uniforms: ArrayLike, logits: ArrayLike) -> Array:
     """
     weights = jnp.exp(logits - logsumexp(logits))
     return platform_dependent(
-        sorted_uniforms, weights, cpu=_inverse_cdf_cpu, default=_inverse_cdf_default
+        sorted_uniforms,
+        stop_gradient(weights),
+        cpu=_inverse_cdf_cpu,
+        default=_inverse_cdf_default,
     )
 
 
@@ -80,3 +84,8 @@ def _inverse_cdf_numba(su, ws, idx):
             j += 1
             s += ws[j]
         idx[n] = j
+
+
+def apply_resampling_indices(positions: ArrayTreeLike, idx: Array) -> ArrayTree:
+    """Apply resampling indices to positions."""
+    return tree_map(lambda x: x[idx], positions)

{cuthbertlib-0.0.8 → cuthbertlib-0.0.10}/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.8 → cuthbertlib-0.0.10}/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.
 
@@ -55,8 +56,9 @@ def simulate(
         idx, x0_res, idx_log_p = carry
         key_prop, key_acc = keys_t
 
-        prop_idx = multinomial.resampling(key_prop, log_weight_x0_all, n_samples)
-        x0_prop = jax.tree.map(lambda z: z[prop_idx], x0_all)
+        prop_idx, _, x0_prop = multinomial.resampling(
+            key_prop, log_weight_x0_all, x0_all, n_samples
+        )
         prop_log_p = jax.vmap(log_density)(x0_prop, x1_all)
 
         log_alpha = prop_log_p - idx_log_p
@@ -70,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.8 → cuthbertlib-0.0.10}/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.8 → cuthbertlib-0.0.10}/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.8 → cuthbertlib-0.0.10}/pyproject.toml

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

cuthbertlib-0.0.8/cuthbertlib/linalg/tria.py

@@ -1,21 +0,0 @@
-"""Implements triangularization operator a matrix via QR decomposition."""
-
-import jax
-
-from cuthbertlib.types import Array
-
-
-def tria(A: Array) -> Array:
-    r"""A triangularization operator using QR decomposition.
-
-    Args:
-        A: The matrix to triangularize.
-
-    Returns:
-        A lower triangular matrix $R$ such that $R R^\top = A A^\top$.
-
-    Reference:
-        [Arasaratnam and Haykin (2008)](https://ieeexplore.ieee.org/document/4524036): Square-Root Quadrature Kalman Filtering
-    """
-    _, R = jax.scipy.linalg.qr(A.T, mode="economic")
-    return R.T

cuthbertlib-0.0.8/cuthbertlib/resampling/README.md

@@ -1,26 +0,0 @@
-# Resampling
-
-This sub-repository provides a unified interface for a variety of resampling
-methods, which convert a set of weighted samples into an unweighted one which
-likely contains duplicates.
-
-A typical call to the library would be:
-
-```python
-sampling_key, resampling_key = jax.random.split(jax.random.key(0))
-particles = jax.random.normal(sampling_key, (100, 2))
-logits = jax.vmap(lambda x: jnp.where(jnp.all(x > 0), 0, -jnp.inf))(particles)
-
-resampled_indices = resampling.multinomial.resampling(resampling_key, logits, 100)
-resampled_particles = particles[resampled_indices]
-```
-
-Or for conditional resampling:
-
-```python
-# Here we resample but keep particle at index 0 fixed
-conditional_resampled_indices = resampling.multinomial.conditional_resampling(
-    resampling_key, logits, 100, pivot_in=0, pivot_out=0
-)
-conditional_resampled_particles = particles[conditional_resampled_indices]
-```

cuthbertlib-0.0.8/cuthbertlib/resampling/__init__.py

@@ -1,3 +0,0 @@
-from cuthbertlib.resampling import killing, multinomial, systematic
-from cuthbertlib.resampling.protocols import ConditionalResampling, Resampling
-from cuthbertlib.resampling.utils import inverse_cdf

cuthbertlib-0.0.8/cuthbertlib/resampling/protocols.py

@@ -1,92 +0,0 @@
-"""Shared protocols for resampling algorithms."""
-
-from typing import Protocol, runtime_checkable
-
-import jax
-
-from cuthbertlib.types import Array, ArrayLike, KeyArray, ScalarArrayLike
-
-
-@runtime_checkable
-class Resampling(Protocol):
-    """Protocol for resampling operations."""
-
-    def __call__(self, key: KeyArray, logits: ArrayLike, n: int) -> Array:
-        """Computes resampling indices according to given logits.
-
-        Args:
-            key: JAX PRNG key.
-            logits: Logits.
-            n: Number of indices to sample.
-
-        Returns:
-            Array of resampling indices.
-        """
-        ...
-
-
-@runtime_checkable
-class ConditionalResampling(Protocol):
-    """Protocol for conditional resampling operations."""
-
-    def __call__(
-        self,
-        key: KeyArray,
-        logits: ArrayLike,
-        n: int,
-        pivot_in: ScalarArrayLike,
-        pivot_out: ScalarArrayLike,
-    ) -> Array:
-        """Conditional resampling.
-
-        Args:
-            key: JAX PRNG key.
-            logits: Log-weights, possibly unnormalized.
-            n: Number of indices to sample.
-            pivot_in: Index of the particle to keep.
-            pivot_out: Value of the output at index `pivot_in`.
-
-        Returns:
-            Array of size n with indices to use for resampling.
-        """
-        ...
-
-
-def resampling_decorator(func: Resampling, name: str, desc: str = "") -> Resampling:
-    """Decorate Resampling function with unified docstring."""
-    doc = f"""
-    {name} resampling. {desc}
-
-    Args:
-        key: PRNGKey to use in resampling
-        logits: Log-weights, possibly unnormalized.
-        n: Number of indices to sample.
-    
-    Returns:
-        Array of size n with indices to use for resampling.
-    """
-
-    func.__doc__ = doc
-    return jax.jit(func, static_argnames=("n",))
-
-
-def conditional_resampling_decorator(
-    func: ConditionalResampling, name: str, desc: str = ""
-) -> ConditionalResampling:
-    """Decorate ConditionalResampling function with unified docstring."""
-    doc = f"""
-    {name} conditional resampling. {desc}
-
-    Args:
-        key: PRNGKey to use in resampling
-        logits: Log-weights, possibly unnormalized.
-        n: Number of indices to sample
-        pivot_in: Index of the particle to keep
-        pivot_out: Value of the output at index `pivot_in`
-    
-    Returns:
-        Array of size n with indices to use for resampling.
-    """
-
-    func.__doc__ = doc
-    return jax.jit(func, static_argnames=("n",))