metrics.py

"""AUROC and sensitivity@k% with exact tie handling.

Both metrics are computed from mid-ranks, so equal scores contribute equally
regardless of the order rows happen to arrive in. That matters here: the models
saturate near 0 and 1, so a validation set contains large blocks of identical
scores, and an ordinal `argsort` would resolve them by row order.
"""
import numpy as np
from scipy.stats import rankdata, t


def auroc(labels, scores):
    """AUROC via the Mann-Whitney U identity, using mid-ranks for ties.

    Returns NaN if either class is absent.
    """
    labels = np.asarray(labels).astype(bool)
    scores = np.asarray(scores, dtype=float)
    n_pos = int(labels.sum())
    n_neg = labels.size - n_pos
    if n_pos == 0 or n_neg == 0:
        return float("nan")
    ranks = rankdata(scores)  # 'average' method -> mid-ranks
    return float((ranks[labels].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg))


def sensitivity_at_top_k(labels, scores, k_percent=2.0):
    """Expected recall within the top k% under uniformly random tie-breaking.

    Scores tied at the cutoff contribute in proportion to how many of the k
    places they can occupy, which makes the result independent of input order.
    """
    labels = np.asarray(labels).astype(bool)
    scores = np.asarray(scores, dtype=float)
    n_pos = int(labels.sum())
    if n_pos == 0:
        return float("nan")
    n_top = max(1, int(round(scores.size * k_percent / 100)))
    threshold = np.partition(scores, -n_top)[-n_top]  # k-th largest score
    above = scores > threshold
    tied = scores == threshold
    hits = float((labels & above).sum())
    slots = n_top - int(above.sum())  # places left for the tied block
    if slots > 0 and tied.any():
        hits += float((labels & tied).sum()) * slots / float(tied.sum())
    return hits / n_pos


def mean_ci(values, confidence=0.95):
    """Mean and half-width of the t-based CI over `values`.

    Returns (mean, half_width, n). half_width is NaN when n < 2.
    """
    v = np.asarray(values, dtype=float)
    v = v[np.isfinite(v)]
    if v.size == 0:
        return float("nan"), float("nan"), 0
    if v.size == 1:
        return float(v[0]), float("nan"), 1
    half = t.ppf(0.5 + confidence / 2, v.size - 1) * v.std(ddof=1) / np.sqrt(v.size)
    return float(v.mean()), float(half), int(v.size)


def evaluate(labels, scores, k_percent=2.0):
    """Both metrics plus set sizes, as a dict."""
    labels = np.asarray(labels)
    return {
        "auroc": auroc(labels, scores),
        "sensitivity_at_2pct": sensitivity_at_top_k(labels, scores, k_percent),
        "n_obs": int(labels.size),
        "n_pos": int(labels.sum()),
    }