#!/usr/bin/env python3 # Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca """Task B ablations (protocol §12), torch family process: DSGFN edge types (no distance / no co-movement / no graph), QANSS without parent drifts, QREFA component removal and an equal-weight combination; plus feature-group ablations of the panel ridge and MLP (no spatial parent features, no volume features, no history). Also stores DSGFN attention summaries and RAREF regime probabilities. Writes results/ablation_b/predictions_ablation.parquet and diagnostics.""" from __future__ import annotations import sys import time from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) import numpy as np import polars as pl from wp23 import config as C from wp23 import data_b as DB from wp23.data import month_add from wp23.models_b import RAREF, BHPF, RidgePanel, TAUS, Base from wp23.models_b_torch import DSGFN, MLPPanel, QANSS, QREFA OUT = C.RESULTS / "ablation_b"; OUT.mkdir(parents=True, exist_ok=True) class DSGFN_NoDist(DSGFN): name = "dsgfn_no_distance" def _adjacency(self, df): A = super()._adjacency(df); import torch 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) / 1000 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-level class DSGFN_NoComove(DSGFN): name = "dsgfn_no_comovement" def _adjacency(self, df): import torch 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() d = np.sqrt((cx[:, None] - cx[None, :]) ** 2 + (cy[:, None] - cy[None, :]) ** 2) / 1000; A = np.exp(-(d / self.bw) ** 2) same = (lvl[:, None] == lvl[None, :]) & (cls[:, None] == cls[None, :]) & (org[:, None] == org[None, :]); return torch.tensor(A * same, dtype=torch.float32) class DSGFN_NoGraph(DSGFN): name = "dsgfn_no_graph" def _adjacency(self, df): import torch return torch.ones((df.height, df.height), dtype=torch.float32) * 1e-6 # uniform (uninformative) prior; attention still learned class QANSS_NoParent(QANSS): name = "qanss_no_parent" def _state(self, df): S = super()._state(df); S[:, 3] = 0.0; S[:, 4] = 0.0; return S class EqualCombo(Base): name = "equal_weight_combo"; prob = True def __init__(self): self.components = [QANSS(), BHPF(), DSGFN(epochs=40), RAREF()] def fit(self, tr): for c in self.components: c.fit(tr) return self._store_res(tr) def predict(self, te): return np.mean([c.predict(te) for c in self.components], axis=0) 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) class FeatureDrop(Base): """Ridge panel with a feature group zeroed out (train and test).""" prob = False def __init__(self, group: str): self.group = group; self.name = f"ridge_no_{group}" 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] def _z(self, df): return df.with_columns([pl.lit(0.0).alias(c) for c in self.cols]) def fit(self, tr): self.m = RidgePanel().fit(self._z(tr)); self._res = self.m._res; return self def predict(self, te): return self.m.predict(self._z(te)) def main() -> None: panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet") models = [DSGFN(), DSGFN_NoDist(), DSGFN_NoComove(), DSGFN_NoGraph(), QANSS(), QANSS_NoParent(), QREFA(), EqualCombo(), FeatureDrop("parent"), FeatureDrop("volume"), FeatureDrop("momentum"), RidgePanel()] for comp in ("QANSS", "BHPF", "DSGFN", "RAREF"): 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) rows, att_rows, reg_rows = [], [], [] for h in (1, 3, 6): ph = panel.filter(pl.col("h") == h); origins = sorted(ph["origin"].unique().to_list()) for k in range(0, len(origins), 3): o = origins[k]; test_o = origins[k:k + 3] tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o) if tr["origin"].n_unique() < 10: continue te = ph.filter(pl.col("origin").is_in(test_o)) for m in models: t0 = time.time() try: m.fit(tr); p = m.predict(te); Q = m.predict_quantiles(te) except Exception as ex: print("FAILED", h, o, m.name, repr(ex), flush=True); continue 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))) rows.append(d); print(h, o, m.name, f"{time.time() - t0:.1f}s", flush=True) if m.name == "dsgfn" and hasattr(m, "last_att"): for (oo, cc), A in m.last_att.items(): sub = te.filter((pl.col("origin") == oo) & (pl.col("class") == cc)) if A.shape[0] != sub.height: continue ids = sub["geo_id"].to_list(); lv = sub["level"].to_list() for i in range(A.shape[0]): 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)))}) if m.name == "raref": P = m._probs(pl.concat([m._train_df.select(te.columns), te], how="diagonal"))[-te.height:] for (oo, cc), pr in zip(te.select(["origin", "class"]).iter_rows(), P): reg_rows.append({"h": h, "origin": oo, "class": cc, "p_state0": float(pr[0]), "p_state1": float(pr[1]), "p_state2": float(pr[2])}) pl.concat(rows).write_parquet(OUT / "predictions_ablation.parquet") pl.DataFrame(att_rows).write_csv(OUT / "dsgfn_attention.csv") if att_rows else None print("done") if __name__ == "__main__": main()