feat: add switching & regime models

Author: Kleyt0n

docs/api/index.rst

@@ -18,9 +18,11 @@ Complete reference for all public classes and functions in dynaris.
 +------------------+-------------------------------------------------------------+
 | :doc:`filters`   | Kalman, EKF, UKF, and Particle filters                     |
 +------------------+-------------------------------------------------------------+
+| :doc:`switching` | Markov-switching models, Hamilton filter, Kim smoother       |
++------------------+-------------------------------------------------------------+
 | :doc:`smoothers` | Rauch-Tung-Striebel backward smoother                       |
 +------------------+-------------------------------------------------------------+
-| :doc:`estimation` | MLE, EM algorithm, diagnostics, transforms                 |
+| :doc:`estimation` | MLE, EM algorithm, diagnostics, model selection             |
 +------------------+-------------------------------------------------------------+
 | :doc:`forecast`  | Multi-step forecasting and batch processing                 |
 +------------------+-------------------------------------------------------------+
@@ -39,6 +41,7 @@ Complete reference for all public classes and functions in dynaris.
    models
    core
    filters
+   switching
    smoothers
    estimation
    forecast

docs/api/switching.rst

@@ -0,0 +1,44 @@
+Switching & Regime Models
+=========================
+
+Markov-switching state-space models with K discrete regimes, each with
+its own linear-Gaussian dynamics. The Hamilton filter tracks filtered
+regime probabilities, and the Kim smoother provides smoothed estimates.
+
+Model
+-----
+
+.. autoclass:: dynaris.core.switching.MarkovSwitchingSSM
+   :members:
+   :show-inheritance:
+
+Hamilton Filter
+---------------
+
+Forward filtering for Markov-switching models. Runs K parallel Kalman
+filters with Kim's moment-matching collapse approximation.
+
+.. autoclass:: dynaris.filters.hamilton.HamiltonFilter
+   :members:
+   :show-inheritance:
+
+.. autofunction:: dynaris.filters.hamilton.hamilton_filter
+
+Kim Smoother
+------------
+
+Backward smoothing for Markov-switching models. Produces smoothed
+state estimates and regime probabilities.
+
+.. autoclass:: dynaris.smoothers.kim.KimSmoother
+   :members:
+   :show-inheritance:
+
+.. autofunction:: dynaris.smoothers.kim.kim_smooth
+
+Model Selection
+---------------
+
+.. autofunction:: dynaris.estimation.model_selection.switching_aic
+
+.. autofunction:: dynaris.estimation.model_selection.switching_bic

src/dynaris/__init__.py

@@ -5,10 +5,13 @@
     FilterProtocol,
     FilterResult,
     GaussianState,
+    MarkovSwitchingSSM,
     NonlinearSSM,
     SmootherProtocol,
     SmootherResult,
     StateSpaceModel,
+    SwitchingFilterResult,
+    SwitchingSmootherResult,
 )
 from dynaris.dlm import (
     DLM,
@@ -21,10 +24,12 @@
 )
 from dynaris.filters import (
     ExtendedKalmanFilter,
+    HamiltonFilter,
     KalmanFilter,
     ParticleFilter,
     UnscentedKalmanFilter,
     ekf_filter,
+    hamilton_filter,
     kalman_filter,
     particle_filter,
     ukf_filter,
@@ -35,7 +40,7 @@
     StochasticVolatility,
     transform_returns,
 )
-from dynaris.smoothers import RTSSmoother, rts_smooth
+from dynaris.smoothers import KimSmoother, RTSSmoother, kim_smooth, rts_smooth
 
 __version__ = "0.1.0"
 
@@ -49,10 +54,13 @@
     "FilterProtocol",
     "FilterResult",
     "GaussianState",
+    "HamiltonFilter",
     "KalmanFilter",
+    "KimSmoother",
     "LocalLevel",
     "LocalLinearTrend",
     "LorenzAttractor",
+    "MarkovSwitchingSSM",
     "NonlinearSSM",
     "ParticleFilter",
     "RTSSmoother",
@@ -62,10 +70,14 @@
     "SmootherResult",
     "StateSpaceModel",
     "StochasticVolatility",
+    "SwitchingFilterResult",
+    "SwitchingSmootherResult",
     "UnscentedKalmanFilter",
     "__version__",
     "ekf_filter",
+    "hamilton_filter",
     "kalman_filter",
+    "kim_smooth",
     "particle_filter",
     "rts_smooth",
     "transform_returns",

src/dynaris/core/__init__.py

@@ -2,18 +2,27 @@
 
 from dynaris.core.nonlinear import NonlinearSSM
 from dynaris.core.protocols import FilterProtocol, SmootherProtocol
