TeX 55.1%
Python 44.2%
1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Heteroskedastic local-level / local-linear-trend state-space models for cell deviations,3with robust (Huber-weighted innovation) filtering and empirical-Bayes pooling of innovation variances.45Observation : ȳ_t = d_t + ε_t, ε_t ~ N(0, σ²/n_t) (n_t = 0 → no observation)6Local level : d_t = d_{t-1} + η_t, η_t ~ N(0, τ²)7Local trend : d_t = d_{t-1} + b_{t-1} + η_t ; b_t = b_{t-1} + ζ_t, ζ_t ~ N(0, τ_b²)89Outputs: filtered (one-sided, real-time) and smoothed (two-sided) paths and variances, the Kalman gain10(share of the period's estimate coming from local data) and the marginal log-likelihood.11"""12from __future__ import annotations1314from dataclasses import dataclass1516import numpy as np1718LOG_TAU2_GRID = np.linspace(np.log(1e-8), np.log(2e-2), 48)19HUBER_K = 1.345202122@dataclass23class SSFit:24 filtered: np.ndarray25 filtered_var: np.ndarray26 smoothed: np.ndarray27 smoothed_var: np.ndarray28 gain: np.ndarray29 tau2: float30 tau2_slope: float31 sigma2: float32 loglik: float33 model: str34 innov_weights: np.ndarray353637def _filter_ll(y, r, tau2, p0):38 T = len(y)39 a = np.empty(T); p = np.empty(T); ap = np.empty(T); pp = np.empty(T); gain = np.zeros(T); v_std = np.zeros(T)40 at, pt, ll = 0.0, p0, 0.041 for t in range(T):42 at_p, pt_p = at, pt + tau243 ap[t], pp[t] = at_p, pt_p44 if np.isfinite(r[t]):45 f = pt_p + r[t]; v = y[t] - at_p; k = pt_p / f46 at = at_p + k * v; pt = pt_p * (1 - k); gain[t] = k; v_std[t] = v / np.sqrt(f)47 ll += -0.5 * (np.log(2 * np.pi * f) + v * v / f)48 else:49 at, pt = at_p, pt_p50 a[t], p[t] = at, pt51 return a, p, ap, pp, gain, ll, v_std525354def _smooth_ll(a, p, ap, pp):55 T = len(a); s = a.copy(); sv = p.copy()56 for t in range(T - 2, -1, -1):57 c = p[t] / pp[t + 1] if pp[t + 1] > 0 else 0.058 s[t] = a[t] + c * (s[t + 1] - ap[t + 1]); sv[t] = p[t] + c * c * (sv[t + 1] - pp[t + 1])59 return s, np.maximum(sv, 0)606162def _filter_llt(y, r, tau2, tau2b, p0):63 """Local linear trend, 2-state Kalman filter + RTS smoother (returns level quantities)."""64 T = len(y)65 Tm = np.array([[1.0, 1.0], [0.0, 1.0]]); Q = np.diag([tau2, tau2b]); Z = np.array([1.0, 0.0])66 a = np.zeros((T, 2)); P = np.zeros((T, 2, 2)); ap = np.zeros((T, 2)); Pp = np.zeros((T, 2, 2)); gain = np.zeros(T); v_std = np.zeros(T)67 at = np.zeros(2); Pt = np.diag([p0, p0 / 10]); ll = 0.068 for t in range(T):69 at_p = Tm @ at; Pt_p = Tm @ Pt @ Tm.T + Q70 ap[t], Pp[t] = at_p, Pt_p71 if np.isfinite(r[t]):72 f = Z @ Pt_p @ Z + r[t]; v = y[t] - Z @ at_p; K = Pt_p @ Z / f73 at = at_p + K * v; Pt = Pt_p - np.outer(K, Z @ Pt_p); gain[t] = K[0]; v_std[t] = v / np.sqrt(f)74 ll += -0.5 * (np.log(2 * np.pi * f) + v * v / f)75 else:76 at, Pt = at_p, Pt_p77 a[t], P[t] = at, Pt78 s = a.copy(); S = P.copy()79 for t in range(T - 2, -1, -1):80 try:81 C = P[t] @ Tm.T @ np.linalg.inv(Pp[t + 1])82 except np.linalg.LinAlgError:83 C = np.zeros((2, 2))84 s[t] = a[t] + C @ (s[t + 1] - ap[t + 1]); S[t] = P[t] + C @ (S[t + 1] - Pp[t + 1]) @ C.T85 return a[:, 0], P[:, 0, 0], s[:, 0], np.maximum(S[:, 0, 0], 0), gain, ll, v_std868788def fit_cell(y: np.ndarray, n: np.ndarray, sigma2: float, model: str = "ll", log_tau2: float | None = None,89 prior: tuple[float, float] | None = None, robust: bool = True, p0: float = 0.25, tau2_slope_ratio: float = 0.05) -> SSFit:90 """Fit one cell. If `log_tau2` is None it is chosen on LOG_TAU2_GRID by maximising91 loglik + log N(log τ²; prior) (prior=None → maximum likelihood). Robust: two passes of Huber92 down-weighting of standardised innovations (observation variance inflated for outlying months)."""93 y = np.asarray(y, float); n = np.asarray(n, float)94 r0 = np.where(n > 0, sigma2 / np.maximum(n, 1), np.inf)95 yy = np.where(np.isfinite(r0), np.nan_to_num(y), 0.0)96 w = np.ones(len(y))9798 def run(lt2, r):99 t2 = np.exp(lt2)100 if model == "llt":101 a, p, s, sv, g, ll, vs = _filter_llt(yy, r, t2, t2 * tau2_slope_ratio, p0)102 else:103 a, p, ap, pp, g, ll, vs = _filter_ll(yy, r, t2, p0); s, sv = _smooth_ll(a, p, ap, pp)104 return a, p, s, sv, g, ll, vs105106 def choose(r):107 if log_tau2 is not None:108 return log_tau2109 lls = np.array([run(lt, r)[5] for lt in LOG_TAU2_GRID])110 if prior is not None:111 m, v = prior; lls = lls - 0.5 * (LOG_TAU2_GRID - m) ** 2 / v112 return float(LOG_TAU2_GRID[int(np.argmax(lls))])113114 r = r0.copy(); lt = choose(r)115 if robust:116 for _ in range(2):117 vs = run(lt, r)[6]118 u = np.abs(vs); w = np.where(u <= HUBER_K, 1.0, HUBER_K / np.maximum(u, 1e-9)); w[~np.isfinite(r0)] = 1.0119 r = np.where(np.isfinite(r0), r0 / w, np.inf); lt = choose(r)120 a, p, s, sv, g, ll, vs = run(lt, r)121 t2 = float(np.exp(lt))122 return SSFit(a, p, s, sv, g, t2, t2 * tau2_slope_ratio if model == "llt" else 0.0, sigma2, ll, model, w)123124125def profile_loglik(y: np.ndarray, n: np.ndarray, sigma2: float, model: str = "ll", p0: float = 0.25) -> np.ndarray:126 """Log-likelihood of a cell on LOG_TAU2_GRID (for empirical-Bayes pooling across cells)."""127 y = np.asarray(y, float); n = np.asarray(n, float)128 r = np.where(n > 0, sigma2 / np.maximum(n, 1), np.inf); yy = np.where(np.isfinite(r), np.nan_to_num(y), 0.0)129 out = np.empty(len(LOG_TAU2_GRID))130 for i, lt in enumerate(LOG_TAU2_GRID):131 t2 = np.exp(lt)132 out[i] = _filter_llt(yy, r, t2, t2 * 0.05, p0)[5] if model == "llt" else _filter_ll(yy, r, t2, p0)[5]133 return out134135136def eb_prior(profiles: list[np.ndarray], n_iter: int = 50, v_floor: float = 0.05) -> tuple[float, float]:137 """Grid EM for a normal prior on log τ² shared by a group of cells: maximise Σ_c log ∫ L_c(τ²) N(log τ²; m, v)."""138 m, v = float(np.mean([LOG_TAU2_GRID[int(np.argmax(p))] for p in profiles])), 1.0139 g = LOG_TAU2_GRID140 for _ in range(n_iter):141 post_m, post_v = [], []142 for p in profiles:143 lp = p - 0.5 * (g - m) ** 2 / v; lp -= lp.max(); w = np.exp(lp); w /= w.sum()144 mu = np.sum(w * g); post_m.append(mu); post_v.append(np.sum(w * (g - mu) ** 2))145 m_new = float(np.mean(post_m)); v_new = float(max(np.mean(post_v) + np.var(post_m), v_floor))146 if abs(m_new - m) < 1e-6 and abs(v_new - v) < 1e-6:147 break148 m, v = m_new, v_new149 return m, v150