Remove init from filter

cuthbert/README.md

@@ -10,7 +10,8 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-model_inputs = ...
+init_model_inputs = ...
+filter_model_inputs = ...
 
 # Load inference method
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
@@ -20,10 +21,10 @@ kalman_filter = cuthbert.gaussian.kalman.build_filter(
 )   # build_filter function takes all inference-specific arguments, swap this out for different inference methods.
 
 # Online inference
-state = kalman_filter.init_prepare(tree.map(lambda x: x[0], model_inputs))
+state = kalman_filter.init_prepare(init_model_inputs)
 
-for t in range(1, T):
-    model_inputs_t = tree.map(lambda x: x[t], model_inputs)
+for t in range(T):
+    model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
     prepare_state = kalman_filter.filter_prepare(model_inputs_t)
     state = kalman_filter.filter_combine(state, prepare_state)
 ```
@@ -33,7 +34,8 @@ Or for offline inference:
 ```python
 kalman_smoother = cuthbert.gaussian.kalman.build_smoother(get_dynamics_params)
 
-filter_states = cuthbert.filter(kalman_filter, model_inputs)
-smoother_states = cuthbert.smoother(kalman_smoother, filter_states, model_inputs)
+init_state = kalman_filter.init_prepare(init_model_inputs)
+filter_states = cuthbert.filter(kalman_filter, filter_model_inputs, init_state)
+smoother_states = cuthbert.smoother(kalman_smoother, filter_states, filter_model_inputs)
 ```
 <!-- --8<-- [end:unified_interface] -->

cuthbert/factorial/README.md

@@ -47,7 +47,8 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-model_inputs = ...
+init_model_inputs = ...
+filter_model_inputs = ...
 
 # Define function to extract the factorial indices from model inputs
 # Here we assume model_inputs is a NamedTuple with a field `factorial_inds`
@@ -64,12 +65,11 @@ kalman_filter = cuthbert.gaussian.kalman.build_filter(
 )
 
 # Online inference
-init_model_inputs = tree.map(lambda x: x[0], model_inputs)
 factorial_state = kalman_filter.init_prepare(init_model_inputs)
 factorial_state = factorializer.factorialize_init_state(factorial_state, init_model_inputs)
 
-for t in range(1, T):
-    model_inputs_t = tree.map(lambda x: x[t], model_inputs)
+for t in range(T):
+    model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
     local_state = factorializer.extract_and_join(prev_factorial_state, model_inputs_t)
     prepare_state = kalman_filter.filter_prepare(model_inputs_t)
     local_joint_filtered_state = kalman_filter.filter_combine(local_state, prepare_state)
@@ -82,9 +82,12 @@ 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, local_filter_states, final_factorial_state = (
+init_factorial_state = kalman_filter.init_prepare(init_model_inputs)
+init_factorial_state = factorializer.factorialize_init_state(init_factorial_state, init_model_inputs)
+
+local_filter_states, final_factorial_state = (
     cuthbert.factorial.filter(
-        kalman_filter, factorializer, model_inputs, output_factorial=False
+        kalman_filter, factorializer, filter_model_inputs, init_factorial_state, output_factorial=False
     )
 )
 ```

cuthbert/factorial/filtering.py

