feat: add bayesian estimation

Author: Kleyt0n

pyproject.toml

@@ -27,6 +27,12 @@ dependencies = [
     "scipy>=1.12",
 ]
 
+[project.optional-dependencies]
+bayesian = [
+    "numpyro>=0.15",
+    "arviz>=0.18",
+]
+
 [project.urls]
 Homepage = "https://github.com/quant-sci/dynaris"
 Repository = "https://github.com/quant-sci/dynaris"
@@ -88,6 +94,8 @@ module = [
     "pandas",
     "pandas.*",
     "scipy.*",
+    "numpyro.*",
+    "arviz.*",
 ]
 ignore_missing_imports = true
 

src/dynaris/__init__.py

@@ -22,6 +22,7 @@
     Regression,
     Seasonal,
 )
+from dynaris.estimation.bayesian import BayesianResult, fit_bayesian
 from dynaris.filters import (
     ExtendedKalmanFilter,
     HamiltonFilter,
@@ -48,6 +49,7 @@
     "DLM",
     "SSM",
     "Autoregressive",
+    "BayesianResult",
     "BearingsTracking",
     "Cycle",
     "ExtendedKalmanFilter",
@@ -75,6 +77,7 @@
     "UnscentedKalmanFilter",
     "__version__",
     "ekf_filter",
+    "fit_bayesian",
     "hamilton_filter",
     "kalman_filter",
     "kim_smooth",

src/dynaris/estimation/__init__.py

@@ -1,22 +1,47 @@
-"""Parameter estimation: MLE, EM, and model diagnostics."""
+"""Parameter estimation: MLE, EM, Bayesian, and model diagnostics."""
 
+from dynaris.estimation.bayesian import BayesianResult, fit_bayesian
+from dynaris.estimation.comparison import compute_loo, compute_waic, to_arviz
 from dynaris.estimation.diagnostics import acf, ljung_box, pacf, standardized_residuals
 from dynaris.estimation.em import EMResult, fit_em
 from dynaris.estimation.mle import MLEResult, fit_mle
 from dynaris.estimation.model_selection import switching_aic, switching_bic
+from dynaris.estimation.predictive import (
+    posterior_predictive_check,
+    posterior_predictive_forecast,
+    prior_predictive,
+)
+from dynaris.estimation.priors import (
+    combine_priors,
+    half_normal_log_prior,
+    inverse_gamma_log_prior,
+    normal_log_prior,
+)
 from dynaris.estimation.transforms import inverse_softplus, softplus
 
 __all__ = [
+    "BayesianResult",
     "EMResult",
     "MLEResult",
     "acf",
+    "combine_priors",
+    "compute_loo",
+    "compute_waic",
+    "fit_bayesian",
     "fit_em",
     "fit_mle",
+    "half_normal_log_prior",
+    "inverse_gamma_log_prior",
     "inverse_softplus",
     "ljung_box",
+    "normal_log_prior",
     "pacf",
+    "posterior_predictive_check",
+    "posterior_predictive_forecast",
+    "prior_predictive",
     "softplus",
     "standardized_residuals",
     "switching_aic",
     "switching_bic",
+    "to_arviz",
 ]

src/dynaris/estimation/bayesian.py

@@ -0,0 +1,176 @@
+"""Bayesian estimation for state-space models via MCMC.
+
+Uses NumPyro's NUTS sampler to draw posterior samples of model parameters,
+with the Kalman filter log-likelihood as the data term.
+
+Requires the ``bayesian`` extra: ``pip install dynaris[bayesian]``
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from dynaris.core.results import FilterResult
+from dynaris.core.state_space import StateSpaceModel
+from dynaris.estimation.priors import LogPriorFn
+from dynaris.filters.kalman import kalman_filter
+
+ModelFactory = Any  # Callable[[Array], StateSpaceModel]
+
+
+def _require_numpyro() -> Any:
+    try:
+        import numpyro
+
+        return numpyro
+    except ImportError:
+        msg = "Bayesian estimation requires numpyro. Install with: pip install dynaris[bayesian]"
+        raise ImportError(msg) from None
+
+
+@dataclass(frozen=True)
+class BayesianResult:
+    """Result of Bayesian MCMC estimation.
+
+    Attributes:
+        samples: Posterior samples (unconstrained), shape (n_samples, n_params).
+        log_likelihood_samples: Log-likelihood at each sample, shape (n_samples,).
+        model: StateSpaceModel at the posterior mean parameters.
+        filter_result: FilterResult from the posterior-mean model.
+        param_names: Optional parameter labels.
+        info: Sampler diagnostics (divergences, acceptance rate, etc.).
+    """
+
+    samples: Array
+    log_likelihood_samples: Array
+    model: StateSpaceModel
+    filter_result: FilterResult
+    param_names: tuple[str, ...] | None = None
+    info: dict[str, Any] | None = None
+
+
+def _flat_prior(params: Array) -> Array:
+    """Flat (improper) prior: constant zero log-density."""
+    return jnp.array(0.0)
+
+
+def fit_bayesian(
+    model_fn: ModelFactory,
+    observations: Array,
+    init_params: Array,
+    log_prior_fn: LogPriorFn | None = None,
+    n_warmup: int = 500,
+    n_samples: int = 1000,
+    key: Array | None = None,
+    param_names: tuple[str, ...] | None = None,
+) -> BayesianResult:
+    """Fit a state-space model via Bayesian MCMC (NUTS).
+
+    Uses NumPyro's NUTS sampler with automatic warmup adaptation.
+    The log-posterior is the Kalman filter log-likelihood plus the
+    log-prior.
+
+    Args:
+        model_fn: Maps unconstrained parameter vector to a
+            :class:`StateSpaceModel`. Same pattern as :func:`fit_mle`.
+        observations: Observation sequence, shape (T, obs_dim).
+        init_params: Initial (unconstrained) parameter vector.
+        log_prior_fn: Log-prior function. Defaults to flat prior.
+        n_warmup: Number of NUTS warmup steps.
+        n_samples: Number of posterior samples to draw.
+        key: JAX PRNG key. Defaults to ``PRNGKey(0)``.
+        param_names: Optional names for each parameter dimension.
+
+    Returns:
+        BayesianResult with posterior samples and fitted model.
+
+    Example::
+
+        import jax.numpy as jnp
+        from dynaris import LocalLevel
+        from dynaris.estimation import fit_bayesian
+        from dynaris.estimation.priors import inverse_gamma_log_prior
+
+        def model_fn(params):
+            return LocalLevel(
+                sigma_level=jnp.exp(params[0]),
+                sigma_obs=jnp.exp(params[1]),
+            )
+
+        result = fit_bayesian(
+            model_fn, observations, jnp.zeros(2),
+            log_prior_fn=inverse_gamma_log_prior(shape=2.0, scale=1.0),
+        )
+    """
+    numpyro = _require_numpyro()
+    from numpyro.infer import MCMC, NUTS
+
+    observations = jnp.asarray(observations)
+    init_params = jnp.asarray(init_params)
+    if key is None:
+        key = jax.random.PRNGKey(0)
+    if log_prior_fn is None:
+        log_prior_fn = _flat_prior
+
+    n_params = init_params.shape[0]
+
+    # Define the log-density for the sampler
+    @jax.jit
+    def _log_density(params: Array) -> Array:
+        model = model_fn(params)
+        fr = kalman_filter(model, observations)
+        return fr.log_likelihood + log_prior_fn(params)
+
+    # NumPyro model: sample unconstrained params, factor by log-density
+    def _numpyro_model() -> None:
+        params = numpyro.sample(
+            "params",
+            numpyro.distributions.Normal(0.0, 100.0).expand([n_params]).to_event(1),
+        )
+        log_density = _log_density(params)
+        numpyro.factor("log_density", log_density)
+
+    # Run MCMC
+    kernel = NUTS(
+        _numpyro_model, init_strategy=numpyro.infer.init_to_value(values={"params": init_params})
+    )
+    mcmc = MCMC(kernel, num_warmup=n_warmup, num_samples=n_samples, progress_bar=False)
+    mcmc.run(key)
+
+    samples = mcmc.get_samples()["params"]  # (n_samples, n_params)
+
+    # Compute per-sample log-likelihoods
+    @jax.jit
+    def _compute_ll(params: Array) -> Array:
+        model = model_fn(params)
+        return kalman_filter(model, observations).log_likelihood
+
+    log_lls = jax.vmap(_compute_ll)(samples)
+
+    # Build posterior-mean model
+    mean_params = jnp.mean(samples, axis=0)
+    mean_model = model_fn(mean_params)
+    mean_fr = kalman_filter(mean_model, observations)
+
+    # Diagnostics
+    info: dict[str, Any] = {}
+    extra_fields = mcmc.get_extra_fields()
+    if "diverging" in extra_fields:
+        info["n_divergences"] = int(np.sum(np.asarray(extra_fields["diverging"])))
+    if "accept_prob" in extra_fields:
+        info["mean_accept_prob"] = float(np.mean(np.asarray(extra_fields["accept_prob"])))
+
+    return BayesianResult(
+        samples=samples,
+        log_likelihood_samples=log_lls,
+        model=mean_model,
+        filter_result=mean_fr,
+        param_names=param_names,
+        info=info,
+    )