1"""Chronological, walk-forward and firm-held-out splits with purge and embargo (RESEARCH_DESIGN §6).23An observation is identified by (entity_id, fiscal_period_end, cutoff, A_q). The *target* is revealed at A_q, so a4training observation is purged if its A_q falls after (split_start - embargo); test observations start at split_start.5"""6from __future__ import annotations78import datetime as dt9import hashlib1011import polars as pl121314def _q_start(d: dt.date) -> dt.date:15 return dt.date(d.year, 3 * ((d.month - 1) // 3) + 1, 1)161718def _shift_quarters(d: dt.date, n: int) -> dt.date:19 m = d.month - 1 + 3 * n20 return dt.date(d.year + m // 12, m % 12 + 1, 1)212223def purge_train(df: pl.DataFrame, boundary: dt.date, embargo_quarters: int = 1) -> pl.DataFrame:24 """Keep training rows whose outcome was revealed before boundary - embargo (quarter-aligned)."""25 limit = _shift_quarters(_q_start(boundary), -embargo_quarters)26 return df.filter(pl.col("A_q") < limit)272829def chronological(df: pl.DataFrame, cfg: dict) -> dict[str, pl.DataFrame]:30 tr0, tr1 = [dt.date.fromisoformat(str(x)) for x in cfg["train"]]31 va0, va1 = [dt.date.fromisoformat(str(x)) for x in cfg["validation"]]32 te0, te1 = [dt.date.fromisoformat(str(x)) for x in cfg["test"]]33 emb = cfg["walk_forward"]["embargo_quarters"]34 train = purge_train(df.filter((pl.col("cutoff") >= tr0) & (pl.col("cutoff") <= tr1)), va0, emb)35 valid = purge_train(df.filter((pl.col("cutoff") >= va0) & (pl.col("cutoff") <= va1)), te0, emb)36 test = df.filter((pl.col("cutoff") >= te0) & (pl.col("cutoff") <= te1))37 return {"train": train, "validation": valid, "test": test}383940def walk_forward(df: pl.DataFrame, cfg: dict, last_year: int) -> list[tuple[int, pl.DataFrame, pl.DataFrame]]:41 """Expanding annual folds: train on cutoffs < Jan 1 of year y (purged, embargo), test on cutoffs in year y."""42 out = []43 emb = cfg["walk_forward"]["embargo_quarters"]44 for y in range(cfg["walk_forward"]["first_test_year"], last_year + 1):45 b = dt.date(y, 1, 1)46 train = purge_train(df.filter(pl.col("cutoff") < b), b, emb)47 test = df.filter((pl.col("cutoff") >= b) & (pl.col("cutoff") < dt.date(y + 1, 1, 1)))48 if train.height and test.height:49 out.append((y, train, test))50 return out515253def firm_holdout_mask(entity_ids: pl.Series, fraction: float = 0.2, salt: str = "wp39") -> pl.Series:54 def h(e: str) -> bool:55 return int(hashlib.sha256(f"{salt}:{e}".encode()).hexdigest()[:8], 16) / 0xFFFFFFFF < fraction56 return pl.Series([h(e) for e in entity_ids.to_list()], dtype=pl.Boolean)575859def overlap_check(train: pl.DataFrame, test: pl.DataFrame) -> dict:60 """Diagnostics used by tests: no training outcome revealed on/after the first test cutoff; no shared observation ids."""61 first_test_cutoff = test["cutoff"].min()62 leak = train.filter(pl.col("A_q") >= first_test_cutoff).height63 key = ["entity_id", "fiscal_period_end", "h"] if "h" in train.columns else ["entity_id", "fiscal_period_end"]64 shared = train.select(key).join(test.select(key), on=key, how="inner").height65 return {"train_outcomes_revealed_in_test": leak, "shared_observations": shared}66