diff --git a/odyssey/inference/uncertainty.py b/odyssey/inference/uncertainty.py new file mode 100644 index 0000000..e655467 --- /dev/null +++ b/odyssey/inference/uncertainty.py @@ -0,0 +1,274 @@ +"""Subject-clustered bootstrap confidence intervals for AUROC. + +Every model/baseline AUROC this project reports today is a bare point +estimate on one held-out draw. That understates how much a reported +"hazard beats GBM by 0.02"-style claim should be trusted: the tuned GBM is +a CONTROL in the value-tail-transform arms (it never reads the value +channel and scores the same held-out shards in every arm), yet its own +AUROC on the same cell moves by up to ~2.9 percentage points across arms +(measured on real data, 2026-08-24: death@8h 0.9452/0.9518/0.9229 across +three arms). Every point-estimate comparison in the registry has been read +against a number that moves by about that much on its own. + +TWO SOURCES OF SPREAD, NOT ONE -- the distinction this module exists to +keep visible, not conflate: + +- FINITE-SAMPLE VARIANCE: the held-out split is one draw from the + population; a cell with 222 positives out of 136,850 rows carries real + sampling error no matter how good the models are. This is what + :func:`bootstrap_auroc` measures. +- REFIT VARIANCE: different fits of the SAME model (different train-time + randomness, different upstream floating-point paths, ...) score + differently even on identical held-out data -- the GBM spread above is + this, not finite-sample variance. NO bootstrap of one fitted model's + predictions can see this; measuring it needs k independent refits with + different seeds, which is out of scope here. + +A bootstrap interval from this module answers "how much would this AUROC +plausibly move if we drew a different held-out sample from the same +population, holding the fitted model fixed" -- it does NOT answer "how +much would this AUROC move on a refit". Reporting a bootstrap interval as +"the" uncertainty on a comparison between two independently-fit models +(hazard head vs. GBM, arm vs. arm) understates the total spread, because it +omits refit variance entirely. Any number built from this module must say +which of the two sources it carries. + +Not yet wired into :class:`~odyssey.inference.alerts.AlertMetrics` or the +alerts pipeline, and does not change any existing reported number -- +landed as a tested, standalone unit first. +""" + +from dataclasses import dataclass +from typing import Optional, Sequence, Union + +import numpy as np +from sklearn.metrics import roc_auc_score + + +@dataclass(frozen=True) +class BootstrapAUROC: + """A subject-clustered bootstrap AUROC summary. + + See module docstring for what this measures (finite-sample variance) + and what it does not (refit variance). + """ + + point_estimate: float + """AUROC on the data exactly as observed, no resampling.""" + mean: Optional[float] + """Mean AUROC across usable bootstrap resamples. None if every + resample was skipped (see ``n_skipped``).""" + std: Optional[float] + """Sample std (ddof=1) across usable resamples. None if fewer than 2 + resamples were usable -- a std from 0 or 1 points is not meaningful.""" + ci_low: Optional[float] + """Lower percentile bound (``100 * alpha / 2``). None under the same + condition as ``mean``.""" + ci_high: Optional[float] + """Upper percentile bound (``100 * (1 - alpha / 2)``). None under the + same condition as ``mean``.""" + n_boot_used: int + """Resamples that produced a usable AUROC.""" + n_boot_skipped: int + """Resamples discarded because the resampled ``y`` was single-class + (AUROC undefined) -- see requirement 3 in the module docstring. A + large count here means the interval rests on fewer resamples than + requested and the cell is sparse; it must stay visible, not get + silently absorbed into a narrower-looking ``n_boot``.""" + + +def bootstrap_auroc( + y: Union[np.ndarray, Sequence[float]], + p: Union[np.ndarray, Sequence[float]], + subject_ids: Union[np.ndarray, Sequence[int]], + *, + n_boot: int = 1000, + seed: int = 0, + alpha: float = 0.05, +) -> Optional[BootstrapAUROC]: + """Subject-clustered bootstrap AUROC: point estimate, mean/std/CI, skip count. + + Returns None if the OBSERVED ``y`` is single-class -- AUROC is + undefined for the cell itself, not just for some resamples, so the + caller should report the cell unscoreable rather than substitute 0.5 + (the same discipline + :class:`~odyssey.inference.survivalpfn_baseline.SurvivalPFNBaselineModel`'s + degeneracy warning already follows for a different kind of undefined + AUROC). + + RESAMPLES SUBJECTS, NOT ROWS -- the one requirement that makes or + breaks this function. Landmark rows are many-per-subject and heavily + correlated (the same patient contributes a row every few hours with + nearly identical features and the same outcome); resampling rows + would treat those as independent draws and produce intervals far too + narrow to trust. Each resample draws ``n_subjects`` subjects WITH + REPLACEMENT and takes every one of a drawn subject's rows, whole. + + A resample whose drawn ``y`` is single-class makes AUROC undefined for + that resample specifically; it is skipped and counted + (``n_boot_skipped``), never silently dropped -- see + :class:`BootstrapAUROC`. + + Rows are grouped by subject ONCE up front (an O(n log n) sort), and + each resample's row selection is a pure index gather via + :func:`numpy.repeat`/cumulative-offset arithmetic -- no per-subject + Python loop and no dataframe operation inside the resample loop. + + ``p`` is ALSO sorted and grouped into tied-value buckets once, up + front, rather than re-sorted inside the loop: a first version called + :func:`sklearn.metrics.roc_auc_score` per resample, which re-sorts the + resampled array every time and measured ~16ms/call on a + 140,000-row cell -- 1000 resamples x 12 cells landed at several + minutes, failing the seconds-not-minutes bar. Each resample instead + computes a WEIGHTED Mann-Whitney U statistic against the one + precomputed sort: a resample only changes how many times each + ORIGINAL row is duplicated (its weight, from + :func:`numpy.bincount` on the resampled row indices), never the + relative order of two distinct row values, so the average-rank-per-tie + bucket only needs computing once and each resample becomes a handful + of O(n) weighted sums instead of a fresh O(n log n) sort. Verified + against :func:`sklearn.metrics.roc_auc_score` directly (see the + module's tests) rather than trusted by derivation alone. Measured + speedup: ~7x on the case above, comfortably seconds for 1000 x 12. + """ + y_arr = np.asarray(y, dtype=np.float64) + p_arr = np.asarray(p, dtype=np.float64) + subject_arr = np.asarray(subject_ids) + if not (len(y_arr) == len(p_arr) == len(subject_arr)): + raise ValueError( + "bootstrap_auroc: y, p, and subject_ids must describe the same " + f"rows; got lengths {len(y_arr)}, {len(p_arr)}, {len(subject_arr)}" + ) + n_rows = len(y_arr) + if len(np.unique(y_arr)) < 2: + return None + + point_estimate = float(roc_auc_score(y_arr, p_arr)) + + # Group rows by subject once: `order` sorted by subject, `boundaries` + # the start offset of each subject's run within `order`. Every + # resample below is then `order[boundaries[s]:boundaries[s+1]]` per + # drawn subject s, vectorized rather than looped. + unique_subjects, inverse = np.unique(subject_arr, return_inverse=True) + n_subjects = len(unique_subjects) + subj_order = np.argsort(inverse, kind="stable") + subj_counts = np.bincount(inverse, minlength=n_subjects) + subj_boundaries = np.concatenate([[0], np.cumsum(subj_counts)]) + + # Group rows by tied p-value once: `p_group_of_row[i]` is which + # ascending-p tie-bucket original row i falls into. + p_group_of_row, n_groups = _group_p_ties(p_arr) + + rng = np.random.default_rng(seed) + scores = [] + n_skipped = 0 + for _ in range(n_boot): + drawn = rng.integers(0, n_subjects, size=n_subjects) + row_idx = _gather_rows_for_drawn_subjects( + drawn, subj_boundaries, subj_order, subj_counts + ) + auc = _weighted_auroc(row_idx, y_arr, p_group_of_row, n_rows, n_groups) + if auc is None: + n_skipped += 1 + continue + scores.append(auc) + + n_used = len(scores) + if n_used == 0: + return BootstrapAUROC(point_estimate, None, None, None, None, 0, n_boot) + + scores_arr = np.array(scores) + mean = float(scores_arr.mean()) + std = float(scores_arr.std(ddof=1)) if n_used > 1 else None + ci_low = float(np.percentile(scores_arr, 100 * alpha / 2)) + ci_high = float(np.percentile(scores_arr, 100 * (1 - alpha / 2))) + return BootstrapAUROC(point_estimate, mean, std, ci_low, ci_high, n_used, n_skipped) + + +def _group_p_ties(p: np.ndarray) -> tuple[np.ndarray, int]: + """Bucket rows by tied ``p`` value: ``(p_group_of_row, n_groups)``. + + ``p_group_of_row[i]`` is which ascending-``p`` tie-bucket original row + ``i`` falls into (0 = smallest ``p``); a group with more than one + member is a genuine tie in the score, which the weighted mid-rank + formula in :func:`_weighted_auroc` must average over correctly even + when the tied rows carry different ``y`` labels -- see that + function's tests for why a mixed-label tie bucket is the case that + actually exercises this. + """ + n_rows = len(p) + if n_rows == 0: + return np.empty(0, dtype=np.int64), 0 + p_sort_order = np.argsort(p, kind="quicksort") + sorted_p = p[p_sort_order] + is_new_group = np.empty(n_rows, dtype=bool) + is_new_group[0] = True + np.not_equal(sorted_p[1:], sorted_p[:-1], out=is_new_group[1:]) + group_id_sorted = np.cumsum(is_new_group) - 1 + n_groups = int(group_id_sorted[-1]) + 1 + p_group_of_row = np.empty(n_rows, dtype=np.int64) + p_group_of_row[p_sort_order] = group_id_sorted + return p_group_of_row, n_groups + + +def _weighted_auroc( + row_idx: np.ndarray, + y: np.ndarray, + p_group_of_row: np.ndarray, + n_rows: int, + n_groups: int, +) -> Optional[float]: + """AUROC of ``y``/``p`` restricted+duplicated to ``row_idx``. + + Uses a precomputed p-value sort rather than a fresh one. A bootstrap + resample only changes how many times each ORIGINAL row + appears (its weight); it never changes the relative order of two rows + with distinct p-values, so the Mann-Whitney U rank-sum formula can be + computed from per-tie-bucket WEIGHT totals instead of a fresh sort: + weight the row's own bucket by its multiplicity, get each bucket's + average rank from its cumulative weight, then apply the ordinary + rank-sum AUROC formula on those weighted ranks. Returns None if the + resample's ``y`` came out single-class (AUROC undefined for that + resample). + """ + weight = np.bincount(row_idx, minlength=n_rows).astype(np.float64) + pos_weight = weight * y + group_weight = np.bincount(p_group_of_row, weights=weight, minlength=n_groups) + group_pos_weight = np.bincount( + p_group_of_row, weights=pos_weight, minlength=n_groups + ) + n_pos = float(group_pos_weight.sum()) + n_neg = float(group_weight.sum() - n_pos) + if n_pos == 0.0 or n_neg == 0.0: + return None + cum_weight_before = np.concatenate([[0.0], np.cumsum(group_weight)[:-1]]) + avg_rank_per_group = cum_weight_before + (group_weight + 1.0) / 2.0 + sum_ranks_pos = float((avg_rank_per_group * group_pos_weight).sum()) + return (sum_ranks_pos - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg) + + +def _gather_rows_for_drawn_subjects( + drawn: np.ndarray, + boundaries: np.ndarray, + order: np.ndarray, + counts: np.ndarray, +) -> np.ndarray: + """Row indices for one resample's drawn subjects, no Python-level per-subject loop. + + ``drawn`` is a subject-position array (may repeat, whole-subject + with-replacement draw). Each drawn subject's row count is variable, so + building the concatenated row-index array is a "ragged repeat": + ``np.repeat`` expands each drawn subject's SLOT into as many copies as + it has rows, then cumulative per-slot offsets recover each output + position's index within its own subject's row run. + """ + drawn_counts = counts[drawn] + total_rows = int(drawn_counts.sum()) + if total_rows == 0: + return np.empty(0, dtype=order.dtype) + slot_of_row = np.repeat(np.arange(len(drawn)), drawn_counts) + slot_start = np.concatenate([[0], np.cumsum(drawn_counts)[:-1]]) + within_subject_pos = np.arange(total_rows) - slot_start[slot_of_row] + subject_start = boundaries[drawn[slot_of_row]] + result: np.ndarray = order[subject_start + within_subject_pos] + return result diff --git a/tests/odyssey/inference/test_uncertainty.py b/tests/odyssey/inference/test_uncertainty.py new file mode 100644 index 0000000..38d1b1d --- /dev/null +++ b/tests/odyssey/inference/test_uncertainty.py @@ -0,0 +1,231 @@ +"""Tests for the subject-clustered bootstrap AUROC helper.""" + +import numpy as np +import pytest +from sklearn.metrics import roc_auc_score + +from odyssey.inference.uncertainty import ( + BootstrapAUROC, + _gather_rows_for_drawn_subjects, + _group_p_ties, + _weighted_auroc, + bootstrap_auroc, +) + + +def _correlated_fixture() -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """19 single-row background subjects + 1 big subject with 100 identical rows. + + Background subjects are perfectly rank-separable (y=1 always scores + higher than y=0), giving a clean high baseline AUROC on their own. The + big subject is a homogeneous block of 100 rows, all y=0 with a very + high score (a false-positive cluster) -- whether or not it appears in + a resample swings the AUROC a lot, and it is an all-or-nothing, + whole-subject event only a SUBJECT bootstrap can reflect: a uniform + ROW bootstrap draws ~84% of its rows from this block on almost every + resample (it is 100 of 119 real rows), so the row-level composition + barely moves resample to resample. + """ + rng = np.random.default_rng(0) + bg_y = np.array([i % 2 for i in range(19)], dtype=float) + bg_p = np.where( + bg_y == 1, 0.6 + rng.uniform(0, 0.1, 19), 0.1 + rng.uniform(0, 0.1, 19) + ) + bg_subjects = np.arange(1, 20) + + big_y = np.zeros(100) + big_p = np.full(100, 0.99) + big_subjects = np.full(100, 1000) + + y = np.concatenate([bg_y, big_y]) + p = np.concatenate([bg_p, big_p]) + subjects = np.concatenate([bg_subjects, big_subjects]) + return y, p, subjects + + +def _row_bootstrap_std(y: np.ndarray, p: np.ndarray, n_boot: int, seed: int) -> float: + """Compute a naive ROW-level bootstrap std, for comparison only. + + Not exposed by the module under test, since resampling rows + independently is exactly the wrong thing to do on correlated + landmark data (see the module docstring). Used here only to prove the + subject-clustered version is not narrower than this on correlated + data. + """ + rng = np.random.default_rng(seed) + n = len(y) + scores = [] + for _ in range(n_boot): + idx = rng.integers(0, n, size=n) + y_b, p_b = y[idx], p[idx] + if len(np.unique(y_b)) < 2: + continue + scores.append(roc_auc_score(y_b, p_b)) + return float(np.std(np.array(scores), ddof=1)) + + +def test_subject_clustered_ci_is_wider_than_row_bootstrap_on_correlated_data() -> None: + y, p, subjects = _correlated_fixture() + + result = bootstrap_auroc(y, p, subjects, n_boot=2000, seed=0) + assert result is not None + assert result.std is not None + + row_std = _row_bootstrap_std(y, p, n_boot=2000, seed=0) + + assert result.std > row_std * 3 # not just marginally wider -- the + # whole point of resampling subjects is that a single dominant, + # homogeneous subject's presence/absence is a coin flip at the + # subject level and nearly invisible at the row level. + + +def test_single_class_observed_y_returns_none() -> None: + y = np.ones(10) + p = np.linspace(0, 1, 10) + subjects = np.arange(10) + + assert bootstrap_auroc(y, p, subjects) is None + + +def test_skip_counter_counts_degenerate_resamples() -> None: + """One positive subject among many negatives. + + Most subject-level resamples that never draw it are single-class and + must be skipped and counted, not silently dropped. + """ + n_negative_subjects = 30 + y = np.array([0.0] * n_negative_subjects + [1.0]) + p = np.linspace(0, 1, n_negative_subjects + 1) + subjects = np.arange(n_negative_subjects + 1) + + result = bootstrap_auroc(y, p, subjects, n_boot=500, seed=0) + + assert result is not None + assert result.n_boot_skipped > 0 + assert result.n_boot_used + result.n_boot_skipped == 500 + + +def test_seeded_reproducibility() -> None: + y, p, subjects = _correlated_fixture() + + a = bootstrap_auroc(y, p, subjects, n_boot=200, seed=42) + b = bootstrap_auroc(y, p, subjects, n_boot=200, seed=42) + + assert a == b + + +def test_different_seeds_can_give_different_results() -> None: + y, p, subjects = _correlated_fixture() + + a = bootstrap_auroc(y, p, subjects, n_boot=200, seed=1) + b = bootstrap_auroc(y, p, subjects, n_boot=200, seed=2) + + assert a is not None and b is not None + assert a.mean != b.mean + + +def test_shape_mismatch_raises() -> None: + y = np.array([0.0, 1.0, 0.0]) + p = np.array([0.1, 0.9]) + subjects = np.array([1, 2, 3]) + + with pytest.raises(ValueError, match="same rows"): + bootstrap_auroc(y, p, subjects) + + +def test_point_estimate_matches_plain_roc_auc_score() -> None: + y, p, subjects = _correlated_fixture() + + result = bootstrap_auroc(y, p, subjects, n_boot=50, seed=0) + + assert result is not None + assert result.point_estimate == pytest.approx(roc_auc_score(y, p)) + + +def test_all_resamples_skipped_returns_none_mean_std_ci_but_a_count() -> None: + # Every subject carries only one class label; with a single positive + # subject far outnumbered, a very small n_boot with an unlucky seed + # can plausibly skip everything -- constructed directly by isolating + # exactly one subject to guarantee a resample that never draws it. + y = np.array([1.0, 0.0]) + p = np.array([0.9, 0.1]) + subjects = np.array([1, 2]) + + result = bootstrap_auroc(y, p, subjects, n_boot=1, seed=7) + + assert result is not None + if result.n_boot_used == 0: + assert result.mean is None + assert result.std is None + assert result.ci_low is None + assert result.ci_high is None + assert result.n_boot_skipped == 1 + else: + # rare with n_boot=1 but not impossible depending on the draw + assert result.n_boot_used == 1 + + +def _subject_grouping( + subjects: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Reproduce the module's own internal subject-grouping, for white-box tests.""" + unique_subjects, inverse = np.unique(subjects, return_inverse=True) + n_subjects = len(unique_subjects) + order = np.argsort(inverse, kind="stable") + counts = np.bincount(inverse, minlength=n_subjects) + boundaries = np.concatenate([[0], np.cumsum(counts)]) + return order, boundaries, counts + + +def test_weighted_auroc_matches_sklearn_exactly_under_ties_and_multiplicity() -> None: + """The load-bearing case for the fast weighted-rank rewrite. + + Real alerts columns are coarse -- SurvivalPFN's vasopressor cells had + 175 distinct probabilities across 111,450 rows -- so ties are the + normal case, not an edge case, and it is exactly the interaction of a + resampled row's MULTIPLICITY (weight > 1, from a subject being drawn + more than once) with a TIE bucket that the mid-rank arithmetic has to + get right. A tie bucket where every row shares one label would let a + mid-rank bug cancel out silently; every bucket here mixes y=0 and + y=1 rows at the identical p value on purpose. + + 6 subjects: subject 0 (2 rows, p=0.5, y=[0,1]) and subject 1 (1 row, + p=0.5, y=1) share one tie bucket across subjects; subject 2 (3 rows, + p=0.2, y=[0,0,1]) is a mixed tie within one subject; subjects 3 (1 + row, p=0.9, y=1) and 4 (2 rows, p=0.9, y=[0,1]) share another + cross-subject tie bucket; subject 5 (1 row, p=0.1, y=0) is untied. + """ + y = np.array([0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0]) + p = np.array([0.5, 0.5, 0.5, 0.2, 0.2, 0.2, 0.9, 0.9, 0.9, 0.1]) + subjects = np.array([0, 0, 1, 2, 2, 2, 3, 4, 4, 5]) + + order, boundaries, counts = _subject_grouping(subjects) + p_group_of_row, n_groups = _group_p_ties(p) + n_rows = len(y) + + # Hand-picked draws (subject POSITIONS, matching how bootstrap_auroc + # itself draws: rng.integers(0, n_subjects, size=n_subjects)), each + # exercising a different multiplicity pattern. The second and third + # each draw a subject 3 times, on purpose, not left to chance. + draws = [ + np.array([0, 1, 2, 3, 4, 5]), # every subject exactly once + np.array([0, 0, 3, 3, 3, 5]), # subject 0 x2, subject 3 x3 + np.array([1, 2, 2, 4, 4, 4]), # subject 2 x2, subject 4 x3 + ] + for drawn in draws: + row_idx = _gather_rows_for_drawn_subjects(drawn, boundaries, order, counts) + actual = _weighted_auroc(row_idx, y, p_group_of_row, n_rows, n_groups) + + y_b, p_b = y[row_idx], p[row_idx] + assert len(np.unique(y_b)) == 2 # this draw must actually exercise AUROC + expected = roc_auc_score(y_b, p_b) + + assert actual == pytest.approx(expected, abs=1e-9) + + +def test_bootstrap_auroc_result_is_frozen() -> None: + y, p, subjects = _correlated_fixture() + result = bootstrap_auroc(y, p, subjects, n_boot=10, seed=0) + assert isinstance(result, BootstrapAUROC) + with pytest.raises(Exception): # noqa: B017, PT011 -- dataclasses.FrozenInstanceError + result.mean = 0.5 # type: ignore[misc]