TeX 62.3%
Python 37.6%
1#!/usr/bin/env python32# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca3"""Task B ablations (protocol §12), torch family process: DSGFN edge types (no distance / no co-movement / no graph), QANSS without4parent drifts, QREFA component removal and an equal-weight combination; plus feature-group ablations of the panel ridge and MLP5(no spatial parent features, no volume features, no history). Also stores DSGFN attention summaries and RAREF regime probabilities.6Writes results/ablation_b/predictions_ablation.parquet and diagnostics."""7from __future__ import annotations89import sys10import time11from pathlib import Path1213sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))14import numpy as np15import polars as pl1617from wp23 import config as C18from wp23 import data_b as DB19from wp23.data import month_add20from wp23.models_b import RAREF, BHPF, RidgePanel, TAUS, Base21from wp23.models_b_torch import DSGFN, MLPPanel, QANSS, QREFA2223OUT = C.RESULTS / "ablation_b"; OUT.mkdir(parents=True, exist_ok=True)242526class DSGFN_NoDist(DSGFN):27 name = "dsgfn_no_distance"28 def _adjacency(self, df):29 A = super()._adjacency(df); import torch30 cx = df["cx"].fill_null(0.0).to_numpy(); cy = df["cy"].fill_null(0.0).to_numpy(); d = np.sqrt((cx[:, None] - cx[None, :]) ** 2 + (cy[:, None] - cy[None, :]) ** 2) / 100031 K = torch.tensor(np.exp(-(d / self.bw) ** 2), dtype=torch.float32); return A / (K + 1e-6) # removes the distance kernel, keeps co-movement × same-level323334class DSGFN_NoComove(DSGFN):35 name = "dsgfn_no_comovement"36 def _adjacency(self, df):37 import torch38 cx = df["cx"].fill_null(0.0).to_numpy(); cy = df["cy"].fill_null(0.0).to_numpy(); lvl = df["level"].to_numpy(); cls = df["class"].to_numpy(); org = df["origin"].to_numpy()39 d = np.sqrt((cx[:, None] - cx[None, :]) ** 2 + (cy[:, None] - cy[None, :]) ** 2) / 1000; A = np.exp(-(d / self.bw) ** 2)40 same = (lvl[:, None] == lvl[None, :]) & (cls[:, None] == cls[None, :]) & (org[:, None] == org[None, :]); return torch.tensor(A * same, dtype=torch.float32)414243class DSGFN_NoGraph(DSGFN):44 name = "dsgfn_no_graph"45 def _adjacency(self, df):46 import torch47 return torch.ones((df.height, df.height), dtype=torch.float32) * 1e-6 # uniform (uninformative) prior; attention still learned484950class QANSS_NoParent(QANSS):51 name = "qanss_no_parent"52 def _state(self, df):53 S = super()._state(df); S[:, 3] = 0.0; S[:, 4] = 0.0; return S545556class EqualCombo(Base):57 name = "equal_weight_combo"; prob = True58 def __init__(self): self.components = [QANSS(), BHPF(), DSGFN(epochs=40), RAREF()]59 def fit(self, tr):60 for c in self.components: c.fit(tr)61 return self._store_res(tr)62 def predict(self, te): return np.mean([c.predict(te) for c in self.components], axis=0)63 def predict_quantiles(self, te, taus=TAUS): return np.sort(np.mean([c.predict_quantiles(te, taus) for c in self.components], axis=0), axis=1)646566class FeatureDrop(Base):67 """Ridge panel with a feature group zeroed out (train and test)."""68 prob = False69 def __init__(self, group: str):70 self.group = group; self.name = f"ridge_no_{group}"71 self.cols = {"parent": ["reg_d3", "reg_d12", "reg_vol12", "prov_d3", "prov_d12", "prov_vol12"], "volume": ["n_last", "n_mean12", "n_ratio", "se_last"], "momentum": ["d1", "d3", "d6", "d12", "accel"]}[group]72 def _z(self, df): return df.with_columns([pl.lit(0.0).alias(c) for c in self.cols])73 def fit(self, tr): self.m = RidgePanel().fit(self._z(tr)); self._res = self.m._res; return self74 def predict(self, te): return self.m.predict(self._z(te))757677def main() -> None:78 panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet")79 models = [DSGFN(), DSGFN_NoDist(), DSGFN_NoComove(), DSGFN_NoGraph(), QANSS(), QANSS_NoParent(), QREFA(), EqualCombo(), FeatureDrop("parent"), FeatureDrop("volume"), FeatureDrop("momentum"), RidgePanel()]80 for comp in ("QANSS", "BHPF", "DSGFN", "RAREF"):81 q = QREFA(components=[c for c in [QANSS(), BHPF(), DSGFN(epochs=40), RAREF()] if type(c).__name__ != comp]); q.name = f"qrefa_without_{comp.lower()}"; models.append(q)82 rows, att_rows, reg_rows = [], [], []83 for h in (1, 3, 6):84 ph = panel.filter(pl.col("h") == h); origins = sorted(ph["origin"].unique().to_list())85 for k in range(0, len(origins), 3):86 o = origins[k]; test_o = origins[k:k + 3]87 tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)88 if tr["origin"].n_unique() < 10: continue89 te = ph.filter(pl.col("origin").is_in(test_o))90 for m in models:91 t0 = time.time()92 try:93 m.fit(tr); p = m.predict(te); Q = m.predict_quantiles(te)94 except Exception as ex:95 print("FAILED", h, o, m.name, repr(ex), flush=True); continue96 d = te.select(["class", "level", "geo_id", "origin", "h", "target", "n_last", "region_code"]).with_columns(pl.lit(m.name).alias("model"), pl.lit(o).alias("fit_origin"), pl.Series("pred", np.asarray(p, float)), pl.Series("q05", np.asarray(Q[:, 0], float)), pl.Series("q95", np.asarray(Q[:, 4], float)))97 rows.append(d); print(h, o, m.name, f"{time.time() - t0:.1f}s", flush=True)98 if m.name == "dsgfn" and hasattr(m, "last_att"):99 for (oo, cc), A in m.last_att.items():100 sub = te.filter((pl.col("origin") == oo) & (pl.col("class") == cc))101 if A.shape[0] != sub.height: continue102 ids = sub["geo_id"].to_list(); lv = sub["level"].to_list()103 for i in range(A.shape[0]):104 j = int(np.argmax(A[i])); att_rows.append({"h": h, "origin": oo, "class": cc, "level": lv[i], "geo_id": ids[i], "top_neighbour": ids[j], "top_weight": float(A[i, j]), "entropy": float(-np.sum(A[i] * np.log(A[i] + 1e-12)))})105 if m.name == "raref":106 P = m._probs(pl.concat([m._train_df.select(te.columns), te], how="diagonal"))[-te.height:]107 for (oo, cc), pr in zip(te.select(["origin", "class"]).iter_rows(), P):108 reg_rows.append({"h": h, "origin": oo, "class": cc, "p_state0": float(pr[0]), "p_state1": float(pr[1]), "p_state2": float(pr[2])})109 pl.concat(rows).write_parquet(OUT / "predictions_ablation.parquet")110 pl.DataFrame(att_rows).write_csv(OUT / "dsgfn_attention.csv") if att_rows else None111 print("done")112113114if __name__ == "__main__":115 main()116