Add factorial docs

cuthbert/factorial/README.md

@@ -1,5 +1,7 @@
 # Factorial State-Space Models
 
+<!-- --8<-- [start:body] -->
+
 A factorial state-space model is a state-space model where the dynamics distribution
 factors into a product of independent distributions across factors
 
@@ -62,7 +64,9 @@ kalman_filter = cuthbert.gaussian.kalman.build_filter(
 )
 
 # Online inference
-factorial_state = kalman_filter.init_prepare(tree.map(lambda x: x[0], model_inputs))
+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)
@@ -130,4 +134,5 @@ Or directly using the `cuthbert.smoother`:
 smoother_states = cuthbert.smoother(
     kalman_smoother, filter_states_single_factor, model_inputs_single_factor
 )
-```
+```
+<!-- --8<-- [end:body] -->

cuthbert/factorial/gaussian.py

@@ -98,7 +98,6 @@ def join(local_factorial_state: KalmanState) -> KalmanState:
     Args:
         local_factorial_state: Factorial Kalman state storing means and chol_covs
             with shape (F, d) and (F, d, d) respectively.
-        factorial_inds: Indices of the factors to extract. Integer array.
 
     Returns:
         Joint local Kalman state with no factorial index dimension.

docs/api_cuthbert/factorial/discrete.md

@@ -0,0 +1,3 @@
+# Discrete Factorial Models
+
+::: cuthbert.factorial.discrete

docs/api_cuthbert/factorial/gaussian.md

@@ -0,0 +1,3 @@
+# Gaussian Factorial Models
+
+::: cuthbert.factorial.gaussian

docs/api_cuthbert/factorial/index.md

@@ -0,0 +1,13 @@
+# Factorial State-Space Models
+
+- [Discrete factorial models](discrete.md) - discrete factorial hidden Markov models.
+- [Gaussian factorial models](gaussian.md) - Gaussian factorial state-space models with factorializers.
+- [Sequential Monte Carlo factorial models](smc.md) - factorial particle filters and smoothers.
+- [Types](types.md) - core types and interfaces for factorial models.
+- [Utilities](utils.md) - helper functions for converting between serial and factorial representations.
+
+## Overview
+
+Factorial state-space models factorize the dynamics distribution into independent factors:
+
+--8<-- "cuthbert/factorial/README.md:body"

docs/api_cuthbert/factorial/smc.md

@@ -0,0 +1,3 @@
+# Sequential Monte Carlo Factorial Models
+
+::: cuthbert.factorial.smc

docs/api_cuthbert/factorial/types.md

@@ -0,0 +1,3 @@
+# Factorial Types
+
+::: cuthbert.factorial.types

docs/api_cuthbert/factorial/utils.md

@@ -0,0 +1,3 @@
+# Factorial Utilities
+
+::: cuthbert.factorial.utils

docs/api_cuthbert/index.md

@@ -4,6 +4,7 @@
 - [General smoothing API](smoothing.md)
 - [Discrete hidden Markov models](discrete.md)
 - [Ensemble Kalman Filter](enkf.md)
+- [Factorial state-space models](factorial/index.md)
 - [Gaussian filters and smoothers](gaussian/index.md)
 - [Sequential Monte Carlo](smc/index.md)
 

zensical.toml

@@ -28,6 +28,14 @@ nav = [
       "api_cuthbert/enkf/ensemble_kalman_filter.md",
       "api_cuthbert/enkf/types.md",
     ] },
+    { "Factorial state-space models" = [
+      "api_cuthbert/factorial/index.md",
+      "api_cuthbert/factorial/discrete.md",
+      "api_cuthbert/factorial/gaussian.md",
+      "api_cuthbert/factorial/smc.md",
+      "api_cuthbert/factorial/types.md",
+      "api_cuthbert/factorial/utils.md",
+    ] },
     { "Gaussian filters and smoothers" = [
       "api_cuthbert/gaussian/index.md",
       { "Exact Kalman" = "api_cuthbert/gaussian/kalman.md" },