1#!/usr/bin/env python2"""Phase 2d/3a: cohorts (event-time per horizon + fixed calendar), PIT joins, exclusion cascade, base rates."""3from __future__ import annotations45import datetime as dt6import json7import sys8from pathlib import Path910import polars as pl1112sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))13from wp39.pipeline import Paths, get_logger, load_config # noqa: E40214from wp39.provenance import now_iso, write_json # noqa: E40215from wp39.samples import (apply_admissibility, calendar_cutoffs, event_cutoffs, exclusion_cascade, join_fundamentals, # noqa: E40216 join_market)171819def hist_features(sample: pl.DataFrame, quarters: pl.DataFrame) -> pl.DataFrame:20 """Information age, last available surprise, mean of last 4 surprises, trailing adverse rate — all from quarters21 available at the cutoff (pit_qidx and earlier)."""22 qq = quarters.select(["entity_id", "qidx", "S_opinc", "y_adverse_opinc", "A_q"]).sort(["entity_id", "qidx"])23 qq = qq.with_columns([24 pl.col("S_opinc").alias("last_surprise"),25 pl.col("S_opinc").rolling_mean(window_size=4, min_samples=2).over("entity_id").alias("mean_surprise4"),26 pl.col("y_adverse_opinc").cast(pl.Float64).rolling_mean(window_size=8, min_samples=2).over("entity_id").alias("hist_adverse_rate"),27 pl.col("A_q").alias("A_pit"),28 ]).select(["entity_id", pl.col("qidx").alias("pit_qidx"), "last_surprise", "mean_surprise4", "hist_adverse_rate", "A_pit"])29 s = sample.join(qq, on=["entity_id", "pit_qidx"], how="left")30 return s.with_columns((pl.col("cutoff") - pl.col("A_pit")).dt.total_days().alias("info_age_td"))313233def build(design: str, cfg: dict, P: Paths, quarters: pl.DataFrame, windows: pl.DataFrame, sym: pl.DataFrame, log) -> pl.DataFrame:34 if design == "event":35 cut = event_cutoffs(quarters, cfg["cutoffs"]["event_time_horizons_td"])36 else:37 s0, s1 = dt.date.fromisoformat(str(cfg["sample"]["start"])), dt.date.fromisoformat(str(cfg["sample"]["end"]))38 cut = calendar_cutoffs(quarters, s0, s1)39 cut = apply_admissibility(cut)40 log.info("%s cutoffs after admissibility: %d", design, cut.height)41 j = join_market(cut, windows, sym)42 j = join_fundamentals(j, quarters)43 # attach the outcome of quarter q (the row keyed by entity/fiscal_period_end)44 ycols = ["S_opinc", "S_netinc", "S_ocf", "S_rev", "ocf_innov", "roa_yoy", "y_adverse_opinc", "y_adverse_netinc", "y_roa_decline",45 "y_ocf_negative", "y_liquidity_deterioration", "y_eps_miss", "y_abs_surprise", "assets_tot_lag1", "_gap_ok8",46 "release_time", "eps_estimate"]47 j = j.join(quarters.select(["entity_id", "fiscal_period_end"] + ycols), on=["entity_id", "fiscal_period_end"], how="left")48 j = hist_features(j, quarters)49 sample, cascade = exclusion_cascade(j, cfg["sample"]["min_price_usd"], cfg["sample"]["min_valid_days_in_window"])50 cascade.write_csv(P.results / f"table_sample_cascade_{design}.csv")51 log.info("%s sample: %d rows; cascade=%s", design, sample.height, cascade.to_dicts())52 return sample535455def main() -> int:56 P = Paths.default()57 log = get_logger("04_labels", P.logs / "04_labels.log")58 cfg = load_config("research")59 quarters = pl.read_parquet(P.processed / "quarters.parquet")60 uni = pl.read_parquet(P.processed / "universe.parquet").filter(pl.col("eligible"))61 sym = uni.select(["entity_id", "ticker", "sic_division", "sic"]).rename({"ticker": "symbol"})62 wins = [p for p in (P.cache / "windows").glob("*.parquet") if p.stem in set(sym["symbol"].to_list())]63 windows = pl.concat([pl.read_parquet(p) for p in wins], how="diagonal_relaxed")64 log.info("windows: %d rows, %d symbols", windows.height, len(wins))65 quarters = quarters.filter(pl.col("entity_id").is_in(sym["entity_id"].to_list()))66 out = {}67 for design in ("event", "calendar"):68 s = build(design, cfg, P, quarters, windows, sym, log)69 s.write_parquet(P.processed / f"sample_{design}.parquet")70 out[design] = {71 "rows": s.height, "firms": s["entity_id"].n_unique(), "quarters": s.select(["entity_id", "fiscal_period_end"]).unique().height,72 "by_h": s.group_by("h").agg([pl.len().alias("n"), pl.col("y_adverse_opinc").mean().alias("base_rate")]).sort("h").to_dicts() if design == "event" else None,73 "base_rates": {y: float(s[y].mean()) for y in ("y_adverse_opinc", "y_adverse_netinc", "y_roa_decline", "y_ocf_negative",74 "y_liquidity_deterioration", "y_abs_surprise")},75 "years": s.with_columns(pl.col("cutoff").dt.year().alias("y")).group_by("y").len().sort("y").to_dicts(),76 }77 write_json(P.results / "sample_audit.json", {"built_at": now_iso(), **out})78 print(json.dumps(out, indent=1, default=str))79 return 0808182if __name__ == "__main__":83 raise SystemExit(main())84