-from dynaris.core.results import FilterResult, SmootherResult
+from dynaris.core.results import (
+    FilterResult,
+    SmootherResult,
+    SwitchingFilterResult,
+    SwitchingSmootherResult,
+)
 from dynaris.core.ssm import SSM
 from dynaris.core.state_space import StateSpaceModel
+from dynaris.core.switching import MarkovSwitchingSSM
 from dynaris.core.types import GaussianState
 
 __all__ = [
     "SSM",
     "FilterProtocol",
     "FilterResult",
     "GaussianState",
+    "MarkovSwitchingSSM",
     "NonlinearSSM",
     "SmootherProtocol",
     "SmootherResult",
     "StateSpaceModel",
+    "SwitchingFilterResult",
+    "SwitchingSmootherResult",
 ]

src/dynaris/core/results.py

@@ -58,3 +58,69 @@ class SmootherResult(NamedTuple):
     predicted_covariances: Array  # (T, n, n)
     log_likelihood: Array  # ()
     observations: Array  # (T, m)
+
+
+class SwitchingFilterResult(NamedTuple):
+    """Output of a Hamilton forward filtering pass for switching models.
+
+    Contains mixture-collapsed state estimates and per-regime quantities.
+
+    Attributes:
+        filtered_states: Mixture-collapsed filtered means, shape (T, n).
+        filtered_covariances: Mixture-collapsed filtered covs, shape (T, n, n).
+        predicted_states: Mixture-collapsed predicted means, shape (T, n).
+        predicted_covariances: Mixture-collapsed predicted covs, shape (T, n, n).
+        log_likelihood: Total log-likelihood scalar, shape ().
+        observations: Input observations, shape (T, m).
+        regime_filtered_probs: Filtered regime probabilities, shape (T, K).
+        regime_predicted_probs: Predicted regime probabilities, shape (T, K).
+        regime_filtered_states: Per-regime filtered means, shape (T, K, n).
+        regime_filtered_covs: Per-regime filtered covs, shape (T, K, n, n).
+        regime_predicted_states: Per-regime predicted means, shape (T, K, n).
+        regime_predicted_covs: Per-regime predicted covs, shape (T, K, n, n).
+    """
+
+    filtered_states: Array  # (T, n)
+    filtered_covariances: Array  # (T, n, n)
+    predicted_states: Array  # (T, n)
+    predicted_covariances: Array  # (T, n, n)
+    log_likelihood: Array  # ()
+    observations: Array  # (T, m)
+    regime_filtered_probs: Array  # (T, K)
+    regime_predicted_probs: Array  # (T, K)
+    regime_filtered_states: Array  # (T, K, n)
+    regime_filtered_covs: Array  # (T, K, n, n)
+    regime_predicted_states: Array  # (T, K, n)
+    regime_predicted_covs: Array  # (T, K, n, n)
+
+
+class SwitchingSmootherResult(NamedTuple):
+    """Output of a Kim backward smoothing pass for switching models.
+
+    Attributes:
+        smoothed_states: Mixture-collapsed smoothed means, shape (T, n).
+        smoothed_covariances: Mixture-collapsed smoothed covs, shape (T, n, n).
+        filtered_states: Mixture-collapsed filtered means, shape (T, n).
+        filtered_covariances: Mixture-collapsed filtered covs, shape (T, n, n).
+        predicted_states: Mixture-collapsed predicted means, shape (T, n).
+        predicted_covariances: Mixture-collapsed predicted covs, shape (T, n, n).
+        log_likelihood: Total log-likelihood scalar, shape ().
+        observations: Input observations, shape (T, m).
+        regime_smoothed_probs: Smoothed regime probabilities, shape (T, K).
+        regime_filtered_probs: Filtered regime probabilities, shape (T, K).
+        regime_smoothed_states: Per-regime smoothed means, shape (T, K, n).
+        regime_smoothed_covs: Per-regime smoothed covs, shape (T, K, n, n).
+    """
+
+    smoothed_states: Array  # (T, n)
+    smoothed_covariances: Array  # (T, n, n)
+    filtered_states: Array  # (T, n)
+    filtered_covariances: Array  # (T, n, n)
+    predicted_states: Array  # (T, n)
+    predicted_covariances: Array  # (T, n, n)
+    log_likelihood: Array  # ()
+    observations: Array  # (T, m)
+    regime_smoothed_probs: Array  # (T, K)
+    regime_filtered_probs: Array  # (T, K)
+    regime_smoothed_states: Array  # (T, K, n)
+    regime_smoothed_covs: Array  # (T, K, n, n)

src/dynaris/core/ssm.py

@@ -28,17 +28,20 @@
 from jax import Array
 
 from dynaris.core.nonlinear import NonlinearSSM
