Unify explicit filter initialization across inference methods

cuthbert/README.md

@@ -10,18 +10,18 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-init_model_inputs = ...
 filter_model_inputs = ...
 
 # Load inference method
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
-    get_init_params=get_init_params,
+    m0=m0,
+    chol_P0=chol_P0,
     get_dynamics_params=get_dynamics_params,
     get_observation_params=get_observation_params,
 )   # build_filter function takes all inference-specific arguments, swap this out for different inference methods.
 
 # Online inference
-state = kalman_filter.init_prepare(init_model_inputs)
+state = kalman_filter.init_prepare()
 
 for t in range(T):
     model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
@@ -34,7 +34,7 @@ Or for offline inference:
 ```python
 kalman_smoother = cuthbert.gaussian.kalman.build_smoother(get_dynamics_params)
 
-init_state = kalman_filter.init_prepare(init_model_inputs)
+init_state = kalman_filter.init_prepare()
 filter_states = cuthbert.filter(kalman_filter, filter_model_inputs, init_state)
 smoother_states = cuthbert.smoother(kalman_smoother, filter_states, filter_model_inputs)
 ```

cuthbert/discrete/filter.py

@@ -13,7 +13,6 @@
 from jax import tree
 
 from cuthbert.discrete.types import (
-    GetInitDist,
     GetObsLogLikelihoods,
     GetTransitionMatrix,
 )
