You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

110 lines
4.2 KiB
Python

#!/usr/bin/env python3
"""Train LightGBM ML scorer for PEAD long, save booster + feature manifest.
Outputs:
- data/ml_models/lgbm_v1_pead_long.txt (LGBM booster, text format)
- data/ml_models/lgbm_v1_pead_long_manifest.json (feature schema + category vocab)
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
from lightgbm import LGBMClassifier
NUMERIC = [
"reaction_day_return", "gap_size", "close_location", "volume_ratio_20d",
"document_quality_score", "parse_confidence_overall", "oneoff_penalty",
"market_cap_proxy", "avg_dollar_volume_20d",
"pre_event_volatility_20d", "pre_event_rsi_14", "pre_event_bb_position",
"pre_event_obv_slope_20d",
"macro_vix", "macro_hy_spread", "macro_t10y2y", "prior_event_fwd5d",
"lm_positive_pct", "lm_negative_pct", "lm_net_sentiment",
"earnings_surprise_pct",
"sue_lag_1_pct", "sue_lag_2_pct", "sue_lag_3_pct",
"sue_hist_mean_4q", "sue_hist_mean_8q", "sue_hist_mean_12q",
"sue_hist_pos_rate_4q", "sue_hist_pos_rate_12q",
"sue_hist_latest_pct", "sue_hist_streak_pos",
"peer_sector_event_count_365d", "peer_sector_surprise_median_365d",
"peer_sector_surprise_mean_365d", "peer_sector_surprise_pos_rate_365d",
"peer_relative_surprise_pct_365d",
"peer_sector_sue_hist_mean_4q_median_365d",
"peer_sector_sue_hist_mean_4q_mean_365d",
"peer_relative_sue_hist_mean_4q_365d",
"peer_sector_sue_hist_pos_rate_4q_mean_365d",
"pre_event_short_ratio", "pre_event_sector_momentum_20d",
"prior_catalyst_count_60d", "prior_catalyst_type_diversity_60d",
"price_vs_sma20", "pre_event_momentum_20d",
]
CAT = ["event_type", "event_direction", "guidance_status", "sector"]
SNAPSHOT = Path("data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_canonical_ftb_fix_v2")
OUT_DIR = Path("data/ml_models")
def load(split: str) -> pd.DataFrame:
schema = pq.read_schema(SNAPSHOT / f"{split}.parquet")
cols = [c for c in NUMERIC + CAT + ["mfe_10d", "mae_10d", "fwd_return_10d"] if c in schema.names]
df = pq.read_table(SNAPSHOT / f"{split}.parquet", columns=cols).to_pandas()
for c in NUMERIC:
if c not in df.columns:
df[c] = np.nan
df[c] = pd.to_numeric(df[c], errors="coerce")
for c in CAT:
if c not in df.columns:
df[c] = "missing"
df[c] = df[c].fillna("missing").astype(str)
df["target"] = ((df["mfe_10d"] - df["mae_10d"]) > 0.08).astype(int)
return df
def main() -> None:
OUT_DIR.mkdir(parents=True, exist_ok=True)
train = load("train")
cat_vocab = {c: sorted(train[c].dropna().unique().tolist()) for c in CAT}
# encode categories as integer codes for stable serialisation
for c in CAT:
codes = {v: i for i, v in enumerate(cat_vocab[c])}
train[c] = train[c].map(codes).fillna(-1).astype("int32")
feats = NUMERIC + CAT
model = LGBMClassifier(
n_estimators=400, learning_rate=0.03, num_leaves=31,
min_child_samples=50, subsample=0.8, colsample_bytree=0.8,
objective="binary", random_state=42, verbosity=-1,
)
model.fit(train[feats], train["target"], categorical_feature=CAT)
booster = model.booster_
model_path = OUT_DIR / "lgbm_v1_pead_long.txt"
booster.save_model(str(model_path))
manifest = {
"numeric_features": NUMERIC,
"categorical_features": CAT,
"categorical_vocab": cat_vocab,
"snapshot": str(SNAPSHOT),
"n_estimators": 400,
"lgbm_version": __import__("lightgbm").__version__,
}
manifest_path = OUT_DIR / "lgbm_v1_pead_long_manifest.json"
manifest_path.write_text(json.dumps(manifest, indent=2))
# quick eval: re-predict on train + valid + test
for split in ["train", "valid", "test"]:
df = load(split)
for c in CAT:
codes = {v: i for i, v in enumerate(cat_vocab[c])}
df[c] = df[c].map(codes).fillna(-1).astype("int32")
prob = model.predict_proba(df[feats])[:, 1]
print(f"{split}: mean={prob.mean():.4f} std={prob.std():.4f} p10={np.percentile(prob,10):.3f} p90={np.percentile(prob,90):.3f}")
print(f"\nSaved booster: {model_path}")
print(f"Saved manifest: {manifest_path}")
if __name__ == "__main__":
main()