-from dynaris.core.results import FilterResult
+from dynaris.core.results import FilterResult, SwitchingFilterResult
 from dynaris.core.state_space import StateSpaceModel
+from dynaris.core.switching import MarkovSwitchingSSM
 from dynaris.core.types import GaussianState
 from dynaris.filters.ekf import ekf_filter
+from dynaris.filters.hamilton import hamilton_filter
 from dynaris.filters.kalman import kalman_filter
 from dynaris.filters.particle import particle_filter
 from dynaris.filters.ukf import ukf_filter
 
-_VALID_FILTERS = {"auto", "kalman", "ekf", "ukf", "particle"}
+_VALID_FILTERS = {"auto", "kalman", "ekf", "ukf", "particle", "hamilton"}
 _LINEAR_FILTERS = {"kalman"}
 _NONLINEAR_FILTERS = {"ekf", "ukf", "particle"}
+_SWITCHING_FILTERS = {"hamilton"}
 
 
 def _to_jax_2d(y: Any) -> tuple[Array, pd.DatetimeIndex | None]:
@@ -92,14 +95,17 @@ class SSM:
 
     def __init__(
         self,
-        model: StateSpaceModel | NonlinearSSM,
+        model: StateSpaceModel | NonlinearSSM | MarkovSwitchingSSM,
         filter: str = "auto",
         *,
         key: Array | None = None,
         **filter_kwargs: Any,
     ) -> None:
-        if not isinstance(model, (StateSpaceModel, NonlinearSSM)):
-            msg = f"model must be a StateSpaceModel or NonlinearSSM, got {type(model).__name__}"
+        if not isinstance(model, (StateSpaceModel, NonlinearSSM, MarkovSwitchingSSM)):
+            msg = (
+                f"model must be a StateSpaceModel, NonlinearSSM, or "
+                f"MarkovSwitchingSSM, got {type(model).__name__}"
+            )
             raise TypeError(msg)
 
         if filter not in _VALID_FILTERS:
@@ -108,29 +114,41 @@ def __init__(
 
         # Resolve auto-selection
         is_linear = isinstance(model, StateSpaceModel)
-        filter_name = ("kalman" if is_linear else "ukf") if filter == "auto" else filter
+        is_switching = isinstance(model, MarkovSwitchingSSM)
+        if filter == "auto":
+            if is_switching:
+                filter_name = "hamilton"
+            elif is_linear:
+                filter_name = "kalman"
+            else:
+                filter_name = "ukf"
+        else:
+            filter_name = filter
 
         # Validate filter/model compatibility
         if filter_name in _LINEAR_FILTERS and not is_linear:
-            msg = f"Filter {filter_name!r} requires a StateSpaceModel, got NonlinearSSM."
+            msg = f"Filter {filter_name!r} requires a StateSpaceModel."
             raise ValueError(msg)
-        if filter_name in _NONLINEAR_FILTERS and is_linear:
-            msg = f"Filter {filter_name!r} requires a NonlinearSSM, got StateSpaceModel."
+        if filter_name in _NONLINEAR_FILTERS and (is_linear or is_switching):
+            msg = f"Filter {filter_name!r} requires a NonlinearSSM."
+            raise ValueError(msg)
+        if filter_name in _SWITCHING_FILTERS and not is_switching:
+            msg = f"Filter {filter_name!r} requires a MarkovSwitchingSSM."
             raise ValueError(msg)
 
         self._model = model
         self._filter_name = filter_name
         self._filter_kwargs = filter_kwargs
         self._key = key
-        self._filter_result: FilterResult | None = None
+        self._filter_result: FilterResult | SwitchingFilterResult | None = None
         self._observations: Array | None = None
         self._index: pd.DatetimeIndex | None = None
         self._is_fitted = False
 
     # --- Properties ---
 
     @property
-    def model(self) -> StateSpaceModel | NonlinearSSM:
+    def model(self) -> StateSpaceModel | NonlinearSSM | MarkovSwitchingSSM:
         """The underlying state-space model."""
         return self._model
 
@@ -140,7 +158,7 @@ def filter_name(self) -> str:
         return self._filter_name
 
     @property
-    def filter_result(self) -> FilterResult:
+    def filter_result(self) -> FilterResult | SwitchingFilterResult:
         """Filter result from the last ``fit()`` call."""
         if self._filter_result is None:
             msg = "Model not fitted. Call .fit() first."
@@ -190,6 +208,10 @@ def fit(
                 initial_state=initial_state,
                 **self._filter_kwargs,
             )
+        elif self._filter_name == "hamilton":
+            self._filter_result = hamilton_filter(
+                self._model, obs, initial_state=initial_state
+            )
 
         self._is_fitted = True
         return self