h2ml 0.1.0__py3-none-any.whl

h2ml/__init__.py

@@ -0,0 +1,2 @@
+from h2ml.pipeline.pipeline import H2MLPipeline, PipelineConfig, PipelineData, PipelineResult  # noqa: F401
+from h2ml.pipeline.base import TaskType  # noqa: F401

h2ml/evaluation/compare.py

@@ -0,0 +1,283 @@
+"""
+h2ml/evaluation/compare.py
+
+compare_results() — compare multiple PipelineResult objects in one DataFrame.
+
+Each row represents one pipeline run. Useful for selecting the best configuration
+across runs with different PipelineConfig settings, feature schemas, or datasets.
+"""
+from __future__ import annotations
+
+from typing import Optional, TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+from loguru import logger
+
+from h2ml.pipeline.pipeline import _MINIMIZE
+
+if TYPE_CHECKING:
+    from h2ml.pipeline.pipeline import PipelineResult
+
+
+def compare_results(
+    results: list["PipelineResult"],
+    labels:  Optional[list[str]] = None,
+    metric:  Optional[str]       = None,
+    n_folds: Optional[int]       = None,
+) -> pd.DataFrame:
+    """
+    Summarise multiple PipelineResult objects as a single comparison DataFrame.
+
+    One row per result. Columns:
+
+        Run                — label (from labels, or Run_0 / Run_1 / …)
+        Metric             — short metric name used for Score_Mean (e.g. "AUC", "R2")
+        Best_Model         — winning model class name
+        Best_Stage         — default | reduced | optimized
+        Y_Transform        — winning y-transform, or None
+        Score_Mean         — mean CV score for the selected metric
+        Score_Std          — fold std for the selected metric
+        Conservative_Bound — variance-penalised score used for ranking. Direction is
+                             derived automatically from the metric name via _MINIMIZE:
+                               higher-is-better: Score_Mean - Score_Std / sqrt(n_folds)
+                               lower-is-better:  Score_Mean + Score_Std / sqrt(n_folds)
+                             A model with high mean but high variance is pulled toward
+                             a worse value, so stable models are preferred over
+                             optimistically-estimated ones.
+        Brier_Mean         — Brier_Test_Mean from the winning stage's agg_df
+                             (classification only; NaN otherwise)
+        OOF_Brier          — Brier score computed on the assembled OOF predictions
+                             (more robust than the fold-averaged Brier_Mean;
+                             classification only; NaN otherwise)
+        N_Features         — feature count in the winning feature stage
+        Completed_Steps    — list of pipeline steps that finished
+
+    Args:
+        results: List of PipelineResult objects to compare.
+        labels:  Optional run labels (same length as results).
+                 Defaults to "Run_0", "Run_1", …
+        metric:  Short metric name to use as the common comparison column
+                 (e.g. "AUC", "R2", "RMSE"). When provided, Score_Mean and
+                 Score_Std are read from the ``<metric>_Test_Mean/Std`` columns
+                 of each result's winning-stage agg DataFrame, making
+                 cross-run scores directly comparable even when runs were
+                 optimised on different metrics. When omitted, Score_Mean falls
+                 back to each result's ``best_model_value``. Minimisation
+                 direction and sort order are derived automatically from
+                 ``_MINIMIZE`` — no need to pass ``ascending`` manually.
+        n_folds: Number of CV folds used in all runs. When provided it overrides
+                 automatic inference from the fold DataFrames, which can fail if
+                 results were loaded from disk or the fold DataFrames are absent.
+                 Conservative_Bound cannot be computed without a fold count and
+                 will fall back to Score_Mean if both this argument and
+                 inference fail.
+
+    Returns:
+        DataFrame sorted by Conservative_Bound (falls back to Score_Mean when
+        fold count cannot be inferred), one row per result.
+
+    Example:
+        >>> r1 = pipeline_a.run(store)
+        >>> r2 = pipeline_b.run(store)
+        >>> compare_results([r1, r2], labels=["baseline", "spatial_cv"], metric="AUC")
+    """
+    if labels is not None and len(labels) != len(results):
+        raise ValueError(
+            f"labels has {len(labels)} entries but results has {len(results)}."
+        )
+
+    if metric is None:
+        unique_metrics = {r.metric for r in results if r.metric is not None}
+        if len(unique_metrics) > 1:
+            raise ValueError(
+                f"compare_results: results use different metrics {sorted(unique_metrics)} — "
+                "Score_Mean values are not directly comparable. "
+                "Pass metric= to select a common comparison metric."
+            )
+
+    rows = []
+    for i, result in enumerate(results):
+        label        = labels[i] if labels is not None else f"Run_{i}"
+        fold_count   = n_folds if n_folds is not None else _infer_n_folds(result)
+
+        if metric is not None:
+            score_mean, score_std = _get_score_for_metric(result, metric)
+            display_metric = metric
+        else:
+            score_mean     = result.best_model_value
+            score_std      = result.best_model_std
+            display_metric = result.metric
+
+        minimize = _MINIMIZE.get(display_metric, False) if display_metric else False
+
+        conservative_bound: Optional[float] = None
+        if score_mean is not None and score_std is not None and fold_count:
+            penalty = score_std / np.sqrt(fold_count)
+            conservative_bound = score_mean + penalty if minimize else score_mean - penalty
+
+        rows.append({
+            "Run":                label,
+            "Metric":             display_metric,
+            "Best_Model":         result.best_model_name,
+            "Best_Stage":         result.best_stage,
+            "Y_Transform":        result.y_transform,
+            "Score_Mean":         score_mean,
+            "Score_Std":          score_std,
+            "Conservative_Bound": conservative_bound,
+            "Brier_Mean":         _get_brier_from_agg(result),
+            "OOF_Brier":          _get_oof_brier(result),
+            "N_Features":         _get_n_features(result),
+            "Completed_Steps":    result.completed_steps,
+        })
+
+    df = pd.DataFrame(rows)
+    if df.empty:
+        return df
+
+    # Sort direction derived from metric, not from an ascending argument
+    if metric is not None:
+        sort_ascending = _MINIMIZE.get(metric, False)
+    else:
+        unique_minimize = {_MINIMIZE.get(r.metric, False) for r in results if r.metric}
+        sort_ascending  = unique_minimize.pop() if len(unique_minimize) == 1 else False
+
+    if df["Conservative_Bound"].isna().all():
+        logger.warning(
+            "compare_results: Conservative_Bound could not be computed for any result "
+            "(fold count unknown and inference failed). Sorting by Score_Mean instead. "
+            "Pass n_folds= to enable variance penalisation."
+        )
+        sort_col = "Score_Mean"
+    else:
+        sort_col = "Conservative_Bound"
+
+    if sort_col in df.columns and df[sort_col].notna().any():
+        df = df.sort_values(sort_col, ascending=sort_ascending, ignore_index=True)
+
+    return df
+
+
+# ---------------------------------------------------------------------------
+# Private helpers
+# ---------------------------------------------------------------------------
+
+def _get_score_for_metric(
+    result: "PipelineResult",
+    metric: str,
+) -> tuple[Optional[float], Optional[float]]:
+    """Return (mean, std) for *metric* from the winning stage's agg DataFrame."""
+    mean_col = f"{metric}_Test_Mean"
+    std_col  = f"{metric}_Test_Std"
+
+    stage = result.best_stage
+    if stage == "optimized":
+        agg_df = result.step4_agg_df
+    elif stage == "reduced":
+        agg_df = result.step3_agg_df
+    else:
+        agg_df = result.step1_agg_df
+
+    if agg_df is None or mean_col not in agg_df.columns or agg_df.empty:
+        logger.warning(
+            f"compare_results: column '{mean_col}' not found in the winning-stage "
+            f"agg DataFrame — Score_Mean will be NaN for this result."
+        )
+        return None, None
+
+    mask = agg_df["Model"] == result.best_model_name
+    if result.y_transform and "Y_transform" in agg_df.columns:
+        mask = mask & (agg_df["Y_transform"] == result.y_transform)
+    if stage == "reduced" and "Stage" in agg_df.columns:
+        mask = mask & (agg_df["Stage"] == "reduced")
+
+    subset = agg_df.loc[mask]
+    if subset.empty:
+        return None, None
+
+    row      = subset.iloc[0]
+    mean_val = float(row[mean_col])
+    std_val  = float(row[std_col]) if std_col in subset.columns else None
+    return mean_val, std_val
+
+
+def _infer_n_folds(result: "PipelineResult") -> Optional[int]:
+    """Infer fold count, preferring the stored splitter to avoid under-counting when folds fail."""
+    if result.splitter is not None and hasattr(result.splitter, "n_splits"):
+        return result.splitter.n_splits
+    for df in (result.step1_fold_df, result.step3_fold_df, result.step4_fold_df):
+        if df is not None and "Fold" in df.columns and not df.empty:
+            return int(df["Fold"].max()) + 1
+    return None
+
+
+def _get_n_features(result: "PipelineResult") -> Optional[int]:
+    """Return the feature count for the winning feature stage."""
+    feature_stage = result.best_feature_stage or result.best_stage
+    if feature_stage == "reduced" and result.features_reduced is not None:
+        return result.features_reduced.n_features
+    if result.features is not None:
+        return result.features.n_features
+    return None
+
+
+def _get_brier_from_agg(result: "PipelineResult") -> Optional[float]:
+    """
+    Extract Brier_Test_Mean for the winning (model, stage, transform) from
+    whichever agg_df corresponds to the best_stage.
+    Returns None when Brier was not computed or the column is absent.
+    """
+    stage = result.best_stage
+
+    if stage == "optimized":
+        df = result.step4_agg_df
+        if df is not None and "Brier_Test_Mean" in df.columns and not df.empty:
+            return float(df.iloc[0]["Brier_Test_Mean"])
+        return None
+
+    df = result.step3_agg_df if stage == "reduced" else result.step1_agg_df
+    if df is None or "Brier_Test_Mean" not in df.columns:
+        return None
+
+    mask = df["Model"] == result.best_model_name
+    if result.y_transform and "Y_transform" in df.columns:
+        mask = mask & (df["Y_transform"] == result.y_transform)
+    if stage == "reduced" and "Stage" in df.columns:
+        mask = mask & (df["Stage"] == "reduced")
+
+    subset = df.loc[mask, "Brier_Test_Mean"]
+    return float(subset.iloc[0]) if not subset.empty else None
+
+
+def _get_oof_brier(result: "PipelineResult") -> Optional[float]:
+    """
+    Compute Brier score on assembled OOF predictions.
+    Only meaningful for classification — returns None for regression or when
+    OOF data is unavailable.
+    """
+    from sklearn.metrics import brier_score_loss
+    from h2ml.pipeline.base import TaskType
+
+    cv = result.best_cv_result
+    if cv is None or cv.task_type != TaskType.CLASSIFICATION:
+        return None
+
+    oof_pred   = result.oof_predictions
+    oof_labels = result.oof_labels
+    if oof_pred is None or oof_labels is None:
+        return None
+    # Drop missing positions (failed folds); use pd.isnull to handle both
+    # float NaN (numeric labels) and None (string/object labels from oof_labels)
+    valid = ~(np.isnan(oof_pred) if oof_pred.ndim == 1 else np.isnan(oof_pred).any(axis=1))
+    valid &= ~pd.isnull(oof_labels)
+    if not valid.any():
+        return None
+    y_true = oof_labels[valid]
+    y_prob = oof_pred[valid]
+    if y_prob.ndim == 2:
+        # Multiclass Brier
+        from sklearn.preprocessing import label_binarize
+        classes = np.unique(y_true.astype(int))
+        y_bin   = label_binarize(y_true, classes=classes)
+        return float(np.mean(np.sum((y_bin - y_prob) ** 2, axis=1)))
+    return float(brier_score_loss(y_true, y_prob))

