drift-or-shift 0.1.0a1__py3-none-any.whl

caliblab/__init__.py

@@ -0,0 +1,159 @@
+"""Simple benchmarking utilities for calibrated classifiers."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+import pandas as pd
+from sklearn.calibration import CalibratedClassifierCV
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import StratifiedKFold, train_test_split
+
+try:  # scikit-learn >= 1.6
+    from sklearn.frozen import FrozenEstimator
+except ImportError:  # pragma: no cover - scikit-learn < 1.6
+    FrozenEstimator = None
+
+_ESTIMATOR_REGISTRY = {
+    "logistic": LogisticRegression,
+    "random_forest": RandomForestClassifier,
+}
+
+
+def _make_estimator(spec: dict, random_state: int | None):
+    """Instantiate a base estimator from spec."""
+    estimator_type = spec.get("type")
+    if estimator_type not in _ESTIMATOR_REGISTRY:
+        raise ValueError(f"Unsupported estimator type: {estimator_type!r}")
+    estimator_cls = _ESTIMATOR_REGISTRY[estimator_type]
+    params = dict(spec.get("params", {}))
+    if random_state is not None:
+        params.setdefault("random_state", random_state)
+    if estimator_type == "logistic":
+        params.setdefault("max_iter", 1000)
+    return estimator_cls(**params)
+
+
+def _calibrate_prefit(estimator, method: str) -> CalibratedClassifierCV:
+    """Wrap an already-fitted estimator in a calibrator.
+
+    ``cv="prefit"`` is deprecated in scikit-learn 1.6 and removed in 1.8, so we
+    prefer ``FrozenEstimator`` where it exists and fall back only for older
+    releases still covered by our dependency floor.
+    """
+    if FrozenEstimator is not None:
+        return CalibratedClassifierCV(FrozenEstimator(estimator), method=method)
+    return CalibratedClassifierCV(  # pragma: no cover - scikit-learn < 1.6
+        estimator=estimator, method=method, cv="prefit"
+    )
+
+
+_MIN_CALIBRATION_PER_CLASS = 5
+
+
+def _validate_calibration_set(y_calib: np.ndarray, holdout: float) -> None:
+    """Fail early and clearly when the held-out split is too small to calibrate.
+
+    ``CalibratedClassifierCV`` cross-validates internally even over a frozen
+    estimator, so it needs at least as many members of each class as its
+    default number of folds. Left to sklearn this surfaces as "n_splits=5
+    cannot be greater than the number of members in each class", which says
+    nothing about the knob the caller actually turned.
+    """
+    counts = np.bincount(np.asarray(y_calib, dtype=int), minlength=2)
+    smallest = int(counts.min())
+    if smallest < _MIN_CALIBRATION_PER_CLASS:
+        raise ValueError(
+            f"calibration_holdout={holdout} leaves only {smallest} sample(s) of "
+            f"the rarest class to calibrate on; at least "
+            f"{_MIN_CALIBRATION_PER_CLASS} are needed. Raise the holdout "
+            f"fraction, use fewer CV folds, or pass more data."
+        )
+
+
+def benchmark(
+    X: np.ndarray,
+    y: np.ndarray,
+    estimators: Iterable[dict],
+    calibrations: Iterable[str],
+    *,
+    cv_params: dict,
+    calibration_holdout: float | None = None,
+) -> pd.DataFrame:
+    """Run calibration benchmarking over a dataset.
+
+    Args:
+        calibration_holdout: Fraction of each fold's training data to hold out
+            for fitting the calibration map. ``None`` (the default) fits the
+            calibrator on the same data the base estimator was trained on.
+
+    On fitting the calibrator on the estimator's own training data:
+
+    Textbook practice is to calibrate on held-out data, because a model's
+    predictions on data it has already seen are over-confident and the map
+    learned from them need not transfer. We measured it on the shipped
+    breast-cancer benchmark, holding the base model and the calibration-set
+    size fixed so that only the overlap differed, and the effect on test Brier
+    score was within noise::
+
+        logistic      sigmoid    -0.0014  (SE 0.0003)
+        logistic      isotonic   +0.0006  (SE 0.0009)
+        random_forest sigmoid    +0.0001  (SE 0.0003)
+        random_forest isotonic   -0.0002  (SE 0.0010)
+
+    The default therefore stays as it was, so published numbers do not move.
+    Pass ``calibration_holdout`` to opt into a held-out split; on a harder
+    dataset or a higher-capacity estimator the difference may well matter.
+
+    ``tests/test_caliblab_leakage.py`` re-runs that comparison, so the claim
+    above is checked rather than merely asserted.
+    """
+    if calibration_holdout is not None and not 0.0 < calibration_holdout < 1.0:
+        raise ValueError("calibration_holdout must lie in the open interval (0, 1)")
+    X = np.asarray(X)
+    y = np.asarray(y)
+    cv = StratifiedKFold(
+        n_splits=cv_params.get("n_splits", 5),
+        shuffle=True,
+        random_state=cv_params.get("random_state"),
+    )
+    results = []
+    for fold, (train_idx, test_idx) in enumerate(cv.split(X, y)):
+        X_train, X_test = X[train_idx], X[test_idx]
+        y_train, y_test = y[train_idx], y[test_idx]
+        X_fit, y_fit = X_train, y_train
+        X_calib, y_calib = X_train, y_train
+        if calibration_holdout is not None:
+            X_fit, X_calib, y_fit, y_calib = train_test_split(
+                X_train,
+                y_train,
+                test_size=calibration_holdout,
+                stratify=y_train,
+                random_state=cv_params.get("random_state"),
+            )
+            _validate_calibration_set(y_calib, calibration_holdout)
+        for spec in estimators:
+            for calibration in calibrations:
+                estimator = _make_estimator(spec, cv_params.get("random_state"))
+                estimator.fit(X_fit, y_fit)
+                predictor = estimator
+                if calibration != "none":
+                    calibrator = _calibrate_prefit(estimator, calibration)
+                    calibrator.fit(X_calib, y_calib)
+                    predictor = calibrator
+                probs = predictor.predict_proba(X_test)[:, 1]
+                for local_idx, prob in enumerate(probs):
+                    results.append(
+                        {
+                            "fold": fold,
+                            "estimator": spec.get("name", spec.get("type")),
+                            "estimator_type": spec.get("type"),
+                            "calibration": calibration,
+                            "sample_index": int(test_idx[local_idx]),
+                            "y_true": int(y_test[local_idx]),
+                            "y_prob": float(prob),
+                        }
+                    )
+    return pd.DataFrame(results)

drift_or_shift/__init__.py

@@ -0,0 +1,113 @@
+"""drift-or-shift: label shift and offset correction utilities."""
+
+__version__ = "0.1.0a1"
+
+from .calibration import (
+    apply_calibrator,
+    find_best_temperature,
+    fit_isotonic_calibrator,
+    temperature_scale,
+)
+from .data_synth import (
+    concept_drift_mu1,
+    default_mu0,
+    default_mu1,
+    make_gaussian_binary,
+    make_multimodal_binary,
+    resample_to_prevalence,
+)
+from .drift_monitor import feature_drift_summary, univariate_feature_stats
+from .drift_variants import (
+    apply_covariance_shift,
+    density_ratio_shift,
+    drift_score_from_ratio,
+    inject_label_noise,
+    shift_mean_vector,
+)
+from .ess import effective_sample_size, ess_fraction
+from .experiments._common import (
+    DEFAULT_COSTS,
+    DRIFT_FEATURE_METRICS,
+    PI_TEST_GRID,
+    PI_TRAIN,
+    SEEDS,
+    ExperimentConfig,
+    aggregate_mean_std,
+    feature_drift_metrics,
+)
+from .io_utils import ensure_dir, save_json, save_table, timestamped_run_dir
+from .metrics import oracle_threshold_min_risk, pr_auc, risk_cost_sensitive, roc_auc
+from .models import fit_logistic_regression, predict_logits
+from .plotting import (
+    plot_auc_pr_vs_prevalence,
+    plot_ess_vs_alpha,
+    plot_risk_vs_prevalence,
+)
+from .reporting import (
+    collect_summary_jsons,
+    format_results_overview,
+    select_latest_summaries,
+    write_results_overview,
+)
+from .shift import (
+    apply_logit_offset,
+    decision_from_logits,
+    logit_offset,
+    odds,
+    threshold_from_costs,
+    validate_prevalence,
+)
+
+__all__ = [
+    "DEFAULT_COSTS",
+    "DRIFT_FEATURE_METRICS",
+    "PI_TEST_GRID",
+    "PI_TRAIN",
+    "SEEDS",
+    "ExperimentConfig",
+    "__version__",
+    "aggregate_mean_std",
+    "apply_calibrator",
+    "apply_covariance_shift",
+    "apply_logit_offset",
+    "collect_summary_jsons",
+    "concept_drift_mu1",
+    "decision_from_logits",
+    "default_mu0",
+    "default_mu1",
+    "density_ratio_shift",
+    "drift_score_from_ratio",
+    "effective_sample_size",
+    "ensure_dir",
+    "ess_fraction",
+    "feature_drift_metrics",
+    "feature_drift_summary",
+    "find_best_temperature",
+    "fit_isotonic_calibrator",
+    "fit_logistic_regression",
+    "format_results_overview",
+    "inject_label_noise",
+    "logit_offset",
+    "make_gaussian_binary",
+    "make_multimodal_binary",
+    "odds",
+    "oracle_threshold_min_risk",
+    "plot_auc_pr_vs_prevalence",
+    "plot_ess_vs_alpha",
+    "plot_risk_vs_prevalence",
+    "pr_auc",
+    "predict_logits",
+    "resample_to_prevalence",
+    "risk_cost_sensitive",
+    "roc_auc",
+    "save_json",
+    "save_table",
+    "select_latest_summaries",
+    "shift_mean_vector",
+    "temperature_scale",
+    "threshold_from_costs",
+    "timestamped_run_dir",
+    "univariate_feature_stats",
+    "validate_prevalence",
+    "write_results_overview",
+]

drift_or_shift/calibration.py

@@ -0,0 +1,92 @@
+"""Calibration helpers for logistic outputs under label shift."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+from numpy.typing import ArrayLike
+from sklearn.isotonic import IsotonicRegression
+
+
+def _sigmoid(logits: np.ndarray) -> np.ndarray:
+    """Return probability after applying the logistic sigmoid function.
+
+    Evaluated piecewise. The direct form ``1 / (1 + exp(-x))`` overflows for
+    large negative x: the result saturates to the correct 0.0, but NumPy emits
+    an ``overflow encountered in exp`` warning on the way, which is noise that
+    can mask a real one. For x < 0 we use ``exp(x) / (1 + exp(x))`` instead,
+    where the exponential underflows harmlessly rather than overflowing.
+    """
+    x = np.asarray(logits, dtype=float)
+    out = np.empty_like(x)
+    positive = x >= 0
+    out[positive] = 1.0 / (1.0 + np.exp(-x[positive]))
+    exp_x = np.exp(x[~positive])
+    out[~positive] = exp_x / (1.0 + exp_x)
+    return out
+
+
+def temperature_scale(logits: ArrayLike, temperature: float) -> np.ndarray:
+    """
+    Adjust logits by a temperature factor.
+
+    Higher temperatures soften the distribution, while temperatures below 1 boost confidence.
+    """
+    if temperature <= 0:
+        raise ValueError("temperature must be positive")
+    logits_arr = np.asarray(logits, dtype=float)
+    return logits_arr / float(temperature)
+
+
+def find_best_temperature(
+    logits: ArrayLike,
+    targets: ArrayLike,
+    *,
+    grid: Iterable[float] | None = None,
+) -> float:
+    """
+    Return the temperature that minimizes logistic cross-entropy on the provided data.
+
+    A default grid spans from 0.05 to 3.0.
+    """
+    if grid is None:
+        grid = list(
+            np.concatenate([np.logspace(-2, -0.5, 5), np.linspace(0.1, 3.0, 50)])
+        )
+    logits_arr = np.asarray(logits, dtype=float)
+    targets_arr = np.asarray(targets, dtype=float)
+    if logits_arr.shape != targets_arr.shape:
+        raise ValueError("logits and targets must share shape")
+    if not (np.all((targets_arr == 0) | (targets_arr == 1))):
+        raise ValueError("targets must be binary (0 or 1)")
+    losses: list[tuple[float, float]] = []
+    eps = 1e-12
+    for temperature in grid:
+        if temperature <= 0:
+            continue
+        scaled = _sigmoid(logits_arr / temperature)
+        loss = -(
+            targets_arr * np.log(scaled + eps)
+            + (1 - targets_arr) * np.log(1 - scaled + eps)
+        ).mean()
+        losses.append((temperature, float(loss)))
+    if not losses:
+        raise ValueError("no valid temperatures found in the grid")
+    return min(losses, key=lambda item: item[1])[0]
+
+
+def fit_isotonic_calibrator(
+    scores: ArrayLike, targets: ArrayLike, *, out_of_bounds: str = "clip"
+) -> IsotonicRegression:
+    """
+    Fit an isotonic regression calibrator mapping scores to calibrated probabilities.
+    """
+    calib = IsotonicRegression(out_of_bounds=out_of_bounds)
+    calib.fit(np.asarray(scores, dtype=float), np.asarray(targets, dtype=float))
+    return calib
+
+
+def apply_calibrator(calibrator: IsotonicRegression, scores: ArrayLike) -> np.ndarray:
+    """Apply the fitted calibrator to a new set of scores."""
+    return calibrator.predict(np.asarray(scores, dtype=float))

