| """
|
| ml/ensemble.py — Ensemble modeling: XGBoost + LightGBM + LSTM/Transformer.
|
|
|
| Combines multiple model families for robust predictions via:
|
| 1. Weighted averaging (simple, robust)
|
| 2. Stacking meta-learner (learns optimal blend)
|
|
|
| GPU-accelerated LSTM for temporal pattern capture.
|
| """
|
| from __future__ import annotations
|
|
|
| import logging
|
| from dataclasses import dataclass
|
|
|
| import numpy as np
|
| import pandas as pd
|
| from sklearn.linear_model import LogisticRegression
|
|
|
| logger = logging.getLogger("trading_system.ml.ensemble")
|
|
|
|
|
| @dataclass
|
| class EnsemblePrediction:
|
| """Output from ensemble model."""
|
| probs: np.ndarray
|
| member_probs: dict[str, np.ndarray]
|
| member_weights: dict[str, float]
|
|
|
|
|
|
|
|
|
|
|
|
|
| def _create_xgboost(n_classes: int = 3, **kwargs):
|
| import xgboost as xgb
|
| defaults = dict(
|
| n_estimators=100, max_depth=4, learning_rate=0.04,
|
| subsample=0.75, colsample_bytree=0.65,
|
| reg_alpha=2.0, reg_lambda=4.0,
|
| objective="multi:softprob", num_class=n_classes,
|
| eval_metric="mlogloss", verbosity=0,
|
| use_label_encoder=False, min_child_weight=15,
|
| )
|
| defaults.update(kwargs)
|
| return xgb.XGBClassifier(**defaults)
|
|
|
|
|
| def _create_lightgbm(n_classes: int = 3, **kwargs):
|
| import lightgbm as lgb
|
| defaults = dict(
|
| n_estimators=100, max_depth=4, learning_rate=0.04,
|
| subsample=0.75, colsample_bytree=0.65,
|
| reg_alpha=2.0, reg_lambda=4.0,
|
| objective="multiclass", num_class=n_classes,
|
| min_child_samples=25, verbose=-1,
|
| )
|
| defaults.update(kwargs)
|
| return lgb.LGBMClassifier(**defaults)
|
|
|
|
|
| class LSTMSignalModel:
|
| """Lightweight LSTM for temporal pattern recognition.
|
|
|
| Architecture: Input → LSTM(hidden) → Dropout → FC → Softmax
|
| Uses GPU if available. Falls back to CPU gracefully.
|
| """
|
|
|
| def __init__(
|
| self,
|
| input_dim: int = 30,
|
| hidden_dim: int = 64,
|
| num_layers: int = 2,
|
| dropout: float = 0.3,
|
| seq_len: int = 20,
|
| n_classes: int = 3,
|
| lr: float = 0.001,
|
| epochs: int = 30,
|
| batch_size: int = 128,
|
| ):
|
| self.input_dim = input_dim
|
| self.hidden_dim = hidden_dim
|
| self.num_layers = num_layers
|
| self.dropout = dropout
|
| self.seq_len = seq_len
|
| self.n_classes = n_classes
|
| self.lr = lr
|
| self.epochs = epochs
|
| self.batch_size = batch_size
|
| self._model = None
|
| self._device = None
|
| self._fitted = False
|
| self.classes_ = np.arange(n_classes)
|
|
|
| def _build_model(self):
|
| import torch
|
| import torch.nn as nn
|
|
|
| self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
| class _LSTMNet(nn.Module):
|
| def __init__(self_, input_dim, hidden_dim, num_layers, dropout, n_classes):
|
| super().__init__()
|
| self_.lstm = nn.LSTM(
|
| input_dim, hidden_dim, num_layers,
|
| batch_first=True, dropout=dropout if num_layers > 1 else 0,
|
| )
|
| self_.dropout = nn.Dropout(dropout)
|
| self_.fc = nn.Linear(hidden_dim, n_classes)
|
|
|
| def forward(self_, x):
|
|
|
| lstm_out, _ = self_.lstm(x)
|
| last_hidden = lstm_out[:, -1, :]
|
| out = self_.dropout(last_hidden)
|
| return self_.fc(out)
|
|
|
| self._model = _LSTMNet(
|
| self.input_dim, self.hidden_dim, self.num_layers,
|
| self.dropout, self.n_classes,
|
| ).to(self._device)
|
|
|
| def _make_sequences(self, X: np.ndarray) -> np.ndarray:
|
| """Convert flat feature matrix into overlapping sequences."""
|
| n = len(X)
|
| if n <= self.seq_len:
|
|
|
| padded = np.zeros((self.seq_len, X.shape[1]))
|
| padded[-n:] = X
|
| return padded[np.newaxis, ...]
|
|
|
| sequences = []
|
| for i in range(self.seq_len, n):
|
| sequences.append(X[i - self.seq_len:i])
|
| return np.array(sequences)
|
|
|
| def fit(self, X: np.ndarray, y: np.ndarray):
|
| """Train LSTM on feature sequences."""
|
| import torch
|
| import torch.nn as nn
|
|
|
| self.input_dim = X.shape[1]
|
| self._build_model()
|
|
|
|
|
| seqs = self._make_sequences(X)
|
| labels = y[self.seq_len:]
|
|
|
| if len(seqs) != len(labels):
|
| min_len = min(len(seqs), len(labels))
|
| seqs = seqs[:min_len]
|
| labels = labels[:min_len]
|
|
|
| if len(seqs) < 50:
|
| logger.warning("Too few sequences for LSTM training (%d)", len(seqs))
|
| self._fitted = False
|
| return
|
|
|
|
|
| X_t = torch.FloatTensor(seqs).to(self._device)
|
| y_t = torch.LongTensor(labels).to(self._device)
|
|
|
|
|
| classes, counts = np.unique(labels, return_counts=True)
|
| weight = torch.FloatTensor(
|
| [len(labels) / (len(classes) * c) for c in counts]
|
| ).to(self._device)
|
|
|
| optimizer = torch.optim.Adam(self._model.parameters(), lr=self.lr,
|
| weight_decay=1e-4)
|
| criterion = nn.CrossEntropyLoss(weight=weight)
|
| scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
| optimizer, patience=5, factor=0.5,
|
| )
|
|
|
| self._model.train()
|
| dataset = torch.utils.data.TensorDataset(X_t, y_t)
|
| loader = torch.utils.data.DataLoader(
|
| dataset, batch_size=self.batch_size, shuffle=True,
|
| )
|
|
|
| best_loss = float('inf')
|
| patience_counter = 0
|
|
|
| for epoch in range(self.epochs):
|
| epoch_loss = 0.0
|
| for batch_X, batch_y in loader:
|
| optimizer.zero_grad()
|
| output = self._model(batch_X)
|
| loss = criterion(output, batch_y)
|
| loss.backward()
|
| torch.nn.utils.clip_grad_norm_(self._model.parameters(), 1.0)
|
| optimizer.step()
|
| epoch_loss += loss.item()
|
|
|
| avg_loss = epoch_loss / len(loader)
|
| scheduler.step(avg_loss)
|
|
|
| if avg_loss < best_loss - 0.001:
|
| best_loss = avg_loss
|
| patience_counter = 0
|
| else:
|
| patience_counter += 1
|
| if patience_counter >= 8:
|
| break
|
|
|
| self._fitted = True
|
|
|
| def predict_proba(self, X) -> np.ndarray:
|
| """Predict class probabilities."""
|
| import torch
|
|
|
| if not self._fitted or self._model is None:
|
| n = len(X) if hasattr(X, '__len__') else 1
|
| return np.full((n, self.n_classes), 1.0 / self.n_classes)
|
|
|
| if isinstance(X, pd.DataFrame):
|
| X = X.values
|
|
|
| seqs = self._make_sequences(X)
|
|
|
| self._model.eval()
|
| with torch.no_grad():
|
| X_t = torch.FloatTensor(seqs).to(self._device)
|
| outputs = self._model(X_t)
|
| probs = torch.softmax(outputs, dim=1).cpu().numpy()
|
|
|
|
|
| n_orig = len(X)
|
| n_pred = len(probs)
|
| if n_pred < n_orig:
|
| pad = np.full((n_orig - n_pred, self.n_classes), 1.0 / self.n_classes)
|
| probs = np.vstack([pad, probs])
|
|
|
| return probs
|
|
|
|
|
|
|
|
|
|
|
|
|
| class EnsembleModel:
|
| """Combine XGBoost + LightGBM + LSTM for robust predictions.
|
|
|
| Combination strategies:
|
| - "average": Weighted average of member probabilities.
|
| - "stacking": Train logistic meta-learner on out-of-fold predictions.
|
|
|
| Walk-forward compatible: can train/predict incrementally.
|
| """
|
|
|
| def __init__(
|
| self,
|
| combination: str = "average",
|
| use_lstm: bool = True,
|
| n_classes: int = 3,
|
| xgb_weight: float = 0.35,
|
| lgb_weight: float = 0.35,
|
| lstm_weight: float = 0.30,
|
| lstm_seq_len: int = 20,
|
| lstm_hidden: int = 64,
|
| lstm_epochs: int = 25,
|
| ):
|
| self.combination = combination
|
| self.use_lstm = use_lstm
|
| self.n_classes = n_classes
|
| self.xgb_weight = xgb_weight
|
| self.lgb_weight = lgb_weight
|
| self.lstm_weight = lstm_weight
|
| self.lstm_seq_len = lstm_seq_len
|
| self.lstm_hidden = lstm_hidden
|
| self.lstm_epochs = lstm_epochs
|
|
|
| self._xgb = None
|
| self._lgb = None
|
| self._lstm = None
|
| self._stacker = None
|
| self._fitted_members: set[str] = set()
|
|
|
| def fit(self, X: np.ndarray, y: np.ndarray):
|
| """Train all ensemble members.
|
|
|
| Args:
|
| X: Feature matrix (n_samples, n_features).
|
| y: Labels (n_samples,), integer class indices.
|
| """
|
| if isinstance(X, pd.DataFrame):
|
| feature_names = list(X.columns)
|
| X_arr = X.values
|
| else:
|
| feature_names = None
|
| X_arr = X
|
|
|
|
|
| mask = ~np.isnan(X_arr).any(axis=1) & ~np.isnan(y)
|
| X_clean = X_arr[mask]
|
| y_clean = y[mask]
|
|
|
| if len(X_clean) < 100 or len(np.unique(y_clean)) < 2:
|
| logger.warning("Insufficient data for ensemble (%d samples)", len(X_clean))
|
| return
|
|
|
|
|
| try:
|
| self._xgb = _create_xgboost(self.n_classes)
|
| if feature_names:
|
| self._xgb.fit(pd.DataFrame(X_clean, columns=feature_names), y_clean)
|
| else:
|
| self._xgb.fit(X_clean, y_clean)
|
| self._fitted_members.add("xgboost")
|
| except Exception as e:
|
| logger.warning("XGBoost training failed: %s", e)
|
|
|
|
|
| try:
|
| self._lgb = _create_lightgbm(self.n_classes)
|
| if feature_names:
|
| self._lgb.fit(pd.DataFrame(X_clean, columns=feature_names), y_clean)
|
| else:
|
| self._lgb.fit(X_clean, y_clean)
|
| self._fitted_members.add("lightgbm")
|
| except Exception as e:
|
| logger.warning("LightGBM training failed: %s", e)
|
|
|
|
|
| if self.use_lstm:
|
| try:
|
| self._lstm = LSTMSignalModel(
|
| input_dim=X_clean.shape[1],
|
| hidden_dim=self.lstm_hidden,
|
| seq_len=self.lstm_seq_len,
|
| n_classes=self.n_classes,
|
| epochs=self.lstm_epochs,
|
| )
|
| self._lstm.fit(X_clean, y_clean)
|
| if self._lstm._fitted:
|
| self._fitted_members.add("lstm")
|
| except Exception as e:
|
| logger.warning("LSTM training failed: %s", e)
|
|
|
|
|
| if self.combination == "stacking" and len(self._fitted_members) >= 2:
|
| self._fit_stacker(X_clean, y_clean, feature_names)
|
|
|
| def _fit_stacker(self, X: np.ndarray, y: np.ndarray, feature_names=None):
|
| """Train stacking meta-learner on in-sample predictions."""
|
| meta_features = self._get_meta_features(X, feature_names)
|
| if meta_features is not None and len(meta_features) > 50:
|
| self._stacker = LogisticRegression(
|
| C=0.5, max_iter=300, multi_class="multinomial",
|
| )
|
| self._stacker.fit(meta_features, y)
|
|
|
| def _get_meta_features(self, X: np.ndarray, feature_names=None) -> np.ndarray | None:
|
| """Get stacked features from member predictions."""
|
| parts = []
|
| if "xgboost" in self._fitted_members:
|
| if feature_names:
|
| p = self._xgb.predict_proba(pd.DataFrame(X, columns=feature_names))
|
| else:
|
| p = self._xgb.predict_proba(X)
|
| parts.append(p)
|
| if "lightgbm" in self._fitted_members:
|
| if feature_names:
|
| p = self._lgb.predict_proba(pd.DataFrame(X, columns=feature_names))
|
| else:
|
| p = self._lgb.predict_proba(X)
|
| parts.append(p)
|
| if "lstm" in self._fitted_members:
|
| p = self._lstm.predict_proba(X)
|
| parts.append(p)
|
|
|
| if not parts:
|
| return None
|
| return np.hstack(parts)
|
|
|
| def predict_proba(self, X) -> EnsemblePrediction:
|
| """Get ensemble probability predictions.
|
|
|
| Returns EnsemblePrediction with combined and per-member probs.
|
| """
|
| if isinstance(X, pd.DataFrame):
|
| feature_names = list(X.columns)
|
| X_arr = X.values
|
| else:
|
| feature_names = None
|
| X_arr = X
|
|
|
| X_clean = np.nan_to_num(X_arr, nan=0.0)
|
| n = len(X_clean)
|
| member_probs = {}
|
|
|
|
|
| if "xgboost" in self._fitted_members:
|
| try:
|
| if feature_names:
|
| p = self._xgb.predict_proba(pd.DataFrame(X_clean, columns=feature_names))
|
| else:
|
| p = self._xgb.predict_proba(X_clean)
|
| p = self._pad_probs(p, self.n_classes)
|
| member_probs["xgboost"] = p
|
| except Exception:
|
| pass
|
|
|
| if "lightgbm" in self._fitted_members:
|
| try:
|
| if feature_names:
|
| p = self._lgb.predict_proba(pd.DataFrame(X_clean, columns=feature_names))
|
| else:
|
| p = self._lgb.predict_proba(X_clean)
|
| p = self._pad_probs(p, self.n_classes)
|
| member_probs["lightgbm"] = p
|
| except Exception:
|
| pass
|
|
|
| if "lstm" in self._fitted_members:
|
| try:
|
| p = self._lstm.predict_proba(X_clean)
|
| p = self._pad_probs(p, self.n_classes)
|
| member_probs["lstm"] = p
|
| except Exception:
|
| pass
|
|
|
| if not member_probs:
|
| uniform = np.full((n, self.n_classes), 1.0 / self.n_classes)
|
| return EnsemblePrediction(
|
| probs=uniform, member_probs={}, member_weights={},
|
| )
|
|
|
|
|
| if self.combination == "stacking" and self._stacker is not None:
|
| meta = np.hstack(list(member_probs.values()))
|
| combined = self._stacker.predict_proba(meta)
|
| weights = {k: 1.0 / len(member_probs) for k in member_probs}
|
| else:
|
|
|
| weight_map = {
|
| "xgboost": self.xgb_weight,
|
| "lightgbm": self.lgb_weight,
|
| "lstm": self.lstm_weight,
|
| }
|
| active = {k: weight_map.get(k, 0.33) for k in member_probs}
|
| total_w = sum(active.values())
|
| weights = {k: v / total_w for k, v in active.items()}
|
|
|
| combined = np.zeros((n, self.n_classes))
|
| for name, w in weights.items():
|
| combined += w * member_probs[name]
|
|
|
| return EnsemblePrediction(
|
| probs=combined,
|
| member_probs=member_probs,
|
| member_weights=weights,
|
| )
|
|
|
| def _pad_probs(self, probs: np.ndarray, n_classes: int) -> np.ndarray:
|
| """Ensure probability matrix has correct number of columns."""
|
| if probs.shape[1] >= n_classes:
|
| return probs[:, :n_classes]
|
| padded = np.full((len(probs), n_classes), 1.0 / n_classes)
|
| padded[:, :probs.shape[1]] = probs
|
| return padded
|
|
|
