tabml-guard 0.1.0__py3-none-any.whl

tabml_guard/__init__.py

@@ -0,0 +1,12 @@
+"""TabML-Guard: Explainable, Drift-Aware AutoML for Time-Ordered Tabular Data."""
+
+from .__version__ import __version__
+from .config import Config, load_config
+from .pipeline import TabMLGuard
+
+__all__ = [
+    "__version__",
+    "TabMLGuard",
+    "Config",
+    "load_config",
+]

tabml_guard/__version__.py

@@ -0,0 +1,2 @@
+__version__ = "0.1.0"
+

tabml_guard/baselines.py

@@ -0,0 +1,416 @@
+"""Stage 1 — Regularized baselines (the bias-variance foundation).
+
+Before any heavy model, establish an honest, interpretable yardstick: linear
+models under different regularizers. This makes two teaching points concrete and
+measurable:
+
+  * **L2 (ridge-style) shrinks** coefficients smoothly toward zero;
+  * **L1 (lasso-style) zeroes** coefficients outright -> built-in feature
+    selection (we report the resulting sparsity);
+  * **ElasticNet** blends the two.
+
+The classification analog of Ridge/Lasso/ElasticNet is Logistic Regression with
+``penalty in {l2, l1, elasticnet}`` — identical regularization mathematics on the
+linear term, just wrapped in a logistic loss. (Be ready to defend this to an
+examiner: "Ridge/Lasso are regression names; the *penalty* is the idea, and it
+transfers directly to logistic regression.")
+
+Draws on Paper 1 (RegStack, ESWA 2025, DOI 10.1016/j.eswa.2025.127766): treating
+regularization as a first-class ingredient, not an afterthought.
+
+    Classification -> LogisticRegression(penalty in {None, l2, l1, elasticnet})
+    Regression     -> LinearRegression / Ridge / Lasso / ElasticNet
+"""
+from __future__ import annotations
+
+import warnings
+from typing import Optional
+
+import numpy as np
+
+from .config import Config
+from .data import DataBundle
+from .metrics import primary_metric_name
+from .modeling import (
+    ModelResult,
+    feature_names_out,
+    fit_evaluate,
+    leaderboard,
+    log_result,
+    pick_best,
+)
+
+
+# --------------------------------------------------------------------------- #
+# Model zoo (built lazily so an import problem surfaces as a clear message)
+# --------------------------------------------------------------------------- #
+def _sklearn_new_api() -> bool:
+    """True on scikit-learn >= 1.8, where ``penalty=`` is deprecated for ``l1_ratio``.
+
+    On 1.8+ a stray ``penalty=`` is silently overridden by ``l1_ratio``'s default
+    of 0.0 (pure L2) — which would make all four regularizers identical L2 fits.
+    Detecting the version lets us speak the right dialect.
+    """
+    import sklearn
+
+    try:
+        major, minor = (int(x) for x in sklearn.__version__.split(".")[:2])
+        return (major, minor) >= (1, 8)
+    except Exception:
+        return False
+
+
+def _make_logreg(kind: str, C, rs: int, new_api: bool):
+    """LogisticRegression for kind in {none, l2, l1, elasticnet} at strength ``C``.
+
+    On 1.8+ the elasticnet spectrum is expressed via ``l1_ratio`` (0=L2, 1=L1,
+    0.5=EN) and ``C=inf`` means unregularized; 'saga' is required once
+    ``l1_ratio > 0`` (lbfgs is fine and faster at 0). ``n_jobs`` is omitted: it has
+    no effect on LogisticRegression and was removed in 1.8.
+    """
+    from sklearn.linear_model import LogisticRegression
+
+    if new_api:
+        cfgs = {
+            "none": dict(C=np.inf, l1_ratio=0.0, solver="lbfgs", max_iter=3000),
+            "l2": dict(C=C, l1_ratio=0.0, solver="lbfgs", max_iter=3000),
+            "l1": dict(C=C, l1_ratio=1.0, solver="saga", max_iter=5000, tol=1e-3, random_state=rs),
+            "elasticnet": dict(C=C, l1_ratio=0.5, solver="saga", max_iter=5000, tol=1e-3, random_state=rs),
+        }
+    else:
+        cfgs = {
+            "none": dict(penalty=None, solver="lbfgs", max_iter=3000),
+            "l2": dict(penalty="l2", C=C, solver="lbfgs", max_iter=3000),
+            "l1": dict(penalty="l1", C=C, solver="saga", max_iter=5000, tol=1e-3, random_state=rs),
+            "elasticnet": dict(penalty="elasticnet", C=C, l1_ratio=0.5, solver="saga",
+                               max_iter=5000, tol=1e-3, random_state=rs),
+        }
+    return LogisticRegression(**cfgs[kind])
+
+
+def _classification_models(cfg: Config) -> dict:
+    """Four logistic-regression regularizers at a common C=1.0: none / L2 / L1 / EN.
+
+    Keeping C fixed across penalties makes the comparison fair. The
+    *regularization path* (below) is what reveals what each penalty actually does
+    as its strength varies — the headline four can look near-identical when the
+    data is abundant and low-dimensional, which is itself a finding.
+    """
+    rs = cfg.random_state
+    new_api = _sklearn_new_api()
+    return {
+        "logreg_none": _make_logreg("none", None, rs, new_api),
+        "logreg_l2": _make_logreg("l2", 1.0, rs, new_api),
+        "logreg_l1": _make_logreg("l1", 1.0, rs, new_api),
+        "logreg_elasticnet": _make_logreg("elasticnet", 1.0, rs, new_api),
+    }
+
+
+def _regression_models(cfg: Config) -> dict:
+    from sklearn.linear_model import ElasticNet, Lasso, LinearRegression, Ridge
+
+    rs = cfg.random_state
+    return {
+        "linear": LinearRegression(),
+        "ridge": Ridge(alpha=1.0, random_state=rs),
+        "lasso": Lasso(alpha=0.001, max_iter=10000, random_state=rs),
+        "elasticnet": ElasticNet(
+            alpha=0.001, l1_ratio=0.5, max_iter=10000, random_state=rs
+        ),
+    }
+
+
+# --------------------------------------------------------------------------- #
+# Coefficient diagnostics — the "regularization actually did something" evidence
+# --------------------------------------------------------------------------- #
+def _coef_diagnostics(res: ModelResult) -> dict:
+    """Sparsity, coefficient norms, and the strongest features for one model."""
+    model = res.estimator.named_steps["model"]
+    coef = getattr(model, "coef_", None)
+    if coef is None:
+        return {}
+    coef = np.asarray(coef, dtype=float)
+    names = feature_names_out(res.estimator)
+
+    flat = coef.ravel()
+    n_total = int(flat.size)
+    n_zero = int(np.sum(np.abs(flat) < 1e-6))
+
+    # Rank features by importance. For multiclass, average |coef| across classes
+    # so each feature gets one score aligned with the feature-name list.
+    if coef.ndim > 1 and coef.shape[0] > 1:
+        strength = np.abs(coef).mean(axis=0)
+    else:
+        strength = np.abs(flat)
+    order = np.argsort(strength)[::-1][:12]
+    signed = flat if flat.size == strength.size else coef.mean(axis=0)
+    top = [
+        [names[i] if i < len(names) else f"f{i}", round(float(signed[i]), 5)]
+        for i in order
+    ]
+
+    return {
+        "n_coefficients": n_total,
+        "n_zero": n_zero,
+        "sparsity": round(n_zero / n_total, 4) if n_total else 0.0,
+        "l1_norm": round(float(np.abs(flat).sum()), 5),
+        "l2_norm": round(float(np.linalg.norm(flat)), 5),
+        "top_features": top,
+    }
+
+
+def _regularization_note(results: list[ModelResult], task: str, path: Optional[dict] = None) -> str:
+    """A plain-English takeaway: why the headline looks flat, and what the path shows."""
+    metric = primary_metric_name(task)
+    vals = [r.metrics.get(metric, float("nan")) for r in results]
+    vals = [v for v in vals if v == v]  # drop nan
+    best = pick_best(results, task)
+
+    parts = []
+    if best is not None and vals:
+        spread = max(vals) - min(vals)
+        parts.append(
+            f"the {len(results)} regularizers span only {spread:.4f} {metric} "
+            f"(best '{best.name}' {best.metrics.get(metric, float('nan')):.4f}) — "
+            f"with abundant rows and few features, penalties barely move the fit at the default strength"
+        )
+
+    if path:
+        F = path.get("n_features", 0)
+        if path.get("type") == "classification" and path.get("l1", {}).get("nonzero"):
+            nz, Cs = path["l1"]["nonzero"], path["C"]
+            i = int(np.argmin(nz))
+            parts.append(
+                f"L1 path: at C={Cs[i]:g} only {nz[i]}/{F} features survive (sparsity), "
+                f"rising to {max(nz)}/{F} as C grows"
+            )
+        elif path.get("type") == "regression" and path.get("lasso", {}).get("nonzero"):
+            nz, al = path["lasso"]["nonzero"], path["alpha"]
+            i = int(np.argmax(al))  # strongest penalty = largest alpha
+            parts.append(
+                f"Lasso path: at alpha={al[i]:g} only {nz[i]}/{F} features survive (sparsity)"
+            )
+
+    return "; ".join(parts) if parts else "regularized linear baselines fitted."
+
+
+# --------------------------------------------------------------------------- #
+# Regularization path — the "what does the penalty actually do?" demonstration
+# --------------------------------------------------------------------------- #
+def _regularization_path(bundle: DataBundle, cfg: Config) -> Optional[dict]:
+    """Sweep penalty strength to visualize L1 sparsity and L2 shrinkage.
+
+    Fits on an INTERNAL train/validation split of the *training* data only (the
+    held-out test set is never touched), so this is a clean diagnostic rather than
+    test-set peeking. Linear fits are cheap, so a 12-point sweep costs seconds.
+
+    Returns per-strength: number of surviving (non-zero) coefficients, the L2 norm
+    of the coefficient vector (shrinkage), and a validation score (the bias-variance
+    sweet spot). Classification sweeps C for L1 & L2; regression sweeps alpha for
+    Lasso & Ridge.
+    """
+    try:
+        from .modeling import make_pipeline
+
+        rs = cfg.random_state
+        n = len(bundle.X_train)
+        if n < 50:
+            return None
+        cut = int(n * 0.8)
+        Xtr, Xval = bundle.X_train.iloc[:cut], bundle.X_train.iloc[cut:]
+        ytr, yval = bundle.y_train[:cut], bundle.y_train[cut:]
+
+        def _fit_coef(estimator):
+            pipe = make_pipeline(bundle.preprocessor, estimator)
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                pipe.fit(Xtr, ytr)
+            coef = np.asarray(pipe.named_steps["model"].coef_, dtype=float).ravel()
+            return pipe, coef
+
+        if bundle.task == "classification":
+            from sklearn.metrics import roc_auc_score
+
+            new_api = _sklearn_new_api()
+            Cs = np.logspace(-3, 2, 12)
+            out: dict = {
+                "type": "classification",
+                "C": [round(float(c), 6) for c in Cs],
+                "l1": {"nonzero": [], "l2norm": [], "val_roc_auc": []},
+                "l2": {"nonzero": [], "l2norm": [], "val_roc_auc": []},
+            }
+            n_features = 0
+            for kind in ("l1", "l2"):
+                for c in Cs:
+                    pipe, coef = _fit_coef(_make_logreg(kind, float(c), rs, new_api))
+                    n_features = coef.size
+                    out[kind]["nonzero"].append(int(np.sum(np.abs(coef) > 1e-6)))
+                    out[kind]["l2norm"].append(round(float(np.linalg.norm(coef)), 5))
+                    try:
+                        proba = pipe.predict_proba(Xval)[:, 1]
+                        out[kind]["val_roc_auc"].append(round(float(roc_auc_score(yval, proba)), 4))
+                    except Exception:
+                        out[kind]["val_roc_auc"].append(float("nan"))
+            out["n_features"] = int(n_features)
+            return out
+
+        # regression
+        from sklearn.linear_model import Lasso, Ridge
+        from sklearn.metrics import mean_squared_error
+
+        alphas = np.logspace(-4, 1, 12)
+        out = {
+            "type": "regression",
+            "alpha": [round(float(a), 6) for a in alphas],
+            "lasso": {"nonzero": [], "l2norm": [], "val_rmse": []},
+            "ridge": {"nonzero": [], "l2norm": [], "val_rmse": []},
+        }
+        n_features = 0
+        for kind in ("lasso", "ridge"):
+            for a in alphas:
+                est = (
+                    Lasso(alpha=float(a), max_iter=10000, random_state=rs)
+                    if kind == "lasso"
+                    else Ridge(alpha=float(a), random_state=rs)
+                )
+                pipe, coef = _fit_coef(est)
+                n_features = coef.size
+                out[kind]["nonzero"].append(int(np.sum(np.abs(coef) > 1e-6)))
+                out[kind]["l2norm"].append(round(float(np.linalg.norm(coef)), 5))
+                pred = pipe.predict(Xval)
+                out[kind]["val_rmse"].append(round(float(np.sqrt(mean_squared_error(yval, pred))), 5))
+        out["n_features"] = int(n_features)
+        return out
+    except Exception as exc:  # a diagnostic must never sink the pipeline
+        print(f"[stage1] regularization path skipped: {type(exc).__name__}: {exc}")
+        return None
+
+
+def _print_path(path: Optional[dict]) -> None:
+    """Compact console view of the sweep so the mechanism is visible during a run."""
+    if not path:
+        return
+    F = path.get("n_features", 0)
+    if path.get("type") == "classification":
+        Cs = path["C"]
+        l1nz, l2n = path["l1"]["nonzero"], path["l2"]["l2norm"]
+        lo, hi = 0, len(Cs) - 1
+        print(
+            f"[stage1] reg-path (of {F} features): "
+            f"L1 keeps {l1nz[lo]}@C={Cs[lo]:g} -> {l1nz[hi]}@C={Cs[hi]:g} (sparsity); "
+            f"L2 |coef| grows {l2n[lo]:g} -> {l2n[hi]:g} (shrinkage)"
+        )
+    elif path.get("type") == "regression":
+        al = path["alpha"]
+        lnz = path["lasso"]["nonzero"]
+        lo, hi = 0, len(al) - 1
+        print(
+            f"[stage1] reg-path (of {F} features): "
+            f"Lasso keeps {lnz[hi]}@alpha={al[hi]:g} -> {lnz[lo]}@alpha={al[lo]:g} (sparsity)"
+        )
+
+
+# --------------------------------------------------------------------------- #
+# Learning curve — bias/variance on the strongest baseline only (cost-bounded)
+# --------------------------------------------------------------------------- #
+def _learning_curve(best: ModelResult, bundle: DataBundle, cfg: Config) -> Optional[dict]:
+    """Train/validation score vs. training-set size, for the best baseline.
+
+    Uses the pipeline (so preprocessing is refit inside each fold — no leakage),
+    a time-aware CV when we have a time column, and only 5 points on 3 folds to
+    stay cheap. ``learning_curve`` uses cross_validate internally, so
+    ``TimeSeriesSplit`` is fine here (unlike stacking's cross_val_predict).
+    """
+    try:
+        from sklearn.base import clone
+        from sklearn.model_selection import (
+            KFold,
+            StratifiedKFold,
+            TimeSeriesSplit,
+            learning_curve,
+        )
+
+        n_classes = len(bundle.target_classes) if bundle.target_classes else 2
+        if bundle.task == "classification":
+            scoring = "roc_auc" if n_classes == 2 else "roc_auc_ovr"
+        else:
+            scoring = "neg_root_mean_squared_error"
+
+        if bundle.has_time:
+            cv = TimeSeriesSplit(n_splits=3)
+        elif bundle.task == "classification":
+            cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=cfg.random_state)
+        else:
+            cv = KFold(n_splits=3, shuffle=True, random_state=cfg.random_state)
+
+        sizes, train_scores, val_scores = learning_curve(
+            clone(best.estimator),
+            bundle.X_train,
+            bundle.y_train,
+            train_sizes=np.linspace(0.1, 1.0, 5),
+            cv=cv,
+            scoring=scoring,
+            n_jobs=-1,
+            shuffle=False,
+        )
+        return {
+            "model": best.name,
+            "metric": scoring,
+            "train_sizes": [int(s) for s in sizes],
+            "train_mean": [round(float(v), 4) for v in train_scores.mean(axis=1)],
+            "val_mean": [round(float(v), 4) for v in val_scores.mean(axis=1)],
+        }
+    except Exception as exc:  # never let a diagnostic sink the pipeline
+        print(f"[stage1] learning curve skipped: {type(exc).__name__}: {exc}")
+        return None
+
+
+# --------------------------------------------------------------------------- #
+# Entry point
+# --------------------------------------------------------------------------- #
+def run(bundle: DataBundle, cfg: Config, ctx: Optional[dict] = None) -> dict:
+    """Fit regularized baselines; return metrics + coefficient diagnostics.
+
+    Side effect: registers every fitted result in ``ctx['results']`` and records
+    ``ctx['baseline_best']`` so later stages (and Stage 2's overall-best pick) can
+    see the baselines.
+    """
+    ctx = ctx if ctx is not None else {}
+    task = bundle.task
+    models = _classification_models(cfg) if task == "classification" else _regression_models(cfg)
+
+    results: list[ModelResult] = []
+    for name, est in models.items():
+        try:
+            res = fit_evaluate(name, est, bundle, cfg)
+            res.extra["coefficients"] = _coef_diagnostics(res)
+            results.append(res)
+            log_result("stage1", res, task, cfg.report.fbeta)
+        except Exception as exc:
+            print(f"[stage1] {name} failed: {type(exc).__name__}: {exc} — skipped.")
+
+    if not results:
+        return {"status": "error", "stage": "baselines", "reason": "no baseline trained"}
+
+    best = pick_best(results, task)
+    reg = ctx.setdefault("results", {})
+    for r in results:
+        reg[r.name] = r
+    ctx["baseline_best"] = best
+
+    lc = _learning_curve(best, bundle, cfg)
+    path = _regularization_path(bundle, cfg)
+    _print_path(path)
+    note = _regularization_note(results, task, path)
+
+    return {
+        "status": "ok",
+        "stage": "baselines",
+        "primary_metric": primary_metric_name(task),
+        "best": best.name,
+        "models": [r.summary() for r in results],
+        "leaderboard": leaderboard(results, task),
+        "learning_curve": lc,
+        "regularization_path": path,
+        "note": note,
+    }

