threshold-opt 0.1.0__py3-none-any.whl

threshold_opt/calibration.py

@@ -0,0 +1,146 @@
+"""Post-hoc calibration methods.
+
+Implements Dirichlet calibration (Kull et al., NeurIPS 2019) as a
+self-contained module, avoiding an external dependency on `dirichletcal`.
+"""
+
+import numpy as np
+import numpy.typing as npt
+from scipy.optimize import minimize
+from scipy.special import softmax
+
+
+class DirichletCalibrator:
+    """Dirichlet calibration (Kull et al., NeurIPS 2019).
+
+    Learns a linear mapping in log-probability space:
+
+        calibrated = softmax(W @ log(p) + b)
+
+    where W is (K x K) and b is (K,), parameterised as a single
+    (K x K+1) weight matrix.  Trained by minimising NLL on a held-out
+    calibration set.
+
+    Parameters
+    ----------
+    reg_lambda : float
+        Weight-matrix L2 regularisation. With ``reg_mu is None`` (the
+        DirL2 baseline in Kull et al., 2019) it penalises every entry of
+        ``W`` with no normalisation. Under DirODIR (``reg_mu`` set) it
+        penalises only the off-diagonal entries of ``W`` with the
+        ``1/(k(k-1))`` factor from the paper.
+    reg_mu : float or None
+        Selects the DirODIR (Off-Diagonal and Intercept Regularised)
+        variant of Kull et al. (2019). When not None, ``reg_lambda``
+        switches to off-diagonal-only and ``reg_mu`` applies an L2
+        penalty ``mu * (1/k) * sum(b**2)`` to the intercept; the
+        diagonal of ``W`` is left free. When None, the plain DirL2
+        baseline is used and the intercept is unregularised.
+    """
+
+    def __init__(self, reg_lambda: float = 0.0, reg_mu: float | None = None):
+        self.reg_lambda = reg_lambda
+        self.reg_mu = reg_mu
+        self.weights_: npt.NDArray[np.float64] | None = None
+
+    def fit(
+        self,
+        probs: npt.NDArray[np.float64],
+        y: npt.NDArray[np.int_],
+    ) -> "DirichletCalibrator":
+        """Fit the calibrator on held-out probabilities and labels."""
+        n_samples, n_classes = probs.shape
+        y = np.asarray(y)
+
+        # Initialise as identity mapping (no calibration).
+        W0 = np.zeros((n_classes, n_classes + 1))
+        W0[:, :n_classes] = np.eye(n_classes)
+
+        log_p = np.log(np.clip(probs, 1e-12, 1.0))
+        X = np.hstack([log_p, np.ones((n_samples, 1))])
+
+        result = minimize(
+            _dirichlet_loss_and_grad,
+            W0.flatten(),
+            args=(X, y, n_classes, self.reg_lambda, self.reg_mu),
+            jac=True,
+            method="L-BFGS-B",
+            options={"maxiter": 1000},
+        )
+        self.weights_ = result.x.reshape(n_classes, n_classes + 1)
+        return self
+
+    def predict_proba(
+        self, probs: npt.NDArray[np.float64]
+    ) -> npt.NDArray[np.float64]:
+        """Return calibrated probabilities."""
+        if self.weights_ is None:
+            raise ValueError("Calibrator not fitted.")
+        n_samples = probs.shape[0]
+        log_p = np.log(np.clip(probs, 1e-12, 1.0))
+        X = np.hstack([log_p, np.ones((n_samples, 1))])
+        logits = X @ self.weights_.T
+        return softmax(logits, axis=1)
+
+
+def _dirichlet_loss_and_grad(
+    w_flat: npt.NDArray[np.float64],
+    X: npt.NDArray[np.float64],
+    y: npt.NDArray[np.int_],
+    n_classes: int,
+    reg_lambda: float,
+    reg_mu: float | None,
+) -> tuple[float, npt.NDArray[np.float64]]:
+    """Dirichlet-calibration NLL (+ L2/ODIR reg) and its analytical gradient.
+
+    `X` is the precomputed `[log_p | 1]` design matrix of shape
+    (n_samples, n_classes + 1). `W` is reshaped from `w_flat` as
+    (n_classes, n_classes + 1).
+    """
+    n_samples = X.shape[0]
+    W = w_flat.reshape(n_classes, n_classes + 1)
+    logits = X @ W.T
+
+    # Softmax and log-softmax in one numerically-stable pass.
+    shifted = logits - logits.max(axis=1, keepdims=True)
+    exp_shifted = np.exp(shifted)
+    sum_exp = exp_shifted.sum(axis=1, keepdims=True)
+    log_probs = shifted - np.log(sum_exp)
+    probs_soft = exp_shifted / sum_exp
+
+    row_idx = np.arange(n_samples)
+    nll = -np.mean(log_probs[row_idx, y])
+
+    # d(nll)/d(logits) = (softmax - onehot(y)) / n_samples
+    # d(nll)/d(W)      = d(logits).T @ X   (since logits = X @ W.T)
+    dlogits = probs_soft
+    dlogits[row_idx, y] -= 1.0
+    dlogits /= n_samples
+    grad_W = dlogits.T @ X
+
+    # Regularisation follows Kull et al. (2019). ``reg_mu`` acts as a
+    # mode switch rather than an additive penalty:
+    #   reg_mu is None -> DirL2 baseline from the paper: plain L2 on the
+    #                     full weight matrix (intercept excluded).
+    #   reg_mu is set  -> DirODIR (eq. in §3): reg_lambda penalises the
+    #                     off-diagonal entries of W with a 1/(k(k-1))
+    #                     factor, reg_mu penalises the intercept b with
+    #                     a 1/k factor, and the diagonal of W is free.
+    loss = nll
+    if reg_mu is None:
+        if reg_lambda > 0:
+            loss += reg_lambda * np.sum(W[:, :n_classes] ** 2)
+            grad_W[:, :n_classes] += 2.0 * reg_lambda * W[:, :n_classes]
+    else:
+        if reg_lambda > 0 and n_classes > 1:
+            off_diag = W[:, :n_classes].copy()
+            np.fill_diagonal(off_diag, 0.0)
+            norm_off = 1.0 / (n_classes * (n_classes - 1))
+            loss += reg_lambda * norm_off * np.sum(off_diag**2)
+            grad_W[:, :n_classes] += 2.0 * reg_lambda * norm_off * off_diag
+        if reg_mu > 0:
+            norm_b = 1.0 / n_classes
+            loss += reg_mu * norm_b * np.sum(W[:, -1] ** 2)
+            grad_W[:, -1] += 2.0 * reg_mu * norm_b * W[:, -1]
+
+    return float(loss), grad_W.ravel()