@@ -43,22 +42,22 @@ def log_normalizing_constant(self) -> Array:
 
 
 def build_filter(
-    get_init_dist: GetInitDist,
+    init_dist: Array,
     get_trans_matrix: GetTransitionMatrix,
     get_obs_lls: GetObsLogLikelihoods,
 ) -> Filter:
     r"""Builds a filter object for discrete hidden Markov models.
 
     Args:
-        get_init_dist: Function to get initial state probabilities $m_i = p(x_0 = i)$.
+        init_dist: Array of initial state probabilities $m_i = p(x_0 = i)$.
         get_trans_matrix: Function to get the transition matrix $A_{ij} = p(x_t = j \mid x_{t-1} = i)$.
         get_obs_lls: Function to get observation log likelihoods $b_i = \log p(y_t | x_t = i)$.
 
     Returns:
         Filter object. Suitable for associative scan.
     """
     return Filter(
-        init_prepare=partial(init_prepare, get_init_dist=get_init_dist),
+        init_prepare=partial(init_prepare, init_dist=init_dist),
         filter_prepare=partial(
             filter_prepare, get_trans_matrix=get_trans_matrix, get_obs_lls=get_obs_lls
         ),
@@ -67,27 +66,23 @@ def build_filter(
     )
 
 
-def init_prepare(
-    model_inputs: ArrayTreeLike, get_init_dist: GetInitDist, key: KeyArray | None = None
-) -> DiscreteFilterState:
+def init_prepare(init_dist: Array, key: KeyArray | None = None) -> DiscreteFilterState:
     """Prepare the initial state for the filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_dist: Function to get initial state probabilities m_i = p(x_0 = i).
+        init_dist: Array of initial state probabilities m_i = p(x_0 = i).
         key: JAX random key - not used.
 
     Returns:
         Prepared state for the filter.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    init_dist = get_init_dist(model_inputs)
+    init_dist = jnp.asarray(init_dist)
     N = init_dist.shape[-1]
     f = init_dist[..., None, :] * jnp.ones((N, N))
     # repeat init_dist N times onto newly added penultimate axis
     log_g = jnp.zeros_like(init_dist)
     return DiscreteFilterState(
-        elem=filtering.FilterScanElement(f, log_g), model_inputs=model_inputs
+        elem=filtering.FilterScanElement(f, log_g), model_inputs=None
     )
 
 

cuthbert/discrete/types.py

@@ -5,22 +5,6 @@
 from cuthbertlib.types import Array, ArrayTreeLike
 
 
-class GetInitDist(Protocol):
-    """Protocol for specifying the initial distribution."""
-
-    def __call__(self, model_inputs: ArrayTreeLike) -> Array:
-        """Get the initial distribution.
-
-        Args:
-            model_inputs: Model inputs.
-
-        Returns:
-            An array $m$ of shape (N,) where N is the number of states,
-                with $m_i = p(x_0 = i)$.
-        """
-        ...
-
-
 class GetTransitionMatrix(Protocol):
     """Protocol for specifying the transition matrix."""
 

cuthbert/ensemble_kalman/ensemble_kalman_filter.py

@@ -81,7 +81,8 @@ def build_filter(
     """Builds an Ensemble Kalman Filter object.
 
     Args:
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, generates a single sample from
+            the initial distribution.
         get_dynamics: Function to get dynamics function (x_t, key) -> x_{t+1} ~ p(x_{t+1} | x_t) from model inputs.
         get_observations: Function to get observation function, chol_R, and y from model inputs.
         n_particles: Number of particles.
@@ -134,7 +135,6 @@ def build_filter(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
     init_sample: InitSample,
     n_particles: int,
     store_predicted_ensemble: bool = False,
@@ -143,8 +143,8 @@ def init_prepare(
     """Prepare the initial state for the EnKF.
 
     Args:
-        model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, generates a single sample from
+            the initial distribution.
         n_particles: Number of particles.
         store_predicted_ensemble: Whether to store incoming forecast ensembles.
         key: JAX random key.
@@ -155,19 +155,18 @@ def init_prepare(
     Raises:
         ValueError: If key is None.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Sample ensemble from initial distribution
     keys = random.split(key, n_particles)
-    ensemble = jax.vmap(init_sample, (0, None))(keys, model_inputs)
+    ensemble = jax.vmap(init_sample)(keys)
     predicted_ensemble = dummy_tree_like(ensemble) if store_predicted_ensemble else None
 
     return EnKFState(
         key=key,
         ensemble=ensemble,
-        model_inputs=model_inputs,
+        model_inputs=None,
         log_normalizing_constant=jnp.array(0.0),
         predicted_ensemble=predicted_ensemble,
     )
@@ -184,7 +183,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, sampling from the initial
+            distribution. Bind initial parameters when building the filter.
         n_particles: Number of particles.
         store_predicted_ensemble: Whether to store incoming forecast ensembles.
         key: JAX random key.
@@ -200,10 +200,11 @@ def filter_prepare(
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Infer state shape from init_sample
-    dummy_particle = jax.eval_shape(init_sample, key, model_inputs)
+    dummy_particle = jax.eval_shape(init_sample, key)
     x_dim = dummy_particle.shape[0]
-    ensemble = jnp.empty((n_particles, x_dim))
-    ensemble = dummy_tree_like(ensemble)
+    ensemble = dummy_tree_like(
+        jax.ShapeDtypeStruct((n_particles, x_dim), dummy_particle.dtype)
+    )
     predicted_ensemble = ensemble if store_predicted_ensemble else None
 
     return EnKFState(

cuthbert/ensemble_kalman/types.py

@@ -12,12 +12,11 @@
 class InitSample(Protocol):
     """Protocol for sampling from the initial distribution."""
 
-    def __call__(self, key: KeyArray, model_inputs: ArrayTreeLike) -> Array:
+    def __call__(self, key: KeyArray) -> Array:
         """Sample from the initial distribution.
 
         Args:
             key: JAX PRNG key.
-            model_inputs: Model inputs.
 
         Returns:
             Sample from the initial distribution, shape (x_dim,).

cuthbert/factorial/README.md

@@ -47,7 +47,6 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-init_model_inputs = ...
 filter_model_inputs = ...
 
 # Define function to extract the factorial indices from model inputs
@@ -59,14 +58,15 @@ factorializer = cuthbert.factorial.gaussian.build_factorializer(get_factorial_in
 
 # Load inference method, with parameter extraction functions defined for factorial inference
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
-    get_init_params=get_init_params,  # Init specified to generate factorial state
+    m0=m0,   # Shape (F, d) i.e. includes factorial dimension
+    chol_P0=chol_P0,  # Shape (F, d, d)
     get_dynamics_params=get_dynamics_params,  # Dynamics specified to act on joint local state
     get_observation_params=get_observation_params,  # Observation specified to act on joint local state
 )
 
 # Online inference
-factorial_state = kalman_filter.init_prepare(init_model_inputs)
-factorial_state = factorializer.factorialize_init_state(factorial_state, init_model_inputs)
+factorial_state = kalman_filter.init_prepare()
+factorial_state = factorializer.factorialize_init_state(factorial_state)
 
 for t in range(T):
     model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
@@ -82,8 +82,8 @@ You can also use `cuthbert.factorial.filter` for convenient offline filtering.
 Note that associative/parallel filtering is not supported for factorial filtering.
 
 ```python
-init_factorial_state = kalman_filter.init_prepare(init_model_inputs)
-init_factorial_state = factorializer.factorialize_init_state(init_factorial_state, init_model_inputs)
+init_factorial_state = kalman_filter.init_prepare()
+init_factorial_state = factorializer.factorialize_init_state(init_factorial_state)
 
 local_filter_states, final_factorial_state = (
     cuthbert.factorial.filter(

cuthbert/factorial/filtering.py

@@ -6,6 +6,7 @@
 
 from cuthbert.factorial.types import Factorializer
 from cuthbert.inference import Filter
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -61,6 +62,7 @@ def filter(
     """
     T = tree.leaves(model_inputs)[0].shape[0]
     prepare_keys = _prepare_keys(key, T)
+    init_state = init_state._replace(model_inputs=dummy_leading_element(model_inputs))
 
     def body_local(prev_factorial_state, prep_inp_and_k):
         prep_inp, k = prep_inp_and_k

cuthbert/factorial/smc.py

@@ -14,7 +14,7 @@
 from cuthbert.smc.marginal_particle_filter import MarginalParticleFilterState
 from cuthbert.smc.particle_filter import ParticleFilterState
 from cuthbertlib.resampling import Resampling, ess_decorator
-from cuthbertlib.types import Array, ArrayLike, ArrayTreeLike
+from cuthbertlib.types import Array, ArrayLike
 
 GeneralParticleFilterState = TypeVar(
     "GeneralParticleFilterState", ParticleFilterState, MarginalParticleFilterState
@@ -58,7 +58,7 @@ def build_factorializer(
 
 
 def factorialize_init_state(
-    init_state: GeneralParticleFilterState, model_inputs: ArrayTreeLike
+    init_state: GeneralParticleFilterState,
 ) -> GeneralParticleFilterState:
     """Convert initial SMC state particles from `(N, F, ...)` to `(F, N, ...)`.
 
@@ -69,7 +69,6 @@ def factorialize_init_state(
 
     Args:
         init_state: Output from particle filter `init_prepare`
-        model_inputs: The model inputs at the first time point - unused.
     """
     particles = tree.map(lambda x: jnp.moveaxis(x, 0, 1), init_state.particles)
     n_factors = tree.leaves(particles)[0].shape[0]

cuthbert/factorial/types.py

@@ -143,14 +143,11 @@ def __call__(
 class FactorializeInitState(Protocol):
     """Protocol for factorial post-processing of `init_prepare`."""
 
-    def __call__(
-        self, init_state: ArrayTreeLike, model_inputs: ArrayTreeLike
-    ) -> ArrayTree:
+    def __call__(self, init_state: ArrayTreeLike) -> ArrayTree:
         """Any processing of the output of `init_prepare` for factorial inference.
 
         Args:
             init_state: Output from base inference method's `init_prepare`
-            model_inputs: The model inputs at the first time point.
         """
         ...
 
@@ -182,9 +179,7 @@ class Factorializer(NamedTuple):
     join: Join
     marginalize: Marginalize
     insert: Insert
-    factorialize_init_state: FactorializeInitState = lambda init_state, model_inputs: (
-        init_state
-    )
+    factorialize_init_state: FactorializeInitState = lambda init_state: init_state
 
     def extract_and_join(
         self, factorial_state: ArrayTreeLike, model_inputs: ArrayTreeLike

cuthbert/factorial/utils.py

@@ -4,6 +4,7 @@
 from jax import tree, vmap
 
 from cuthbert.factorial.types import Extract
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayLike, ArrayTree, ArrayTreeLike
 
 
@@ -30,7 +31,9 @@ def serial_to_factorial(
             maximum index, preserving the behaviour of earlier versions.
         init_factorial_tree: Optional initial factorial tree to use, as the first
             element of each returned tree.
-            Leaves with shape (F, ...)
+            Leaves with shape (F, ...). Subtrees that are None in the initial
+            tree (such as model_inputs) are padded with dummy values matching
+            the history.
 
     Returns:
         A single tree when a single index is selected; otherwise, a list of trees in the
@@ -72,12 +75,23 @@ def extract_single_factor(factorial_index: ArrayLike) -> ArrayTree:
             return factor_states
 
         initial_state = extract(init_factorial_tree, factorial_index)
+
+        def prepend_initial(initial_leaf, history_subtree):
+            if initial_leaf is None:
+                # Subtrees absent from the initial tree (such as model_inputs)
+                # are padded with dummy values matching the history.
+                initial_leaf = dummy_leading_element(history_subtree)
+            return tree.map(
+                lambda initial, history: jnp.concatenate([initial[None], history]),
+                initial_leaf,
+                history_subtree,
+            )
+
         return tree.map(
-            lambda initial_leaf, history_leaf: jnp.concatenate(
-                [initial_leaf[None], history_leaf]
-            ),
+            prepend_initial,
             initial_state,
             factor_states,
+            is_leaf=lambda x: x is None,
         )
 
     factorial_trees = [

cuthbert/filtering.py

@@ -7,7 +7,7 @@
 from jax.lax import associative_scan, scan
 
 from cuthbert.inference import Filter
-from cuthbert.utils import dummy_tree_like
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -54,8 +54,7 @@ def filter(
     else:
         prepare_keys = random.split(key, T)
 
-    sample_model_inputs = tree.map(lambda x: x[0], model_inputs)
-    init_state = init_state._replace(model_inputs=dummy_tree_like(sample_model_inputs))
+    init_state = init_state._replace(model_inputs=dummy_leading_element(model_inputs))
 
     if parallel:
         other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(