feat: dynamic factor models and new examples

Author: Kleyt0n

docs/api/factor.rst

@@ -0,0 +1,35 @@
+Dynamic Factor Models
+=====================
+
+Dynamic Factor Models (DFMs) reduce high-dimensional multivariate time
+series to a small number of latent factors. The observation matrix
+(loading matrix) maps latent factors to observed variables.
+
+High-Level API
+--------------
+
+.. autoclass:: dynaris.models.dfm_api.DFMModel
+   :members:
+   :show-inheritance:
+
+Model Factory
+-------------
+
+.. autofunction:: dynaris.models.factor.DynamicFactorModel
+
+Estimation
+----------
+
+.. autofunction:: dynaris.estimation.dfm.fit_dfm_em
+
+.. autoclass:: dynaris.estimation.dfm.DFMResult
+   :members:
+
+Utilities
+---------
+
+.. autofunction:: dynaris.models.factor.initialize_loadings_pca
+
+.. autofunction:: dynaris.models.factor.rotate_loadings
+
+.. autofunction:: dynaris.models.factor.apply_identification_constraints

docs/api/index.rst

@@ -14,6 +14,8 @@ Complete reference for all public classes and functions in dynaris.
 +------------------+-------------------------------------------------------------+
 | :doc:`models`    | Built-in nonlinear models (stochastic vol, tracking, etc.)  |
 +------------------+-------------------------------------------------------------+
+| :doc:`factor`    | Dynamic Factor Models (DFMModel, loadings, rotation)        |
++------------------+-------------------------------------------------------------+
 | :doc:`core`      | ``StateSpaceModel``, ``GaussianState``, result containers   |
 +------------------+-------------------------------------------------------------+
 | :doc:`filters`   | Kalman, EKF, UKF, and Particle filters                     |
@@ -39,6 +41,7 @@ Complete reference for all public classes and functions in dynaris.
    dlm
    components
    models
+   factor
    core
    filters
    switching

examples/bayesian_nile.py