drift_or_shift/data_synth.py

@@ -0,0 +1,170 @@
+"""Synthetic data and resampling helpers for label shift experiments."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import cast
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+def default_mu0(d: int) -> np.ndarray:
+    """Return the default negative-class mean (zeros)."""
+    return np.zeros(d)
+
+
+def default_mu1(d: int) -> np.ndarray:
+    """Return the default positive-class mean (first five dims at 1)."""
+    mu = np.zeros(d)
+    mu[: min(5, d)] = 1.0
+    return mu
+
+
+def concept_drift_mu1(
+    mu1: ArrayLike, shift_dim: int = 5, delta: float = 0.5
+) -> np.ndarray:
+    """Return mu1 shifted along one dimension to simulate concept drift."""
+    mu_arr = np.asarray(mu1, dtype=float)
+    if shift_dim >= mu_arr.shape[0]:
+        raise IndexError("shift_dim must be within the dimensionality of mu1.")
+    mu_arr = mu_arr.copy()
+    mu_arr[shift_dim] += delta
+    return mu_arr
+
+
+def _validate_prevalence(pi: float) -> None:
+    if not 0 < pi < 1:
+        raise ValueError("pi must lie in the open interval (0, 1).")
+
+
+def _build_covariance(sigma: float | np.ndarray, d: int) -> np.ndarray:
+    if np.isscalar(sigma):
+        scalar = float(cast(float, np.asarray(sigma, dtype=float)))
+        return np.eye(d, dtype=float) * scalar
+    cov = np.asarray(sigma, dtype=float)
+    if cov.shape != (d, d):
+        raise ValueError("Covariance must be scalar or a d-by-d array.")
+    return cov
+
+
+def make_gaussian_binary(
+    n: int,
+    d: int,
+    pi: float,
+    mu0: ArrayLike,
+    mu1: ArrayLike,
+    sigma: float | np.ndarray,
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate binary Gaussian data for label shift experiments."""
+    if n <= 0:
+        raise ValueError("n must be positive.")
+    if d <= 0:
+        raise ValueError("d must be positive.")
+    _validate_prevalence(pi)
+    mu0_arr = np.asarray(mu0, dtype=float)
+    mu1_arr = np.asarray(mu1, dtype=float)
+    if mu0_arr.shape != (d,) or mu1_arr.shape != (d,):
+        raise ValueError("Mean vectors must match dimensionality d.")
+    cov = _build_covariance(sigma, d)
+
+    probabilities = [1 - pi, pi]
+    labels = rng.choice([0, 1], size=n, p=probabilities)
+    X = np.empty((n, d), dtype=float)
+    if np.any(labels == 0):
+        count0 = np.sum(labels == 0)
+        X[labels == 0] = rng.multivariate_normal(mu0_arr, cov, size=count0)
+    if np.any(labels == 1):
+        count1 = np.sum(labels == 1)
+        X[labels == 1] = rng.multivariate_normal(mu1_arr, cov, size=count1)
+    return X, labels
+
+
+def resample_to_prevalence(
+    X: np.ndarray, y: np.ndarray, pi_target: float, rng: np.random.Generator
+) -> tuple[np.ndarray, np.ndarray]:
+    """Resample a fixed pool to a target prevalence via stratified sampling."""
+    _validate_prevalence(pi_target)
+    if X.shape[0] != y.shape[0]:
+        raise ValueError("X and y must have the same number of samples.")
+    n = len(y)
+    target_pos = round(pi_target * n)
+    target_neg = n - target_pos
+    pos_idx = np.where(y == 1)[0]
+    neg_idx = np.where(y == 0)[0]
+    replace_pos = target_pos > len(pos_idx)
+    replace_neg = target_neg > len(neg_idx)
+    chosen_pos = rng.choice(pos_idx, size=target_pos, replace=replace_pos)
+    chosen_neg = rng.choice(neg_idx, size=target_neg, replace=replace_neg)
+    indices = np.concatenate([chosen_pos, chosen_neg])
+    rng.shuffle(indices)
+    return X[indices], y[indices]
+
+
+def make_multimodal_binary(
+    n: int,
+    d: int,
+    pi: float,
+    components_neg: Sequence[tuple[ArrayLike, float]],
+    components_pos: Sequence[tuple[ArrayLike, float]],
+    sigma: float | np.ndarray,
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate a binary dataset where each class is a mixture of Gaussians."""
+    if n <= 0:
+        raise ValueError("n must be positive.")
+    if d <= 0:
+        raise ValueError("d must be positive.")
+    _validate_prevalence(pi)
+    neg_means, neg_weights = _prep_mixture(components_neg, d)
+    pos_means, pos_weights = _prep_mixture(components_pos, d)
+    cov = _build_covariance(sigma, d)
+    labels = rng.choice([0, 1], size=n, p=[1 - pi, pi])
+    X = np.empty((n, d), dtype=float)
+    if np.any(labels == 0):
+        X[labels == 0] = _sample_mixture(
+            cov, neg_means, neg_weights, rng, size=np.sum(labels == 0)
+        )
+    if np.any(labels == 1):
+        X[labels == 1] = _sample_mixture(
+            cov, pos_means, pos_weights, rng, size=np.sum(labels == 1)
+        )
+    return X, labels
+
+
+def _prep_mixture(
+    components: Sequence[tuple[ArrayLike, float]], d: int
+) -> tuple[list[np.ndarray], np.ndarray]:
+    if not components:
+        raise ValueError("At least one mixture component is required.")
+    means: list[np.ndarray] = []
+    weights = []
+    for mean, weight in components:
+        if weight < 0:
+            raise ValueError("Mixture weights must be non-negative.")
+        arr = np.asarray(mean, dtype=float)
+        if arr.shape != (d,):
+            raise ValueError("Component means must match dimensionality d.")
+        means.append(arr)
+        weights.append(weight)
+    weights_arr = np.asarray(weights, dtype=float)
+    total = float(weights_arr.sum())
+    if total <= 0:
+        raise ValueError("Mixture weights must sum to a positive value.")
+    return means, weights_arr / total
+
+
+def _sample_mixture(
+    cov: np.ndarray,
+    means: Sequence[np.ndarray],
+    weights: np.ndarray,
+    rng: np.random.Generator,
+    *,
+    size: int,
+) -> np.ndarray:
+    if size == 0:
+        return np.empty((0, cov.shape[0]))
+    indices = rng.choice(len(means), size=size, p=weights)
+    samples = [rng.multivariate_normal(means[idx], cov, size=1)[0] for idx in indices]
+    return np.vstack(samples)

drift_or_shift/drift_monitor.py

@@ -0,0 +1,124 @@
+"""Lightweight helpers for monitoring feature drift."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+def _ks_statistic(reference: np.ndarray, target: np.ndarray) -> float:
+    if reference.size == 0 or target.size == 0:
+        return 0.0
+    reference_sorted = np.sort(reference)
+    target_sorted = np.sort(target)
+    combined = np.sort(np.concatenate([reference_sorted, target_sorted]))
+    cdf_ref = (
+        np.searchsorted(reference_sorted, combined, side="right")
+        / reference_sorted.size
+    )
+    cdf_target = (
+        np.searchsorted(target_sorted, combined, side="right") / target_sorted.size
+    )
+    return float(np.max(np.abs(cdf_ref - cdf_target)))
+
+
+def _covariance_matrix(data: np.ndarray) -> np.ndarray:
+    return np.atleast_2d(np.cov(data, rowvar=False))
+
+
+def _correlation_matrix(data: np.ndarray) -> np.ndarray:
+    if data.shape[1] == 0:
+        return np.zeros((0, 0))
+    corr = np.corrcoef(data, rowvar=False)
+    if corr.ndim == 0:
+        corr = np.array([[corr]])
+    return np.nan_to_num(corr, nan=0.0, posinf=0.0, neginf=0.0)
+
+
+def covariance_frobenius_diff(reference: np.ndarray, target: np.ndarray) -> float:
+    """Return the Frobenius difference between training and target covariance matrices."""
+    cov_ref = _covariance_matrix(reference)
+    cov_target = _covariance_matrix(target)
+    return float(np.linalg.norm(cov_ref - cov_target, ord="fro"))
+
+
+def correlation_mean_diff(reference: np.ndarray, target: np.ndarray) -> float:
+    """Compute average absolute deviation between correlation matrices (off-diagonal)."""
+    d = reference.shape[1]
+    if d <= 1:
+        return 0.0
+    corr_ref = _correlation_matrix(reference)
+    corr_target = _correlation_matrix(target)
+    mask = ~np.eye(d, dtype=bool)
+    diffs = np.abs(corr_ref - corr_target)
+    return float(np.mean(diffs[mask]))
+
+
+def multivariate_projection_ks(
+    reference: np.ndarray,
+    target: np.ndarray,
+    *,
+    projections: int = 12,
+    rng: np.random.Generator | None = None,
+) -> list[float]:
+    """Approximate multivariate drift via KS on random projections."""
+    rng = rng or np.random.default_rng(0)
+    d = reference.shape[1]
+    if d == 0 or projections <= 0:
+        return [0.0]
+    stats: list[float] = []
+    for _ in range(projections):
+        direction = np.asarray(rng.normal(size=d), dtype=float)
+        norm = np.linalg.norm(direction)
+        if norm == 0:
+            continue
+        direction /= norm
+        ref_proj = reference @ direction
+        target_proj = target @ direction
+        stats.append(_ks_statistic(ref_proj, target_proj))
+    return stats or [0.0]
+
+
+def univariate_feature_stats(
+    X_ref: ArrayLike,
+    X_target: ArrayLike,
+) -> list[dict]:
+    """Return per-feature drift statistics between reference and target arrays."""
+    ref_arr = np.asarray(X_ref, dtype=float)
+    target_arr = np.asarray(X_target, dtype=float)
+    if ref_arr.ndim != 2 or target_arr.ndim != 2:
+        raise ValueError("input arrays must be two-dimensional")
+    if ref_arr.shape[1] != target_arr.shape[1]:
+        raise ValueError("feature dimensionality must match")
+    if ref_arr.shape[0] == 0 or target_arr.shape[0] == 0:
+        raise ValueError("input datasets must contain samples")
+    stats: list[dict] = []
+    for idx in range(ref_arr.shape[1]):
+        ref_col = ref_arr[:, idx]
+        target_col = target_arr[:, idx]
+        stats.append(
+            {
+                "feature": idx,
+                "mean_diff": float(np.abs(ref_col.mean() - target_col.mean())),
+                "std_diff": float(np.abs(ref_col.std(ddof=0) - target_col.std(ddof=0))),
+                "ks_stat": _ks_statistic(ref_col, target_col),
+            }
+        )
+    return stats
+
+
+def feature_drift_summary(stats: Sequence[dict]) -> dict[str, float]:
+    """Aggregate univariate stats into summary metrics suitable for tracking."""
+    if not stats:
+        raise ValueError("stats cannot be empty")
+    mean_diffs = [entry["mean_diff"] for entry in stats]
+    std_diffs = [entry["std_diff"] for entry in stats]
+    ks_stats = [entry["ks_stat"] for entry in stats]
+    return {
+        "feature_max_mean_diff": float(max(mean_diffs)),
+        "feature_max_std_diff": float(max(std_diffs)),
+        "feature_max_ks": float(max(ks_stats)),
+        "feature_mean_ks": float(np.mean(ks_stats)),
+    }