threshold_opt/cv.py

@@ -0,0 +1,51 @@
+"""Cross-validation fold-count helpers shared by strategies and the
+experiments pipeline.
+
+Kept out of :mod:`threshold_opt.base` (which holds base-class content)
+because this is a stand-alone utility: the CV-tuning strategies
+(``EmpiricalThreshold``, ``ThresholdRatioMulticlass``, the calibration
+strategies) and the pipeline's fold precompute all need the same
+"largest StratifiedKFold this label vector can honour" decision, and
+all need to react to "CV is impossible here" identically.
+"""
+
+import numpy as np
+import numpy.typing as npt
+
+
+class InsufficientCVData(ValueError):
+    """Raised when a class has too few members for the requested K-fold CV.
+
+    Subclasses ``ValueError`` so existing ``except ValueError`` handlers
+    still catch it, while callers that want to treat "CV is impossible on
+    this split" distinctly (e.g. the experiments pipeline, which logs it
+    as a countable ``CV_SKIP`` rather than a generic warning) can catch
+    it by type first.
+    """
+
+
+def safe_n_splits(y: npt.ArrayLike, max_splits: int = 5) -> int:
+    """Largest ``StratifiedKFold`` ``n_splits`` honourable on *y*.
+
+    Returns ``min(max_splits, smallest_class_count)`` (never below 2).
+    Raises :class:`InsufficientCVData` when the smallest class has fewer
+    than 2 members, because no ``n_splits >= 2`` stratified split exists
+    in that case -- the old ``max(2, ...)`` clamp returned 2 here and
+    then crashed deep inside sklearn with a misleading ``ValueError`` /
+    ``IndexError``.
+
+    The result is a deterministic function of *y* alone, so the pipeline
+    precompute and a strategy's own lazy refit (which key off the same
+    fit-set labels) still derive the same K and the shared fold cache
+    stays interchangeable with the lazy path.
+    """
+    counts = np.unique(np.asarray(y), return_counts=True)[1]
+    if len(counts) == 0:
+        raise InsufficientCVData("Cannot run CV: no samples.")
+    min_count = int(counts.min())
+    if min_count < 2:
+        raise InsufficientCVData(
+            f"Cannot run StratifiedKFold: smallest class has "
+            f"{min_count} member(s), need >= 2."
+        )
+    return min(max_splits, min_count)