tabml_guard/boosting.py

@@ -0,0 +1,260 @@
+"""Stage 2 — Boosting ensembles + a stacking meta-model (the workhorses).
+
+Three gradient-boosting libraries (XGBoost, LightGBM, CatBoost), a RandomForest
+as an older-ensemble reference point, and a stacking model whose meta-learner is
+a regularized linear model — RegStack in spirit.
+
+Draws on Paper 1 (RegStack, ESWA 2025 — regularized stacking) and Paper 2
+(BO-XGBoost, ESWA 2026, DOI 10.1016/j.eswa.2026.131260 — XGBoost as the core
+learner, later tuned in Stage 3).
+
+Robustness choices worth knowing:
+  * **Optional dependencies.** Each booster is imported lazily; if a library
+    isn't installed (or fails to construct), that model is skipped with a message
+    and the pipeline carries on. The tool degrades gracefully instead of dying.
+  * **Stacking CV must be a partition.** scikit-learn's StackingClassifier builds
+    meta-features with ``cross_val_predict``, which requires every sample to fall
+    in exactly one validation fold. ``TimeSeriesSplit`` is *not* a partition
+    (its earliest rows never appear in a validation fold) and would raise
+    "cross_val_predict only works for partitions". So the meta-feature CV uses
+    stratified/plain k-fold. Time-honesty is preserved because the held-out
+    *temporal* test set is never part of this CV and is only scored at the end.
+  * **No nested oversubscription.** Base learners use all cores (n_jobs=-1);
+    stacking itself runs folds sequentially (n_jobs=1) so we don't spawn
+    cores^2 threads on a laptop.
+  * **Uniform preprocessing.** Trees don't need scaling/one-hot, but every model
+    shares the one leakage-safe preprocessor so the comparison is clean and the
+    system stays simple. (Scaling is harmless to trees.)
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+from .config import Config
+from .data import DataBundle
+from .metrics import primary_metric_name
+from .modeling import (
+    ModelResult,
+    feature_names_out,
+    fit_evaluate,
+    leaderboard,
+    log_result,
+    pick_best,
+)
+
+
+# --------------------------------------------------------------------------- #
+# Model builders (lazy imports: a missing library skips only that model)
+# --------------------------------------------------------------------------- #
+def _build_xgboost(task: str, rs: int):
+    import xgboost as xgb  # noqa: F401  (import may raise ImportError -> skipped)
+
+    params = dict(
+        n_estimators=300, learning_rate=0.1, max_depth=6,
+        subsample=0.9, colsample_bytree=0.9,
+        tree_method="hist",          # histogram splits -> scales to large data
+        random_state=rs, n_jobs=-1,
+    )
+    if task == "classification":
+        return xgb.XGBClassifier(eval_metric="logloss", **params)
+    return xgb.XGBRegressor(**params)
+
+
+def _build_lightgbm(task: str, rs: int):
+    import lightgbm as lgb
+
+    params = dict(
+        n_estimators=300, learning_rate=0.1, num_leaves=31, max_depth=-1,
+        subsample=0.9, colsample_bytree=0.9,
+        random_state=rs, n_jobs=-1, verbose=-1,   # verbose=-1 silences LightGBM
+    )
+    if task == "classification":
+        return lgb.LGBMClassifier(**params)
+    return lgb.LGBMRegressor(**params)
+
+
+def _build_catboost(task: str, rs: int):
+    from catboost import CatBoostClassifier, CatBoostRegressor
+
+    params = dict(
+        iterations=300, learning_rate=0.1, depth=6, random_seed=rs,
+        verbose=0,
+        allow_writing_files=False,   # don't litter a catboost_info/ dir (Windows!)
+        thread_count=-1,
+    )
+    if task == "classification":
+        return CatBoostClassifier(**params)
+    return CatBoostRegressor(**params)
+
+
+def _build_random_forest(task: str, rs: int):
+    from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+    params = dict(n_estimators=300, n_jobs=-1, random_state=rs)
+    if task == "classification":
+        return RandomForestClassifier(**params)
+    return RandomForestRegressor(**params)
+
+
+_BUILDERS = {
+    "xgboost": _build_xgboost,
+    "lightgbm": _build_lightgbm,
+    "catboost": _build_catboost,
+    "random_forest": _build_random_forest,
+}
+
+
+# --------------------------------------------------------------------------- #
+# Stacking meta-model
+# --------------------------------------------------------------------------- #
+def _build_stacking(task: str, available: list[str], rs: int, cv_splits: int):
+    """A stacking ensemble over the available boosters with a regularized meta-learner.
+
+    See the module docstring for why the internal CV is k-fold, not TimeSeriesSplit.
+    """
+    from sklearn.ensemble import StackingClassifier, StackingRegressor
+    from sklearn.linear_model import LogisticRegression, RidgeCV
+    from sklearn.model_selection import KFold, StratifiedKFold
+
+    # Fresh, unfitted base estimators (never reuse ones already fitted above).
+    estimators = [(name, _BUILDERS[name](task, rs)) for name in available]
+
+    if task == "classification":
+        cv = StratifiedKFold(n_splits=cv_splits, shuffle=True, random_state=rs)
+        # Regularized (L2 by default) linear meta-learner. No penalty= arg: it is
+        # the default on old sklearn and is deprecated on 1.8+ (would warn).
+        final = LogisticRegression(max_iter=2000)
+        return StackingClassifier(
+            estimators=estimators, final_estimator=final, cv=cv,
+            n_jobs=1, passthrough=False,
+        )
+    cv = KFold(n_splits=cv_splits, shuffle=True, random_state=rs)
+    final = RidgeCV()  # regularized linear meta-learner for regression
+    return StackingRegressor(
+        estimators=estimators, final_estimator=final, cv=cv,
+        n_jobs=1, passthrough=False,
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Feature-importance diagnostics (tree models expose feature_importances_)
+# --------------------------------------------------------------------------- #
+def _importance_diagnostics(res: ModelResult) -> dict:
+    model = res.estimator.named_steps["model"]
+    imp = getattr(model, "feature_importances_", None)
+    if imp is None:
+        return {}
+    imp = np.asarray(imp, dtype=float)
+    names = feature_names_out(res.estimator)
+    order = np.argsort(imp)[::-1][:15]
+    top = [
+        [names[i] if i < len(names) else f"f{i}", round(float(imp[i]), 5)]
+        for i in order
+    ]
+    return {"top_features": top}
+
+
+# --------------------------------------------------------------------------- #
+# Entry point
+# --------------------------------------------------------------------------- #
+def run(bundle: DataBundle, cfg: Config, ctx: Optional[dict] = None) -> dict:
+    """Train ensembles + stacking; return a leaderboard and stash the best models.
+
+    Side effects on ``ctx``:
+      * ``ctx['results']``       — every fitted ModelResult, keyed by name
+      * ``ctx['ensemble_best']`` — best Stage-2 model
+      * ``ctx['best_booster']``  / ``ctx['best_booster_name']`` — best *single*
+        booster (excludes stacking) -> what Stage 3 will tune
+      * ``ctx['best_overall']``  / ``ctx['best_model']`` — best across ALL stages
+        so far (baselines + ensembles) -> what Stages 4/5 explain & monitor
+    """
+    ctx = ctx if ctx is not None else {}
+    task = bundle.task
+    rs = cfg.random_state
+    requested = list(cfg.models.boosters)
+
+    results: list[ModelResult] = []
+    available: list[str] = []
+
+    for name in requested:
+        builder = _BUILDERS.get(name)
+        if builder is None:
+            print(f"[stage2] unknown booster '{name}' — skipped.")
+            continue
+        try:
+            est = builder(task, rs)
+        except ImportError as exc:
+            print(f"[stage2] '{name}' unavailable (library not installed): {exc} — skipped.")
+            continue
+        except Exception as exc:
+            print(f"[stage2] '{name}' could not be constructed: {exc} — skipped.")
+            continue
+        try:
+            res = fit_evaluate(name, est, bundle, cfg)
+            res.extra["importance"] = _importance_diagnostics(res)
+            results.append(res)
+            available.append(name)
+            log_result("stage2", res, task, cfg.report.fbeta)
+        except Exception as exc:
+            print(f"[stage2] '{name}' failed during fit/eval: {type(exc).__name__}: {exc} — skipped.")
+
+    # Stacking needs at least two working base learners to be meaningful.
+    if cfg.models.stacking:
+        if len(available) >= 2:
+            try:
+                stack = _build_stacking(task, available, rs, cfg.tuning.cv_splits)
+                res = fit_evaluate("stacking", stack, bundle, cfg)
+                results.append(res)
+                log_result("stage2", res, task, cfg.report.fbeta)
+            except Exception as exc:
+                print(f"[stage2] stacking failed: {type(exc).__name__}: {exc} — skipped.")
+        else:
+            print(f"[stage2] stacking skipped: needs >=2 base learners, have {len(available)}.")
+
+    if not results:
+        # Every booster was skipped (e.g. none of the libraries are installed).
+        # Don't starve downstream stages: fall back to the best baseline that
+        # Stage 1 already registered in the shared context.
+        reg = ctx.setdefault("results", {})
+        fallback = pick_best(list(reg.values()), task) if reg else None
+        if fallback is not None:
+            ctx["best_overall"] = fallback
+            ctx["best_model"] = fallback.estimator
+        return {
+            "status": "error",
+            "stage": "boosting",
+            "reason": "no ensemble model trained — are xgboost/lightgbm/catboost installed?",
+            "requested": requested,
+            "fell_back_to": fallback.name if fallback else None,
+        }
+
+    # Register in shared context.
+    reg = ctx.setdefault("results", {})
+    for r in results:
+        reg[r.name] = r
+
+    best_ensemble = pick_best(results, task)
+    single = [r for r in results if r.name in _BUILDERS]     # exclude stacking
+    best_booster = pick_best(single, task) if single else best_ensemble
+    best_overall = pick_best(list(reg.values()), task)        # incl. baselines
+
+    ctx["ensemble_best"] = best_ensemble
+    ctx["best_booster"] = best_booster
+    ctx["best_booster_name"] = best_booster.name if best_booster else None
+    ctx["best_overall"] = best_overall
+    ctx["best_model"] = best_overall.estimator if best_overall else None
+
+    return {
+        "status": "ok",
+        "stage": "boosting",
+        "primary_metric": primary_metric_name(task),
+        "available_boosters": available,
+        "skipped_boosters": [n for n in requested if n not in available],
+        "best": best_ensemble.name if best_ensemble else None,
+        "best_booster": best_booster.name if best_booster else None,
+        "best_overall": best_overall.name if best_overall else None,
+        "models": [r.summary() for r in results],
+        "leaderboard": leaderboard(list(reg.values()), task),
+    }