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.
88 lines
3.8 KiB
Python
88 lines
3.8 KiB
Python
"""Enrich snapshot with Tier 3 features: OU theta, Gravitational Pull, Market Temperature.
|
|
All computed from price bars (no extra API needed beyond what Tier 2 already fetches).
|
|
|
|
Usage:
|
|
PYTHONUNBUFFERED=1 uv run python3 scripts/enrich_tier3_features.py \
|
|
--input data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_tier2 \
|
|
--output data/parquet/midlarge-liquid-long-v1_bucketfix_full_audit_tier3
|
|
"""
|
|
from __future__ import annotations
|
|
import argparse, json, sys
|
|
from pathlib import Path
|
|
import pyarrow as pa, pyarrow.parquet as pq
|
|
import requests
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from libs.features.market_features import pre_event_ou_theta, pre_event_gravitational_pull, pre_event_market_temperature
|
|
from libs.oracle_client.models import PriceBar
|
|
|
|
ORACLE_URL = "http://localhost:18001"
|
|
FEATURES = ["pre_event_ou_theta_60d", "pre_event_gravitational_pull", "pre_event_market_temperature"]
|
|
|
|
|
|
def fetch_bars(ticker: str, event_date: str) -> list[PriceBar]:
|
|
from datetime import datetime, timedelta
|
|
end_dt = datetime.strptime(event_date, "%Y-%m-%d")
|
|
start_dt = end_dt - timedelta(days=120)
|
|
resp = requests.get(f"{ORACLE_URL}/api/v1/price/data/{ticker}",
|
|
params={"start_date": start_dt.strftime("%Y-%m-%d"), "end_date": event_date}, timeout=30)
|
|
if resp.status_code != 200: return []
|
|
data = resp.json()
|
|
bars = []
|
|
for b in (data.get("bars") or data.get("data") or []):
|
|
try:
|
|
bars.append(PriceBar(date=b["date"], open=float(b.get("open",0)), high=float(b.get("high",0)),
|
|
low=float(b.get("low",0)), close=float(b.get("close",0)), volume=int(b.get("volume",0))))
|
|
except: continue
|
|
return bars
|
|
|
|
|
|
def enrich_split(input_path: Path, output_path: Path):
|
|
table = pq.read_table(input_path)
|
|
n = len(table)
|
|
existing = set(table.column_names)
|
|
if all(f in existing for f in FEATURES):
|
|
nulls = {f: table.column(f).null_count for f in FEATURES}
|
|
if all(v < n * 0.3 for v in nulls.values()):
|
|
print(f" Already enriched, copying"); pq.write_table(table, output_path); return
|
|
|
|
tickers = table.column("ticker").to_pylist()
|
|
dates = table.column("event_date").to_pylist()
|
|
results = {f: [None]*n for f in FEATURES}
|
|
ok = 0
|
|
for i in range(n):
|
|
if i % 200 == 0: print(f" {i}/{n} ({ok} ok)...")
|
|
try:
|
|
bars = fetch_bars(tickers[i], str(dates[i]))
|
|
if bars:
|
|
results["pre_event_ou_theta_60d"][i] = pre_event_ou_theta(bars, str(dates[i]), 60)
|
|
results["pre_event_gravitational_pull"][i] = pre_event_gravitational_pull(bars, str(dates[i]))
|
|
results["pre_event_market_temperature"][i] = pre_event_market_temperature(bars, str(dates[i]))
|
|
if results["pre_event_ou_theta_60d"][i] is not None: ok += 1
|
|
except: pass
|
|
print(f" Done: {ok}/{n}")
|
|
for f in FEATURES:
|
|
arr = pa.array(results[f], type=pa.float64())
|
|
table = table.set_column(table.column_names.index(f), f, arr) if f in existing else table.append_column(f, arr)
|
|
pq.write_table(table, output_path)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", required=True); parser.add_argument("--output", required=True)
|
|
args = parser.parse_args()
|
|
inp, out = Path(args.input), Path(args.output)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
m = inp / "manifest.json"
|
|
if m.exists():
|
|
manifest = json.loads(m.read_text()); manifest["snapshot_id"] = out.name
|
|
(out / "manifest.json").write_text(json.dumps(manifest, indent=2))
|
|
for split in ["train","valid","test"]:
|
|
p = inp / f"{split}.parquet"
|
|
if not p.exists(): continue
|
|
print(f"Enriching {split}...")
|
|
enrich_split(p, out / f"{split}.parquet")
|
|
print("Done!")
|
|
|
|
if __name__ == "__main__": main()
|