Simplify init_prepare

cuthbert/filtering.py

@@ -7,6 +7,7 @@
 from jax.lax import associative_scan, scan
 
 from cuthbert.inference import Filter
+from cuthbert.utils import dummy_tree_like
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -53,6 +54,9 @@ 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))
+
     if parallel:
         other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
             model_inputs, prepare_keys

cuthbert/gaussian/kalman.py

@@ -9,11 +9,7 @@
 from jax import numpy as jnp
 from jax import tree
 
-from cuthbert.gaussian.types import (
-    GetDynamicsParams,
-    GetInitParams,
-    GetObservationParams,
-)
+from cuthbert.gaussian.types import GetDynamicsParams, GetObservationParams
 from cuthbert.inference import Filter, Smoother
 from cuthbert.utils import dummy_tree_like
 from cuthbertlib.kalman import filtering, smoothing
@@ -62,15 +58,17 @@ def chol_cov(self) -> Array:
 
 
 def build_filter(
-    get_init_params: GetInitParams,
+    m0: Array,
+    chol_P0: Array,
     get_dynamics_params: GetDynamicsParams,
     get_observation_params: GetObservationParams,
 ) -> Filter:
     """Builds an exact Kalman filter object for linear Gaussian SSMs.
 
     Args:
-        get_init_params: Function to get m0, chol_P0 to initialize filter state,
-            given model inputs sufficient to define p(x_0) = N(m0, chol_P0 @ chol_P0^T).
+        m0: The mean vector of the initial state.
+        chol_P0: The generalized Cholesky factor of the covariance matrix of
+            the initial state.
         get_dynamics_params: Function to get dynamics parameters, F, c, chol_Q
             given model inputs sufficient to define
             p(x_t | x_{t-1}) = N(F @ x_{t-1} + c, chol_Q @ chol_Q^T).
@@ -82,10 +80,7 @@ def build_filter(
         Filter object for exact Kalman filter. Suitable for associative scan.
     """
     return Filter(
-        init_prepare=partial(
-            init_prepare,
-            get_init_params=get_init_params,
-        ),
+        init_prepare=partial(init_prepare, m0=m0, chol_P0=chol_P0),
         filter_prepare=partial(
             filter_prepare,
             get_dynamics_params=get_dynamics_params,
@@ -132,23 +127,21 @@ def build_smoother(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
-    get_init_params: GetInitParams,
+    m0: Array,
+    chol_P0: Array,
     key: KeyArray | None = None,
 ) -> KalmanFilterState:
     """Prepare the initial state for the Kalman filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_params: Function to get m0, chol_P0 from model inputs.
+        m0: The mean vector of the initial state.
+        chol_P0: The generalized Cholesky factor of the covariance matrix of
+            the initial state.
         key: JAX random key - not used.
 
     Returns:
         State for the Kalman filter.
-            Contains mean and chol_cov (generalised Cholesky factor of covariance).
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    m0, chol_P0 = get_init_params(model_inputs)
     elem = filtering.FilterScanElement(
         A=jnp.zeros_like(chol_P0),
         b=m0,
@@ -157,7 +150,7 @@ def init_prepare(
         Z=jnp.zeros_like(chol_P0),
         ell=jnp.array(0.0),
     )
-    return KalmanFilterState(elem=elem, model_inputs=model_inputs)
+    return KalmanFilterState(elem=elem, model_inputs=None)
 
 
 def filter_prepare(

cuthbert/gaussian/types.py

@@ -7,14 +7,6 @@
 
 
 ### Kalman types
-class GetInitParams(Protocol):
-    """Protocol for defining the initial distribution of a linear Gaussian SSM."""
-
-    def __call__(self, model_inputs: ArrayTreeLike) -> tuple[Array, Array]:
-        """Get initial parameters (m0, chol_P0) from model inputs."""
-        ...
-
-
 class GetDynamicsParams(Protocol):
     """Protocol for defining the dynamics model of a linear Gaussian SSM."""
 

cuthbert/inference.py

@@ -8,18 +8,12 @@
 class InitPrepare(Protocol):
     """Protocol for preparing the initial state for the inference."""
 
-    def __call__(
-        self,
-        model_inputs: ArrayTreeLike,
-        *,
-        key: KeyArray | None = None,
-    ) -> ArrayTree:
+    def __call__(self, *, key: KeyArray | None = None) -> ArrayTree:
         """Prepare the initial state for the inference.
 
         The state at the first time point, prior to any observations.
 
         Args:
-            model_inputs: The model inputs at the first time point.
             key: The key for the random number generator.
                 Optional, as only used for stochastic inference methods