1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Shared helpers: frame coercion, truth join, derived strata, arm selection.34Prediction record contract (docs/architecture.md §1):5 tx_id, split, arm, committee, aggregation, point, q05, q25, q50, q75, q95, cost_usd, latency_s, tokens, run_id6Truth contract (data/processed/analytic.parquet):7 tx_id, price, date, region_group, region_name, mun_code, property_class, subtype, urban_class, price_tercile,8 year_built, floor_area_m2 (+ optional n_sales_3km_12m)9"""10from __future__ import annotations1112from typing import Any1314import numpy as np15import pandas as pd1617try: # polars is a hard dependency of the package, but the evaluation code only needs it for coercion18 import polars as pl19except ImportError: # pragma: no cover20 pl = None # type: ignore[assignment]2122PRED_COLUMNS = ["tx_id", "split", "arm", "committee", "aggregation", "point", "q05", "q25", "q50", "q75", "q95",23 "cost_usd", "latency_s", "tokens", "run_id"]24QUANTILE_COLUMNS = ["q05", "q25", "q50", "q75", "q95"]25TAUS = np.array([0.05, 0.25, 0.50, 0.75, 0.95])26ARM_KEY = ["arm", "committee", "aggregation"]27TRUTH_REQUIRED = ["tx_id", "price"]28TRUTH_OPTIONAL = ["date", "region_group", "region_name", "mun_code", "property_class", "subtype", "urban_class",29 "price_tercile", "year_built", "floor_area_m2", "n_sales_3km_12m", "split"]3031# Region groups used in the paper (preregistration §Sample). Any region not listed maps to "other".32REGION_GROUP_MAP = {33 "Montréal": "Montréal", "Montreal": "Montréal", "Laval": "Montréal",34 "Capitale-Nationale": "Québec", "Québec": "Québec", "Quebec": "Québec",35 "Outaouais": "Outaouais",36}37REGION_GROUP_ORDER = ["Montréal", "Québec", "Outaouais", "other"]383940def to_pandas(df: Any) -> pd.DataFrame:41 """Accept a polars or pandas frame and return a pandas copy."""42 if pl is not None and isinstance(df, pl.DataFrame):43 return df.to_pandas()44 if isinstance(df, pd.DataFrame):45 return df.copy()46 raise TypeError(f"expected a polars or pandas DataFrame, got {type(df)!r}")474849def validate_predictions(preds: pd.DataFrame) -> pd.DataFrame:50 """Check the prediction contract; fill optional cost columns with NaN when absent."""51 missing = [c for c in ["tx_id", "arm", "point"] if c not in preds.columns]52 if missing:53 raise ValueError(f"prediction table is missing required columns {missing}")54 out = preds.copy()55 for c in ["committee", "aggregation", "split", "run_id"]:56 if c not in out.columns:57 out[c] = "none" if c in {"committee"} else ("model" if c == "aggregation" else "unknown")58 for c in QUANTILE_COLUMNS + ["cost_usd", "latency_s", "tokens"]:59 if c not in out.columns:60 out[c] = np.nan61 out["tx_id"] = out["tx_id"].astype(str)62 out["point"] = out["point"].astype(float)63 for c in QUANTILE_COLUMNS + ["cost_usd", "latency_s"]:64 out[c] = pd.to_numeric(out[c], errors="coerce").astype(float)65 out["tokens"] = pd.to_numeric(out["tokens"], errors="coerce")66 if (out["point"] <= 0).any():67 raise ValueError("point estimates must be strictly positive (log error undefined)")68 return out697071def derive_region_group(region_name: pd.Series) -> pd.Series:72 return region_name.map(lambda r: REGION_GROUP_MAP.get(str(r), "other")).astype("string")737475def derive_price_tercile(price: pd.Series) -> pd.Series:76 """Tercile of the sale price over the supplied truth table (labels low/mid/high)."""77 q1, q2 = np.nanquantile(price.astype(float), [1 / 3, 2 / 3])78 return pd.Series(np.where(price <= q1, "low", np.where(price <= q2, "mid", "high")), index=price.index, dtype="string")798081def prepare_truth(truth: Any) -> pd.DataFrame:82 """Coerce the truth table, derive `region_group` and `price_tercile` when absent, parse `date`."""83 t = to_pandas(truth)84 missing = [c for c in TRUTH_REQUIRED if c not in t.columns]85 if missing:86 raise ValueError(f"truth table is missing required columns {missing}")87 t["tx_id"] = t["tx_id"].astype(str)88 t["price"] = t["price"].astype(float)89 if (t["price"] <= 0).any():90 raise ValueError("truth prices must be strictly positive")91 if "region_group" not in t.columns and "region_name" in t.columns:92 t["region_group"] = derive_region_group(t["region_name"])93 if "price_tercile" not in t.columns:94 t["price_tercile"] = derive_price_tercile(t["price"])95 if "date" in t.columns:96 t["date"] = pd.to_datetime(t["date"])97 if t["tx_id"].duplicated().any():98 raise ValueError("truth table has duplicated tx_id")99 return t100101102def join_truth(preds: Any, truth: Any, *, how: str = "inner") -> pd.DataFrame:103 """Join predictions to truth on tx_id (truth columns suffixed only on collision: `split` → `split_truth`)."""104 p = validate_predictions(to_pandas(preds))105 t = prepare_truth(truth)106 keep = [c for c in TRUTH_REQUIRED + TRUTH_OPTIONAL if c in t.columns]107 t = t[keep].rename(columns={"split": "split_truth"})108 joined = p.merge(t, on="tx_id", how=how, validate="many_to_one")109 if how == "inner" and len(joined) < len(p):110 joined.attrs["n_unmatched"] = len(p) - len(joined)111 return joined112113114def arm_filter(preds: pd.DataFrame, spec: str | dict[str, Any]) -> pd.DataFrame:115 """Select the rows of one arm. `spec` is an arm name or a dict with keys among arm/committee/aggregation/run_id."""116 if isinstance(spec, str):117 spec = {"arm": spec}118 mask = np.ones(len(preds), dtype=bool)119 for k, v in spec.items():120 mask &= (preds[k].astype(str) == str(v)).to_numpy()121 sub = preds.loc[mask]122 if sub.empty:123 raise ValueError(f"no predictions match {spec}")124 if sub["tx_id"].duplicated().any():125 combos = sub[ARM_KEY + ["run_id"]].drop_duplicates().to_dict("records")126 raise ValueError(f"arm spec {spec} is ambiguous (several rows per tx_id); disambiguate with committee/aggregation/run_id: {combos[:6]}")127 return sub128129130def arm_label(spec: str | dict[str, Any]) -> str:131 if isinstance(spec, str):132 return spec133 return "/".join(str(spec[k]) for k in ["arm", "committee", "aggregation", "run_id"] if k in spec)134135136def block_ids(df: pd.DataFrame) -> pd.Series:137 """Block = mun_code × calendar month when both are present, else mun_code, else tx_id (i.i.d. fallback)."""138 if "mun_code" in df.columns and df["mun_code"].notna().any():139 mun = df["mun_code"].astype(str)140 if "date" in df.columns and df["date"].notna().any():141 month = pd.to_datetime(df["date"]).dt.strftime("%Y-%m")142 return (mun + "|" + month).rename("block")143 return mun.rename("block")144 return df["tx_id"].astype(str).rename("block")145