TeX 62.3%
Python 37.6%
1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Task B models: forecast the h-month change of a cell's quality-adjusted log index from information at the origin.34Interface: fit(panel_train) ; predict(panel_test) → point forecast of `target`; predict_quantiles(panel_test, taus) optional.5`panel_*` are polars frames from data_b.build_panel restricted to one horizon h (train rows: origins whose target is observed6before the test origin — enforced by the caller).7Simple benchmarks: no-change (0), mean growth (historical mean h-change of the cell), momentum (last 3-m change scaled),8seasonal naive (same-calendar change last year), random walk with drift (12-m drift × h/12).9Time-series benchmarks on each cell's own history: AR(2) on monthly changes, local linear trend (Kalman).10Pooled panel benchmarks: ridge, LightGBM, MLP, LSTM / GRU / TCN on the 13-month history, TFT-lite (variable-selection + attention).11Proposed: DSGFN (graph attention over cells), RAREF (HMM regimes + gated experts), BHPF (hierarchical Bayesian ridge, posterior predictive),12QANSS (Kalman latent state per cell with learned nonlinear drift), QREFA (integration)."""13from __future__ import annotations1415import os1617import numpy as np18import polars as pl19from sklearn.linear_model import Ridge2021from wp23.data_b import FEATURES_B2223TAUS = np.array([0.05, 0.25, 0.5, 0.75, 0.95])242526def _X(df: pl.DataFrame) -> np.ndarray:27 return np.nan_to_num(df.select(FEATURES_B).to_numpy().astype(float))282930def _keys(df: pl.DataFrame, cols: list[str]) -> np.ndarray:31 return df.select(pl.concat_str([pl.col(c).cast(pl.Utf8).fill_null("na") for c in cols], separator="|").alias("k"))["k"].to_numpy().astype(str)323334def _H(df: pl.DataFrame) -> np.ndarray:35 H = np.array(df["hist"].to_list(), float); return H - H[:, [-1]] # history relative to the origin level363738class Base:39 name = "base"; prob = False40 def fit(self, tr): return self41 def predict(self, te): raise NotImplementedError42 def predict_quantiles(self, te, taus=TAUS):43 """Default: empirical residual quantiles from training (homoskedastic)."""44 p = self.predict(te); return p[:, None] + np.quantile(self._res, taus)[None, :]45 def _store_res(self, tr):46 self._res = tr["target"].to_numpy() - self.predict(tr); return self474849class NoChange(Base):50 name = "no_change"51 def fit(self, tr): return self._store_res(tr)52 def predict(self, te): return np.zeros(te.height)535455class MeanGrowth(Base):56 name = "mean_growth"57 def fit(self, tr):58 self.h = int(tr["h"][0]); m = tr.group_by(["class", "level", "geo_id"]).agg(pl.col("target").mean().alias("mg")); self.map = {tuple(k): v for *k, v in m.iter_rows()}59 self.glob = float(tr["target"].mean()); return self._store_res(tr)60 def predict(self, te): return np.array([self.map.get((c, l, g), self.glob) for c, l, g in te.select(["class", "level", "geo_id"]).iter_rows()])616263class Momentum(Base):64 name = "momentum"65 def fit(self, tr):66 self.h = int(tr["h"][0]); x = tr["d3"].to_numpy() / 3 * self.h; y = tr["target"].to_numpy(); self.b = float(np.sum(x * y) / max(np.sum(x * x), 1e-12)); return self._store_res(tr)67 def predict(self, te): return self.b * te["d3"].to_numpy() / 3 * self.h686970class RWDrift(Base):71 name = "rw_drift"72 def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)73 def predict(self, te): return te["d12"].to_numpy() / 12 * self.h747576class SeasonalNaive(Base):77 name = "seasonal_naive"78 def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)79 def predict(self, te):80 H = np.array(te["hist"].to_list(), float); h = self.h81 # change over the same months one year earlier, when the 13-month history covers it (h ≤ 12 → months -12..-12+h)82 if h <= 12 and H.shape[1] >= 13:83 return H[:, h] - H[:, 0]84 return np.zeros(te.height)858687class AR2(Base):88 name = "ar2_cell"89 def fit(self, tr):90 self.h = int(tr["h"][0]); D = np.diff(_H(tr), axis=1); y = tr["target"].to_numpy()91 X = np.column_stack([np.ones(len(D)), D[:, -1], D[:, -2]]); self.b = np.linalg.lstsq(X, y, rcond=None)[0]; return self._store_res(tr)92 def predict(self, te): D = np.diff(_H(te), axis=1); return np.column_stack([np.ones(len(D)), D[:, -1], D[:, -2]]) @ self.b939495class LocalLinearTrend(Base):96 name = "local_linear_trend"97 def fit(self, tr): self.h = int(tr["h"][0]); return self._store_res(tr)98 def predict(self, te):99 import sys100 from wp23 import config as C101 sys.path.insert(0, str(C.WP22_SRC)); from sharp_hpi.statespace import _filter_llt102 out = np.empty(te.height); H = np.array(te["hist"].to_list(), float)103 for i, y in enumerate(H):104 r = np.full(len(y), 1e-4); a, p, s, sv, g, ll, vs = _filter_llt(y, r, 1e-5, 2e-6, 1.0)105 # slope at the end from the filtered state (re-run with 2-state output): approximate slope = last filtered level change106 out[i] = (a[-1] - a[-2]) * self.h107 return out108109110class RidgePanel(Base):111 name = "ridge_panel"112 def fit(self, tr):113 X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; self.m = Ridge(alpha=3.0).fit((X - self.mu) / self.sd, tr["target"].to_numpy()); return self._store_res(tr)114 def predict(self, te): return self.m.predict((_X(te) - self.mu) / self.sd)115116117class LGBMPanel(Base):118 name = "lightgbm_panel"; prob = True119 def fit(self, tr):120 import lightgbm as lgb121 X = _X(tr); y = tr["target"].to_numpy(); kw = dict(n_estimators=300, learning_rate=0.03, num_leaves=15, min_child_samples=40, subsample=0.8, subsample_freq=1, colsample_bytree=0.8, verbose=-1, n_jobs=4, random_state=0)122 self.m = lgb.LGBMRegressor(objective="l2", **kw).fit(X, y); self.q = {t: lgb.LGBMRegressor(objective="quantile", alpha=float(t), **kw).fit(X, y) for t in TAUS}; return self._store_res(tr)123 def predict(self, te): return self.m.predict(_X(te))124 def predict_quantiles(self, te, taus=TAUS): return np.sort(np.column_stack([self.q[t].predict(_X(te)) for t in taus]), axis=1)125126127class GaussianHMM:128 """Two-state Gaussian HMM on a univariate series, EM (Baum–Welch); used on the provincial/regional real-time growth at the origin."""129 def __init__(self, K=3, iters=60, seed=0): self.K, self.iters, self.rng = K, iters, np.random.default_rng(seed)130 def fit(self, x):131 x = np.asarray(x, float); K = self.K; q = np.quantile(x, np.linspace(0.2, 0.8, K)); mu = q.copy(); sd = np.full(K, x.std() + 1e-6); A = np.full((K, K), 0.1) + np.eye(K) * (0.9 - 0.1 * (K - 1)); pi = np.full(K, 1 / K)132 for _ in range(self.iters):133 B = np.exp(-0.5 * ((x[:, None] - mu) / sd) ** 2) / sd; T = len(x); al = np.zeros((T, K)); be = np.zeros((T, K)); c = np.zeros(T)134 al[0] = pi * B[0]; c[0] = al[0].sum() + 1e-300; al[0] /= c[0]135 for t in range(1, T):136 al[t] = (al[t - 1] @ A) * B[t]; c[t] = al[t].sum() + 1e-300; al[t] /= c[t]137 be[-1] = 1138 for t in range(T - 2, -1, -1):139 be[t] = (A @ (B[t + 1] * be[t + 1])) / c[t + 1]140 g = al * be; g /= g.sum(1, keepdims=True) + 1e-300141 xi = np.zeros((K, K))142 for t in range(T - 1):143 m = al[t][:, None] * A * (B[t + 1] * be[t + 1])[None, :]; xi += m / (m.sum() + 1e-300)144 pi = g[0]; A = xi / (xi.sum(1, keepdims=True) + 1e-300); mu = (g * x[:, None]).sum(0) / (g.sum(0) + 1e-300); sd = np.sqrt((g * (x[:, None] - mu) ** 2).sum(0) / (g.sum(0) + 1e-300)) + 1e-4145 self.mu, self.sd, self.A, self.pi = mu, sd, A, pi; self.filtered = al; return self146 def filter_last(self, x):147 x = np.asarray(x, float); B = np.exp(-0.5 * ((x[:, None] - self.mu) / self.sd) ** 2) / self.sd; al = self.pi * B[0]; al /= al.sum() + 1e-300148 for t in range(1, len(x)):149 al = (al @ self.A) * B[t]; al /= al.sum() + 1e-300150 return al151152153class RAREF(Base):154 """Regime-adaptive forecasting: a 3-state Gaussian HMM is estimated at each origin on the provincial real-time monthly growth155 (information ≤ origin only); the filtered regime probabilities gate a mixture of ridge experts (one per regime) estimated on156 training rows weighted by their own origin's regime probabilities. Forecast = probability-weighted expert forecast."""157 name = "raref"; prob = True158 def __init__(self, K=3): self.K = K159 def _probs(self, df):160 # per origin: filtered regime probabilities of the provincial growth series up to that origin161 out = np.zeros((df.height, self.K)); cache = {}162 for i, (o, c) in enumerate(df.select(["origin", "class"]).iter_rows()):163 if (o, c) not in cache:164 hist = df.filter((pl.col("origin") == o) & (pl.col("class") == c) & (pl.col("level") == "province"))165 if hist.height == 0:166 cache[(o, c)] = np.full(self.K, 1 / self.K)167 else:168 g = np.diff(np.array(hist["hist"][0], float)); key = (o, c)169 series = self._growth_series(df, o, c)170 cache[key] = self.hmm_for(series).filter_last(series) if len(series) > 10 else np.full(self.K, 1 / self.K)171 out[i] = cache[(o, c)]172 return out173 def _growth_series(self, df, o, c):174 # concatenate provincial growth across origins ≤ o from the stored histories (real-time at each origin): use the origin's own 13-month history175 # extended backward with the earliest available origins (all ≤ o → no leakage)176 hs = df.filter((pl.col("class") == c) & (pl.col("level") == "province") & (pl.col("origin") <= o)).sort("origin")177 if hs.height == 0:178 return np.array([])179 seq = []180 for o2, hh in zip(hs["origin"].to_list(), hs["hist"].to_list()):181 seq.append(np.diff(np.array(hh, float))[-1]) # last monthly growth known at each past origin182 return np.array(seq)183 def hmm_for(self, series):184 return GaussianHMM(self.K).fit(series)185 def fit(self, tr):186 self.P = self._probs(tr); X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; Xs = (X - self.mu) / self.sd; y = tr["target"].to_numpy()187 self.experts = []188 for k in range(self.K):189 w = self.P[:, k] + 1e-3; self.experts.append(Ridge(alpha=3.0).fit(Xs, y, sample_weight=w))190 self._train_df = tr; return self._store_res(tr)191 def predict(self, te):192 both = pl.concat([self._train_df.select(te.columns), te], how="diagonal") if te.height else te193 P = self._probs(both)[-te.height:]; Xs = (_X(te) - self.mu) / self.sd194 return sum(P[:, k] * self.experts[k].predict(Xs) for k in range(self.K))195196197class BHPF(Base):198 """Bayesian hierarchical probabilistic forecasting: hierarchical linear model with cell-level random intercepts (shrunk toward199 their region, then the class mean) and common slopes on the standardized features; conjugate Gaussian posterior (ridge-type)200 with variance components estimated by REML-like iteration; predictive distribution = posterior predictive Gaussian with201 cell-specific variance. Laplace/conjugate rather than MCMC: exact for this Gaussian model."""202 name = "bhpf"; prob = True203 def fit(self, tr):204 X = _X(tr); self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; Xs = (X - self.mu) / self.sd; y = tr["target"].to_numpy()205 cells = _keys(tr, ["class", "level", "geo_id"]); regs = _keys(tr, ["class", "region_code"])206 self.cells = sorted(set(cells)); self.regs = sorted(set(regs)); ci = {c: i for i, c in enumerate(self.cells)}; ri = {r: i for i, r in enumerate(self.regs)}207 Zc = np.zeros((len(y), len(self.cells))); Zc[np.arange(len(y)), [ci[c] for c in cells]] = 1208 Zr = np.zeros((len(y), len(self.regs))); Zr[np.arange(len(y)), [ri[r] for r in regs]] = 1209 self.ci, self.ri = ci, ri210 D = np.column_stack([np.ones(len(y)), Xs, Zr, Zc]); p = Xs.shape[1]211 tau_b, tau_r, tau_c, s2 = 1.0, 0.01, 0.01, float(y.var())212 for _ in range(15):213 prec = np.concatenate([[1e-6], np.full(p, 1 / tau_b), np.full(len(self.regs), 1 / tau_r), np.full(len(self.cells), 1 / tau_c)])214 Ainv = np.linalg.inv(D.T @ D / s2 + np.diag(prec)); beta = Ainv @ (D.T @ y / s2); r = y - D @ beta215 s2 = float(np.mean(r ** 2)) + 1e-9216 br = beta[1 + p: 1 + p + len(self.regs)]; bc = beta[1 + p + len(self.regs):]217 tau_r = max(float(np.mean(br ** 2) + np.mean(np.diag(Ainv)[1 + p: 1 + p + len(self.regs)])), 1e-6); tau_c = max(float(np.mean(bc ** 2) + np.mean(np.diag(Ainv)[1 + p + len(self.regs):])), 1e-6)218 self.beta, self.Ainv, self.s2, self.p = beta, Ainv, s2, p; self.vc = {"s2": s2, "tau_region": tau_r, "tau_cell": tau_c}; return self._store_res(tr)219 def _design(self, te):220 Xs = (_X(te) - self.mu) / self.sd; cells = _keys(te, ["class", "level", "geo_id"]); regs = _keys(te, ["class", "region_code"])221 Zc = np.zeros((te.height, len(self.cells))); Zr = np.zeros((te.height, len(self.regs)))222 for i, (c, r) in enumerate(zip(cells, regs)):223 if c in self.ci: Zc[i, self.ci[c]] = 1224 if r in self.ri: Zr[i, self.ri[r]] = 1225 return np.column_stack([np.ones(te.height), Xs, Zr, Zc])226 def predict(self, te): return self._design(te) @ self.beta227 def predict_quantiles(self, te, taus=TAUS):228 from scipy.stats import norm229 D = self._design(te); m = D @ self.beta; v = np.einsum("ij,jk,ik->i", D, self.Ainv, D) + self.s2230 return m[:, None] + np.sqrt(v)[:, None] * norm.ppf(taus)[None, :]231232233def registry(family: str = "classic", fast: bool = False) -> list:234 if family == "torch":235 from wp23.models_b_torch import DSGFN, GRUPanel, LSTMPanel, MLPPanel, QANSS, QREFA, TCNPanel, TFTLite236 base = [MLPPanel(), LSTMPanel(), GRUPanel(), TCNPanel(), TFTLite(), DSGFN(), QANSS()]237 if not fast:238 base.append(QREFA())239 return base240 if family == "trees":241 return [LGBMPanel()]242 return [NoChange(), MeanGrowth(), Momentum(), RWDrift(), SeasonalNaive(), AR2(), LocalLinearTrend(), RidgePanel(), RAREF(), BHPF()]243