"""Prediction model — XGBoost classifier + regression for price direction & %."""
import json
import logging
import os
import pickle
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import (
accuracy_score, precision_score, recall_score, f1_score,
mean_absolute_error, mean_squared_error, r2_score,
)
from backend.config import config
logger = logging.getLogger(__name__)
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
class Predictor:
"""XGBoost-based price prediction model."""
def __init__(self):
self.model_class = None # Classification model (up/down)
self.model_reg = None # Regression model (% change)
self.feature_names: List[str] = []
self.is_trained = False
self.train_metrics: Dict[str, Any] = {}
self.trained_at: Optional[str] = None
self.ticker: Optional[str] = None
self.bar_count: int = 0
@property
def model_path(self) -> str:
model_dir = config.storage.get("model_dir", "models/saved")
path = os.path.join(PROJECT_ROOT, model_dir)
os.makedirs(path, exist_ok=True)
return path
def train(
self,
X: np.ndarray,
y_class: np.ndarray,
y_reg: np.ndarray,
feature_names: List[str],
ticker: str = "unknown",
) -> Dict[str, Any]:
"""Train both classification and regression models.
Args:
X: Feature matrix
y_class: Binary labels (1=up, 0=down)
y_reg: Continuous labels (% change)
feature_names: Column names
ticker: Ticker symbol for labeling
Returns:
Training metrics dict
"""
from xgboost import XGBClassifier, XGBRegressor
mc = config.model
# Split
X_train, X_test, yc_train, yc_test, yr_train, yr_test = train_test_split(
X, y_class, y_reg,
test_size=mc.get("train_test_split", 0.2),
random_state=mc.get("random_state", 42),
stratify=y_class,
)
# --- Classification model ---
self.model_class = XGBClassifier(
n_estimators=mc.get("n_estimators", 200),
max_depth=mc.get("max_depth", 6),
learning_rate=mc.get("learning_rate", 0.1),
random_state=mc.get("random_state", 42),
use_label_encoder=False,
eval_metric="logloss",
n_jobs=-1,
)
self.model_class.fit(X_train, yc_train)
# --- Regression model ---
self.model_reg = XGBRegressor(
n_estimators=mc.get("n_estimators", 200),
max_depth=mc.get("max_depth", 6),
learning_rate=mc.get("learning_rate", 0.1),
random_state=mc.get("random_state", 42),
n_jobs=-1,
)
self.model_reg.fit(X_train, yr_train)
# --- Evaluate ---
yc_pred = self.model_class.predict(X_test)
yr_pred = self.model_reg.predict(X_test)
# Check if we have both classes in test set
has_both_classes = len(np.unique(yc_pred)) > 1
self.train_metrics = {
"classification": {
"accuracy": float(accuracy_score(yc_test, yc_pred)),
"precision": float(precision_score(yc_test, yc_pred, zero_division=0)) if has_both_classes else 0.0,
"recall": float(recall_score(yc_test, yc_pred, zero_division=0)) if has_both_classes else 0.0,
"f1": float(f1_score(yc_test, yc_pred, zero_division=0)) if has_both_classes else 0.0,
},
"regression": {
"mae": float(mean_absolute_error(yr_test, yr_pred)),
"rmse": float(np.sqrt(mean_squared_error(yr_test, yr_pred))),
"r2": float(r2_score(yr_test, yr_pred)),
},
"train_size": int(len(X_train)),
"test_size": int(len(X_test)),
"features": len(feature_names),
}
self.feature_names = feature_names
self.is_trained = True
self.trained_at = datetime.now().isoformat()
self.ticker = ticker
self.bar_count = len(X)
logger.info(
f"Trained {ticker}: acc={self.train_metrics['classification']['accuracy']:.3f}, "
f"r2={self.train_metrics['regression']['r2']:.3f}, "
f"n={self.bar_count}, features={len(feature_names)}"
)
return self.train_metrics
def predict(self, X: np.ndarray) -> Dict[str, Any]:
"""Predict on new features.
Returns:
{
'direction': 'up' | 'down' | 'neutral',
'confidence': float, # max(class_prob)
'prob_up': float,
'prob_down': float,
'pct_change': float, # predicted % change
'features': {name: value}, # for debugging
}
"""
if not self.is_trained:
raise RuntimeError("Model not trained yet")
X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
# Classification
proba = self.model_class.predict_proba(X)[0]
prob_up = float(proba[0][1]) # Probability of class 1 (up)
prob_down = float(proba[0][0])
direction_pred = self.model_class.predict(X)[0]
# Regression
pct_pred = float(self.model_reg.predict(X)[0])
# Confidence
confidence = max(prob_up, prob_down)
# Direction with threshold
threshold = config.model.get("classification_threshold", 0.55)
if confidence < threshold:
direction = "neutral"
elif direction_pred == 1:
direction = "up"
else:
direction = "down"
return {
"direction": direction,
"confidence": round(confidence, 4),
"prob_up": round(prob_up, 4),
"prob_down": round(prob_down, 4),
"pct_change": round(pct_pred, 4),
"features": {name: round(float(val), 4) for name, val in zip(self.feature_names, X[0])},
}
def feature_importance(self) -> List[Dict[str, Any]]:
"""Get feature importance from classification model."""
if not self.is_trained or self.model_class is None:
return []
importances = self.model_class.feature_importances_
indexed = list(zip(self.feature_names, importances))
indexed.sort(key=lambda x: x[1], reverse=True)
return [
{"feature": name, "importance": round(float(imp), 4)}
for name, imp in indexed[:20] # Top 20
]
def save(self, ticker: Optional[str] = None) -> str:
"""Save model to disk."""
target = ticker or self.ticker or "model"
path = os.path.join(self.model_path, f"{target}.pkl")
with open(path, "wb") as f:
pickle.dump({
"model_class": self.model_class,
"model_reg": self.model_reg,
"feature_names": self.feature_names,
"is_trained": self.is_trained,
"trained_at": self.trained_at,
"ticker": self.ticker,
"bar_count": self.bar_count,
"metrics": self.train_metrics,
}, f)
logger.info(f"Model saved: {path}")
return path
def load(self, ticker: str) -> bool:
"""Load model from disk."""
path = os.path.join(self.model_path, f"{ticker}.pkl")
if not os.path.exists(path):
return False
with open(path, "rb") as f:
data = pickle.load(f)
self.model_class = data["model_class"]
self.model_reg = data["model_reg"]
self.feature_names = data["feature_names"]
self.is_trained = data["is_trained"]
self.trained_at = data["trained_at"]
self.ticker = data["ticker"]
self.bar_count = data["bar_count"]
self.train_metrics = data["metrics"]
logger.info(f"Model loaded: {path} (trained {self.trained_at})")
return True
def info(self) -> Dict[str, Any]:
"""Get model info summary."""
return {
"is_trained": self.is_trained,
"ticker": self.ticker,
"trained_at": self.trained_at,
"bar_count": self.bar_count,
"feature_count": len(self.feature_names),
"metrics": self.train_metrics,
"feature_importance": self.feature_importance() if self.is_trained else [],
}