"""Reusable statistical algorithms for the released NOHARM analysis. The notebook keeps each cohort, estimand, contrast, and validation visible. This module contains only algorithms reused across panels: specialty-stratified case-bootstrap intervals, crossed random-effects models, and multiple comparisons with the best. Nothing here imports the private MAST codebase. """ from __future__ import annotations import warnings from math import erfc, sqrt from typing import Any, TypeAlias import numpy as np import pandas as pd __all__ = [ "N_BOOT", "PCTILE", "SEED", "cluster_boot", "crossed_re", "model_ci", "top_tier", ] N_BOOT = 10_000 SEED = 0 PCTILE = (2.5, 97.5) TableRow: TypeAlias = dict[str, Any] DataFrame: TypeAlias = pd.DataFrame def _holm(p_values: dict[str, float]) -> dict[str, float]: """Holm step-down multiple-comparison adjustment (monotone in rank).""" items = sorted(p_values.items(), key=lambda x: x[1]) m = len(items) out: dict[str, float] = {} running = 0.0 for k, (model, p) in enumerate(items): running = max(min(1.0, p * (m - k)), running) out[model] = running return out def top_tier( table: list[TableRow], metric: str, cohort: list[str] | None = None, higher_is_better: bool = True, stratum_col: str = "stratum", alpha: float = 0.05, n_boot: int = N_BOOT, seed: int = SEED, ) -> dict[str, Any]: """Reproduce the published "top-tier" bracket: the set of models that cannot be statistically separated from the leader. This mirrors the production MCB (multiple comparisons with the best): a specialty-stratified paired cluster bootstrap over the intersection of the cohort's base cases (one shared resample per stratum and draw), a one-sided achieved significance level per model (fraction of draws the leader is not ahead, floored at 1/B), then a Holm step-down adjustment. Top tier = models whose Holm-adjusted p exceeds alpha. Pass `cohort` (e.g. the keys of donoharm-top-tier.json's p_holm block) to match a specific published view; default is every model in the table. Returns {leader, top_tier, p_holm}. """ mcv: dict[str, dict[str, float]] = {} case_strata: dict[str, str] = {} for r in table: if cohort is not None and r["model"] not in cohort: continue v = r[metric] if v == v: case_id = r["case_id"] stratum = r.get(stratum_col) if stratum is None or pd.isna(stratum): raise ValueError( f"top_tier requires {stratum_col!r} for case {case_id!r}" ) stratum = str(stratum) previous = case_strata.setdefault(case_id, stratum) if previous != stratum: raise ValueError( f"case {case_id!r} maps to multiple strata: " f"{previous!r}, {stratum!r}" ) mcv.setdefault(r["model"], {})[case_id] = v if not mcv: raise ValueError("top_tier requires at least one model with valid values") common = sorted(set.intersection(*(set(v) for v in mcv.values()))) if not common: raise ValueError("top_tier models have no common cases") models = sorted(mcv) n = len(common) M = np.array([[mcv[m][c] for c in common] for m in models]) means = M.mean(axis=1) leader_idx = int(np.argmax(means) if higher_is_better else np.argmin(means)) leader = models[leader_idx] rng = np.random.default_rng(seed) M_boot_sum = np.zeros((len(models), n_boot), dtype=float) positions_by_stratum: dict[str, list[int]] = {} for index, case_id in enumerate(common): positions_by_stratum.setdefault(case_strata[case_id], []).append(index) for stratum in sorted(positions_by_stratum): positions = np.asarray(positions_by_stratum[stratum]) sampled = rng.choice( positions, size=(n_boot, len(positions)), replace=True ) M_boot_sum += M[:, sampled].sum(axis=2) M_boot = M_boot_sum / n diff = M_boot[leader_idx][None, :] - M_boot # leader minus each model if not higher_is_better: diff = -diff p_one = { m: float(max((diff[i] <= 0).mean(), 1.0 / n_boot)) for i, m in enumerate(models) if i != leader_idx } p_holm = _holm(p_one) p_holm[leader] = 1.0 return { "leader": leader, "top_tier": sorted(m for m, p in p_holm.items() if p > alpha), "p_holm": p_holm, } def model_ci( df: DataFrame, metric: str, *, case_col: str = "case_id", stratum_col: str = "stratum", n_boot: int = N_BOOT, seed: int = SEED, ) -> pd.Series: """Return a mean and stratified case-bootstrap interval for one model. ``df`` should contain rows for a single model. Duplicate case rows are collapsed before resampling so the case, not the source row, remains the sampling unit. The returned series has ``mean``, ``lo``, and ``hi`` fields, which makes this function convenient with ``DataFrameGroupBy.apply``. """ columns = [case_col, stratum_col, metric] complete_values = df[columns].dropna(subset=[case_col, metric]) if complete_values[stratum_col].isna().any(): raise ValueError(f"model_ci requires {stratum_col!r} for every case") stratum_counts = complete_values.groupby(case_col)[stratum_col].nunique() inconsistent = stratum_counts[stratum_counts != 1] if not inconsistent.empty: raise ValueError( "model_ci cases map to multiple strata: " + ", ".join(str(case) for case in inconsistent.index[:5]) ) cases = complete_values.drop_duplicates(case_col) if cases.empty: return pd.Series({"mean": np.nan, "lo": np.nan, "hi": np.nan}) strata = [ group[metric].to_numpy(dtype=float) for _, group in cases.groupby(stratum_col, sort=False) ] rng = np.random.default_rng(seed) draws = np.zeros(n_boot, dtype=float) for values in strata: indices = rng.integers(0, len(values), size=(n_boot, len(values))) draws += values[indices].sum(axis=1) draws /= sum(len(values) for values in strata) lo, hi = np.percentile(draws, PCTILE) return pd.Series( {"mean": float(cases[metric].mean()), "lo": float(lo), "hi": float(hi)} ) def crossed_re( df: DataFrame, group_col: str, levels: list[str], *, metric: str = "F1_weighted", clinician_col: str = "clinician", case_col: str = "case_id", ) -> dict[str, Any]: """Fit adjusted group means with clinician and case random intercepts. The model is ``metric ~ 0 + group`` with crossed random intercepts for the clinician and case. The first item in ``levels`` is the reference arm. Returns per-arm estimates and two-sided normal contrasts against that arm. The notebook does not call this function: participant-level human-study rows are withheld from this bundle (aggregate-only release under the study IRB), so the human-study panels plot the released fits. It is kept as the exact specification of the estimator behind those released estimates. """ import statsmodels.api as sm required = [metric, group_col, clinician_col, case_col] subset = df.dropna(subset=required).copy() if subset.empty: raise ValueError("crossed_re requires at least one complete observation") missing = [level for level in levels if not (subset[group_col] == level).any()] if missing: raise ValueError(f"crossed_re levels absent from data: {missing}") design = np.column_stack( [(subset[group_col] == level).astype(float).to_numpy() for level in levels] ) model_data = pd.DataFrame( design, columns=[f"x{i}" for i in range(len(levels))] ) model_data["y"] = subset[metric].to_numpy() model_data["clinician"] = subset[clinician_col].to_numpy() model_data["case"] = subset[case_col].to_numpy() model_data["all"] = 1 formula = "y ~ 0 + " + " + ".join(model_data.columns[: len(levels)]) model = sm.MixedLM.from_formula( formula, data=model_data, re_formula="0", vc_formula={ "clinician": "0 + C(clinician)", "case": "0 + C(case)", }, groups="all", ) with warnings.catch_warnings(): warnings.simplefilter("ignore", RuntimeWarning) warnings.simplefilter("ignore", UserWarning) try: fitted = model.fit(reml=True, method=["lbfgs"]) except Exception: fitted = model.fit(reml=True, method=["bfgs"]) beta = np.asarray(fitted.fe_params).ravel() standard_errors = np.asarray(fitted.bse_fe).ravel() covariance = np.asarray(fitted.cov_params())[: len(levels), : len(levels)] arms = [ { "key": level, "mean": float(beta[index]), "se": float(standard_errors[index]), "lo": float(beta[index] - 1.96 * standard_errors[index]), "hi": float(beta[index] + 1.96 * standard_errors[index]), "n": int((subset[group_col] == level).sum()), } for index, level in enumerate(levels) ] contrasts: dict[str, dict[str, float]] = {} for index, level in enumerate(levels[1:], start=1): delta = float(beta[index] - beta[0]) standard_error = float( np.sqrt( covariance[0, 0] + covariance[index, index] - 2 * covariance[0, index] ) ) z_score = abs(delta / standard_error) contrasts[level] = { "delta": delta, "se": standard_error, "p": erfc(z_score / sqrt(2)), } return {"arms": arms, "contrasts": contrasts, "n_obs": len(subset)} def cluster_boot( df: DataFrame, cluster: str, value: str, *, stratum: str = "stratum", weight: str | None = None, n_boot: int = N_BOOT, seed: int = 11, ) -> tuple[float, float, float]: """Bootstrap a weighted mean by resampling whole clusters within strata. Returns ``(estimate, lower, upper)``. When ``weight`` is omitted, every row contributes equally. Missing cluster, stratum, value, or weight fields are excluded. Each cluster must belong to exactly one stratum. """ required_values = [cluster, value] + ([weight] if weight else []) subset = df.dropna(subset=required_values) if subset.empty: return (float("nan"),) * 3 if subset[stratum].isna().any(): raise ValueError(f"cluster_boot requires {stratum!r} for every cluster") weights = ( subset[weight].to_numpy(dtype=float) if weight else np.ones(len(subset), dtype=float) ) weighted_values = subset[value].to_numpy(dtype=float) * weights cluster_rows = pd.DataFrame( { "cluster": subset[cluster].astype(str).to_numpy(), "stratum": subset[stratum].astype(str).to_numpy(), "numerator": weighted_values, "denominator": weights, } ) grouped = cluster_rows.groupby("cluster", sort=True).agg( stratum=("stratum", "first"), n_strata=("stratum", "nunique"), numerator=("numerator", "sum"), denominator=("denominator", "sum"), ) inconsistent = grouped[grouped.n_strata != 1] if not inconsistent.empty: raise ValueError( "cluster_boot clusters map to multiple strata: " + ", ".join(inconsistent.index[:5]) ) if grouped.denominator.sum() == 0: raise ValueError("cluster_boot weights sum to zero") rng = np.random.default_rng(seed) draw_numerators = np.zeros(n_boot, dtype=float) draw_denominators = np.zeros(n_boot, dtype=float) for _, stratum_rows in grouped.groupby("stratum", sort=True): numerators = stratum_rows.numerator.to_numpy() denominators = stratum_rows.denominator.to_numpy() indices = rng.integers( 0, len(stratum_rows), size=(n_boot, len(stratum_rows)) ) draw_numerators += numerators[indices].sum(axis=1) draw_denominators += denominators[indices].sum(axis=1) if np.any(draw_denominators == 0): raise ValueError("cluster_boot produced a zero-weight resample") draws = draw_numerators / draw_denominators lo, hi = np.percentile(draws, PCTILE) return ( float(weighted_values.sum() / weights.sum()), float(lo), float(hi), )