1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Chronological partitions, spatial-holdout flag and stratified LLM evaluation subsamples.34* ``assign_splits``: every sale is assigned to train / validation / test / holdout by its sale date using the5 closed windows of ``configs/data.yaml``; sales outside every window get ``split = "unused"`` (reported).6 ``spatial_holdout_region`` flags the regions withheld in the spatial-generalisation experiment.7* ``stratified_subsample``: deterministic stratified random subsample of each evaluation partition8 (strata = region_group × property_class × price_tercile, proportional allocation with largest-remainder9 rounding, seeded draws). Baselines are evaluated on the full partitions; the LLM arms on these subsamples.10"""11from __future__ import annotations1213import hashlib14import json15import time16from datetime import date17from pathlib import Path18from typing import Any1920import numpy as np21import polars as pl2223from ..config import PROCESSED, RESULTS, SEED, SYNTHETIC, load_config2425SPLIT_ORDER = ("train", "validation", "test", "holdout")262728def split_windows(cfg: dict[str, Any]) -> dict[str, tuple[date, date]]:29 out: dict[str, tuple[date, date]] = {}30 for name in SPLIT_ORDER:31 w = cfg["splits"][name]32 out[name] = (date.fromisoformat(w["start"]), date.fromisoformat(w["end"]))33 return out343536def check_windows(cfg: dict[str, Any]) -> None:37 """Windows must be ordered, non-overlapping and contiguous (no sale can fall in two partitions)."""38 w = split_windows(cfg)39 prev_end: date | None = None40 for name in SPLIT_ORDER:41 start, end = w[name]42 if start > end:43 raise ValueError(f"split {name}: start {start} after end {end}")44 if prev_end is not None and start <= prev_end:45 raise ValueError(f"split {name} overlaps the previous window (starts {start} <= {prev_end})")46 prev_end = end474849def assign_splits(analytic: pl.DataFrame, cfg: dict[str, Any] | None = None) -> pl.DataFrame:50 """→ (tx_id, split, spatial_holdout_region)."""51 cfg = cfg or load_config("data")52 check_windows(cfg)53 w = split_windows(cfg)54 expr: pl.Expr = pl.lit("unused")55 for name in reversed(SPLIT_ORDER):56 start, end = w[name]57 expr = pl.when((pl.col("date") >= start) & (pl.col("date") <= end)).then(pl.lit(name)).otherwise(expr)58 regions = list(cfg.get("spatial_holdout", {}).get("regions", []))59 return analytic.select(60 "tx_id",61 expr.alias("split"),62 pl.col("region_name").is_in(regions).alias("spatial_holdout_region"),63 )646566def _stratum_label(cols: list[str]) -> pl.Expr:67 return pl.concat_str([pl.col(c).cast(pl.String) for c in cols], separator="|")686970def proportional_allocation(counts: dict[str, int], n: int) -> dict[str, int]:71 """Largest-remainder (Hamilton) allocation of ``n`` draws across strata proportional to their sizes, capped at72 the stratum size; the shortfall of capped strata is redistributed among the others."""73 keys = sorted(counts)74 total = sum(counts.values())75 if total == 0 or n <= 0:76 return {k: 0 for k in keys}77 n = min(n, total)78 alloc = {k: 0 for k in keys}79 remaining = n80 open_keys = list(keys)81 while remaining > 0 and open_keys:82 pool = sum(counts[k] - alloc[k] for k in open_keys)83 quotas = {k: remaining * (counts[k] - alloc[k]) / pool for k in open_keys}84 floors = {k: int(np.floor(quotas[k])) for k in open_keys}85 left = remaining - sum(floors.values())86 order = sorted(open_keys, key=lambda k: (-(quotas[k] - floors[k]), k))87 for k in order[:left]:88 floors[k] += 189 for k in open_keys:90 alloc[k] += floors[k]91 # cap and reopen92 over = {k: alloc[k] - counts[k] for k in open_keys if alloc[k] > counts[k]}93 for k, o in over.items():94 alloc[k] = counts[k]95 remaining = n - sum(alloc.values())96 open_keys = [k for k in open_keys if alloc[k] < counts[k]]97 return alloc9899100def _derive_seed(seed: int, split: str, stratum: str) -> int:101 h = hashlib.sha256(f"{seed}|{split}|{stratum}".encode()).hexdigest()102 return int(h[:8], 16)103104105def stratified_subsample(106 analytic: pl.DataFrame,107 splits: pl.DataFrame,108 cfg: dict[str, Any] | None = None,109 n_by_split: dict[str, int] | None = None,110 seed: int | None = None,111) -> pl.DataFrame:112 """→ (tx_id, split, stratum, draw_seed). Deterministic given (table, config, seed)."""113 cfg = cfg or load_config("data")114 sub = cfg["llm_eval_subsample"]115 strata_cols = list(sub["strata"])116 n_by_split = n_by_split or {"validation": sub["validation_n"], "test": sub["test_n"], "holdout": sub["holdout_n"]}117 seed = cfg.get("seed", SEED) if seed is None else seed118 df = analytic.select(["tx_id", *strata_cols]).join(splits.select("tx_id", "split"), on="tx_id", how="inner")119 df = df.with_columns(_stratum_label(strata_cols).alias("stratum"))120 frames: list[pl.DataFrame] = []121 for split, n in n_by_split.items():122 part = df.filter(pl.col("split") == split)123 counts = {k: int(v) for k, v in part.group_by("stratum").len().iter_rows()}124 alloc = proportional_allocation(counts, int(n))125 for stratum in sorted(alloc):126 k = alloc[stratum]127 if k == 0:128 continue129 ids = np.array(sorted(part.filter(pl.col("stratum") == stratum)["tx_id"].to_list()))130 s = _derive_seed(seed, split, stratum)131 rng = np.random.default_rng(s)132 chosen = rng.choice(ids, size=k, replace=False)133 frames.append(pl.DataFrame({"tx_id": sorted(chosen.tolist()), "split": split, "stratum": stratum, "draw_seed": s}, schema={"tx_id": pl.String, "split": pl.String, "stratum": pl.String, "draw_seed": pl.Int64}))134 if not frames:135 return pl.DataFrame(schema={"tx_id": pl.String, "split": pl.String, "stratum": pl.String, "draw_seed": pl.Int64})136 out = pl.concat(frames)137 return out.sort(["split", "stratum", "tx_id"])138139140def run(synthetic: bool = False) -> dict[str, Any]:141 cfg = load_config("data")142 t0 = time.perf_counter()143 base = SYNTHETIC if synthetic else PROCESSED144 analytic = pl.read_parquet(base / "analytic.parquet")145 splits = assign_splits(analytic, cfg)146 n_by = {"validation": 40, "test": 60, "holdout": 30} if synthetic else None147 subs = stratified_subsample(analytic, splits, cfg, n_by_split=n_by)148 splits.write_parquet(base / "splits.parquet")149 subs.write_parquet(base / "llm_subsamples.parquet")150 counts = {k: int(v) for k, v in splits.group_by("split").len().iter_rows()}151 sub_counts = {k: int(v) for k, v in subs.group_by("split").len().iter_rows()}152 summary = {153 "synthetic": synthetic,154 "splits": {k: counts.get(k, 0) for k in (*SPLIT_ORDER, "unused")},155 "spatial_holdout_rows": int(splits["spatial_holdout_region"].sum()),156 "spatial_holdout_regions": cfg["spatial_holdout"]["regions"],157 "windows": {k: [a.isoformat(), b.isoformat()] for k, (a, b) in split_windows(cfg).items()},158 "llm_subsamples": sub_counts,159 "llm_subsample_strata": int(subs["stratum"].n_unique()),160 "seed": cfg.get("seed", SEED),161 "seconds": round(time.perf_counter() - t0, 2),162 }163 out = RESULTS / "data" / ("splits_summary_synthetic.json" if synthetic else "splits_summary.json")164 out.parent.mkdir(parents=True, exist_ok=True)165 out.write_text(json.dumps(summary, indent=1, ensure_ascii=False))166 return summary167