cuthbertlib 0.0.10__tar.gz → 0.0.12__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuthbertlib
-Version: 0.0.10
+Version: 0.0.12
 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.10 → cuthbertlib-0.0.12}/cuthbertlib/enkf/filtering.py

@@ -13,7 +13,7 @@ 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
+from cuthbertlib.types import Array, KeyArray, ScalarArray
 
 ObservationFn = Callable[[Array], Array]
 DynamicsFn = Callable[[Array, KeyArray], Array]

{cuthbertlib-0.0.10 → cuthbertlib-0.0.12}/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.10 → cuthbertlib-0.0.12}/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.10 → cuthbertlib-0.0.12}/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.12/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.10 → cuthbertlib-0.0.12}/pyproject.toml

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