1import datetime as dt23import polars as pl45from wp39.splits import chronological, firm_holdout_mask, overlap_check, walk_forward67CFG = {"train": ["2012-01-01", "2018-12-31"], "validation": ["2019-01-01", "2020-12-31"], "test": ["2021-01-01", "2026-10-30"],8 "walk_forward": {"first_test_year": 2016, "embargo_quarters": 1}, "firm_holdout_fraction": 0.2}91011def _panel():12 rows = []13 d = dt.date(2012, 1, 15)14 for i in range(140):15 c = d + dt.timedelta(days=int(i * 36.5))16 rows.append({"entity_id": f"frd:E{i % 7}", "fiscal_period_end": c - dt.timedelta(days=30), "h": 20, "cutoff": c,17 "A_q": c + dt.timedelta(days=28), "y": i % 3 == 0})18 return pl.DataFrame(rows)192021def test_chronological_purge_and_embargo():22 s = chronological(_panel(), CFG)23 assert s["train"]["A_q"].max() < dt.date(2018, 10, 1) # validation starts 2019-01-01, embargo one quarter24 assert s["validation"]["A_q"].max() < dt.date(2020, 10, 1)25 assert overlap_check(s["train"], s["test"])["shared_observations"] == 0262728def test_walk_forward_folds():29 folds = walk_forward(_panel(), CFG, 2025)30 for y, tr, te in folds:31 assert tr["A_q"].max() < dt.date(y - 1, 10, 1)32 assert te["cutoff"].min().year == y333435def test_firm_holdout_is_deterministic():36 e = pl.Series(["frd:A", "frd:B", "frd:C"])37 assert firm_holdout_mask(e).to_list() == firm_holdout_mask(e).to_list()38