TeX 62.3%
Python 37.6%
1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Task B neural models (PyTorch): panel MLP / LSTM / GRU / TCN / TFT-lite benchmarks and the proposed DSGFN, QANSS and QREFA.3Never import lightgbm / xgboost in this process (OpenMP conflict on macOS)."""4from __future__ import annotations56import os78import numpy as np9import polars as pl10import torch11import torch.nn as nn1213from wp23.data_b import FEATURES_B14from wp23.models_b import TAUS, Base, BHPF, RAREF, _H, _X1516torch.set_num_threads(int(os.environ.get("WP23_TORCH_THREADS", "6")))171819def _pinball(pred, y, taus):20 d = y[:, None] - pred; return torch.mean(torch.maximum(taus * d, (taus - 1) * d))212223class TorchPanel(Base):24 """Shared training loop for neural panel models: inputs = standardized features + relative history; output = 5 quantiles."""25 prob = True26 def __init__(self, epochs=60, lr=2e-3, batch=512, seed=0, wd=1e-4):27 self.epochs, self.lr, self.batch, self.seed, self.wd = epochs, lr, batch, seed, wd28 def _inputs(self, df, fit=False):29 X = _X(df); H = _H(df)30 if fit:31 self.mu, self.sd = X.mean(0), X.std(0) + 1e-9; self.hs = H.std() + 1e-932 return torch.tensor((X - self.mu) / self.sd, dtype=torch.float32), torch.tensor(H / self.hs, dtype=torch.float32)33 def build(self, d_x, T): raise NotImplementedError34 def forward(self, x, hist): return self.net(x, hist)35 def fit(self, tr):36 torch.manual_seed(self.seed); x, hist = self._inputs(tr, fit=True); y = torch.tensor(tr["target"].to_numpy(), dtype=torch.float32)37 self.ys = float(y.std()) + 1e-9; y = y / self.ys; taus = torch.tensor(TAUS, dtype=torch.float32)38 self.build(x.shape[1], hist.shape[1]); self._extra_fit(tr)39 opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.wd); n = len(y)40 for ep in range(self.epochs):41 perm = torch.randperm(n)42 for i in range(0, n, self.batch):43 b = perm[i:i + self.batch]; opt.zero_grad(); loss = _pinball(self.forward(x[b], hist[b]), y[b], taus); loss.backward(); opt.step()44 self.eval_mode(); return self._store_res(tr)45 def parameters(self): return self.net.parameters()46 def eval_mode(self): self.net.eval()47 def _extra_fit(self, tr): pass48 def predict_quantiles(self, te, taus=TAUS):49 with torch.no_grad():50 x, hist = self._inputs(te); Q = self.forward(x, hist).numpy() * self.ys51 return np.sort(Q, axis=1)52 def predict(self, te): return self.predict_quantiles(te)[:, 2]535455class MLPNet(nn.Module):56 def __init__(self, d_x, T, hidden=64):57 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))58 def forward(self, x, h): return self.f(torch.cat([x, h], 1))596061class MLPPanel(TorchPanel):62 name = "mlp_panel"63 def build(self, d_x, T): self.net = MLPNet(d_x, T)646566class RNNNet(nn.Module):67 def __init__(self, d_x, T, kind="lstm", hidden=32):68 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))69 def forward(self, x, h):70 out, _ = self.rnn(h.unsqueeze(-1)); return self.f(torch.cat([out[:, -1], x], 1))717273class LSTMPanel(TorchPanel):74 name = "lstm"75 def build(self, d_x, T): self.net = RNNNet(d_x, T, "lstm")767778class GRUPanel(TorchPanel):79 name = "gru"80 def build(self, d_x, T): self.net = RNNNet(d_x, T, "gru")818283class TCNNet(nn.Module):84 def __init__(self, d_x, T, ch=16):85 super().__init__()86 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)87 self.f = nn.Sequential(nn.Linear(ch + d_x, 64), nn.GELU(), nn.Linear(64, 5))88 def forward(self, x, h):89 z = h.unsqueeze(1); T = z.shape[-1]90 z = torch.relu(self.c1(z))[..., :T]; z = torch.relu(self.c2(z))[..., :T]; z = torch.relu(self.c3(z))[..., :T] # causal crops91 return self.f(torch.cat([z[:, :, -1], x], 1))929394class TCNPanel(TorchPanel):95 name = "tcn"96 def build(self, d_x, T): self.net = TCNNet(d_x, T)979899class TFTLiteNet(nn.Module):100 """Temporal-fusion-style: gated variable selection over static/known features, self-attention over the history embedding."""101 def __init__(self, d_x, T, d=32):102 super().__init__(); self.vs = nn.Sequential(nn.Linear(d_x, d_x), nn.Softmax(dim=-1)); self.xe = nn.Linear(d_x, d)103 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)104 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))105 def forward(self, x, h):106 xs = self.xe(x * self.vs(x)); hs = self.he(h.unsqueeze(-1)) + self.pos107 a, _ = self.att(hs, hs, hs, need_weights=False); ctx = a[:, -1]108 g = self.gate(torch.cat([xs, ctx], 1)); return self.f(torch.cat([g * xs, (1 - g) * ctx], 1))109110111class TFTLite(TorchPanel):112 name = "tft_lite"113 def build(self, d_x, T): self.net = TFTLiteNet(d_x, T)114115116# ---- proposed architectures -----------------------------------------------------------------------117class DSGFNNet(nn.Module):118 """Dynamic spatial graph forecasting network: per-node encoding of history + features, then graph attention over the cells119 present at the same origin with an adjacency prior (distance kernel + co-movement + composition similarity) and learned120 attention; the head reads the node's own encoding and its attention-aggregated neighbourhood."""121 def __init__(self, d_x, T, d=32):122 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)123 self.prior_w = nn.Parameter(torch.tensor(1.0)); self.f = nn.Sequential(nn.Linear(2 * d, 64), nn.GELU(), nn.Linear(64, 5))124 def forward(self, x, h, A):125 z = self.enc(torch.cat([x, h], 1)); q, k, v = self.q(z), self.k(z), self.v(z)126 logits = q @ k.T / np.sqrt(q.shape[1]) + self.prior_w * torch.log(A + 1e-6)127 logits = logits - torch.eye(len(z)) * 1e9 # no self-attention: neighbours only128 att = torch.softmax(logits, dim=1); nb = att @ v129 return self.f(torch.cat([z, nb], 1)), att130131132class DSGFN(TorchPanel):133 name = "dsgfn"134 def __init__(self, bandwidth_km=60.0, **kw):135 super().__init__(**kw); self.bw = bandwidth_km136 if bandwidth_km != 60.0: self.name = f"dsgfn_bw{int(bandwidth_km)}"137 def build(self, d_x, T): self.net = DSGFNNet(d_x, T)138 def _adjacency(self, df: pl.DataFrame) -> torch.Tensor:139 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()140 d = np.sqrt((cx[:, None] - cx[None, :]) ** 2 + (cy[:, None] - cy[None, :]) ** 2) / 1000; A = np.exp(-(d / self.bw) ** 2)141 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))142 same = (lvl[:, None] == lvl[None, :]) & (cls[:, None] == cls[None, :]) & (org[:, None] == org[None, :]); A = A * same143 return torch.tensor(A, dtype=torch.float32)144 def fit(self, tr):145 torch.manual_seed(self.seed); self.build(len(FEATURES_B), 13); taus = torch.tensor(TAUS, dtype=torch.float32)146 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-9147 groups = tr.with_row_index("_i").group_by(["origin", "class"]).agg(pl.col("_i"))148 idxs = [torch.tensor(g) for g in groups["_i"].to_list()]; As = [self._adjacency(tr[g]) for g in groups["_i"].to_list()]149 opt = torch.optim.AdamW(self.net.parameters(), lr=self.lr, weight_decay=self.wd)150 for ep in range(self.epochs):151 for j in torch.randperm(len(idxs)).tolist():152 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()153 self.net.eval(); return self._store_res(tr)154 def predict_quantiles(self, te, taus=TAUS):155 Q = np.zeros((te.height, 5)); self.last_att = {}156 with torch.no_grad():157 x, h = self._inputs(te)158 for o, c, g in te.with_row_index("_i").group_by(["origin", "class"]).agg(pl.col("_i")).iter_rows():159 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()160 return np.sort(Q, axis=1)161162163class QANSS(Base):164 """Quality-adjusted neural state-space forecaster: the cell's quality-adjusted log level (already separated from composition by165 the WP22 stage-1 hedonic) follows a latent level + latent drift; the drift's transition is a small network of the cell's recent166 state (drift, volatility, volume ratio, parent drift) estimated on training origins; the Kalman filter at the origin gives the167 state and its variance; the forecast is level + Σ_{j≤h} drift_j with drift propagated by the learned transition and the168 predictive variance from the filter plus the learned drift noise."""169 name = "qanss"; prob = True170 def __init__(self, epochs=200, lr=5e-3, seed=0): self.epochs, self.lr, self.seed = epochs, lr, seed171 def _state(self, df):172 H = np.array(df["hist"].to_list(), float); D = np.diff(H, axis=1)173 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() / 3174 return np.column_stack([drift, vol, nr - 1, pd_, rd, D[:, -1], D[:, -2]])175 def fit(self, tr):176 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-9177 x = torch.tensor((S - self.mu) / self.sd, dtype=torch.float32); y = torch.tensor(tr["target"].to_numpy(), dtype=torch.float32)178 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 change179 opt = torch.optim.Adam(self.net.parameters(), lr=self.lr)180 for _ in range(self.epochs):181 opt.zero_grad(); out = self.net(x); m = out[:, 0] * self.h; ls = out[:, 1]182 loss = torch.mean(0.5 * ((y - m) / torch.exp(ls)) ** 2 + ls); loss.backward(); opt.step()183 self.net.eval(); return self._store_res(tr)184 def _mv(self, te):185 with torch.no_grad():186 out = self.net(torch.tensor((self._state(te) - self.mu) / self.sd, dtype=torch.float32)).numpy()187 se = te["se_last"].fill_null(0.0).to_numpy()188 return out[:, 0] * self.h, np.sqrt(np.exp(2 * out[:, 1]) + se ** 2)189 def predict(self, te): return self._mv(te)[0]190 def predict_quantiles(self, te, taus=TAUS):191 from scipy.stats import norm192 m, s = self._mv(te); return m[:, None] + s[:, None] * norm.ppf(taus)[None, :]193194195class QREFA(Base):196 """Integrated architecture: stacking of the components with the strongest out-of-sample evidence (QANSS drift state,197 DSGFN neighbourhood signal, BHPF hierarchical intercepts, RAREF regime probabilities) through a constrained linear198 combination estimated on the training origins (non-negative weights summing to one, no intercept), so that every component199 is ablatable by removing it from the stack."""200 name = "qrefa"; prob = True201 def __init__(self, components=None): self.components = components or [QANSS(), BHPF(), DSGFN(epochs=40), RAREF()]202 def fit(self, tr):203 from scipy.optimize import nnls204 F = []205 for c in self.components:206 c.fit(tr); F.append(c.predict(tr))207 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))208 return self._store_res(tr)209 def predict(self, te): return np.column_stack([c.predict(te) for c in self.components]) @ self.w210 def predict_quantiles(self, te, taus=TAUS):211 Q = sum(w * c.predict_quantiles(te, taus) for w, c in zip(self.w, self.components)); return np.sort(Q, axis=1)212213214215216def registry_torch_fast(family: str = "torch", fast: bool = True) -> list:217 return [MLPPanel(epochs=30), LSTMPanel(epochs=30), TCNPanel(epochs=30), TFTLite(epochs=30), DSGFN(epochs=30), QANSS(epochs=150)]218