1"""Negative controls, placebos and subgroup stress tests (Phase 5)."""2from __future__ import annotations34import numpy as np5import polars as pl678def shuffle_labels_within_quarter(df: pl.DataFrame, y: str, seed: int) -> pl.DataFrame:9 """Permute the label within calendar quarter of the cutoff: any 'gain' left is overfitting/leakage."""10 rng = np.random.default_rng(seed)11 q = df.with_columns((pl.col("cutoff").dt.year() * 4 + (pl.col("cutoff").dt.month() - 1) // 3).alias("_cq"))12 parts = []13 for _, g in q.group_by("_cq", maintain_order=True):14 v = g[y].to_numpy().copy(); rng.shuffle(v)15 parts.append(g.with_columns(pl.Series(y, v)))16 return pl.concat(parts).drop("_cq")171819def subgroup_masks(df: pl.DataFrame) -> dict[str, pl.Expr]:20 """Pre-specified subgroups (terciles are computed on the TRAINING rows and passed as cut-points by the caller)."""21 return {22 "all": pl.lit(True),23 "size_small": pl.col("tercile_size") == 1, "size_large": pl.col("tercile_size") == 3,24 "illiq_high": pl.col("tercile_illiq") == 3, "illiq_low": pl.col("tercile_illiq") == 1,25 "lev_high": pl.col("tercile_lev") == 3, "lev_low": pl.col("tercile_lev") == 1,26 "loss_firms": pl.col("f_loss") == 1, "profit_firms": pl.col("f_loss") == 0,27 "covid_2020": (pl.col("cutoff").dt.year() == 2020), "ex_covid": (pl.col("cutoff").dt.year() != 2020),28 "after_close_release": pl.col("release_time") == "after_close", "before_open_release": pl.col("release_time") == "before_open",29 "sector_finance": pl.col("sic_division") == "H_finance_insurance_realestate",30 "sector_nonfinance": pl.col("sic_division") != "H_finance_insurance_realestate",31 }323334def add_terciles(train: pl.DataFrame, others: list[pl.DataFrame]) -> list[pl.DataFrame]:35 """Tercile cut-points for size, illiquidity and leverage computed on training rows only."""36 cuts = {}37 for name, col in (("size", "f_log_mcap"), ("illiq", "wd_amihud60"), ("lev", "f_leverage")):38 s = train[col].drop_nulls().drop_nans()39 cuts[name] = (float(s.quantile(1 / 3)), float(s.quantile(2 / 3)))40 out = []41 for d in [train] + others:42 for name, col in (("size", "f_log_mcap"), ("illiq", "wd_amihud60"), ("lev", "f_leverage")):43 lo, hi = cuts[name]44 d = d.with_columns(pl.when(pl.col(col) <= lo).then(1).when(pl.col(col) <= hi).then(2).otherwise(3).alias(f"tercile_{name}"))45 out.append(d)46 return out47