threshold_opt/metrics.py

@@ -0,0 +1,322 @@
+"""Core mathematical functions for risk estimation and threshold derivation."""
+
+import dataclasses
+from typing import Callable, Union
+
+import numpy as np
+import numpy.typing as npt
+
+CostMatrix = npt.ArrayLike
+CostMatrixFactory = Callable[[int], CostMatrix]
+
+
+@dataclasses.dataclass
+class MetricObjective:
+    metric_fn: Callable[[npt.ArrayLike, npt.ArrayLike], float]
+    higher_is_better: bool
+    vectorized_key: str | None = None
+    beta: float = 1.0
+
+
+@dataclasses.dataclass
+class DataDependentCostBuilder:
+    """Cost matrix that depends on the training-set class priors.
+
+    The pipeline materializes the builder once per task by calling
+    ``build_fn(y_train)`` and then passes the resulting cost matrix to
+    strategies just like any other ``CostMatrix`` ndarray. This lets a
+    single objective definition adapt to each dataset's empirical
+    class distribution without duplicating cost-matrix entries.
+    """
+
+    build_fn: Callable[[npt.NDArray], np.ndarray]
+
+
+def balanced_error_cost(y: npt.ArrayLike) -> np.ndarray:
+    """Cost matrix whose expected value equals macro balanced error.
+
+    Macro balanced error rate equals
+    ``(1/K) * sum_y (1 - recall_y) = E_{(X,Y)}[cost(Y, A(X))]`` for
+    ``cost(y, a) = 1 / (K * pi(y))`` when ``a != y`` and 0 on the
+    diagonal, where ``pi(y)`` is the marginal class prior. Strategies
+    that minimize this empirical risk on a held-out set therefore
+    minimize balanced error directly.
+
+    Empty classes (no samples in ``y``) get cost 0; this avoids
+    division by zero but means strategies cannot be punished for
+    failing to predict classes absent at fit time.
+    """
+    y_arr = np.asarray(y)
+    classes, counts = np.unique(y_arr, return_counts=True)
+    K = int(classes.max()) + 1 if len(classes) else 0
+    priors = np.zeros(K)
+    priors[classes.astype(int)] = counts / counts.sum()
+    M = np.zeros((K, K), dtype=np.float64)
+    for i in range(K):
+        if priors[i] <= 0:
+            continue
+        M[i, :] = 1.0 / (K * priors[i])
+        M[i, i] = 0.0
+    return M
+
+
+BALANCED_ERROR_COST = DataDependentCostBuilder(build_fn=balanced_error_cost)
+
+
+OptimizationObjective = Union[
+    CostMatrix, MetricObjective, CostMatrixFactory, DataDependentCostBuilder
+]
+
+
+def get_objective_score(
+    y_true: npt.ArrayLike,
+    y_pred: npt.ArrayLike,
+    optimization_objective: "OptimizationObjective",
+) -> float:
+    """Score a prediction under an OptimizationObjective.
+
+    Higher is always better: metric signs are flipped to minimization
+    form, and cost-matrix objectives return negative empirical risk.
+    """
+    if isinstance(optimization_objective, MetricObjective):
+        sign = 1.0 if optimization_objective.higher_is_better else -1.0
+        return sign * optimization_objective.metric_fn(y_true, y_pred)
+
+    cost_matrix = np.asarray(optimization_objective)
+    return -calculate_empirical_risk(y_true, y_pred, cost_matrix)
+
+
+def calculate_optimal_threshold(cost_matrix: CostMatrix) -> float:
+    """Derives the theoretical optimal threshold p* for a binary cost matrix.
+
+    The optimal threshold is derived by minimizing the expected cost.
+    The formula is:
+        p* = (FP - TN) / ( (FP - TN) + (FN - TP) )
+           = (C_01 - C_00) / ( (C_01 - C_00) + (C_10 - C_11) )
+
+    Args:
+        cost_matrix: 2x2 matrix [[TN_cost, FP_cost], [FN_cost, TP_cost]].
+                     Note: This follows the standard confusion matrix
+                     layout where
+                       C_ij is the cost of predicting class j when truth is i.
+
+    Returns:
+        float: The probability threshold in [0, 1].
+
+    Raises:
+        ValueError: If the cost matrix is not 2x2.
+    """
+    cost_matrix = np.asarray(cost_matrix)
+    if cost_matrix.shape != (2, 2):
+        raise ValueError(
+            f"Binary threshold requires 2x2 matrix, got {cost_matrix.shape}"
+        )
+
+    c_00, c_01 = cost_matrix[0]
+    c_10, c_11 = cost_matrix[1]
+
+    numerator = c_01 - c_00
+    denominator = (c_01 - c_00) + (c_10 - c_11)
+
+    if np.isclose(denominator, 0):
+        return 0.5
+
+    return float(numerator / denominator)
+
+
+def calculate_empirical_risk(
+    y_true: npt.ArrayLike,
+    y_pred: npt.ArrayLike,
+    cost_matrix: CostMatrix,
+) -> float:
+    """Computes average realized cost (Empirical Risk).
+
+    Args:
+        y_true: The true labels.
+        y_pred: The predicted labels.
+        cost_matrix: The cost matrix where `cost_matrix[i, j]`
+          represents the cost
+          of predicting class `j` when the true class is `i`.
+
+    Returns:
+        The average empirical risk.
+
+    Raises:
+        ValueError: If the cost matrix is not square.
+    """
+    cost_matrix = np.asarray(cost_matrix)
+
+    if cost_matrix.ndim != 2 or cost_matrix.shape[0] != cost_matrix.shape[1]:
+        raise ValueError("Cost matrix must be square.")
+
+    y_true = _as_integer_labels(y_true, "y_true")
+    y_pred = _as_integer_labels(y_pred, "y_pred")
+
+    n = y_true.shape[0]
+    if n == 0:
+        return 0.0
+
+    if y_pred.shape[0] != n:
+        raise ValueError(
+            f"y_true and y_pred must have the same length, "
+            f"got {n} and {y_pred.shape[0]}."
+        )
+
+    n_classes = cost_matrix.shape[0]
+    for name, arr in (("y_true", y_true), ("y_pred", y_pred)):
+        if arr.min() < 0 or arr.max() >= n_classes:
+            raise ValueError(
+                f"{name} contains labels outside [0, {n_classes - 1}]."
+            )
+
+    total_cost = cost_matrix[y_true, y_pred].sum()
+    return float(total_cost / n)
+
+
+def _as_integer_labels(
+    labels: npt.ArrayLike, name: str
+) -> npt.NDArray[np.intp]:
+    arr = np.asarray(labels)
+    if np.issubdtype(arr.dtype, np.integer) or arr.dtype == bool:
+        return arr.astype(np.intp, copy=False)
+    if np.issubdtype(arr.dtype, np.floating):
+        cast = arr.astype(np.intp)
+        if not np.array_equal(cast, arr):
+            raise ValueError(
+                f"{name} must be integer-valued, got non-integer floats."
+            )
+        return cast
+    raise ValueError(f"{name} must be integer labels, got dtype {arr.dtype}.")
+
+
+# ---------------------------------------------------------------------------
+# Vectorized scoring for binary threshold grid search
+# ---------------------------------------------------------------------------
+
+
+def binary_confusion_counts_vectorized(
+    y_true: npt.NDArray[np.int_],
+    probs: npt.NDArray[np.float64],
+    thresholds: npt.NDArray[np.float64],
+) -> tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray]:
+    """Compute TP, FP, FN, TN for all thresholds simultaneously.
+
+    Uses numpy broadcasting to avoid per-threshold Python loops.
+
+    Args:
+        y_true: shape (n_samples,), values in {0, 1}.
+        probs: shape (n_samples,), predicted probability of class 1.
+        thresholds: shape (n_thresholds,).
+
+    Returns:
+        (tp, fp, fn, tn) each of shape (n_thresholds,).
+    """
+    # pred_positive has shape (n_thresholds, n_samples)
+    pred_positive = probs[np.newaxis, :] >= thresholds[:, np.newaxis]
+    positives = y_true.astype(bool)
+    negatives = ~positives
+
+    tp = (pred_positive & positives).sum(axis=1)
+    fp = (pred_positive & negatives).sum(axis=1)
+    fn = (~pred_positive & positives).sum(axis=1)
+    tn = (~pred_positive & negatives).sum(axis=1)
+    return tp, fp, fn, tn
+
+
+def vectorized_scores_from_counts(
+    tp: npt.NDArray,
+    fp: npt.NDArray,
+    fn: npt.NDArray,
+    tn: npt.NDArray,
+    metric_key: str,
+    *,
+    beta: float = 1.0,
+) -> npt.NDArray[np.float64]:
+    """Compute metric scores for all thresholds from confusion counts.
+
+    All inputs have shape (n_thresholds,). Returns shape (n_thresholds,)
+    holding the metric's natural value (e.g. F1 or balanced accuracy
+    higher-is-better, balanced error lower-is-better); the caller is
+    responsible for sign conventions via ``MetricObjective.higher_is_better``.
+    """
+    tp = tp.astype(np.float64)
+    fp = fp.astype(np.float64)
+    fn = fn.astype(np.float64)
+    tn = tn.astype(np.float64)
+
+    if np.any(np.isnan(tp) | np.isnan(fp) | np.isnan(fn) | np.isnan(tn)):
+        raise ValueError(
+            "Confusion counts contain NaN values. Check input arrays for NaN."
+        )
+
+    if metric_key == "f1":
+        num = 2.0 * tp
+        den = 2.0 * tp + fp + fn
+        return np.divide(num, den, out=np.zeros_like(num), where=den > 0)
+
+    if metric_key == "f_beta":
+        beta2 = beta**2
+        num = (1.0 + beta2) * tp
+        den = (1.0 + beta2) * tp + beta2 * fn + fp
+        return np.divide(num, den, out=np.zeros_like(num), where=den > 0)
+
+    if metric_key == "mcc":
+        num = tp * tn - fp * fn
+        product = (tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)
+        den = np.sqrt(np.maximum(product, 0.0))
+        return np.divide(num, den, out=np.zeros_like(num), where=den > 0)
+
+    if metric_key == "balanced_accuracy":
+        tpr = np.divide(
+            tp, tp + fn, out=np.zeros_like(tp), where=(tp + fn) > 0
+        )
+        tnr = np.divide(
+            tn, tn + fp, out=np.zeros_like(tn), where=(tn + fp) > 0
+        )
+        return (tpr + tnr) / 2.0
+
+    if metric_key == "balanced_error":
+        # 1 - balanced accuracy = (FPR + FNR) / 2; the natural cost form
+        # of BA under prior-weighted misclassification costs (the BER /
+        # balanced error rate of Menon et al. 2013). Returned as-is so the
+        # caller can negate via ``higher_is_better=False``.
+        tpr = np.divide(
+            tp, tp + fn, out=np.zeros_like(tp), where=(tp + fn) > 0
+        )
+        tnr = np.divide(
+            tn, tn + fp, out=np.zeros_like(tn), where=(tn + fp) > 0
+        )
+        return 1.0 - (tpr + tnr) / 2.0
+
+    if metric_key == "gmean":
+        tpr = np.divide(
+            tp, tp + fn, out=np.zeros_like(tp), where=(tp + fn) > 0
+        )
+        tnr = np.divide(
+            tn, tn + fp, out=np.zeros_like(tn), where=(tn + fp) > 0
+        )
+        return np.sqrt(tpr * tnr)
+
+    raise ValueError(f"Unknown vectorized metric key: {metric_key!r}")
+
+
+def vectorized_cost_scores(
+    tp: npt.NDArray,
+    fp: npt.NDArray,
+    fn: npt.NDArray,
+    tn: npt.NDArray,
+    cost_matrix: npt.NDArray,
+) -> npt.NDArray[np.float64]:
+    """Compute negative empirical risk for all thresholds at once.
+
+    Returns scores where higher is better (negative cost).
+    """
+    cost_matrix = np.asarray(cost_matrix, dtype=np.float64)
+    n = (tp + fp + fn + tn).astype(np.float64)
+    total_cost = (
+        tn * cost_matrix[0, 0]
+        + fp * cost_matrix[0, 1]
+        + fn * cost_matrix[1, 0]
+        + tp * cost_matrix[1, 1]
+    )
+    return -(total_cost / n)