UQO Working Paper No. 3 — Hedonic housing price models for the US: parametric, quantile, and machine-learning approaches.
1#!/usr/bin/env python32# Author: Simon-Pierre Boucher — contact@spboucher.ai3#4"""Geographic validation designs on the corrected (v2) design matrix.56November 2026 review, Stage D batch N1 (items WP03-EDIT-01, ADVE-01, ADVE-02,7ADVE-03, ADVE-04, ADVE-16, ADVE-17, ECON-01, ECON-05, ECON-06, EDIT-03, EDIT-06,8THEO-02, THEO-03, THEO-04, THEO-11, THEO-13, ADVE-21). Refactor of905_revision_geo_validation.py and of the ablation of 01_spatial_robustness.py.1011Designs (trimmed sample, 783,872 listings; the v1.1 random split is drawn on12the full index first, then implausible rows are removed):13 random 80/20, seed 42 1 fold14 size_matched random split, geographic training size 1 fold15 ten_state CA NY TX FL OH CO NC WA IL GA held out 1 fold16 gkf_state GroupKFold(5) by state 5 folds17 gkf_zip3 GroupKFold(5) by 3-digit ZIP 5 folds18 gkf_zip5 GroupKFold(5) by 5-digit ZIP 5 folds19 loso leave-one-state-out 51 folds2021Models: OLS (62 regressors + lot indicator), OLS without region dummies,22OLS + state fixed effects (where the test states are in the training set),23XGBoost (full / + 6 missingness indicators / no region / no school / no school24no region), LightGBM (full). For every boosted outer fit the number of rounds25is (i) selected by an inner GroupKFold(3) over the TRAINING groups on the grid26{25..200 by 25, 300..1000 by 100} ("nested"; the test fold is never scored27during selection), (ii) fixed at 100 and (iii) fixed at 1,000 with early28stopping on 10% of the training rows ("fixed" rounds, for comparability with29v1.2). Seeds 42-46 on the outer fits of the full and +missingness variants30(--loso-seeds controls LOSO); one seed elsewhere.3132Factorial ablation (random and ten-state, 5 seeds, fixed rounds): six nested33stages, region x school, stage 3 + tax rate only, stage 3 + lat/lon, stage 3 +34state dummies, stage 4 with the school block replaced by a state identifier,35stage 4 without Walk/Bike/Transit.3637Outputs (results/geo/): metrics.csv, rounds_selection.csv, trees_curves.csv,38predictions/<design>_<model>_<variant>_seed<k>.parquet (out-of-fold39predictions: nested / 100 / 1000 rounds), fold_membership.csv, ablation.csv,40n_corrupt_heldout.csv, rf_geo.csv (--rf), run_meta.json. Resumable: every41(design, fold, model, variant, seed) unit is cached under predictions/parts/.4243Usage: python scripts/05_geo_designs.py [--designs ...] [--seeds 42 43 44 45 46]44 [--loso-seeds 5] [--workers 16] [--resume] [--subsample N] [--skip-ablation] [--rf]45"""4647import argparse48import json49import sys50import time51from pathlib import Path5253import numpy as np54import pandas as pd55from sklearn.metrics import mean_squared_error, r2_score56from sklearn.model_selection import GroupKFold, KFold, train_test_split5758sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))59from wp3 import config, design # noqa: E40260from wp3.models import XGB_PARAMS # noqa: E4026162OUT = config.RESULTS_DIR / "geo"63PRED = OUT / "predictions"64PARTS = PRED / "parts"65SCHOOL = ["avg_school_rating", "school_count", "nearest_school_distance"]66GRID = list(range(25, 201, 25)) + list(range(300, 1001, 100))67CURVE_GRID = [25, 50, 75, 100, 125, 150, 175, 200, 300, 400, 500, 600, 700, 800, 900, 1000]68FINE_GRID = list(range(25, 201, 5)) + list(range(300, 1001, 100))69DESIGNS = ["random", "size_matched", "ten_state", "gkf_state", "gkf_zip3", "gkf_zip5", "loso"]70LGB_PARAMS = dict(n_estimators=1000, max_depth=8, num_leaves=31, learning_rate=0.05, subsample=0.8,71 subsample_freq=0, colsample_bytree=0.8, min_child_samples=20, reg_alpha=0.1,72 reg_lambda=1.0, verbose=-1)73LOG = None747576def log(msg):77 print(msg, flush=True)787980# ----------------------------------------------------------------------------- metrics81def metrics(y, yhat, state=None, group=None):82 out = {"r2": r2_score(y, yhat), "rmse": float(np.sqrt(mean_squared_error(y, yhat))),83 "mdape": float(np.median(np.abs(np.exp(yhat - y) - 1)) * 100),84 "mean_error": float(np.mean(yhat - y))}85 for name, g in [("state", state), ("group", group)]:86 if g is None:87 continue88 gs = pd.Series(g)89 yd = y - pd.Series(y).groupby(gs.values).transform("mean").values90 pdm = yhat - pd.Series(yhat).groupby(gs.values).transform("mean").values91 out[f"r2_within_{name}"] = r2_score(yd, pdm) # oracle level correction92 out[f"between_share_{name}"] = float(1 - yd.var() / y.var())93 # SSE decomposition of the raw prediction: within + between (level) components94 e = yhat - y95 eb = pd.Series(e).groupby(gs.values).transform("mean").values96 sst = float(((y - y.mean()) ** 2).sum())97 out[f"sse_within_share_{name}"] = float(((e - eb) ** 2).sum() / sst)98 out[f"sse_between_share_{name}"] = float((eb ** 2).sum() / sst)99 return out100101102# ----------------------------------------------------------------------------- fits103def fit_xgb(Xtr, ytr, seed, workers, n_estimators=1000):104 from xgboost import XGBRegressor105 p = dict(XGB_PARAMS); p["random_state"] = seed; p["n_jobs"] = workers; p["n_estimators"] = n_estimators106 X_fit, X_eval, y_fit, y_eval = train_test_split(Xtr, ytr, test_size=0.1, random_state=config.RANDOM_STATE)107 m = XGBRegressor(**p).fit(X_fit, y_fit, eval_set=[(X_eval, y_eval)], verbose=0)108 n_trees = m.get_booster().num_boosted_rounds()109 return m, n_trees, (X_eval, y_eval)110111112def fit_lgb(Xtr, ytr, seed, workers):113 import lightgbm as lgb114 p = dict(LGB_PARAMS); p["random_state"] = seed; p["n_jobs"] = workers115 X_fit, X_eval, y_fit, y_eval = train_test_split(Xtr, ytr, test_size=0.1, random_state=config.RANDOM_STATE)116 m = lgb.LGBMRegressor(**p).fit(X_fit, y_fit, eval_set=[(X_eval, y_eval)],117 callbacks=[lgb.early_stopping(50, verbose=False)])118 n_trees = m.booster_.num_trees()119 return m, n_trees, (X_eval, y_eval)120121122def predict_k(m, X, k, n_trees, kind):123 k = int(min(k, n_trees))124 if kind == "xgb":125 return m.predict(X, iteration_range=(0, k))126 return m.predict(X, num_iteration=k)127128129def select_rounds(X, y, tr, groups, seed, workers, kind):130 """Inner GroupKFold(3) over the training groups (KFold if no groups); returns k*, grid table."""131 if groups is None:132 splitter = KFold(n_splits=3, shuffle=True, random_state=seed).split(tr)133 else:134 splitter = GroupKFold(n_splits=3).split(tr, groups=groups[tr])135 scores = []136 for j, (a, b) in enumerate(splitter):137 ia, ib = tr[a], tr[b]138 if kind == "xgb":139 m, n_trees, _ = fit_xgb(X[ia], y[ia], seed, workers)140 else:141 m, n_trees, _ = fit_lgb(X[ia], y[ia], seed, workers)142 scores.append([r2_score(y[ib], predict_k(m, X[ib], k, n_trees, kind)) for k in GRID])143 s = np.array(scores)144 mean = s.mean(axis=0)145 k_star = GRID[int(np.argmax(mean))]146 return k_star, mean, s147148149def ols_fit_predict(Xtr, ytr, Xte):150 """OLS via centred least squares (sklearn LinearRegression): numerically robust to the151 corrupt parking column of the legacy coding (a raw constant + lstsq gives 0.613 instead152 of the stored 0.630 on the random split)."""153 from sklearn.linear_model import LinearRegression154 # float64 is mandatory: on float32 input LinearRegression's lstsq returns rank 1 (R2 = 0.06)155 return LinearRegression().fit(np.asarray(Xtr, np.float64), ytr).predict(np.asarray(Xte, np.float64))156157158def fe_fit_predict(Xtr, ytr, gtr, Xte, gte):159 Xtr = np.asarray(Xtr, np.float64); Xte = np.asarray(Xte, np.float64)160 gtr = pd.Series(gtr); gte = pd.Series(gte)161 Xd = Xtr - pd.DataFrame(Xtr).groupby(gtr.values).transform("mean").values162 yd = ytr - pd.Series(ytr).groupby(gtr.values).transform("mean").values163 beta, *_ = np.linalg.lstsq(Xd, yd, rcond=None)164 alpha = pd.Series(ytr - Xtr @ beta).groupby(gtr.values).mean()165 return Xte @ beta + gte.map(alpha).fillna(alpha.mean()).values166167168# ----------------------------------------------------------------------------- designs169def make_folds(info, idx_full_random, geo_states):170 """Return {design: (group_name, [(train_idx, test_idx, label, groups_array_or_None)])}."""171 n = len(info)172 state = info["address_state"].values; z3 = info["zip3"].values; z5 = info["zip5"].values173 ok = info["implausible"].values == 0174 all_idx = np.arange(n)175 folds = {}176 # random: v1.1 split on the full index, trimmed177 idx_tr, idx_te = idx_full_random178 folds["random"] = ("none", [(idx_tr[ok[idx_tr]], idx_te[ok[idx_te]], "random", None)])179 # size matched: permutation on the full index, first n_geo_train rows train, trimmed180 geo_mask = np.isin(state, geo_states)181 rng = np.random.RandomState(config.RANDOM_STATE)182 perm = rng.permutation(n)183 n_geo_train = int((~geo_mask).sum())184 i_tr, i_te = perm[:n_geo_train], perm[n_geo_train:]185 folds["size_matched"] = ("none", [(i_tr[ok[i_tr]], i_te[ok[i_te]], "size_matched", None)])186 folds["ten_state"] = ("state", [(all_idx[~geo_mask & ok], all_idx[geo_mask & ok], "ten_state", state)])187 for name, g in [("gkf_state", state), ("gkf_zip3", z3), ("gkf_zip5", z5)]:188 tidx = all_idx[ok]189 fl = []190 for k, (a, b) in enumerate(GroupKFold(n_splits=5).split(tidx, groups=g[tidx])):191 fl.append((tidx[a], tidx[b], f"fold{k + 1}", g))192 folds[name] = (name.replace("gkf_", ""), fl)193 fl = []194 for st in sorted(np.unique(state)):195 te = (state == st) & ok196 fl.append((all_idx[~te & ok], all_idx[te], st, state))197 folds["loso"] = ("state", fl)198 return folds199200201def variant_columns(cols, region_cols):202 return {203 "full": [c for c in cols if c not in design.MISS_BLOCK],204 "plus_missing": list(cols),205 "no_region": [c for c in cols if c not in design.MISS_BLOCK and c not in region_cols],206 "no_school": [c for c in cols if c not in design.MISS_BLOCK and c not in SCHOOL],207 "no_school_no_region": [c for c in cols if c not in design.MISS_BLOCK and c not in SCHOOL208 and c not in region_cols],209 }210211212def unit_key(dsg, fold, model, variant, seed):213 return f"{dsg}__{fold}__{model}__{variant}__seed{seed}"214215216def run_unit(X, y, info, cols_all, dsg, fold_label, tr, te, groups, model, variant, vcols, seed,217 workers, state, curves, selection_cache):218 """One (design, fold, model, variant, seed) unit; returns (metric rows, prediction frame)."""219 key = unit_key(dsg, fold_label, model, variant, seed)220 part = PARTS / f"{key}.parquet"221 mpart = PARTS / f"{key}.json"222 if part.exists() and mpart.exists():223 return json.loads(mpart.read_text()), None224 t0 = time.time()225 ci = [cols_all.index(c) for c in vcols]226 Xv = X[:, ci]227 st_te = state[te]228 g_te = groups[te] if groups is not None and dsg.startswith("gkf_zip") else None229 base = {"design": dsg, "fold": fold_label, "model": model, "variant": variant, "seed": seed,230 "n_train": int(len(tr)), "n_test": int(len(te)), "n_features": len(vcols),231 "n_states_test": int(len(np.unique(st_te)))}232 rows, pred = [], {"zpid": info["zpid"].values[te], "y": y[te], "fold": fold_label}233 if model.startswith("ols"):234 if model == "ols_state_fe":235 yhat = fe_fit_predict(Xv[tr], y[tr], state[tr], Xv[te], st_te)236 else:237 yhat = ols_fit_predict(Xv[tr], y[tr], Xv[te])238 rows.append({**base, "rounds_rule": "ols", "rounds": None, **metrics(y[te], yhat, st_te, g_te)})239 pred["pred"] = yhat240 else:241 kind = "xgb" if model == "xgb" else "lgb"242 # nested selection (one seed per fold x variant x model; cached)243 skey = f"{dsg}__{fold_label}__{model}__{variant}"244 if skey not in selection_cache:245 k_star, mean, s = select_rounds(Xv, y, tr, groups, config.RANDOM_STATE, workers, kind)246 selection_cache[skey] = (k_star, mean.tolist(), s.tolist())247 pd.DataFrame({"design": dsg, "fold": fold_label, "model": model, "variant": variant,248 "rounds": GRID, "inner_r2_mean": mean, "inner_r2_fold1": s[0],249 "inner_r2_fold2": s[1], "inner_r2_fold3": s[2], "selected": [k == k_star for k in GRID]}) \250 .to_csv(PARTS / f"{skey}__selection.csv", index=False)251 k_star = selection_cache[skey][0]252 if kind == "xgb":253 m, n_trees, (Xe, ye) = fit_xgb(Xv[tr], y[tr], seed, workers)254 best = int(m.best_iteration) + 1255 else:256 m, n_trees, (Xe, ye) = fit_lgb(Xv[tr], y[tr], seed, workers)257 best = int(m.best_iteration_ or n_trees)258 for rule, k in [("nested", k_star), ("fixed100", 100), ("fixed1000", 1000)]:259 yhat = predict_k(m, Xv[te], k, n_trees, kind)260 rows.append({**base, "rounds_rule": rule, "rounds": int(min(k, n_trees)),261 "selected_rounds": k_star, "best_iteration": best, "n_trees": n_trees,262 **metrics(y[te], yhat, st_te, g_te)})263 pred[f"pred_{rule}"] = yhat264 # curve along the grid (held-out fold and early-stopping set)265 grid = FINE_GRID if (dsg in ("ten_state", "random") and variant in ("full", "plus_missing")) else CURVE_GRID266 for k in grid:267 if k > n_trees:268 continue269 curves.append({"design": dsg, "fold": fold_label, "model": model, "variant": variant, "seed": seed,270 "rounds": k, "r2_test": r2_score(y[te], predict_k(m, Xv[te], k, n_trees, kind)),271 "r2_eval_within_train": r2_score(ye, predict_k(m, Xe, k, n_trees, kind))})272 pd.DataFrame(pred).to_parquet(part, index=False)273 for r in rows:274 r["elapsed_s"] = round(time.time() - t0, 1)275 mpart.write_text(json.dumps(rows))276 r0 = rows[0]277 log(f" {key:60s} r2={r0['r2']:.4f} within={r0.get('r2_within_state', float('nan')):.4f} "278 f"{'k*=' + str(rows[0].get('selected_rounds')) if 'selected_rounds' in r0 else ''} ({time.time() - t0:.0f}s)")279 return rows, None280281282# ----------------------------------------------------------------------------- ablation283def ablation_specs(cols, region_cols, state_dummies_cols, state_id_col, latlon_cols):284 """`cols` must be the 62 regressors + lot_missing (not the extra state/lat-lon columns)."""285 S, L, A, N, M, I = (config.STRUCTURAL_FEATS, config.LOT_FEATS + ["lot_missing"], config.AMENITY_FEATS,286 config.NEIGHBORHOOD_FEATS, config.MARKET_FEATS, config.INTERACTION_FEATS)287 full = [c for c in cols if c not in design.MISS_BLOCK and c not in state_dummies_cols288 and c != state_id_col and c not in latlon_cols]289 assert len(full) == 63, len(full)290 return {291 "1_structural": S,292 "2_plus_lot": S + L,293 "3_plus_amenities": S + L + A,294 "4_plus_neighborhood": S + L + A + N,295 "5_plus_market": S + L + A + N + M,296 "6_full_model": full,297 "4a_stage3_plus_tax_only": S + L + A + ["property_tax_rate"],298 "4b_stage3_plus_latlon": S + L + A + latlon_cols,299 "4c_stage3_plus_state_dummies": S + L + A + state_dummies_cols,300 "4d_stage4_school_to_state_id": S + L + A + [c for c in N if c not in SCHOOL] + [state_id_col],301 "4e_stage4_no_walk_bike_transit": S + L + A + [c for c in N if c not in ("walk_score", "bike_score", "transit_score")],302 "6_full_no_region": [c for c in full if c not in region_cols],303 "6_full_no_school": [c for c in full if c not in SCHOOL],304 "6_full_no_school_no_region": [c for c in full if c not in SCHOOL and c not in region_cols],305 "6_full_plus_latlon": full + latlon_cols,306 }307308309# ----------------------------------------------------------------------------- main310def main():311 ap = argparse.ArgumentParser(description=__doc__)312 ap.add_argument("--designs", nargs="+", default=DESIGNS)313 ap.add_argument("--seeds", type=int, nargs="+", default=[42, 43, 44, 45, 46])314 ap.add_argument("--loso-seeds", type=int, default=5)315 ap.add_argument("--workers", type=int, default=16)316 ap.add_argument("--resume", action="store_true")317 ap.add_argument("--subsample", type=int, default=0, help="debug: rows")318 ap.add_argument("--skip-ablation", action="store_true")319 ap.add_argument("--skip-lgb", action="store_true")320 ap.add_argument("--rf", action="store_true", help="Random-Forest row under ten_state")321 args = ap.parse_args()322 t_start = time.time()323 for d in (OUT, PRED, PARTS):324 d.mkdir(parents=True, exist_ok=True)325326 Xdf, y, info, ex = design.load_design("v2", trimmed=False, standardize=False, missing_block=True)327 if args.subsample:328 keep = np.sort(np.random.RandomState(0).choice(len(y), args.subsample, replace=False))329 Xdf, y, info = Xdf.iloc[keep].reset_index(drop=True), y[keep], info.iloc[keep].reset_index(drop=True)330 cols_all = list(Xdf.columns)331 region_cols = [c for c in cols_all if c.startswith("region_")]332 # extra columns for the ablation333 state = info["address_state"].values334 sd = pd.get_dummies(pd.Series(state), prefix="state", drop_first=True).astype(np.float32)335 Xdf = pd.concat([Xdf, sd, pd.DataFrame({"state_id": pd.factorize(state)[0].astype(np.float32),336 "latitude": info["latitude"].values,337 "longitude": info["longitude"].values})], axis=1)338 cols_all = list(Xdf.columns)339 X = Xdf.values.astype(np.float32)340 idx_tr, idx_te = train_test_split(np.arange(len(y)), test_size=0.2, random_state=config.RANDOM_STATE)341 folds = make_folds(info, (idx_tr, idx_te), config.HOLDOUT_STATES)342 variants = variant_columns([c for c in cols_all if c in ex["meta"]["regressors"] + ["lot_missing"] + design.MISS_BLOCK],343 region_cols)344 log(f"design: {X.shape}, trimmed {int((info['implausible'] == 0).sum()):,}; "345 f"designs {args.designs}; seeds {args.seeds}; workers {args.workers}")346347 # fold membership (EDIT-13)348 fm = []349 for dsg, (gname, fl) in folds.items():350 for tr, te, label, g in fl:351 fm.append({"design": dsg, "fold": label, "n_train": len(tr), "n_test": len(te),352 "n_states_test": len(np.unique(state[te])),353 "states_test": " ".join(sorted(np.unique(state[te]))) if gname == "state" or dsg == "loso" else "",354 "n_groups_test": int(len(np.unique(g[te]))) if g is not None else None})355 pd.DataFrame(fm).to_csv(OUT / "fold_membership.csv", index=False)356357 # THEO-11: legacy untrimmed OLS under ten-state358 Xl, yl, il, _ = design.load_design("legacy", trimmed=False, standardize=False, lot_indicator=False)359 if args.subsample:360 Xl, yl, il = Xl.iloc[keep].reset_index(drop=True), yl[keep], il.iloc[keep].reset_index(drop=True)361 geo = np.isin(il["address_state"].values, config.HOLDOUT_STATES)362 yhat = ols_fit_predict(Xl.values[~geo], yl[~geo], Xl.values[geo])363 err = yhat - yl[geo]364 big = np.abs(err) > 2365 sse = float((err ** 2).sum())366 pd.DataFrame([{"specification": "OLS, v1.1 coding, all observations, ten-state holdout",367 "n_test": int(geo.sum()), "r2": r2_score(yl[geo], yhat),368 "n_abs_error_gt_2": int(big.sum()), "sse_share_abs_error_gt_2": float((err[big] ** 2).sum() / sse),369 "n_abs_error_gt_1": int((np.abs(err) > 1).sum()),370 "n_implausible_in_test": int((il["implausible"].values[geo] == 1).sum()),371 "n_abs_error_gt_2_and_implausible": int((big & (il["implausible"].values[geo] == 1)).sum()),372 "r2_without_abs_error_gt_2": r2_score(yl[geo][~big], yhat[~big])}]) \373 .to_csv(OUT / "n_corrupt_heldout.csv", index=False)374 del Xl375376 # ---------------------------------------------------------------- unit list (priority order)377 units = [] # (design, fold_i, model, variant, seed)378 def add(dsg, model, variant, seeds):379 if dsg not in args.designs:380 return381 for i in range(len(folds[dsg][1])):382 for s in seeds:383 units.append((dsg, i, model, variant, s))384 one = [args.seeds[0]]385 loso_seeds = args.seeds[:args.loso_seeds]386 for dsg in DESIGNS: # OLS everywhere (cheap)387 add(dsg, "ols", "full", one); add(dsg, "ols", "no_region", one)388 if dsg in ("random", "size_matched", "gkf_zip3", "gkf_zip5"):389 add(dsg, "ols_state_fe", "no_region", one)390 for dsg in DESIGNS:391 add(dsg, "xgb", "full", loso_seeds if dsg == "loso" else args.seeds)392 for dsg in DESIGNS:393 add(dsg, "xgb", "plus_missing", loso_seeds if dsg == "loso" else args.seeds)394 for dsg in ("ten_state", "gkf_state", "loso"):395 for v in ("no_region", "no_school", "no_school_no_region"):396 add(dsg, "xgb", v, one)397 for dsg in ("random", "gkf_zip3", "gkf_zip5"):398 add(dsg, "xgb", "no_region", one)399 if not args.skip_lgb:400 for dsg in DESIGNS:401 add(dsg, "lgb", "full", one)402 log(f"{len(units)} units")403404 all_rows, curves, selection_cache = [], [], {}405 for n_done, (dsg, i, model, variant, seed) in enumerate(units):406 gname, fl = folds[dsg]407 tr, te, label, g = fl[i]408 rows, _ = run_unit(X, y, info, cols_all, dsg, label, tr, te, g, model, variant, variants[variant],409 seed, args.workers, state, curves, selection_cache)410 all_rows.extend(rows)411 if n_done % 20 == 0 or n_done == len(units) - 1:412 pd.DataFrame(all_rows).to_csv(OUT / "metrics_partial.csv", index=False)413 if curves:414 pd.DataFrame(curves).to_csv(OUT / "trees_curves_partial.csv", index=False)415 log(f"--- {n_done + 1}/{len(units)} units, {(time.time() - t_start) / 3600:.2f} h")416417 # ---------------------------------------------------------------- assemble418 met = pd.DataFrame(all_rows)419 met.to_csv(OUT / "metrics.csv", index=False)420 sel = [pd.read_csv(p) for p in sorted(PARTS.glob("*__selection.csv"))]421 if sel:422 pd.concat(sel).to_csv(OUT / "rounds_selection.csv", index=False)423 if curves: # cached (--resume) units do not recompute the curves: keep the existing file424 pd.DataFrame(curves).to_csv(OUT / "trees_curves.csv", index=False)425 # out-of-fold prediction files: one per design x model x variant x seed426 for (dsg, model, variant, seed), grp in met.groupby(["design", "model", "variant", "seed"]):427 parts = [pd.read_parquet(PARTS / f"{unit_key(dsg, f, model, variant, seed)}.parquet")428 for f in grp["fold"].unique()]429 pd.concat(parts, ignore_index=True).to_parquet(PRED / f"{dsg}_{model}_{variant}_seed{seed}.parquet", index=False)430 # pooled metrics over folds for the multi-fold designs (out-of-fold predictions)431 pooled = []432 for (dsg, model, variant, seed), grp in met.groupby(["design", "model", "variant", "seed"]):433 if grp["fold"].nunique() < 2:434 continue435 p = pd.read_parquet(PRED / f"{dsg}_{model}_{variant}_seed{seed}.parquet")436 st = info.set_index("zpid").loc[p["zpid"].values, "address_state"].values437 for rule in [c.replace("pred_", "") for c in p.columns if c.startswith("pred")]:438 col = "pred" if rule == "pred" else f"pred_{rule}"439 pooled.append({"design": dsg, "model": model, "variant": variant, "seed": seed,440 "rounds_rule": "ols" if rule == "pred" else rule, "n_test": len(p),441 **metrics(p["y"].values, p[col].values, st)})442 pd.DataFrame(pooled).to_csv(OUT / "metrics_pooled.csv", index=False)443444 # ---------------------------------------------------------------- ablation445 if not args.skip_ablation:446 log("Factorial ablation (random + ten_state, 5 seeds, fixed rounds)")447 specs = ablation_specs(cols_all, region_cols, list(sd.columns), "state_id", ["latitude", "longitude"])448 ab_rows = []449 for dsg in ("random", "ten_state"):450 if dsg not in args.designs:451 continue452 tr, te, label, g = folds[dsg][1][0]453 for name, feats in specs.items():454 ci = [cols_all.index(c) for c in feats]455 for seed in args.seeds:456 key = f"ablation__{dsg}__{name}__seed{seed}"457 cache = PARTS / f"{key}.json"458 if cache.exists():459 ab_rows.extend(json.loads(cache.read_text())); continue460 t0 = time.time()461 m, n_trees, (Xe, ye) = fit_xgb(X[tr][:, ci], y[tr], seed, args.workers)462 rows = []463 for rule, k in [("fixed1000", 1000), ("fixed100", 100)]:464 yhat = predict_k(m, X[te][:, ci], k, n_trees, "xgb")465 rows.append({"design": dsg, "spec": name, "seed": seed, "n_features": len(feats),466 "rounds_rule": rule, "rounds": int(min(k, n_trees)),467 "best_iteration": int(m.best_iteration) + 1,468 **metrics(y[te], yhat, state[te]), "elapsed_s": round(time.time() - t0, 1)})469 cache.write_text(json.dumps(rows)); ab_rows.extend(rows)470 log(f" {key:55s} r2={rows[0]['r2']:.4f} r2@100={rows[1]['r2']:.4f} ({time.time() - t0:.0f}s)")471 pd.DataFrame(ab_rows).to_csv(OUT / "ablation.csv", index=False)472473 # ---------------------------------------------------------------- RF under ten_state (ADVE-21)474 if args.rf and "ten_state" in args.designs and not (OUT / "rf_geo.csv").exists():475 from sklearn.ensemble import RandomForestRegressor476 tr, te, label, g = folds["ten_state"][1][0]477 ci = [cols_all.index(c) for c in variants["full"]]478 t0 = time.time()479 rf = RandomForestRegressor(n_estimators=500, max_features=1 / 3, min_samples_leaf=5, random_state=42,480 n_jobs=args.workers).fit(X[tr][:, ci], y[tr])481 yhat = rf.predict(X[te][:, ci])482 pd.DataFrame([{"design": "ten_state", "model": "rf", "variant": "full", "seed": 42,483 "spec": "Random Forest (re-specified: 500 trees, 1/3 features, leaf 5)",484 **metrics(y[te], yhat, state[te]), "elapsed_s": round(time.time() - t0, 1)}]) \485 .to_csv(OUT / "rf_geo.csv", index=False)486 pd.DataFrame({"zpid": info["zpid"].values[te], "y": y[te], "pred": yhat}).to_parquet(487 PRED / "ten_state_rf_full_seed42.parquet", index=False)488 log(f" RF ten_state r2={r2_score(y[te], yhat):.4f} ({time.time() - t0:.0f}s)")489490 import platform491 meta_path = OUT / "run_meta.json"492 meta = {"node": platform.node(), "python": platform.python_version(), "workers": args.workers,493 "seeds": args.seeds, "loso_seeds": args.loso_seeds, "designs": args.designs, "grid": GRID,494 "n_units": len(units), "elapsed_h": round((time.time() - t_start) / 3600, 3),495 "subsample": args.subsample, "resume": bool(args.resume), "finished": time.strftime("%Y-%m-%dT%H:%M:%S%z")}496 if meta_path.exists(): # never overwrite the record of an earlier (e.g. full) batch: keep it under previous_runs497 try:498 prev = json.loads(meta_path.read_text())499 meta["previous_runs"] = prev.pop("previous_runs", []) + [prev]500 except (OSError, ValueError):501 pass502 meta_path.write_text(json.dumps(meta, indent=2))503 log(f"Done in {(time.time() - t_start) / 3600:.2f} h")504505506if __name__ == "__main__":507 main()508