# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca """Task B neural models (PyTorch): panel MLP / LSTM / GRU / TCN / TFT-lite benchmarks and the proposed DSGFN, QANSS and QREFA. Never import lightgbm / xgboost in this process (OpenMP conflict on macOS).""" from __future__ import annotations import os import numpy as np import polars as pl import torch import torch.nn as nn from wp23.data_b import FEATURES_B from wp23.models_b import TAUS, Base, BHPF, RAREF, _H, _X torch.set_num_threads(int(os.environ.get("WP23_TORCH_THREADS", "6"))) def _pinball(pred, y, taus): d = y[:, None] - pred; return torch.mean(torch.maximum(taus * d, (taus - 1) * d)) class TorchPanel(Base): """Shared training loop for neural panel models: inputs = standardized features + relative history; output = 5 quantiles.""" prob = True def __init__(self, epochs=60, lr=2e-3, batch=512, seed=0, wd=1e-4): self.epochs, self.lr, self.batch, self.seed, self.wd = epochs, lr, batch, seed, wd def _inputs(self, df, fit=False): X = _X(df); H = _H(df) if fit: self.mu, self.sd = X.mean(0), X.std(0) + 1e-9; self.hs = H.std() + 1e-9 return torch.tensor((X - self.mu) / self.sd, dtype=torch.float32), torch.tensor(H / self.hs, dtype=torch.float32) def build(self, d_x, T): raise NotImplementedError def forward(self, x, hist): return self.net(x, hist) def fit(self, tr): torch.manual_seed(self.seed); x, hist = self._inputs(tr, fit=True); y = torch.tensor(tr["target"].to_numpy(), dtype=torch.float32) self.ys = float(y.std()) + 1e-9; y = y / self.ys; taus = torch.tensor(TAUS, dtype=torch.float32) self.build(x.shape[1], hist.shape[1]); self._extra_fit(tr) opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.wd); n = len(y) for ep in range(self.epochs): perm = torch.randperm(n) for i in range(0, n, self.batch): b = perm[i:i + self.batch]; opt.zero_grad(); loss = _pinball(self.forward(x[b], hist[b]), y[b], taus); loss.backward(); opt.step() self.eval_mode(); return self._store_res(tr) def parameters(self): return self.net.parameters() def eval_mode(self): self.net.eval() def _extra_fit(self, tr): pass def predict_quantiles(self, te, taus=TAUS): with torch.no_grad(): x, hist = self._inputs(te); Q = self.forward(x, hist).numpy() * self.ys return np.sort(Q, axis=1) def predict(self, te): return self.predict_quantiles(te)[:, 2] class MLPNet(nn.Module): def __init__(self, d_x, T, hidden=64): super().__init__(); self.f = nn.Sequential(nn.Linear(d_x + T, hidden), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, 5)) def forward(self, x, h): return self.f(torch.cat([x, h], 1)) class MLPPanel(TorchPanel): name = "mlp_panel" def build(self, d_x, T): self.net = MLPNet(d_x, T) class RNNNet(nn.Module): def __init__(self, d_x, T, kind="lstm", hidden=32): super().__init__(); cell = nn.LSTM if kind == "lstm" else nn.GRU; self.rnn = cell(1, hidden, batch_first=True); self.f = nn.Sequential(nn.Linear(hidden + d_x, 64), nn.GELU(), nn.Linear(64, 5)) def forward(self, x, h): out, _ = self.rnn(h.unsqueeze(-1)); return self.f(torch.cat([out[:, -1], x], 1)) class LSTMPanel(TorchPanel): name = "lstm" def build(self, d_x, T): self.net = RNNNet(d_x, T, "lstm") class GRUPanel(TorchPanel): name = "gru" def build(self, d_x, T): self.net = RNNNet(d_x, T, "gru") class TCNNet(nn.Module): def __init__(self, d_x, T, ch=16): super().__init__() self.c1 = nn.Conv1d(1, ch, 3, padding=2, dilation=1); self.c2 = nn.Conv1d(ch, ch, 3, padding=4, dilation=2); self.c3 = nn.Conv1d(ch, ch, 3, padding=8, dilation=4) self.f = nn.Sequential(nn.Linear(ch + d_x, 64), nn.GELU(), nn.Linear(64, 5)) def forward(self, x, h): z = h.unsqueeze(1); T = z.shape[-1] z = torch.relu(self.c1(z))[..., :T]; z = torch.relu(self.c2(z))[..., :T]; z = torch.relu(self.c3(z))[..., :T] # causal crops return self.f(torch.cat([z[:, :, -1], x], 1)) class TCNPanel(TorchPanel): name = "tcn" def build(self, d_x, T): self.net = TCNNet(d_x, T) class TFTLiteNet(nn.Module): """Temporal-fusion-style: gated variable selection over static/known features, self-attention over the history embedding.""" def __init__(self, d_x, T, d=32): super().__init__(); self.vs = nn.Sequential(nn.Linear(d_x, d_x), nn.Softmax(dim=-1)); self.xe = nn.Linear(d_x, d) self.he = nn.Linear(1, d); self.pos = nn.Parameter(torch.randn(T, d) * 0.02); self.att = nn.MultiheadAttention(d, 4, batch_first=True) self.gate = nn.Sequential(nn.Linear(2 * d, d), nn.Sigmoid()); self.f = nn.Sequential(nn.Linear(2 * d, 64), nn.GELU(), nn.Linear(64, 5)) def forward(self, x, h): xs = self.xe(x * self.vs(x)); hs = self.he(h.unsqueeze(-1)) + self.pos a, _ = self.att(hs, hs, hs, need_weights=False); ctx = a[:, -1] g = self.gate(torch.cat([xs, ctx], 1)); return self.f(torch.cat([g * xs, (1 - g) * ctx], 1)) class TFTLite(TorchPanel): name = "tft_lite" def build(self, d_x, T): self.net = TFTLiteNet(d_x, T) # ---- proposed architectures ----------------------------------------------------------------------- class DSGFNNet(nn.Module): """Dynamic spatial graph forecasting network: per-node encoding of history + features, then graph attention over the cells present at the same origin with an adjacency prior (distance kernel + co-movement + composition similarity) and learned attention; the head reads the node's own encoding and its attention-aggregated neighbourhood.""" def __init__(self, d_x, T, d=32): super().__init__(); self.enc = nn.Sequential(nn.Linear(d_x + T, d), nn.GELU(), nn.Linear(d, d)); self.q = nn.Linear(d, d); self.k = nn.Linear(d, d); self.v = nn.Linear(d, d) self.prior_w = nn.Parameter(torch.tensor(1.0)); self.f = nn.Sequential(nn.Linear(2 * d, 64), nn.GELU(), nn.Linear(64, 5)) def forward(self, x, h, A): z = self.enc(torch.cat([x, h], 1)); q, k, v = self.q(z), self.k(z), self.v(z) logits = q @ k.T / np.sqrt(q.shape[1]) + self.prior_w * torch.log(A + 1e-6) logits = logits - torch.eye(len(z)) * 1e9 # no self-attention: neighbours only att = torch.softmax(logits, dim=1); nb = att @ v return self.f(torch.cat([z, nb], 1)), att class DSGFN(TorchPanel): name = "dsgfn" def __init__(self, bandwidth_km=60.0, **kw): super().__init__(**kw); self.bw = bandwidth_km if bandwidth_km != 60.0: self.name = f"dsgfn_bw{int(bandwidth_km)}" def build(self, d_x, T): self.net = DSGFNNet(d_x, T) def _adjacency(self, df: pl.DataFrame) -> torch.Tensor: cx = df["cx"].fill_null(0.0).to_numpy(); cy = df["cy"].fill_null(0.0).to_numpy(); H = _H(df); lvl = df["level"].to_numpy(); cls = df["class"].to_numpy(); org = df["origin"].to_numpy() d = np.sqrt((cx[:, None] - cx[None, :]) ** 2 + (cy[:, None] - cy[None, :]) ** 2) / 1000; A = np.exp(-(d / self.bw) ** 2) D = np.diff(H, axis=1); D = (D - D.mean(1, keepdims=True)) / (D.std(1, keepdims=True) + 1e-9); A = A * (0.5 + 0.5 * np.clip(D @ D.T / D.shape[1], -1, 1)) same = (lvl[:, None] == lvl[None, :]) & (cls[:, None] == cls[None, :]) & (org[:, None] == org[None, :]); A = A * same return torch.tensor(A, dtype=torch.float32) def fit(self, tr): torch.manual_seed(self.seed); self.build(len(FEATURES_B), 13); taus = torch.tensor(TAUS, dtype=torch.float32) x_all, h_all = self._inputs(tr, fit=True); y_all = torch.tensor(tr["target"].to_numpy(), dtype=torch.float32); self.ys = float(y_all.std()) + 1e-9 groups = tr.with_row_index("_i").group_by(["origin", "class"]).agg(pl.col("_i")) idxs = [torch.tensor(g) for g in groups["_i"].to_list()]; As = [self._adjacency(tr[g]) for g in groups["_i"].to_list()] opt = torch.optim.AdamW(self.net.parameters(), lr=self.lr, weight_decay=self.wd) for ep in range(self.epochs): for j in torch.randperm(len(idxs)).tolist(): b = idxs[j]; opt.zero_grad(); out, _ = self.net(x_all[b], h_all[b], As[j]); loss = _pinball(out, y_all[b] / self.ys, taus); loss.backward(); opt.step() self.net.eval(); return self._store_res(tr) def predict_quantiles(self, te, taus=TAUS): Q = np.zeros((te.height, 5)); self.last_att = {} with torch.no_grad(): x, h = self._inputs(te) for o, c, g in te.with_row_index("_i").group_by(["origin", "class"]).agg(pl.col("_i")).iter_rows(): b = torch.tensor(g); out, att = self.net(x[b], h[b], self._adjacency(te[g])); Q[g] = out.numpy() * self.ys; self.last_att[(o, c)] = att.numpy() return np.sort(Q, axis=1) class QANSS(Base): """Quality-adjusted neural state-space forecaster: the cell's quality-adjusted log level (already separated from composition by the WP22 stage-1 hedonic) follows a latent level + latent drift; the drift's transition is a small network of the cell's recent state (drift, volatility, volume ratio, parent drift) estimated on training origins; the Kalman filter at the origin gives the state and its variance; the forecast is level + Σ_{j≤h} drift_j with drift propagated by the learned transition and the predictive variance from the filter plus the learned drift noise.""" name = "qanss"; prob = True def __init__(self, epochs=200, lr=5e-3, seed=0): self.epochs, self.lr, self.seed = epochs, lr, seed def _state(self, df): H = np.array(df["hist"].to_list(), float); D = np.diff(H, axis=1) lvl = H[:, -1]; drift = D[:, -3:].mean(1); vol = D.std(1); nr = df["n_ratio"].fill_null(1.0).to_numpy(); pd_ = df["prov_d3"].fill_null(0.0).to_numpy() / 3; rd = df["reg_d3"].fill_null(0.0).to_numpy() / 3 return np.column_stack([drift, vol, nr - 1, pd_, rd, D[:, -1], D[:, -2]]) def fit(self, tr): torch.manual_seed(self.seed); self.h = int(tr["h"][0]); S = self._state(tr); self.mu = S.mean(0); self.sd = S.std(0) + 1e-9 x = torch.tensor((S - self.mu) / self.sd, dtype=torch.float32); y = torch.tensor(tr["target"].to_numpy(), dtype=torch.float32) self.net = nn.Sequential(nn.Linear(S.shape[1], 32), nn.Tanh(), nn.Linear(32, 2)) # mean drift per month and log-sd of the h-step change opt = torch.optim.Adam(self.net.parameters(), lr=self.lr) for _ in range(self.epochs): opt.zero_grad(); out = self.net(x); m = out[:, 0] * self.h; ls = out[:, 1] loss = torch.mean(0.5 * ((y - m) / torch.exp(ls)) ** 2 + ls); loss.backward(); opt.step() self.net.eval(); return self._store_res(tr) def _mv(self, te): with torch.no_grad(): out = self.net(torch.tensor((self._state(te) - self.mu) / self.sd, dtype=torch.float32)).numpy() se = te["se_last"].fill_null(0.0).to_numpy() return out[:, 0] * self.h, np.sqrt(np.exp(2 * out[:, 1]) + se ** 2) def predict(self, te): return self._mv(te)[0] def predict_quantiles(self, te, taus=TAUS): from scipy.stats import norm m, s = self._mv(te); return m[:, None] + s[:, None] * norm.ppf(taus)[None, :] class QREFA(Base): """Integrated architecture: stacking of the components with the strongest out-of-sample evidence (QANSS drift state, DSGFN neighbourhood signal, BHPF hierarchical intercepts, RAREF regime probabilities) through a constrained linear combination estimated on the training origins (non-negative weights summing to one, no intercept), so that every component is ablatable by removing it from the stack.""" name = "qrefa"; prob = True def __init__(self, components=None): self.components = components or [QANSS(), BHPF(), DSGFN(epochs=40), RAREF()] def fit(self, tr): from scipy.optimize import nnls F = [] for c in self.components: c.fit(tr); F.append(c.predict(tr)) F = np.column_stack(F); y = tr["target"].to_numpy(); w, _ = nnls(F, y); self.w = w / (w.sum() + 1e-12) if w.sum() > 0 else np.full(len(w), 1 / len(w)) return self._store_res(tr) def predict(self, te): return np.column_stack([c.predict(te) for c in self.components]) @ self.w def predict_quantiles(self, te, taus=TAUS): Q = sum(w * c.predict_quantiles(te, taus) for w, c in zip(self.w, self.components)); return np.sort(Q, axis=1) def registry_torch_fast(family: str = "torch", fast: bool = True) -> list: return [MLPPanel(epochs=30), LSTMPanel(epochs=30), TCNPanel(epochs=30), TFTLite(epochs=30), DSGFN(epochs=30), QANSS(epochs=150)]