h2ml/evaluation/metrics.py

@@ -0,0 +1,359 @@
+"""
+Consumes CVResult objects and produces metrics DataFrames.
+Responsibilities:
+    - Compute per-fold metrics (classification or regression)
+    - Aggregate fold metrics (mean, std across folds)
+    - Attach domain metadata (species, setting, batch, schema)
+    - Detect overfitting gap between train and test metrics
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+from sklearn.metrics import (
+    # Classification
+    average_precision_score,
+    brier_score_loss,
+    f1_score,
+    log_loss,
+    roc_auc_score,
+    # Regression
+    mean_absolute_error,
+    mean_squared_error,
+    r2_score
+)
+
+from h2ml.pipeline.base import TaskType
+from h2ml.pipeline.cv import CVResult, FoldResult
+
+# ---------------------------------------------------------------------------
+# Metadata container — replaces PipelineSpecs fields in results
+# ---------------------------------------------------------------------------
+
+@dataclass
+class RunMetadata:
+    """
+    Domain-specific context attached to results after CV.
+    Keeps metadata out of the CV engine.
+
+    Args:
+        schema: top-level folder / experiment name
+        target:  target column name.
+        stage:
+            - "default":   step 1 (all features, default params)
+            - "reduced":   step 3 (reduced features, default params)
+            - "optimized": step 4 (reduced or full features, optimized params)
+        batch:    Optional batch/run identifier.
+        schema:   Optional feature schema label.
+    """
+    # Experiment ID
+    schema:  Optional[str] = None
+    target: Optional[str] = None
+
+    # Pipeline stage — set automatically by pipeline._metadata_with()
+    stage: str = "default"  # "default" | "reduced" | "optimized"
+
+    # Optional grouping
+    batch:       Optional[str] = None   # run ID, date, or any grouping label
+    y_transform: Optional[str] = None   # y-transform name, set by run_transform_sweep()
+    notes:       Optional[str] = None   # free text — useful for experiments log
+
+    def to_dict(self) -> dict:
+        return {k.capitalize(): v for k, v in self.__dict__.items() if v is not None}
+
+
+# ---------------------------------------------------------------------------
+# Per-fold metric rows
+# ---------------------------------------------------------------------------
+
+def _classification_metrics(fold: FoldResult) -> dict:
+    """Compute classification metrics for a single fold.
+
+    Supports both binary and multiclass targets. Detection is based on the shape
+    of y_prob_test: 1-D → binary, 2-D → multiclass (OvR weighted averages).
+
+    Degenerate folds (only one class in y_test or y_train) arise with spatial CV
+    on imbalanced targets. Rank-based metrics (AUC, AUC-PR) return NaN for those
+    folds; other metrics use explicit class labels to avoid ValueError.
+    """
+    if fold.y_prob_test is None or fold.y_prob_train is None:
+        raise ValueError(
+            f"[{fold.model_name}] Fold {fold.fold_id} has no probabilities. "
+            "Did you use a BaseClassifier?"
+        )
+
+    is_multiclass = fold.y_prob_test.ndim == 2
+
+    # All classes seen across both splits — needed to keep metric calls stable
+    # even when one split is missing a class (common with spatial CV + imbalance).
+    all_classes = sorted(np.unique(np.concatenate([fold.y_train, fold.y_test])).tolist())
+    single_class_test  = len(np.unique(fold.y_test))  < 2
+    single_class_train = len(np.unique(fold.y_train)) < 2
+
+    if is_multiclass:
+        from sklearn.preprocessing import label_binarize
+        y_test_bin  = label_binarize(fold.y_test,  classes=all_classes)
+        y_train_bin = label_binarize(fold.y_train, classes=all_classes)
+        auc_test  = (np.nan if single_class_test  else
+                     roc_auc_score(fold.y_test,  fold.y_prob_test,  multi_class="ovr", average="weighted"))
+        auc_train = (np.nan if single_class_train else
+                     roc_auc_score(fold.y_train, fold.y_prob_train, multi_class="ovr", average="weighted"))
+        auc_pr_test  = (np.nan if single_class_test  else
+                        average_precision_score(y_test_bin,  fold.y_prob_test,  average="weighted"))
+        auc_pr_train = (np.nan if single_class_train else
+                        average_precision_score(y_train_bin, fold.y_prob_train, average="weighted"))
+        logloss_test  = log_loss(fold.y_test,  fold.y_prob_test,  labels=all_classes)
+        logloss_train = log_loss(fold.y_train, fold.y_prob_train, labels=all_classes)
+        f1_test  = f1_score(fold.y_test,  fold.y_pred_test,  average="weighted", zero_division=0)
+        f1_train = f1_score(fold.y_train, fold.y_pred_train, average="weighted", zero_division=0)
+        # Multiclass Brier: mean over samples of sum-of-squared probability errors
+        brier_test  = float(np.mean(np.sum((y_test_bin  - fold.y_prob_test)**2, axis=1)))
+        brier_train = float(np.mean(np.sum((y_train_bin - fold.y_prob_train)**2, axis=1)))
+    else:
+        auc_test  = (np.nan if single_class_test  else roc_auc_score(fold.y_test,  fold.y_prob_test))
+        auc_train = (np.nan if single_class_train else roc_auc_score(fold.y_train, fold.y_prob_train))
+        auc_pr_test  = (np.nan if single_class_test  else average_precision_score(fold.y_test,  fold.y_prob_test))
+        auc_pr_train = (np.nan if single_class_train else average_precision_score(fold.y_train, fold.y_prob_train))
+        logloss_test  = log_loss(fold.y_test,  fold.y_prob_test,  labels=all_classes)
+        logloss_train = log_loss(fold.y_train, fold.y_prob_train, labels=all_classes)
+        f1_test  = f1_score(fold.y_test,  fold.y_pred_test,  zero_division=0, labels=all_classes)
+        f1_train = f1_score(fold.y_train, fold.y_pred_train, zero_division=0, labels=all_classes)
+        brier_test  = float(brier_score_loss(fold.y_test,  fold.y_prob_test))
+        brier_train = float(brier_score_loss(fold.y_train, fold.y_prob_train))
+
+    return {
+        "AUC_Test":      auc_test,
+        "AUC_Train":     auc_train,
+        "AUC_PR_Test":   auc_pr_test,
+        "AUC_PR_Train":  auc_pr_train,
+        "LogLoss_Test":  logloss_test,
+        "LogLoss_Train": logloss_train,
+        "F1_Test":       f1_test,
+        "F1_Train":      f1_train,
+        "Brier_Test":    brier_test,
+        "Brier_Train":   brier_train,
+        "Fit_Time":      fold.fit_time,
+    }
+
+
+def _regression_metrics(fold: FoldResult) -> dict:
+    """Compute regression metrics for a single fold."""
+    mse_test  = mean_squared_error(fold.y_test,  fold.y_pred_test)
+    mse_train = mean_squared_error(fold.y_train, fold.y_pred_train)
+    return {
+        "R2_Test":    r2_score(fold.y_test,  fold.y_pred_test),
+        "R2_Train":   r2_score(fold.y_train, fold.y_pred_train),
+        "MAE_Test":   mean_absolute_error(fold.y_test,  fold.y_pred_test),
+        "MAE_Train":  mean_absolute_error(fold.y_train, fold.y_pred_train),
+        "RMSE_Test":  np.sqrt(mse_test),
+        "RMSE_Train": np.sqrt(mse_train),
+        "Fit_Time":   fold.fit_time,
+    }
+
+
+# Dispatcher — routes to the right metric function by task type
+_METRIC_FN = {
+    TaskType.CLASSIFICATION: _classification_metrics,
+    TaskType.REGRESSION:     _regression_metrics,
+}
+
+
+# ---------------------------------------------------------------------------
+# Overfitting gap
+# ---------------------------------------------------------------------------
+
+# Metric triples (test_col, train_col, minimize).
+# Gap convention: positive always means the model does worse on test than train (overfitting).
+#   maximize metrics (AUC, R2, F1):  gap = train - test   → positive = overfitting
+#   minimize metrics (LogLoss, MAE):  gap = test  - train  → positive = overfitting
+_OVERFIT_PAIRS = {
+    TaskType.CLASSIFICATION: [
+        ("AUC_Test",     "AUC_Train",     False),
+        ("AUC_PR_Test",  "AUC_PR_Train",  False),
+        ("F1_Test",      "F1_Train",      False),
+        ("LogLoss_Test", "LogLoss_Train", True),
+        ("Brier_Test",   "Brier_Train",   True),
+    ],
+    TaskType.REGRESSION: [
+        ("R2_Test",   "R2_Train",   False),
+        ("MAE_Test",  "MAE_Train",  True),
+        ("RMSE_Test", "RMSE_Train", True),
+    ],
+}
+
+
+def _add_overfit_gap(df: pd.DataFrame, task_type: TaskType) -> pd.DataFrame:
+    """
+    Adds Gap columns for key metrics. A positive gap always means the model
+    performs worse on test than train, regardless of metric direction.
+    """
+    for test_col, train_col, minimize in _OVERFIT_PAIRS.get(task_type, []):
+        if test_col in df.columns and train_col in df.columns:
+            metric_name = test_col.replace("_Test", "")
+            df[f"{metric_name}_Gap"] = (
+                df[test_col] - df[train_col] if minimize
+                else df[train_col] - df[test_col]
+            )
+    return df
+
+
+# ---------------------------------------------------------------------------
+# Main public API
+# ---------------------------------------------------------------------------
+
+def compute_metrics(
+    cv_result: CVResult,
+    metadata:  Optional[RunMetadata] = None,
+) -> pd.DataFrame:
+    """
+    Compute per-fold metrics for a single CVResult.
+
+    Args:
+        cv_result: Output of CrossValidator.run() for one model.
+        metadata:  Optional RunMetadata to attach to every row.
+
+    Returns:
+        DataFrame with one row per fold.
+    """
+    metric_fn = _METRIC_FN.get(cv_result.task_type)
+    if metric_fn is None:
+        raise ValueError(f"Unsupported task type: {cv_result.task_type}")
+
+    rows = []
+    for fold in cv_result.folds:
+        row = {"Model": fold.model_name, "Fold": fold.fold_id}
+        row.update(metric_fn(fold))
+        if metadata:
+            row.update(metadata.to_dict())
+        rows.append(row)
+
+    df = pd.DataFrame(rows)
+    df = _add_overfit_gap(df, cv_result.task_type)
+    return df
+
+
+def compute_metrics_all(
+    cv_results: list[CVResult],
+    metadata:   Optional[RunMetadata] = None,
+) -> pd.DataFrame:
+    """
+    Compute per-fold metrics for a list of CVResults (all models).
+
+    Args:
+        cv_results: Output of CrossValidator.run_all().
+        metadata:   Optional RunMetadata to attach to every row.
+
+    Returns:
+        DataFrame with one row per fold per model.
+    """
+    frames = [compute_metrics(r, metadata) for r in cv_results if r.folds]
+    if not frames:
+        raise RuntimeError("All CV results are empty — every model/fold failed.")
+    return pd.concat(frames, ignore_index=True)
+
+
+def aggregate_metrics(
+        fold_df: pd.DataFrame,
+        group_by_stage: bool = True
+        ) -> pd.DataFrame:
+    """
+    Aggregate per-fold results into mean ± std per model.
+
+    Args:
+        fold_df: Output of compute_metrics() or compute_metrics_all().
+        group_by_stage: If True (default), group by Stage alongside Model and other
+                        metadata columns so each (model, stage) pair gets its own row.
+                        Set False to collapse all stages per model — used in step 4
+                        when the optimized run's folds should be averaged without a
+                        Stage breakdown.
+
+    Returns:
+        DataFrame with one row per model, columns suffixed _Mean and _Std.
+    """
+    # Columns to group by — model name + any metadata fields present in the df
+    group_cols = ["Model"]
+    meta_fields = ["Schema", "Target", "Batch", "Y_transform", "Stage", "Notes"]
+
+    if group_by_stage:
+        group_cols += [c for c in meta_fields if c in fold_df.columns]
+    else:
+        # Exclude Stage — pool all folds across stages per model/transform
+        group_cols += [c for c in meta_fields if c in fold_df.columns
+                       and c != "Stage"]
+
+    numeric_cols = fold_df.select_dtypes(include=np.number).columns.tolist()
+    # Exclude fold id from aggregation
+    numeric_cols = [c for c in numeric_cols if c != "Fold"]
+
+    mean_df = fold_df.groupby(group_cols)[numeric_cols].mean().add_suffix("_Mean")
+    std_df  = fold_df.groupby(group_cols)[numeric_cols].std().add_suffix("_Std")
+
+    agg_df = pd.concat([mean_df, std_df], axis=1).reset_index()
+
+    # Interleave Mean/Std columns for readability: AUC_Test_Mean, AUC_Test_Std, ...
+    metric_cols = []
+    for col in numeric_cols:
+        metric_cols += [f"{col}_Mean", f"{col}_Std"]
+
+    ordered_cols = group_cols + [c for c in metric_cols if c in agg_df.columns]
+    return agg_df[ordered_cols]
+
+
+def select_best(
+    agg_df:   pd.DataFrame,
+    metric:   str           = "AUC_Test_Mean",
+    minimize: bool          = False,
+    n_folds:  Optional[int] = None,
+) -> dict:
+    """
+    Select the best model from an aggregated metrics DataFrame.
+
+    When n_folds is provided, ranking uses a lower/upper confidence bound
+    (LCB for maximize metrics, UCB for minimize metrics):
+
+        maximize:  score = mean - std / sqrt(n_folds)   → idxmax
+        minimize:  score = mean + std / sqrt(n_folds)   → idxmin
+
+    This prefers models whose conservative-bound performance is highest,
+    penalising fold variance in proportion to 1/sqrt(n_folds).
+    When n_folds is None the raw mean is used (original behaviour).
+
+    Args:
+        agg_df:   Output of aggregate_metrics().
+        metric:   Column name to rank by (must end in _Mean, e.g. AUC_Test_Mean).
+        minimize: If True, select the row with the lowest value (e.g. RMSE, MAE).
+        n_folds:  Number of CV folds used to produce agg_df. When provided,
+                  selection uses confidence-bound scoring instead of raw mean.
+
+    Returns:
+        Dict with keys model_name, metric, value, row.
+    """
+    if metric not in agg_df.columns:
+        raise ValueError(
+            f"Metric '{metric}' not found. Available: {list(agg_df.columns)}"
+        )
+
+    std_col = metric.replace("_Mean", "_Std")
+    if n_folds is not None and std_col in agg_df.columns:
+        penalty = agg_df[std_col] / np.sqrt(n_folds)
+        score   = agg_df[metric] + penalty if minimize else agg_df[metric] - penalty
+    else:
+        score = agg_df[metric]
+
+    if score.isna().all():
+        raise RuntimeError(
+            f"select_best: all values in column '{metric}' are NaN. "
+            "Every model may have failed all CV folds. Check cv_warnings on the result."
+        )
+    idx      = score.idxmin() if minimize else score.idxmax()
+    best_row = agg_df.loc[idx]
+
+    return {
+        "model_name": best_row["Model"],
+        "metric":     metric,
+        "value":      best_row[metric],
+        "row":        best_row,
+    }