spb/wp23_uqo

1#!/usr/bin/env python32# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca3"""Task B — pseudo-real-time index forecasting on the cell × origin panel.4For each horizon h and each quarterly re-estimation origin o: training rows are panel rows whose target was observable at o5(origin' + h ≤ o); models are fitted once and used for the forecasts issued at origins o, o+1, o+2 (features at those origins).6Usage: python3 scripts/20_task_b.py [--fast] [--models a,b] [--refit-every 3]"""7from __future__ import annotations89import argparse10import sys11import time12from pathlib import Path1314sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))15import numpy as np16import polars as pl1718from wp23 import config as C19from wp23.data import month_add20from wp23.models_b import TAUS, registry2122OUT = C.RESULTS / "task_b"; OUT.mkdir(parents=True, exist_ok=True)232425def main() -> None:26    ap = argparse.ArgumentParser(); ap.add_argument("--fast", action="store_true"); ap.add_argument("--family", type=str, default="trees"); ap.add_argument("--models", type=str, default=""); ap.add_argument("--refit-every", type=int, default=3)27    ap.add_argument("--min-train-origins", type=int, default=10)28    a = ap.parse_args()29    panel = pl.read_parquet(C.PROCESSED / "task_b_panel.parquet")30    models = registry(a.family, a.fast)31    if a.models:32        keep = set(a.models.split(",")); models = [m for m in models if m.name in keep]33    rows, timing = [], []34    for h in C.HORIZONS:35        ph = panel.filter(pl.col("h") == h); origins = sorted(ph["origin"].unique().to_list())36        for k in range(0, len(origins), a.refit_every):37            o = origins[k]; test_origins = origins[k:k + a.refit_every]38            tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)39            if tr["origin"].n_unique() < a.min_train_origins:40                continue41            te = ph.filter(pl.col("origin").is_in(test_origins))42            for m in models:43                t0 = time.time()44                try:45                    m.fit(tr); p = m.predict(te); Q = m.predict_quantiles(te)46                except Exception as ex:47                    print("FAILED", h, o, m.name, repr(ex), flush=True); continue48                d = te.select(["class", "level", "geo_id", "origin", "h", "target", "n_last", "grade_target", "region_code"]).with_columns(pl.lit(m.name).alias("model"), pl.lit(o).alias("fit_origin"), pl.Series("pred", np.asarray(p, dtype=np.float64)))49                for j, t in enumerate(TAUS):50                    d = d.with_columns(pl.Series(f"q{int(t*100):02d}", np.asarray(Q[:, j], dtype=np.float64)))51                rows.append(d); timing.append({"h": h, "fit_origin": o, "model": m.name, "seconds": round(time.time() - t0, 1), "n_train": tr.height, "n_test": te.height})52                print(h, o, m.name, f"{time.time() - t0:.1f}s", f"MAE={np.mean(np.abs(p - te['target'].to_numpy())):.4f}", flush=True)53            pl.concat(rows).write_parquet(OUT / f"predictions_{a.family}.parquet"); pl.DataFrame(timing).write_csv(OUT / f"timing_{a.family}.csv")54    print("done")555657if __name__ == "__main__":58    main()59