#!/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()