trading-bot/src/trading_bot/data/train_pipeline.py

#!/usr/bin/env python3
"""Generate training dataset and train an XGBoost model using a JSON config.

This script combines the logic from notebooks/spy_direction_dataset.ipynb
and notebooks/train-xboost.ipynb into a single runnable script.
"""
from pathlib import Path
import json
import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.metrics import classification_report, accuracy_score


def find_project_root(start: Path | None = None) -> Path:
    current = (start or Path.cwd()).resolve()
    for candidate in [current, *current.parents]:
        if (candidate / "pyproject.toml").exists():
            return candidate
    raise RuntimeError("Could not find project root containing pyproject.toml")


def load_close(raw_dir: Path, symbol: str, alias: str | None = None) -> pd.DataFrame:
    alias = alias or symbol
    path = raw_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")
    if "close" not in frame.columns:
        raise ValueError(f"{path} does not contain a close column")

    frame = frame.copy()
    frame.index = pd.to_datetime(frame.index)
    frame.index.name = "date"
    frame = frame.sort_index()
    return frame[["close"]].rename(columns={"close": f"{alias}_close"})


def assign_chronological_splits(frame: pd.DataFrame, train_fraction: float, validation_fraction: float, test_fraction: float) -> pd.Series:
    total_fraction = train_fraction + validation_fraction + test_fraction
    if not np.isclose(total_fraction, 1.0):
        raise ValueError(f"Split fractions must sum to 1.0, got {total_fraction}")

    n_rows = len(frame)
    train_end = int(n_rows * train_fraction)
    validation_end = train_end + int(n_rows * validation_fraction)

    split = pd.Series(index=frame.index, dtype="object")
    split.iloc[:train_end] = "train"
    split.iloc[train_end:validation_end] = "validation"
    split.iloc[validation_end:] = "test"
    return split


def build_dataset(raw_data_dir: Path, output_path: Path, symbols: dict, fractions: dict) -> pd.DataFrame:
    prices = pd.concat(
        [
            load_close(raw_data_dir, symbols["SPY"], "SPY"),
            load_close(raw_data_dir, symbols["VIX"] , "VIX"),
            load_close(raw_data_dir, symbols["TLT"] , "TLT"),
            load_close(raw_data_dir, symbols["USO"] , "USO"),
        ],
        axis=1,
        join="inner",
    )

    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)

    spy_forward_close = df["SPY_close"].shift(-5)
    df["spy_up_5d"] = np.nan
    df.loc[spy_forward_close > df["SPY_close"], "spy_up_5d"] = 1.0
    df.loc[spy_forward_close < df["SPY_close"], "spy_up_5d"] = 0.0
    df.loc[spy_forward_close == df["SPY_close"], "spy_up_5d"] = 0.5

    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",
    ]
    TARGET_COLUMN = "spy_up_5d"

    df_model = df[FEATURE_COLUMNS + [TARGET_COLUMN]].dropna().copy()

    # Drop unchanged targets (0.5) to keep binary classification
    df_model = df_model[df_model[TARGET_COLUMN] != 0.5].copy()
    df_model[TARGET_COLUMN] = df_model[TARGET_COLUMN].astype(int)

    df_model["split"] = assign_chronological_splits(
        df_model,
        fractions["train"],
        fractions["validation"],
        fractions["test"],
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    dataset_to_save = df_model.reset_index()
    dataset_to_save.to_parquet(output_path, index=False)

    return df_model, FEATURE_COLUMNS, TARGET_COLUMN


def train_model(df_model: pd.DataFrame, feature_cols: list, target_col: str, config: dict, models_dir: Path):
    train_df = df_model[df_model["split"] == "train"]
    val_df = df_model[df_model["split"] == "validation"]
    test_df = df_model[df_model["split"] == "test"]

    X_train, y_train = train_df[feature_cols], train_df[target_col]
    X_val, y_val = val_df[feature_cols], val_df[target_col]
    X_test, y_test = test_df[feature_cols], test_df[target_col]

    clf_params = dict(config)
    # Ensure early_stopping_rounds present as int

    model = xgb.XGBClassifier(**clf_params)

    eval_set = [(X_train, y_train), (X_val, y_val)]
    fit_kwargs = {"eval_set": eval_set, "verbose": False}

    model.fit(X_train, y_train, **fit_kwargs)

    print(f"Best iteration: {getattr(model, 'best_iteration', None)}")
    evals = model.evals_result()
    if "validation_1" in evals and "logloss" in evals["validation_1"]:
        val_loss = evals["validation_1"]["logloss"]
        print(f"Starting Validation Loss: {val_loss[0]:.4f}")
        print(f"Final Validation Loss:    {val_loss[-1]:.4f}")

    # Evaluate on test set
    y_pred = model.predict(X_test)
    print("\n--- Final Holdout Test Performance ---")
    print(f"Test Set Accuracy: {accuracy_score(y_test, y_pred):.2%}\n")
    print(classification_report(y_test, y_pred))

    # Save model and metadata
    models_dir.mkdir(parents=True, exist_ok=True)
    model_path = models_dir / "spy_xgb_v1.json"
    meta_path = models_dir / "spy_xgb_v1_meta.json"
    model.save_model(str(model_path))

    metadata = {
        "p_base": float(y_train.mean()),
        "feature_cols": feature_cols,
        "last_trained_date": str(df_model.index.max().date()),
        "config": config,
    }
    with open(meta_path, "w") as f:
        json.dump(metadata, f, indent=2)

    print(f"Saved model to {model_path}")


def main():
    project_root = find_project_root()

    # Defaults
    raw_data_dir = project_root / "data" / "alpaca" / "daily"
    output_path = project_root / "data" / "training" / "spy_direction_5d.parquet"
    models_dir = project_root / "models"
    config_path = project_root / "config" / "train_config.json"

    # Load config
    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")
    with open(config_path, "r") as f:
        config = json.load(f)

    symbols = config.get("symbols", {"SPY": "SPY", "VIX": config.get("vix_symbol", "VIXY"), "TLT": "TLT", "USO": "USO"})
    fractions = config.get("fractions", {"train": 0.7, "validation": 0.15, "test": 0.15})

    print("Building dataset...")
    df_model, feature_cols, target_col = build_dataset(raw_data_dir, output_path, symbols, fractions)

    print(f"Rows: {len(df_model):,}")
    print(f"Feature columns: {feature_cols}")
    print(f"Target column: {target_col}")

    split_counts = df_model.groupby("split").size()
    print("Split counts:")
    print(split_counts.to_string())

    print("\nTraining model...")
    model_config = config.get("model", {})
    train_model(df_model, feature_cols, target_col, model_config, models_dir)


if __name__ == "__main__":
    main()