1"""Cohort construction: cutoffs, window aggregation of session features, PIT fundamentals join, exclusions.23Event-time: c = A_q - h trading days (h in horizons). Fixed-calendar: c = month-end sessions, target = next release.4All joins are explicit as-of joins on `session_date <= c` (bars) and `available_date <= c` (fundamentals).5"""6from __future__ import annotations78import datetime as dt910import numpy as np11import polars as pl1213from .calendar import session_array, shift_trading_days, month_end_sessions14from .features import FREQS, INTRADAY_FAMILIES15from .pit_fundamentals import METRICS, pit_feature_block1617WINDOW_MEAN = ["rv", "bpv", "rj", "rsneg_share", "sjv", "rskew", "rkurt", "max_share", "ac1", "vol_first30", "vol_last30",18 "vol_gini", "down_vol_share", "present_share", "rv_mktdown_share"]192021def complete_session_index(feat: pl.DataFrame) -> pl.DataFrame:22 """Reindex a symbol's session features to every XNYS session between its first and last session (rows for missing23 sessions have null features) so that rolling windows are measured in trading days."""24 arr = session_array()25 d0, d1 = feat["session_date"].min(), feat["session_date"].max()26 s = arr[(arr >= np.datetime64(d0)) & (arr <= np.datetime64(d1))]27 idx = pl.DataFrame({"session_date": pl.Series(s.astype("datetime64[D]")).cast(pl.Date)})28 return idx.join(feat, on="session_date", how="left").sort("session_date")293031def rolling_window_features(feat: pl.DataFrame, short: int = 20, long: int = 60, min_valid: int = 15) -> pl.DataFrame:32 """Rolling (trailing, inclusive) aggregates per session; the row at session c is the feature vector for cutoff c."""33 f = complete_session_index(feat)34 exprs = []35 valid = pl.col("f1m_n_bars").is_not_null() & (pl.col("f1m_n_bars") >= 100)36 f = f.with_columns(valid.cast(pl.Int8).alias("_valid"))37 exprs.append(pl.col("_valid").rolling_sum(window_size=short, min_samples=1).alias("w_valid_days"))38 # intraday families: window means (RV-type quantities averaged -> per-day average, same across frequencies)39 for name in FREQS:40 p = f"f{name}_"41 for fam in WINDOW_MEAN:42 exprs.append(pl.col(p + fam).rolling_mean(window_size=short, min_samples=min_valid).alias(f"w{name}_{fam}"))43 exprs.append(pl.col(p + "drift").rolling_sum(window_size=short, min_samples=min_valid).alias(f"w{name}_drift"))44 exprs.append((pl.col(p + "sxy").rolling_sum(window_size=short, min_samples=min_valid) /45 (pl.col(p + "syy").rolling_sum(window_size=short, min_samples=min_valid) + 1e-18)).alias(f"w{name}_beta"))46 exprs.append(pl.col(p + "rv").rolling_std(window_size=short, min_samples=min_valid).alias(f"w{name}_volofvol"))47 # variance ratio: mean intraday RV vs variance of daily close-to-close returns over the window48 exprs.append((pl.col(p + "rv").rolling_mean(window_size=short, min_samples=min_valid) /49 (pl.col("d_ret_cc").rolling_var(window_size=short, min_samples=min_valid) + 1e-18)).alias(f"w{name}_vratio"))50 # daily block51 r = pl.col("d_ret_cc")52 exprs += [53 ((1.0 + r).log().rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_ret20"),54 ((1.0 + r).log().rolling_sum(window_size=long, min_samples=3 * min_valid)).alias("wd_ret60"),55 r.rolling_std(window_size=short, min_samples=min_valid).alias("wd_vol20"),56 r.rolling_std(window_size=long, min_samples=3 * min_valid).alias("wd_vol60"),57 pl.col("d_parkinson").rolling_mean(window_size=short, min_samples=min_valid).sqrt().alias("wd_parkinson20"),58 r.rolling_skew(window_size=long).alias("wd_skew60"),59 r.rolling_kurtosis(window_size=long).alias("wd_kurt60"),60 (pl.col("d_ret_co").rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_overnight20"),61 (pl.col("d_ret_oc").rolling_sum(window_size=short, min_samples=min_valid)).alias("wd_intraday20"),62 (pl.col("d_volume").rolling_mean(window_size=short, min_samples=min_valid) /63 (pl.col("d_volume").rolling_median(window_size=long, min_samples=3 * min_valid) + 1.0)).alias("wd_volratio20"),64 (pl.col("d_volume").rolling_mean(window_size=5, min_samples=3) /65 (pl.col("d_volume").rolling_median(window_size=long, min_samples=3 * min_valid) + 1.0)).alias("wd_volratio5"),66 ((r.abs() / (pl.col("d_dollar_vol") + 1.0)) * 1e6).rolling_mean(window_size=long, min_samples=3 * min_valid).alias("wd_amihud60"),67 (r == 0).cast(pl.Float64).rolling_mean(window_size=long, min_samples=3 * min_valid).alias("wd_zeroret60"),68 pl.col("d_dollar_vol").rolling_mean(window_size=long, min_samples=3 * min_valid).log1p().alias("wd_log_dollar_vol60"),69 pl.col("d_close").alias("wd_close"),70 pl.col("f1m_n_bars").rolling_mean(window_size=short, min_samples=min_valid).alias("wd_bars20"),71 ]72 f = f.with_columns(exprs)73 # max drawdown over 60 sessions from cumulative log returns74 f = f.with_columns((1.0 + r).log().fill_null(0.0).cum_sum().alias("_cum"))75 f = f.with_columns((pl.col("_cum") - pl.col("_cum").rolling_max(window_size=long, min_samples=3 * min_valid)).rolling_min(76 window_size=long, min_samples=3 * min_valid).alias("wd_maxdd60"))77 # daily beta / residual vol vs SPY daily (from 1m family sums at daily horizon are not available -> use 5m beta proxy)78 cols = ["session_date", "symbol", "w_valid_days"] + [c for c in f.columns if c.startswith("w1m_") or c.startswith("w5m_")79 or c.startswith("w30m_") or c.startswith("wd_")]80 return f.select(cols)818283def event_cutoffs(quarters: pl.DataFrame, horizons: list[int]) -> pl.DataFrame:84 """One row per (entity, quarter, h) with cutoff c = A_q - h trading days (requires release_match == ok)."""85 q = quarters.filter(pl.col("release_match") == "ok").select(86 ["entity_id", "fiscal_period_end", "A_q", "filing_date", "available_date", "qidx"])87 rows = []88 a_prev = q.with_columns(pl.col("A_q").shift(1).over("entity_id").alias("A_prev"))89 for r in a_prev.iter_rows(named=True):90 for h in horizons:91 try:92 c = shift_trading_days(r["A_q"], -h)93 except ValueError:94 continue95 rows.append({"entity_id": r["entity_id"], "fiscal_period_end": r["fiscal_period_end"], "h": h, "cutoff": c,96 "A_q": r["A_q"], "F_q": r["filing_date"], "A_prev": r["A_prev"], "qidx": r["qidx"]})97 return pl.DataFrame(rows)9899100def calendar_cutoffs(quarters: pl.DataFrame, start: dt.date, end: dt.date, max_td: int = 65) -> pl.DataFrame:101 """Month-end cutoffs; target = next release after c within max_td trading days."""102 me = pl.DataFrame({"cutoff": month_end_sessions(start, end)})103 q = quarters.filter(pl.col("release_match") == "ok").select(["entity_id", "fiscal_period_end", "A_q", "filing_date", "qidx"])104 ents = q.select("entity_id").unique()105 grid = ents.join(me, how="cross").sort(["entity_id", "cutoff"])106 qq = q.sort(["entity_id", "A_q"]).with_columns(pl.col("A_q").alias("_a"))107 j = grid.join_asof(qq, left_on="cutoff", right_on="_a", by="entity_id", strategy="forward")108 j = j.filter(pl.col("A_q") > pl.col("cutoff"))109 arr = session_array()110 c = np.searchsorted(arr, j["cutoff"].to_numpy().astype("datetime64[D]"), side="right")111 a = np.searchsorted(arr, j["A_q"].to_numpy().astype("datetime64[D]"), side="right")112 j = j.with_columns(pl.Series("h", (a - c).astype(int))).filter(pl.col("h") <= max_td)113 return j.rename({"filing_date": "F_q"}).with_columns(pl.lit(None, dtype=pl.Date).alias("A_prev"))114115116def apply_admissibility(cut: pl.DataFrame) -> pl.DataFrame:117 """c < min(A_q, F_q) and c > A_prev (previous reveal already happened)."""118 return cut.filter((pl.col("cutoff") < pl.col("A_q")) & (pl.col("cutoff") < pl.col("F_q")) &119 (pl.col("A_prev").is_null() | (pl.col("cutoff") > pl.col("A_prev"))))120121122def join_market(cut: pl.DataFrame, win: pl.DataFrame, symbol_of: pl.DataFrame) -> pl.DataFrame:123 """As-of join of the rolling window features at session_date <= cutoff (exact session normally)."""124 c = cut.join(symbol_of, on="entity_id", how="left").sort(["symbol", "cutoff"])125 w = win.sort(["symbol", "session_date"])126 j = c.join_asof(w, left_on="cutoff", right_on="session_date", by="symbol", strategy="backward", tolerance="5d")127 return j.rename({"session_date": "market_asof_date"})128129130def join_fundamentals(j: pl.DataFrame, wide: pl.DataFrame) -> pl.DataFrame:131 """Latest available quarter at cutoff + yoy growth from the quarter 4 back (also available by construction)."""132 w = wide.select(["entity_id", "fiscal_period_end", "available_date", "qidx", "A_q"] + METRICS).sort(["entity_id", "available_date"])133 w4 = wide.select(["entity_id", pl.col("qidx").alias("pit_qidx4"), pl.col("rev").alias("rev_l4"), pl.col("op_inc").alias("op_inc_l4"),134 pl.col("shares").alias("shares_l4"), pl.col("available_date").alias("avail_l4")])135 wr = w.rename({c: c for c in METRICS}).rename({"fiscal_period_end": "pit_period_end", "available_date": "pit_available_date",136 "qidx": "pit_qidx", "A_q": "pit_A"})137 out = j.sort(["entity_id", "cutoff"]).join_asof(wr, left_on="cutoff", right_on="pit_available_date", by="entity_id", strategy="backward")138 out = out.with_columns((pl.col("pit_qidx") - 4).alias("pit_qidx4")).join(w4, on=["entity_id", "pit_qidx4"], how="left")139 out = out.with_columns([140 (pl.col("pit_available_date") <= pl.col("cutoff")).alias("pit_ok"),141 ((pl.col("rev") - pl.col("rev_l4")) / (pl.col("rev_l4").abs() + 1e-9)).alias("f_rev_growth"),142 ((pl.col("shares") - pl.col("shares_l4")) / (pl.col("shares_l4").abs() + 1e-9)).alias("f_net_issuance"),143 (pl.col("qidx") - pl.col("pit_qidx")).alias("quarters_ahead"), # 1 = we predict the quarter right after the latest available144 ])145 out = pit_feature_block(out)146 out = out.with_columns([147 (pl.col("wd_close") * pl.col("shares").abs()).log1p().alias("f_log_mcap"),148 ])149 return out150151152def exclusion_cascade(df: pl.DataFrame, min_price: float, min_valid_days: int) -> tuple[pl.DataFrame, pl.DataFrame]:153 """Apply frozen exclusions in order; return (sample, cascade table)."""154 steps = [155 ("cutoff_admissible", pl.lit(True)),156 ("market_window_available", pl.col("market_asof_date").is_not_null() & (pl.col("w_valid_days") >= min_valid_days)),157 ("price_ge_min", pl.col("wd_close") >= min_price),158 ("fundamentals_available", pl.col("pit_ok").fill_null(False) & pl.col("pit_qidx").is_not_null()),159 ("predicting_next_quarter", pl.col("quarters_ahead") == 1),160 ("expectation_defined", pl.col("S_opinc").is_not_null() & pl.col("assets_tot_lag1").is_not_null()),161 ("regular_fiscal_calendar", pl.col("_gap_ok8") == 8),162 ]163 rows, cur = [], df164 for name, cond in steps:165 before = cur.height166 cur = cur.filter(cond)167 rows.append({"step": name, "remaining": cur.height, "dropped": before - cur.height})168 return cur, pl.DataFrame(rows)169170171def feature_columns(df: pl.DataFrame) -> dict[str, list[str]]:172 fund = [c for c in df.columns if c.startswith("f_")] + ["info_age_td", "last_surprise", "mean_surprise4", "hist_adverse_rate"]173 daily = [c for c in df.columns if c.startswith("wd_") and c != "wd_close"]174 return {175 "fundamentals": [c for c in fund if c in df.columns],176 "daily": daily,177 "intraday_30m": [c for c in df.columns if c.startswith("w30m_")],178 "intraday_5m": [c for c in df.columns if c.startswith("w5m_")],179 "intraday_1m": [c for c in df.columns if c.startswith("w1m_")],180 }181