spb/wp23_uqo

1#!/usr/bin/env python32# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca3"""Task B robustness (protocol §12): rolling 24-origin training windows vs expanding, three seeds for the neural models,4exclusion of thin cell-months (n_last < 10), DSGFN bandwidths 30 / 120 km, winsorised targets. Torch-family process.5Writes results/robustness_b/predictions_robustness.parquet and summary.csv."""6from __future__ import annotations78import sys9import time10from pathlib import Path1112sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))13import numpy as np14import polars as pl1516from wp23 import config as C17from wp23.data import month_add18from wp23.models_b import AR2, Momentum, NoChange, RidgePanel, TAUS19from wp23.models_b_torch import DSGFN, QANSS, TFTLite2021OUT = C.RESULTS / "robustness_b"; OUT.mkdir(parents=True, exist_ok=True)222324def run(panel, models, tag, h=3, window=None, min_n=None, winsor=False, refit=3):25    rows = []26    ph = panel.filter(pl.col("h") == h)27    if min_n: ph = ph.filter(pl.col("n_last") >= min_n)28    if winsor:29        q1, q99 = ph["target"].quantile(0.01), ph["target"].quantile(0.99); ph = ph.with_columns(pl.col("target").clip(q1, q99))30    origins = sorted(ph["origin"].unique().to_list())31    for k in range(0, len(origins), refit):32        o = origins[k]; test_o = origins[k:k + refit]33        tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)34        if window: tr = tr.filter(pl.col("origin") > month_add(o, -window))35        if tr["origin"].n_unique() < 10: continue36        te = ph.filter(pl.col("origin").is_in(test_o))37        for m in models:38            t0 = time.time()39            try: m.fit(tr); p = m.predict(te); Q = m.predict_quantiles(te)40            except Exception as ex: print("FAILED", tag, o, m.name, repr(ex), flush=True); continue41            rows.append(te.select(["class", "level", "geo_id", "origin", "h", "target", "n_last"]).with_columns(pl.lit(tag).alias("variant"), pl.lit(m.name).alias("model"), pl.Series("pred", np.asarray(p, float)), pl.Series("q05", np.asarray(Q[:, 0], float)), pl.Series("q95", np.asarray(Q[:, 4], float))))42            print(tag, o, m.name, f"{time.time() - t0:.1f}s", flush=True)43    return pl.concat(rows) if rows else None444546def main() -> None:47    panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet"); parts = []48    base = lambda: [NoChange(), Momentum(), AR2(), RidgePanel(), QANSS(), TFTLite(), DSGFN()]49    parts.append(run(panel, base(), "expanding"))50    parts.append(run(panel, base(), "rolling_24", window=24))51    parts.append(run(panel, base(), "min_n_10", min_n=10))52    parts.append(run(panel, base(), "winsorised", winsor=True))53    for s in (1, 2):54        parts.append(run(panel, [QANSS(seed=s), TFTLite(seed=s), DSGFN(seed=s)], f"seed_{s}"))55    parts.append(run(panel, [DSGFN(bandwidth_km=30.0), DSGFN(bandwidth_km=120.0)], "bandwidth"))56    # label the bandwidth variants57    out = []58    for p in parts:59        if p is not None: out.append(p)60    allp = pl.concat(out); allp.write_parquet(OUT / "predictions_robustness.parquet")61    s = allp.with_columns((pl.col("pred") - pl.col("target")).abs().alias("ae"), ((pl.col("target") >= pl.col("q05")) & (pl.col("target") <= pl.col("q95"))).cast(pl.Float64).alias("cov90")).group_by(["variant", "model", "level"]).agg(pl.len().alias("n"), pl.col("ae").mean().alias("mae"), pl.col("cov90").mean().alias("cov90")).sort(["level", "variant", "mae"])62    s.write_csv(OUT / "summary.csv"); pl.Config.set_tbl_rows(100); print(s.filter(pl.col("level") == "municipality")); print("done")636465if __name__ == "__main__":66    main()67