""" Stacking Ensemble for fraud detection. Combines predictions from multiple base models (XGBoost, LightGBM, deep learning) using a meta-learner (logistic regression) trained on out-of-fold predictions to avoid data leakage and overfitting. Supports two stacking strategies: - "stack" (default): True out-of-fold stacking. Base models are retrained per fold so OOF predictions never see their own training data. - "blend": Legacy mode. Pre-trained base models generate predictions directly (no per-fold retraining). Faster but prone to leakage. """ import copy import logging import time from typing import Any, Callable, Dict, List, Optional, Tuple, Union import numpy as np from sklearn.linear_model import LogisticRegression from sklearn.metrics import average_precision_score, roc_auc_score from sklearn.model_selection import StratifiedKFold logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Type alias for a model factory: () -> model_instance # --------------------------------------------------------------------------- ModelFactory = Callable[[], Any] class StackingEnsemble: """Stacking ensemble that combines multiple fraud detection models. Uses cross-validated out-of-fold predictions as meta-features to train a second-level logistic regression model. Two construction patterns are supported: 1. **True stacking (recommended)** -- pass *model_factories* so the ensemble can create fresh model instances for each CV fold:: ensemble = StackingEnsemble( model_factories={ "xgb": lambda: FraudXGBoost(xgb_config), "lgb": lambda: FraudLightGBM(lgb_config), }, ) ensemble.train(X_train, y_train, X_val, y_val, method="stack") 2. **Blend mode (legacy)** -- pass pre-trained *base_models* and call ``train(..., method="blend")``:: ensemble = StackingEnsemble(base_models={"xgb": trained_xgb, ...}) ensemble.train(X_train, y_train, X_val, y_val, method="blend") Both patterns expose the same ``predict_proba`` / ``predict`` interface after training. """ def __init__( self, base_models: Optional[Dict[str, Any]] = None, model_factories: Optional[Dict[str, ModelFactory]] = None, config: Optional[dict] = None, ): """Initialize the stacking ensemble. Args: base_models: Dict mapping model name to a *pre-trained* model instance. Each must implement ``predict_proba(X)``. Used in blend mode or as a fallback when *model_factories* is not provided. model_factories: Dict mapping model name to a callable that returns a **new, untrained** model instance. Each returned instance must implement ``train(X_train, y_train, X_val, y_val)`` and ``predict_proba(X)``. Required for true stacking mode. config: Optional ensemble configuration. Recognised keys: * ``cv_folds`` (int, default 5): number of stratified folds. * ``meta_C`` (float, default 1.0): regularisation strength for the logistic-regression meta-learner. * ``random_state`` (int, default 42): random seed. """ if base_models is None and model_factories is None: raise ValueError( "At least one of 'base_models' or 'model_factories' must be provided." ) self.base_models: Dict[str, Any] = base_models or {} self.model_factories: Dict[str, ModelFactory] = model_factories or {} self.config = config or {} self.meta_learner: Optional[LogisticRegression] = None self.cv_folds: int = self.config.get("cv_folds", 5) self._random_state: int = self.config.get("random_state", 42) self._meta_C: float = self.config.get("meta_C", 1.0) # Populated after train(); holds the final base models used at # inference time (retrained on full training set for stacking mode, # or the original pre-trained models for blend mode). self._inference_models: Dict[str, Any] = {} # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def train( self, X_train: np.ndarray, y_train: np.ndarray, X_val: np.ndarray, y_val: np.ndarray, method: str = "stack", feature_names: Optional[list] = None, ) -> None: """Train the stacking ensemble. Args: X_train: Training feature matrix. y_train: Training labels (binary). X_val: Validation feature matrix (used for early-stopping inside base models and for diagnostic evaluation of the ensemble). y_val: Validation labels. method: ``"stack"`` for true out-of-fold stacking (default), ``"blend"`` for legacy pre-trained-model mode. feature_names: Optional feature column names forwarded to base model ``train()`` calls. Ignored in blend mode. """ method = method.lower() if method not in ("stack", "blend"): raise ValueError(f"Unknown method '{method}'. Choose 'stack' or 'blend'.") if method == "stack": if not self.model_factories: raise ValueError( "True stacking (method='stack') requires 'model_factories'. " "Pass pre-trained 'base_models' and use method='blend', or " "provide model_factories." ) oof_predictions = self._train_stack( X_train, y_train, X_val, y_val, feature_names, ) else: if not self.base_models: raise ValueError( "Blend mode requires pre-trained 'base_models'." ) oof_predictions = self._train_blend(X_train, y_train) # ----- Fit the meta-learner on OOF predictions ----- logger.info("Fitting meta-learner on OOF predictions of shape %s", oof_predictions.shape) self.meta_learner = LogisticRegression( C=self._meta_C, max_iter=1000, class_weight="balanced", random_state=self._random_state, ) self.meta_learner.fit(oof_predictions, y_train) # Log meta-learner coefficients (model importance) model_names = list(self._inference_models.keys()) coefs = self.meta_learner.coef_[0] for name, coef in zip(model_names, coefs): logger.info(" Meta-learner weight for %s: %.4f", name, coef) # ----- Evaluate on held-out validation set ----- val_meta_features = self._get_meta_features(X_val) val_proba = self.meta_learner.predict_proba(val_meta_features)[:, 1] val_auc = roc_auc_score(y_val, val_proba) val_ap = average_precision_score(y_val, val_proba) logger.info( "Ensemble validation -- AUC-ROC: %.4f, Avg Precision: %.4f", val_auc, val_ap, ) def predict_proba(self, X: np.ndarray) -> np.ndarray: """Predict fraud probability using the stacking ensemble. Args: X: Feature matrix. Returns: 1-D array of fraud probability scores. Raises: RuntimeError: If the ensemble has not been trained yet. """ if self.meta_learner is None: raise RuntimeError("Ensemble has not been trained. Call train() first.") meta_features = self._get_meta_features(X) return self.meta_learner.predict_proba(meta_features)[:, 1] def predict(self, X: np.ndarray, threshold: float = 0.5) -> np.ndarray: """Predict binary fraud labels. Args: X: Feature matrix. threshold: Decision threshold for the positive class. Returns: Binary prediction array (0 or 1). """ return (self.predict_proba(X) >= threshold).astype(int) # ------------------------------------------------------------------ # Internal: true stacking # ------------------------------------------------------------------ def _train_stack( self, X_train: np.ndarray, y_train: np.ndarray, X_val: np.ndarray, y_val: np.ndarray, feature_names: Optional[list], ) -> np.ndarray: """Run K-fold stacking: retrain base models per fold. Returns the (n_samples, n_models) OOF prediction matrix. After the OOF loop, retrains each base model on the full training set so that ``_inference_models`` are ready for ``predict_proba``. """ model_names = list(self.model_factories.keys()) n_models = len(model_names) n_samples = len(X_train) logger.info( "Training stacking ensemble (method=stack) with %d base models, " "%d CV folds, %d training samples", n_models, self.cv_folds, n_samples, ) oof_predictions = np.zeros((n_samples, n_models)) kf = StratifiedKFold( n_splits=self.cv_folds, shuffle=True, random_state=self._random_state, ) for fold_idx, (train_idx, val_idx) in enumerate(kf.split(X_train, y_train)): fold_start = time.perf_counter() logger.info( " Fold %d/%d -- train: %d samples, val: %d samples", fold_idx + 1, self.cv_folds, len(train_idx), len(val_idx), ) X_fold_train, y_fold_train = X_train[train_idx], y_train[train_idx] X_fold_val, y_fold_val = X_train[val_idx], y_train[val_idx] for model_idx, name in enumerate(model_names): model = self.model_factories[name]() logger.info(" Training '%s' on fold %d ...", name, fold_idx + 1) model.train( X_fold_train, y_fold_train, X_fold_val, y_fold_val, feature_names=feature_names, ) oof_predictions[val_idx, model_idx] = model.predict_proba(X_fold_val) fold_elapsed = time.perf_counter() - fold_start logger.info(" Fold %d complete in %.1fs", fold_idx + 1, fold_elapsed) # Retrain each base model on the FULL training set for inference logger.info("Retraining all base models on the full training set for inference ...") self._inference_models = {} for name in model_names: model = self.model_factories[name]() logger.info(" Retraining '%s' on %d samples ...", name, n_samples) model.train( X_train, y_train, X_val, y_val, feature_names=feature_names, ) self._inference_models[name] = model return oof_predictions # ------------------------------------------------------------------ # Internal: blend (legacy) mode # ------------------------------------------------------------------ def _train_blend( self, X_train: np.ndarray, y_train: np.ndarray, ) -> np.ndarray: """Generate OOF-style predictions from pre-trained models (blend). In blend mode the base models are already trained on the full training set, so the OOF predictions are *not* truly out-of-fold. This is kept for backward compatibility; prefer ``method='stack'`` for proper stacking. Returns the (n_samples, n_models) prediction matrix. """ model_names = list(self.base_models.keys()) n_models = len(model_names) logger.info( "Training stacking ensemble (method=blend) with %d pre-trained base models, " "%d CV folds", n_models, self.cv_folds, ) oof_predictions = np.zeros((len(X_train), n_models)) kf = StratifiedKFold( n_splits=self.cv_folds, shuffle=True, random_state=self._random_state, ) for fold_idx, (train_idx, val_idx) in enumerate(kf.split(X_train, y_train)): logger.info(" Fold %d/%d", fold_idx + 1, self.cv_folds) X_fold_val = X_train[val_idx] for model_idx, (name, model) in enumerate(self.base_models.items()): oof_predictions[val_idx, model_idx] = model.predict_proba(X_fold_val) # In blend mode the pre-trained models are used directly at inference. self._inference_models = dict(self.base_models) return oof_predictions # ------------------------------------------------------------------ # Shared helpers # ------------------------------------------------------------------ def _get_meta_features(self, X: np.ndarray) -> np.ndarray: """Generate meta-features from the inference base-model predictions. Args: X: Raw feature matrix. Returns: (n_samples, n_models) array of base-model probability predictions. """ if not self._inference_models: raise RuntimeError( "No inference models available. Call train() first." ) meta_features = np.column_stack( [model.predict_proba(X) for model in self._inference_models.values()] ) return meta_features