@@ -0,0 +1,173 @@
+"""Bayesian estimation — airline passengers with trend + seasonality.
+
+Demonstrates full Bayesian inference for a composed DLM (local linear
+trend + seasonal) using NumPyro's NUTS sampler. The airline passenger
+data is log-transformed to convert multiplicative seasonality to additive.
+
+This example:
+1. Fits a trend + seasonal model to log(airline passengers) via MLE
+2. Fits the same model via Bayesian MCMC with a normal prior
+3. Compares posterior mean with MLE point estimates
+4. Generates posterior predictive forecasts with credible intervals
+5. Runs prior predictive checks
+"""
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+from dynaris import LocalLinearTrend, Seasonal
+from dynaris.core.state_space import StateSpaceModel
+from dynaris.datasets import load_airline
+from dynaris.estimation import fit_bayesian, fit_mle
+from dynaris.estimation.predictive import posterior_predictive_forecast, prior_predictive
+from dynaris.estimation.priors import normal_log_prior
+
+# --- Data (log-transform for additive seasonality) ---
+y = load_airline()
+obs = jnp.log(jnp.array(y.values, dtype=jnp.float32)).reshape(-1, 1)
+print(f"Airline passengers: {len(obs)} monthly observations (log-transformed)")
+
+
+# --- Model factory: trend + seasonal (period=12) ---
+# params: [log_sigma_level, log_sigma_slope, log_sigma_seasonal, log_sigma_obs]
+def model_fn(params: jax.Array) -> StateSpaceModel:
+    trend = LocalLinearTrend(
+        sigma_level=jnp.exp(params[0]),
+        sigma_slope=jnp.exp(params[1]),
+        sigma_obs=0.0,
+    )
+    seasonal = Seasonal(
+        period=12,
+        sigma_seasonal=jnp.exp(params[2]),
+        sigma_obs=jnp.exp(params[3]),
+    )
+    return trend + seasonal
+
+
+# --- MLE fit (baseline) ---
+print("\nFitting MLE...")
+init_params = jnp.array([0.0, -2.0, -2.0, -1.0])
+mle_result = fit_mle(model_fn, obs, init_params)
+print(f"  MLE params: {np.round(np.asarray(mle_result.params), 3)}")
+print(f"  MLE log-likelihood: {mle_result.log_likelihood:.2f}")
+
+# --- Bayesian fit ---
+print("\nRunning Bayesian MCMC (NUTS)...")
+prior = normal_log_prior(loc=0.0, scale=5.0)
+bayes_result = fit_bayesian(
+    model_fn,
+    obs,
+    init_params=mle_result.params,
+    log_prior_fn=prior,
+    n_warmup=500,
+    n_samples=1000,
+    key=jax.random.PRNGKey(42),
+    param_names=(
+        "log_sigma_level",
+        "log_sigma_slope",
+        "log_sigma_seasonal",
+        "log_sigma_obs",
+    ),
+)
+
+samples = bayes_result.samples
+names = bayes_result.param_names or ()
+print("  Posterior means:")
+for j, name in enumerate(names):
+    post_mean = float(jnp.mean(samples[:, j]))
+    mle_val = float(mle_result.params[j])
+    print(f"    {name}: {post_mean:.3f} (MLE: {mle_val:.3f})")
+mean_ll = float(jnp.mean(bayes_result.log_likelihood_samples))
+print(f"  Posterior mean log-likelihood: {mean_ll:.2f}")
+if bayes_result.info:
+    print(f"  Divergences: {bayes_result.info.get('n_divergences', 'N/A')}")
+
+# --- Posterior predictive forecast (24 months ahead) ---
+print("\nGenerating posterior predictive forecast (24 months)...")
+fc = posterior_predictive_forecast(
+    bayes_result,
+    model_fn,
+    obs,
+    steps=24,
+    n_posterior_samples=200,
+)
+
+# --- Prior predictive ---
+print("Running prior predictive check...")
+
+
+def prior_sample_fn(key: jax.Array) -> jax.Array:
+    return jax.random.normal(key, (4,)) * 1.0 + mle_result.params
+
+
+prior_sims = prior_predictive(model_fn, prior_sample_fn, n_steps=144, n_samples=50)
+
+# --- Plots ---
+fig, axes = plt.subplots(2, 2, figsize=(14, 9))
+
+# Posterior parameter distributions
+ax = axes[0, 0]
+for j, name in enumerate(names):
+    ax.hist(np.asarray(samples[:, j]), bins=25, alpha=0.6, label=name)
+    ax.axvline(float(mle_result.params[j]), color=f"C{j}", linestyle="--", linewidth=1)
+ax.set_xlabel("Parameter value (log scale)")
+ax.set_ylabel("Count")
+ax.set_title("Posterior Distributions (dashed = MLE)")
+ax.legend(frameon=False, fontsize=7)
+
+# Posterior predictive forecast (back-transformed to passenger scale)
+ax = axes[0, 1]
+n_obs = len(obs)
+time_obs = np.arange(n_obs)
+time_fc = np.arange(n_obs, n_obs + 24)
+ax.plot(
+    time_obs[-48:],
+    np.exp(np.asarray(obs[-48:, 0])),
+    "k.-",
+    markersize=2,
+    label="Observed",
+)
+ax.plot(
+    time_fc,
+    np.exp(np.asarray(fc["mean"][:, 0])),
+    "C0-",
+    linewidth=1.5,
+    label="Forecast mean",
+)
+ax.fill_between(
+    time_fc,
+    np.exp(np.asarray(fc["lower"][:, 0])),
+    np.exp(np.asarray(fc["upper"][:, 0])),
+    alpha=0.3,
+    color="C0",
+    label="95% credible interval",
+)
+ax.set_xlabel("Month index")
+ax.set_ylabel("Passengers")
+ax.set_title("Posterior Predictive Forecast (24 months)")
+ax.legend(frameon=False, fontsize=8)
+
+# Prior predictive (log scale — simulations start from zero)
+ax = axes[1, 0]
+for i in range(min(20, prior_sims.shape[0])):
+    ax.plot(np.asarray(prior_sims[i, :, 0]), alpha=0.4, linewidth=0.5, color="gray")
+ax.axhline(float(jnp.mean(obs)), color="k", linestyle="--", linewidth=0.8, label="Data mean")
+ax.set_xlabel("Time step")
+ax.set_ylabel("log(passengers)")
+ax.set_title("Prior Predictive Samples (log scale)")
+ax.legend(frameon=False, fontsize=8)
+
+# Trace plot
+ax = axes[1, 1]
+for j, name in enumerate(names):
+    ax.plot(np.asarray(samples[:, j]), alpha=0.6, linewidth=0.4, label=name)
+ax.set_xlabel("Sample index")
+ax.set_ylabel("Value")
+ax.set_title("MCMC Traces")
+ax.legend(frameon=False, fontsize=7)
+
+fig.suptitle("Bayesian Estimation — Airline Passengers (Trend + Seasonality)")
+fig.tight_layout()
+plt.show()

