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.
97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
"""Build a lightweight bucket-prior ranking model from snapshot labels."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
from libs.backtest.ranking_models import _resolve_feature_key
|
|
|
|
|
|
FEATURE_SPECS = [
|
|
("event_type", 0.08),
|
|
("direction_guidance_combo", 0.27),
|
|
("reaction_bucket", 0.22),
|
|
("close_bucket", 0.15),
|
|
("volume_bucket", 0.12),
|
|
("gap_bucket", 0.08),
|
|
("document_bucket", 0.05),
|
|
("confidence_bucket", 0.03),
|
|
]
|
|
|
|
|
|
def build_model(
|
|
snapshot_path: Path,
|
|
cutoff_event_date: str,
|
|
min_bucket_count: int,
|
|
) -> dict:
|
|
df = pd.read_parquet(snapshot_path)
|
|
df["event_date"] = pd.to_datetime(df["event_date"]).dt.date
|
|
df["reaction_date"] = pd.to_datetime(df["reaction_date"], errors="coerce").dt.date
|
|
cutoff_date = pd.to_datetime(cutoff_event_date).date()
|
|
df = df[
|
|
(df["event_date"] <= cutoff_date)
|
|
& (df["reaction_date"] == df["event_date"])
|
|
& (df["reaction_day_return"] > 0)
|
|
& df["event_type"].isin(["earnings_release", "guidance_update"])
|
|
& df["fwd_return_20d"].notna()
|
|
].copy()
|
|
global_mean = float(df["fwd_return_20d"].mean()) if not df.empty else 0.0
|
|
|
|
features = []
|
|
for name, weight in FEATURE_SPECS:
|
|
bucket_map: dict[str, list[float]] = {}
|
|
for row in df.to_dict("records"):
|
|
key = _resolve_feature_key(name, row)
|
|
if key is None:
|
|
continue
|
|
bucket_map.setdefault(key, []).append(float(row["fwd_return_20d"]))
|
|
values = {
|
|
key: sum(values) / len(values)
|
|
for key, values in bucket_map.items()
|
|
if len(values) >= min_bucket_count
|
|
}
|
|
features.append(
|
|
{
|
|
"name": name,
|
|
"weight": weight,
|
|
"values": values,
|
|
}
|
|
)
|
|
|
|
return {
|
|
"model_type": "bucket_blend_v1",
|
|
"target": "fwd_return_20d",
|
|
"snapshot_id": snapshot_path.parent.name,
|
|
"cutoff_event_date": cutoff_event_date,
|
|
"min_bucket_count": min_bucket_count,
|
|
"training_rows": int(len(df)),
|
|
"global_mean": global_mean,
|
|
"features": features,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Build a bucket-prior ranking model")
|
|
parser.add_argument("--snapshot", required=True, help="Path to the parquet snapshot")
|
|
parser.add_argument("--output", required=True, help="Output JSON path")
|
|
parser.add_argument("--cutoff-event-date", default="2024-12-31")
|
|
parser.add_argument("--min-bucket-count", type=int, default=8)
|
|
args = parser.parse_args()
|
|
|
|
model = build_model(
|
|
snapshot_path=Path(args.snapshot),
|
|
cutoff_event_date=args.cutoff_event_date,
|
|
min_bucket_count=args.min_bucket_count,
|
|
)
|
|
output = Path(args.output)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
output.write_text(json.dumps(model, indent=2))
|
|
print(output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|