@@ -13,16 +13,21 @@ def filter(
     filter_obj: Filter,
     factorializer: Factorializer,
     model_inputs: ArrayTreeLike,
+    init_state: ArrayTree,
+    *,
     output_factorial: bool = False,
     key: KeyArray | None = None,
 ) -> (
-    ArrayTree | tuple[ArrayTree, ArrayTree, ArrayTree]
+    ArrayTree | tuple[ArrayTree, ArrayTree]
 ):  # TODO: Can overload this function so the type checker knows that the output is an ArrayTree if output_factorial is True and a tuple[ArrayTree, ArrayTree, ArrayTree] if output_factorial is False
     """Applies offline factorial filtering for given model inputs.
 
-    `model_inputs` should have leading temporal dimension of length T + 1,
+    `model_inputs` should have leading temporal dimension of length T,
     where T is the number of time steps excluding the initial state.
 
+    The output will have leading temporal dimension of length T + 1, and
+    include the initial state as the first element.
+
     Parallel associative filtering is not supported for factorial filtering.
 
     Note that if output_factorial is True, this function will output a factorial state
@@ -33,42 +38,35 @@ def filter(
     Args:
         filter_obj: The filter inference object.
         factorializer: The factorializer object for the inference method.
-        model_inputs: The model inputs (with leading temporal dimension of length T + 1).
+        model_inputs: The model inputs for filtering
+            (with leading temporal dimension of length T).
+        init_state: The initial factorial state
+            (with factorial dimension but no temporal dimension).
+            Generated by `filter_obj.init_prepare` and then factorialized by
+            `factorializer.factorialize_init_state`.
         output_factorial: If True, return a single state with first temporal dimension
             of length T + 1 and second factorial dimension of length F.
-            If False, return a tuple of states. The first being the initial state
-            with first dimension of length F and no temporal dimension.
-            The second being the local states for each time step, i.e. first
+            If False, return a tuple of states.
+            The first being the local states for each time step, i.e. first
             dimension of length T and no factorial dimension.
-            The third being the final factorial state with first dimension of
+            The second being the final factorial state with first dimension of
             length F and no temporal dimension.
         key: The key for the random number generator.
 
     Returns:
         If output_factorial is True, returns a single state with first temporal dimension
             of length T + 1 and second factorial dimension of length F.
         If output_factorial is False, returns a tuple of
-            (initial factorial state, local states for each time step,
-            final factorial state).
+            (local states for each time step, final factorial state).
     """
-    T = tree.leaves(model_inputs)[0].shape[0] - 1
+    T = tree.leaves(model_inputs)[0].shape[0]
 
     if key is None:
         # This will throw error if used as a key, which is desired behavior
         # (albeit not a useful error, we could improve this)
-        prepare_keys = jnp.empty(T + 1)
+        prepare_keys = jnp.empty(T)
     else:
-        prepare_keys = random.split(key, T + 1)
-
-    init_model_input = tree.map(lambda x: x[0], model_inputs)
-    init_factorial_state = filter_obj.init_prepare(
-        init_model_input, key=prepare_keys[0]
-    )
-    init_factorial_state = factorializer.factorialize_init_state(
-        init_factorial_state, init_model_input
-    )
-
-    prep_model_inputs = tree.map(lambda x: x[1:], model_inputs)
+        prepare_keys = random.split(key, T)
 
     def body_local(prev_factorial_state, prep_inp_and_k):
         prep_inp, k = prep_inp_and_k
@@ -99,12 +97,12 @@ def body_factorial(prev_factorial_state, prep_inp_and_k):
 
         _, factorial_states = scan(
             body_factorial,
-            init_factorial_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            init_state,
+            (model_inputs, prepare_keys),
         )
         factorial_states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]),
-            init_factorial_state,
+            init_state,
             factorial_states,
         )
 
@@ -113,7 +111,7 @@ def body_factorial(prev_factorial_state, prep_inp_and_k):
     else:
         final_factorial_state, local_states = scan(
             body_local,
-            init_factorial_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            init_state,
+            (model_inputs, prepare_keys),
         )
-        return init_factorial_state, local_states, final_factorial_state
+        return local_states, final_factorial_state

cuthbert/filtering.py