examples/bearings_tracking.py

@@ -0,0 +1,89 @@
+"""Bearings-only tracking — EKF on a nonlinear observation model.
+
+A target moves with near-constant velocity in 2D, observed only via
+bearing (angle) from a fixed sensor. The observation function is
+nonlinear (atan2), making this a classic EKF/UKF benchmark.
+
+This example:
+1. Simulates a target on a curved trajectory
+2. Generates noisy bearing measurements from a sensor at the origin
+3. Tracks the 4D state (x, vx, y, vy) using the EKF
+4. Plots the reconstructed trajectory vs ground truth
+"""
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+from dynaris.core.types import GaussianState
+from dynaris.filters.ekf import ekf_filter
+from dynaris.models import BearingsTracking
+
+# --- Model ---
+model = BearingsTracking(sensor_pos=(0.0, 0.0), dt=1.0, sigma_accel=0.01, sigma_bearing=0.05)
+
+# --- Simulate target trajectory (gentle curve) ---
+key = jax.random.PRNGKey(7)
+n_steps = 100
+
+true_state = jnp.array([5.0, 0.3, 5.0, 0.2])
+true_states = []
+observations = []
+for _ in range(n_steps):
+    key, k_obs = jax.random.split(key)
+    true_state = model.f(true_state)
+    bearing = model.h(true_state)
+    noisy_bearing = bearing + jax.random.normal(k_obs, (1,)) * 0.05
+    true_states.append(true_state)
+    observations.append(noisy_bearing)
+true_states = jnp.stack(true_states)
+observations = jnp.stack(observations)
+
+# --- Track with EKF ---
+init = GaussianState(
+    mean=jnp.array([5.0, 0.3, 5.0, 0.2]),
+    cov=jnp.diag(jnp.array([1.0, 0.5, 1.0, 0.5])),
+)
+result = ekf_filter(model, observations, initial_state=init)
+
+# --- Plot trajectory ---
+fig, axes = plt.subplots(1, 2, figsize=(14, 5))
+
+# 2D trajectory
+ax = axes[0]
+ax.plot(
+    np.asarray(true_states[:, 0]),
+    np.asarray(true_states[:, 2]),
+    "k-",
+    linewidth=1.5,
+    label="True path",
+)
+ax.plot(
+    np.asarray(result.filtered_states[:, 0]),
+    np.asarray(result.filtered_states[:, 2]),
+    "o-",
+    markersize=2,
+    linewidth=0.8,
+    label="EKF estimate",
+)
+ax.plot(0, 0, "r^", markersize=10, label="Sensor")
+ax.set_xlabel("x")
+ax.set_ylabel("y")
+ax.set_title("2D Target Trajectory")
+ax.legend(frameon=False)
+ax.set_aspect("equal")
+
+# Bearing observations
+ax = axes[1]
+ax.plot(np.asarray(observations[:, 0]), "k.", markersize=2, alpha=0.5, label="Observed bearings")
+true_bearings = np.asarray(jax.vmap(model.h)(true_states)[:, 0])
+ax.plot(true_bearings, "r-", linewidth=1, label="True bearing")
+ax.set_xlabel("Time step")
+ax.set_ylabel("Bearing (rad)")
+ax.set_title("Bearing Observations")
+ax.legend(frameon=False)
+
+fig.suptitle("Bearings-Only Tracking with EKF")
+fig.tight_layout()
+plt.show()

