Update to cuthbert v0.0.14 (#300)

Cuthbert cut a new release which had some breaking changes (in particular, filters no longer call init_prepare themselves, that's now our jobs). There was also a changing in the filenaming for EnKF (mia culpa). This pins an exact version, since I think cuthbert is still undergoing breaking changes every now and then.

Author: DanWaxman

dynestyx/inference/integrations/cuthbert/discrete_filter.py

@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 import numpyro.distributions as dist
 from cuthbert import filter as cuthbert_filter
-from cuthbert.enkf import ensemble_kalman_filter
+from cuthbert.ensemble_kalman import ensemble_kalman_filter
 from cuthbert.gaussian import kalman, taylor
 from cuthbert.smc import particle_filter
 from cuthbertlib.resampling import (
@@ -250,11 +250,21 @@ def compute_cuthbert_filter(
             f"filter is not associative: {type(filter_config).__name__}."
         )
 
+    init_inputs = jax.tree.map(lambda leaf: leaf[0], cuthbert_inputs)
+    filter_inputs = jax.tree.map(lambda leaf: leaf[1:], cuthbert_inputs)
+    if key is None:
+        init_state = filter_obj.init_prepare(init_inputs)
+        filter_key = None
+    else:
+        init_key, filter_key = jax.random.split(key)
+        init_state = filter_obj.init_prepare(init_inputs, key=init_key)
+
     raw_states = cuthbert_filter(
         filter_obj,
-        cuthbert_inputs,
+        filter_inputs,
+        init_state,
         parallel=cast(bool, parallel),
-        key=key,
+        key=filter_key,
     )
     marginal_loglik = raw_states.log_normalizing_constant[-1]
     states = (

pyproject.toml

@@ -31,8 +31,8 @@ classifiers = [
 ]
 dependencies = [
     "effectful>=0.4.0",
-    "cuthbert>=0.0.10",
-    "cuthbertlib>=0.0.10",
+    "cuthbert==0.0.14",
+    "cuthbertlib==0.0.14",
     "cd-dynamax>=0.3.3",
     "matplotlib>=3.10.7",
     "numpyro>=0.19.0",