Bot first implementation

2026-07-28 21:04:16 +03:00
parent dde7e81bf8
commit a699bfc0fb
6 changed files with 307 additions and 0 deletions
@@ -0,0 +1,11 @@
+"""Model-related modules for the trading bot."""
+
+__all__ = [
+    "DEFAULT_SYMBOLS",
+    "FEATURE_COLUMNS",
+    "PredictionResult",
+    "compute_features",
+    "get_latest_inference_features",
+    "load_raw_close_prices",
+    "predict_latest_probability",
+]
@@ -0,0 +1,244 @@
+"""Prediction pipeline for the initial trading bot prototype.
+
+This module mirrors the notebook prototype by loading raw daily OHLCV parquet
+files, deriving the engineered feature set used during training, and using the
+trained XGBoost model's predict_proba() method to emit a probability for the
+latest available observation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import xgboost as xgb
+
+FEATURE_COLUMNS = [
+    "SPY_ret_5",
+    "SPY_ret_20",
+    "SPY_dist_sma50",
+    "VIX_change_5",
+    "VIX_rank_20",
+    "TLT_ret_10",
+    "USO_ret_5",
+    "SPY_TLT_ratio_ret",
+]
+
+DEFAULT_SYMBOLS = {
+    "SPY": "SPY",
+    "VIXY": "VIX",
+    "TLT": "TLT",
+    "USO": "USO",
+}
+
+
+@dataclass(frozen=True)
+class PredictionResult:
+    """Container for the latest model prediction."""
+
+    prediction_date: pd.Timestamp
+    probability: float
+    base_probability: float
+    feature_columns: list[str]
+
+
+def load_raw_close_prices(
+    raw_data_dir: Path | str, symbols: dict[str, str] | None = None
+) -> pd.DataFrame:
+    """Load raw parquet close prices and merge them into a single frame."""
+
+    symbols = symbols or DEFAULT_SYMBOLS
+    raw_data_dir = Path(raw_data_dir)
+
+    frames: list[pd.DataFrame] = []
+    for symbol, alias in symbols.items():
+        path = raw_data_dir / f"{symbol}.parquet"
+        if not path.exists():
+            raise FileNotFoundError(f"Missing raw data file: {path}")
+
+        frame = pd.read_parquet(path)
+        if "date" in frame.columns:
+            frame = frame.set_index("date")
+
+        frame.index = pd.to_datetime(frame.index)
+        frame = frame.sort_index()
+        frames.append(frame[["close"]].rename(columns={"close": f"{alias}_close"}))
+
+    return pd.concat(frames, axis=1, join="inner")
+
+
+def compute_features(prices: pd.DataFrame) -> pd.DataFrame:
+    """Compute the engineered feature matrix used for inference."""
+
+    df = prices.copy()
+
+    df["SPY_ret_5"] = df["SPY_close"].pct_change(5)
+    df["SPY_ret_20"] = df["SPY_close"].pct_change(20)
+
+    sma_50 = df["SPY_close"].rolling(50).mean()
+    df["SPY_dist_sma50"] = (df["SPY_close"] - sma_50) / sma_50
+
+    df["VIX_change_5"] = df["VIX_close"].pct_change(5)
+    df["VIX_rank_20"] = df["VIX_close"].rolling(20).rank(pct=True)
+
+    df["TLT_ret_10"] = df["TLT_close"].pct_change(10)
+    df["USO_ret_5"] = df["USO_close"].pct_change(5)
+    df["SPY_TLT_ratio_ret"] = (df["SPY_close"] / df["TLT_close"]).pct_change(5)
+
+    return df[FEATURE_COLUMNS]
+
+
+def get_latest_inference_features(
+    raw_data_dir: Path | str,
+    symbols: dict[str, str] | None = None,
+    max_age_days: int = 1,
+) -> tuple[pd.Timestamp, pd.DataFrame]:
+    """Return the latest row of features and its observation date."""
+
+    prices = load_raw_close_prices(raw_data_dir, symbols)
+    features = compute_features(prices).dropna()
+    if features.empty:
+        raise ValueError(
+            "Not enough historical rows to compute 50-day rolling window features."
+        )
+
+    latest_row = features.iloc[[-1]]
+    latest_date = latest_row.index[0]
+
+    now = pd.Timestamp.now()
+    latest_date_naive = (
+        latest_date.tz_localize(None)
+        if latest_date.tz is not None
+        else latest_date
+    )
+    days_old = (now.floor("D") - latest_date_naive.floor("D")).days
+
+    if days_old > max_age_days:
+        warnings.warn(
+            f"STALE DATA WARNING: Latest feature row is from {latest_date.strftime('%Y-%m-%d')} "
+            f"({days_old} day(s) old). Update raw Parquet files before executing trades.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    return latest_date, latest_row
+
+
+def load_model(model_path: Path | str) -> xgb.XGBClassifier:
+    """Load the serialized XGBoost classifier from disk."""
+
+    model = xgb.XGBClassifier()
+    model.load_model(str(model_path))
+    return model
+
+
+def load_feature_metadata(metadata_path: Path | str) -> dict[str, Any]:
+    """Load the model metadata JSON file."""
+
+    with Path(metadata_path).open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def predict_latest_probability(
+    model_path: Path | str,
+    metadata_path: Path | str,
+    raw_data_dir: Path | str,
+    symbols: dict[str, str] | None = None,
+    max_age_days: int = 1,
+) -> PredictionResult:
+    """Return the latest predicted probability for the next SPY move."""
+
+    model = load_model(model_path)
+    metadata = load_feature_metadata(metadata_path)
+    latest_date, latest_features = get_latest_inference_features(
+        raw_data_dir=raw_data_dir,
+        symbols=symbols,
+        max_age_days=max_age_days,
+    )
+
+    feature_columns = list(metadata.get("feature_cols", FEATURE_COLUMNS))
+    model_input = latest_features[feature_columns]
+    probabilities = model.predict_proba(model_input)
+    probability = float(probabilities[0, 1])
+    base_probability = float(metadata.get("p_base", 0.5))
+
+    return PredictionResult(
+        prediction_date=latest_date,
+        probability=probability,
+        base_probability=base_probability,
+        feature_columns=feature_columns,
+    )
+
+
+def get_target_exposure(
+    p_pred: float, p_base: float, sensitivity: float = 5.0
+) -> float:
+    """Maps predicted probability to a target portfolio equity allocation (0.0 to 1.0).
+
+    - p_pred == p_base  --> 50% Target Exposure (Neutral)
+    - p_pred > p_base   --> Scale up toward 100% (Bullish)
+    - p_pred < p_base   --> Scale down toward 0% (Bearish / Cash)
+    """
+    # Calculate deviation from the historical average
+    delta = p_pred - p_base
+
+    # Base target allocation is 50% equity / 50% cash
+    base_allocation = 0.50
+
+    # Sensitivity controls how aggressively probability changes alter allocation
+    # e.g., a +0.08 delta * 5.0 = +0.40 -> 90% Equity Allocation
+    target_allocation = base_allocation + (delta * sensitivity)
+
+    # Clamp bounds strictly between 0% (full cash) and 100% (full SPY)
+    return float(np.clip(target_allocation, 0.0, 1.0))
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments for local prediction runs."""
+
+    parser = argparse.ArgumentParser(description="Predict the latest SPY direction probability.")
+    parser.add_argument(
+        "--model-path",
+        type=Path,
+        default=Path("notebooks/models/spy_xgb_v1.json"),
+        help="Path to the XGBoost model artifact.",
+    )
+    parser.add_argument(
+        "--metadata-path",
+        type=Path,
+        default=Path("notebooks/models/spy_xgb_v1_meta.json"),
+        help="Path to the model metadata JSON file.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        type=Path,
+        default=Path("data/ibkr/daily"),
+        help="Directory containing the raw Parquet market data files.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Run the latest prediction from the command line."""
+
+    args = parse_args()
+    result = predict_latest_probability(
+        model_path=args.model_path,
+        metadata_path=args.metadata_path,
+        raw_data_dir=args.data_dir,
+    )
+    print(
+        f"Date: {result.prediction_date.date()} | Prob: {result.probability:.4f} | "
+        f"Base: {result.base_probability:.4f}"
+    )
+    print(f"Exposure: {get_target_exposure(result.probability, result.base_probability, sensitivity=5.0):.2f}")
+
+
+if __name__ == "__main__":
+    main()