@@ -13,17 +13,25 @@
 def filter(
     filter_obj: Filter,
     model_inputs: ArrayTreeLike,
+    init_state: ArrayTree,
+    *,
     parallel: bool = False,
     key: KeyArray | None = None,
 ) -> ArrayTree:
     """Applies offline filtering given a filter object and model inputs.
 
-    `model_inputs` should have leading temporal dimension of length T + 1,
+    `model_inputs` should have leading temporal dimension of length T,
     where T is the number of time steps excluding the initial state.
 
+    The output will have leading temporal dimension of length T + 1, and
+    include the initial state as the first element.
+
     Args:
         filter_obj: The filter inference object.
-        model_inputs: The model inputs (with leading temporal dimension of length T + 1).
+        model_inputs: The model inputs for filtering
+            (with leading temporal dimension of length T).
+        init_state: The initial state (with no temporal dimension).
+            Generated by `filter_obj.init_prepare`.
         parallel: Whether to run the filter in parallel.
             Requires `filter.associative_filter` to be `True`.
         key: The key for the random number generator.
@@ -36,23 +44,18 @@ def filter(
             f"Parallel filtering attempted but filter.associative is False for {filter_obj}"
         )
 
-    T = tree.leaves(model_inputs)[0].shape[0] - 1
+    T = tree.leaves(model_inputs)[0].shape[0]
 
     if key is None:
         # This will throw error if used as a key, which is desired behavior
         # (albeit not a useful error, we could improve this)
-        prepare_keys = jnp.empty(T + 1)
+        prepare_keys = jnp.empty(T)
     else:
-        prepare_keys = random.split(key, T + 1)
-
-    init_model_input = tree.map(lambda x: x[0], model_inputs)
-    init_state = filter_obj.init_prepare(init_model_input, key=prepare_keys[0])
-
-    prep_model_inputs = tree.map(lambda x: x[1:], model_inputs)
+        prepare_keys = random.split(key, T)
 
     if parallel:
         other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
-            prep_model_inputs, prepare_keys[1:]
+            model_inputs, prepare_keys
         )
         prep_states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]), init_state, other_prep_states
@@ -72,7 +75,7 @@ def body(prev_state, prep_inp_and_k):
         _, states = scan(
             body,
             init_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            (model_inputs, prepare_keys),
         )
         states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]), init_state, states

docs/examples/diff_pf_resampling.md

@@ -123,7 +123,8 @@ def kalman_mll(theta, ys, *, m0, chol_P0):
 
     kf = kalman.build_filter(get_init_params, get_dynamics_params, get_observation_params)
     model_inputs = jnp.arange(ys.shape[0] + 1)
-    states = run_filter(kf, model_inputs, parallel=False)
+    init_state = kf.init_prepare(model_inputs[0])
+    states = run_filter(kf, model_inputs[1:], init_state, parallel=False)
     return states.log_normalizing_constant[-1]
 
 
@@ -196,7 +197,11 @@ def pf_mll(
     )
 
     model_inputs = jnp.arange(ys.shape[0] + 1)
-    states = run_filter(filt, model_inputs, parallel=False, key=key)
+    init_key, filter_key = jax.random.split(key)
+    init_state = filt.init_prepare(model_inputs[0], key=init_key)
+    states = run_filter(
+        filt, model_inputs[1:], init_state, parallel=False, key=filter_key
+    )
     return states.log_normalizing_constant[-1]
 
 
@@ -536,7 +541,11 @@ Finally, we should be sure that the forward pass is not actually being modified
         )
 
         model_inputs = jnp.arange(ys.shape[0] + 1)
-        states = run_filter(filt, model_inputs, parallel=False, key=key)
+        init_key, filter_key = jax.random.split(key)
+        init_state = filt.init_prepare(model_inputs[0], key=init_key)
+        states = run_filter(
+            filt, model_inputs[1:], init_state, parallel=False, key=filter_key
+        )
         return states.particles[-1]
 
 

docs/examples/dynamax_integration.md

@@ -171,7 +171,8 @@ Now we can run the `cuthbert` Kalman filter to obtain the filtering distribution
 
 ```{.python #dynamax-run-filter}
 # Run Kalman filtering
-filtered_states = cuthbert.filter(filter_obj, model_inputs)
+init_state = filter_obj.init_prepare(model_inputs[0])
+filtered_states = cuthbert.filter(filter_obj, model_inputs[1:], init_state)
 
 # Extract filtering results - remove initial time step
 filtered_means = filtered_states.mean[1:]
@@ -327,4 +328,3 @@ filters.
 <<dynamax-visualize>>
 ```
 -->
-

docs/examples/enkf_comparison.md

@@ -164,14 +164,16 @@ ekf = taylor.build_filter(
 )
 
 jitted_filter = jit(run_filter, static_argnames=["filter_obj"])
+filter_model_inputs = model_inputs[1:]
 
 n_timing = 20
-ekf_states = jitted_filter(ekf, model_inputs)  # warm up
+ekf_init_state = ekf.init_prepare(model_inputs[0])
+ekf_states = jitted_filter(ekf, filter_model_inputs, ekf_init_state)  # warm up
 jax.block_until_ready(ekf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(ekf, model_inputs))
+    jax.block_until_ready(jitted_filter(ekf, filter_model_inputs, ekf_init_state))
     _times.append(time.perf_counter() - _t0)
 ekf_time = float(jnp.median(jnp.array(_times)))
 
@@ -194,13 +196,18 @@ enkf = ensemble_kalman_filter.build_filter(
     perturbed_obs=True,
 )
 
-key, enkf_key = random.split(key)
-enkf_states = jitted_filter(enkf, model_inputs, key=enkf_key)  # warm up
+key, enkf_init_key, enkf_filter_key = random.split(key, 3)
+enkf_init_state = enkf.init_prepare(model_inputs[0], key=enkf_init_key)
+enkf_states = jitted_filter(
+    enkf, filter_model_inputs, enkf_init_state, key=enkf_filter_key
+)  # warm up
 jax.block_until_ready(enkf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(enkf, model_inputs, key=enkf_key))
+    jax.block_until_ready(
+        jitted_filter(enkf, filter_model_inputs, enkf_init_state, key=enkf_filter_key)
+    )
     _times.append(time.perf_counter() - _t0)
 enkf_time = float(jnp.median(jnp.array(_times)))
 
@@ -225,13 +232,18 @@ pf = particle_filter.build_filter(
     resampling_fn=adaptive_systematic,
 )
 
-key, pf_key = random.split(key)
-pf_states = jitted_filter(pf, model_inputs, key=pf_key)  # warm up
+key, pf_init_key, pf_filter_key = random.split(key, 3)
+pf_init_state = pf.init_prepare(model_inputs[0], key=pf_init_key)
+pf_states = jitted_filter(
+    pf, filter_model_inputs, pf_init_state, key=pf_filter_key
+)  # warm up
 jax.block_until_ready(pf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(pf, model_inputs, key=pf_key))
+    jax.block_until_ready(
+        jitted_filter(pf, filter_model_inputs, pf_init_state, key=pf_filter_key)
+    )
     _times.append(time.perf_counter() - _t0)
 pf_time = float(jnp.median(jnp.array(_times)))
 

docs/examples/kalman_tracking.md

@@ -142,7 +142,8 @@ function with our constructed filter object and model inputs:
 
 ```{.python #kalman-run-filter}
 # Run the filter
-filtered_states = filter(filter_obj, model_inputs, parallel=True)
+init_state = filter_obj.init_prepare(model_inputs[0])
+filtered_states = filter(filter_obj, model_inputs[1:], init_state, parallel=True)
 
 # Extract results
 means = filtered_states.mean  # Posterior state means
@@ -165,7 +166,9 @@ log_normalizing_constant = filtered_states.log_normalizing_constant  # Log margi
 
     ```{.python #kalman-jit-example}
     jitted_filter = jit(filter, static_argnames=['filter_obj', 'parallel'])
-    filtered_states = jitted_filter(filter_obj, model_inputs, parallel=True)
+    filtered_states = jitted_filter(
+        filter_obj, model_inputs[1:], init_state, parallel=True
+    )
     ```
 
 ### Visualize the results
@@ -213,7 +216,10 @@ car's position.
     )
 
     # Run filtering - cuthbert handles NaNs automatically
-    filtered_states_missing = filter(filter_obj_missing, model_inputs_missing)
+    init_state_missing = filter_obj_missing.init_prepare(model_inputs_missing[0])
+    filtered_states_missing = filter(
+        filter_obj_missing, model_inputs_missing[1:], init_state_missing
+    )
 
     print("Successfully handled missing observations!")
     print("Missing observation period: steps 20-30 during the car's turn")