jarno/trading-bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Bot first implementation

Browse Source
This commit is contained in:
jarno
2026-07-28 21:04:16 +03:00
parent dde7e81bf8
commit a699bfc0fb
6 changed files with 307 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/bot-prototype.ipynb
View File
pyproject.toml
+1
View File
@@ -9,6 +9,7 @@ dependencies = [
"matplotlib>=3.11.1",
"pandas>=2.3.0",
"pyarrow>=20.0.0",
"pytest>=9.1.1",
"scikit-learn>=1.9.0",
"xgboost>=3.2.0",
]
src/trading_bot/models/__init__.py
+11
View File
@@ -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",
]
src/trading_bot/models/prediction.py
+244
View File
@@ -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()
tests/test_prediction.py
+15
View File
@@ -0,0 +1,15 @@
from pathlib import Path
from trading_bot.models.prediction import FEATURE_COLUMNS, predict_latest_probability
def test_predict_latest_probability_returns_probability_between_zero_and_one() -> None:
result = predict_latest_probability(
model_path=Path("notebooks/models/spy_xgb_v1.json"),
metadata_path=Path("notebooks/models/spy_xgb_v1_meta.json"),
raw_data_dir=Path("data/ibkr/daily"),
)
assert set(result.feature_columns).issubset(FEATURE_COLUMNS)
assert 0.0 <= result.probability <= 1.0
assert result.probability > 0.0
uv.lock
Generated
+36
View File
@@ -176,6 +176,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/8f/f3/28ea87be30570f4d6b8fd24380d12fa74e59467ee003755e76aeb29082b8/ib_insync-0.9.86-py3-none-any.whl", hash = "sha256:a61fbe56ff405d93d211dad8238d7300de76dd6399eafc04c320470edec9a4a4", size = 72980, upload-time = "2023-07-02T12:43:29.928Z" },
]
[[package]]
name = "iniconfig"
version = "2.3.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
]
[[package]]
name = "ipykernel"
version = "7.3.0"
@@ -499,6 +508,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/7d/68/d8d58938dfb1370b266a1a729e6d77a985be23689a0496498ee17b2cbf90/platformdirs-4.11.0-py3-none-any.whl", hash = "sha256:360ccded2b7fce0af0ff80cc8f5942a1c5d99b0e856033acb030bfc634709e74", size = 23247, upload-time = "2026-07-21T13:09:35.422Z" },
]
[[package]]
name = "pluggy"
version = "1.6.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
]
[[package]]
name = "prompt-toolkit"
version = "3.0.52"
@@ -587,6 +605,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" },
]
[[package]]
name = "pytest"
version = "9.1.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "colorama", marker = "sys_platform == 'win32'" },
{ name = "iniconfig" },
{ name = "packaging" },
{ name = "pluggy" },
{ name = "pygments" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" },
]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"
@@ -735,6 +769,7 @@ dependencies = [
{ name = "matplotlib" },
{ name = "pandas" },
{ name = "pyarrow" },
{ name = "pytest" },
{ name = "scikit-learn" },
{ name = "xgboost" },
]
@@ -750,6 +785,7 @@ requires-dist = [
{ name = "matplotlib", specifier = ">=3.11.1" },
{ name = "pandas", specifier = ">=2.3.0" },
{ name = "pyarrow", specifier = ">=20.0.0" },
{ name = "pytest", specifier = ">=9.1.1" },
{ name = "scikit-learn", specifier = ">=1.9.0" },
{ name = "xgboost", specifier = ">=3.2.0" },
]