examples/lorenz_tracking.py

@@ -0,0 +1,78 @@
+"""Lorenz attractor tracking — nonlinear filter comparison.
+
+Demonstrates EKF, UKF, and Particle Filter on the chaotic Lorenz system.
+The Lorenz attractor is a classic benchmark for nonlinear filtering,
+with three coupled differential equations producing deterministic chaos.
+
+This example:
+1. Simulates a trajectory from the Lorenz system
+2. Generates noisy partial observations (x and y only)
+3. Tracks the full 3D state using EKF, UKF, and Particle Filter
+4. Compares filter accuracy via correlation with the true state
+"""
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+from dynaris.core.types import GaussianState
+from dynaris.filters.ekf import ekf_filter
+from dynaris.filters.particle import particle_filter
+from dynaris.filters.ukf import ukf_filter
+from dynaris.models import LorenzAttractor
+
+# --- Model ---
+model = LorenzAttractor(dt=0.01, process_noise=0.5, obs_noise=2.0, obs_dims=2)
+
+# --- Simulate true trajectory ---
+key = jax.random.PRNGKey(42)
+k1, k2, k_pf = jax.random.split(key, 3)
+n_steps = 500
+
+state = jnp.array([1.0, 1.0, 1.0])
+true_states = []
+observations = []
+for _ in range(n_steps):
+    k1, k_state, k_obs = jax.random.split(k1, 3)
+    state = model.f(state) + jax.random.normal(k_state, (3,)) * 0.5
+    obs = model.h(state) + jax.random.normal(k_obs, (2,)) * 2.0
+    true_states.append(state)
+    observations.append(obs)
+true_states = jnp.stack(true_states)
+observations = jnp.stack(observations)
+
+# --- Initial state ---
+init = GaussianState(mean=jnp.array([1.0, 1.0, 1.0]), cov=jnp.eye(3) * 10.0)
+
+# --- Run filters ---
+print("Running EKF...")
+ekf_result = ekf_filter(model, observations, initial_state=init)
+
+print("Running UKF...")
+ukf_result = ukf_filter(model, observations, initial_state=init, alpha=1.0)
+
+print("Running Particle Filter (1000 particles)...")
+pf_result = particle_filter(model, observations, n_particles=1000, key=k_pf, initial_state=init)
+
+# --- Compare accuracy ---
+print("\nFilter accuracy (correlation with true x-component):")
+for name, result in [("EKF", ekf_result), ("UKF", ukf_result), ("PF", pf_result)]:
+    corr = float(jnp.corrcoef(jnp.stack([result.filtered_states[:, 0], true_states[:, 0]]))[0, 1])
+    print(f"  {name}: r = {corr:.4f}")
+
+# --- Plot ---
+fig, axes = plt.subplots(3, 1, figsize=(12, 8), sharex=True)
+labels = ["x", "y", "z"]
+for i, ax in enumerate(axes):
+    ax.plot(np.asarray(true_states[:, i]), "k-", alpha=0.4, linewidth=0.8, label="True")
+    ax.plot(np.asarray(ekf_result.filtered_states[:, i]), label="EKF", linewidth=1.0)
+    ax.plot(np.asarray(ukf_result.filtered_states[:, i]), label="UKF", linewidth=1.0)
+    ax.plot(np.asarray(pf_result.filtered_states[:, i]), label="PF", linewidth=1.0)
+    ax.set_ylabel(labels[i])
+    if i == 0:
+        ax.legend(frameon=False, ncol=4)
+axes[-1].set_xlabel("Time step")
+fig.suptitle("Lorenz Attractor Tracking — EKF vs UKF vs Particle Filter")
+fig.tight_layout()
+plt.show()