feat: add non-linear and non-gaussian models

Author: Kleyt0n

docs/api/index.rst

@@ -6,13 +6,17 @@ Complete reference for all public classes and functions in dynaris.
 +------------------+-------------------------------------------------------------+
 | Module           | Description                                                 |
 +==================+=============================================================+
+| :doc:`ssm`       | Unified ``SSM`` class with automatic filter selection       |
++------------------+-------------------------------------------------------------+
 | :doc:`dlm`       | High-level ``DLM`` class (fit, smooth, forecast, plot)      |
 +------------------+-------------------------------------------------------------+
 | :doc:`components` | Six composable building blocks (``LocalLevel``, etc.)      |
 +------------------+-------------------------------------------------------------+
+| :doc:`models`    | Built-in nonlinear models (stochastic vol, tracking, etc.)  |
++------------------+-------------------------------------------------------------+
 | :doc:`core`      | ``StateSpaceModel``, ``GaussianState``, result containers   |
 +------------------+-------------------------------------------------------------+
-| :doc:`filters`   | Kalman filter (predict, update, full forward pass)          |
+| :doc:`filters`   | Kalman, EKF, UKF, and Particle filters                     |
 +------------------+-------------------------------------------------------------+
 | :doc:`smoothers` | Rauch-Tung-Striebel backward smoother                       |
 +------------------+-------------------------------------------------------------+
@@ -29,8 +33,10 @@ Complete reference for all public classes and functions in dynaris.
    :maxdepth: 2
    :hidden:
 
+   ssm
    dlm
    components
+   models
    core
    filters
    smoothers

docs/api/models.rst

@@ -0,0 +1,22 @@
+Built-in Nonlinear Models
+=========================
+
+Factory functions that return :class:`~dynaris.core.nonlinear.NonlinearSSM`
+instances for well-known nonlinear state-space models.
+
+Stochastic Volatility
+---------------------
+
+.. autofunction:: dynaris.models.StochasticVolatility
+
+.. autofunction:: dynaris.models.transform_returns
+
+Bearings-Only Tracking
+----------------------
+
+.. autofunction:: dynaris.models.BearingsTracking
+
+Lorenz Attractor
+----------------
+
+.. autofunction:: dynaris.models.LorenzAttractor

docs/api/ssm.rst

@@ -0,0 +1,9 @@
+SSM (Unified Interface)
+=======================
+
+A unified high-level interface for both linear and nonlinear state-space
+models, with automatic filter selection.
+
+.. autoclass:: dynaris.core.ssm.SSM
+   :members:
+   :show-inheritance:

src/dynaris/__init__.py

@@ -1,6 +1,7 @@
 """dynaris: A JAX-powered Python library for Dynamic Linear Models (DLMs)."""
 
 from dynaris.core import (
+    SSM,
     FilterProtocol,
     FilterResult,
     GaussianState,
@@ -28,13 +29,21 @@
     particle_filter,
     ukf_filter,
 )
+from dynaris.models import (
+    BearingsTracking,
+    LorenzAttractor,
+    StochasticVolatility,
+    transform_returns,
+)
 from dynaris.smoothers import RTSSmoother, rts_smooth
 
 __version__ = "0.1.0"
 
 __all__ = [
     "DLM",
+    "SSM",
     "Autoregressive",
+    "BearingsTracking",
     "Cycle",
     "ExtendedKalmanFilter",
     "FilterProtocol",
@@ -43,6 +52,7 @@
     "KalmanFilter",
     "LocalLevel",
     "LocalLinearTrend",
+    "LorenzAttractor",
     "NonlinearSSM",
     "ParticleFilter",
     "RTSSmoother",
@@ -51,11 +61,13 @@
     "SmootherProtocol",
     "SmootherResult",
     "StateSpaceModel",
+    "StochasticVolatility",
     "UnscentedKalmanFilter",
     "__version__",
     "ekf_filter",
     "kalman_filter",
     "particle_filter",
     "rts_smooth",
+    "transform_returns",
     "ukf_filter",
 ]

src/dynaris/core/__init__.py

@@ -3,10 +3,12 @@
 from dynaris.core.nonlinear import NonlinearSSM
 from dynaris.core.protocols import FilterProtocol, SmootherProtocol
 from dynaris.core.results import FilterResult, SmootherResult
+from dynaris.core.ssm import SSM
 from dynaris.core.state_space import StateSpaceModel
 from dynaris.core.types import GaussianState
 
 __all__ = [
+    "SSM",
     "FilterProtocol",
     "FilterResult",
     "GaussianState",