From c2119a6fd6933b11b3ff33024e6d7f13ea43728d Mon Sep 17 00:00:00 2001
From: ipezygj <231574583+ipezygj@users.noreply.github.com>
Date: Sat, 8 Aug 2026 16:07:59 +0300
Subject: [PATCH] Add bootstrap CI and age adjustment to calculate_c_index

A C-index of 0.75 from 100 deaths and from 5,000 deaths support very
different conclusions, and chronological age alone predicts mortality -
so an unadjusted point estimate both hides its uncertainty and rewards a
clock merely for correlating with age.

- ci_bootstrap_samples=N adds seeded percentile-bootstrap 95% CI columns
  (subject-level resampling; resamples with zero events are skipped)
- adjust_for_age=True computes the C-index on residuals after regressing
  each predictor on chronological age, mirroring the standardization
  already used in calculate_mortality_hazard_ratios
- defaults unchanged: existing callers get the exact same output
- tests: parity with direct lifelines concordance_index, CI brackets and
  determinism, narrowing with n, and an age-only synthetic clock whose
  adjusted C-index collapses to 0.5 while a genuinely informative clock
  retains signal
---
 biolearn/mortality.py                   | 64 ++++++++++++++++--
 biolearn/test/test_mortality_c_index.py | 86 +++++++++++++++++++++++++
 2 files changed, 145 insertions(+), 5 deletions(-)
 create mode 100644 biolearn/test/test_mortality_c_index.py

diff --git a/biolearn/mortality.py b/biolearn/mortality.py
index 42e8959..81ea992 100644
--- a/biolearn/mortality.py
+++ b/biolearn/mortality.py
@@ -42,18 +42,39 @@ def run_predictions(data, predictors_dict):
     return results_df
 
 
-def calculate_c_index(data, predictor_results):
+def calculate_c_index(
+    data,
+    predictor_results,
+    ci_bootstrap_samples=0,
+    seed=42,
+    adjust_for_age=False,
+):
     """
-    Calculates the C-index for each predictor in the predictor_results DataFrame without adjusting for age.
+    Calculates the C-index for each predictor in the predictor_results DataFrame.
 
     Args:
         data (Dataset): A Dataset object containing metadata with columns:
             'dead' - boolean indicating if the subject is dead
             'years_until_death' - time until death or censoring
         predictor_results (pd.DataFrame): DataFrame containing predictor results. Columns are the names of the predictors, and rows are IDs from data.
+        ci_bootstrap_samples (int): When > 0, also report a 95% percentile-bootstrap
+            confidence interval for each C-index (columns 'CI95_low'/'CI95_high'),
+            resampling subjects with replacement this many times. A C-index of 0.75
+            from 100 deaths and from 5,000 deaths support very different conclusions;
+            the interval makes that difference visible.
+        seed (int): Seed for the bootstrap resampling, so reported intervals are
+            reproducible.
+        adjust_for_age (bool): When True, the C-index is computed on the residuals
+            of each predictor after regressing out chronological age (metadata
+            column 'age'), mirroring the standardization used in
+            calculate_mortality_hazard_ratios. Chronological age alone predicts
+            mortality, so an unadjusted C-index rewards a clock merely for
+            correlating with age; the adjusted value asks what the clock adds
+            beyond it.
 
     Returns:
-        pd.DataFrame: A DataFrame containing C-index values for each predictor.
+        pd.DataFrame: A DataFrame containing C-index values for each predictor,
+        plus CI columns when ci_bootstrap_samples > 0.
     """
     # Merge predictor results with metadata
     analysis_df = pd.merge(
@@ -62,12 +83,23 @@ def calculate_c_index(data, predictor_results):
 
     # Remove rows with missing 'dead' or 'years_until_death' values
     analysis_df = analysis_df.dropna(subset=["dead", "years_until_death"])
+    if adjust_for_age:
+        analysis_df = analysis_df.dropna(subset=["age"])
+
+    def compute_scores(frame, clock):
+        predictor_values = frame[clock].astype(float)
+        if adjust_for_age:
+            age = frame["age"].astype(float)
+            slope, intercept = np.polyfit(age, predictor_values, 1)
+            predictor_values = predictor_values - (slope * age + intercept)
+        return predictor_values
 
     c_index_values = []
+    ci_lows = []
+    ci_highs = []
 
     for clock in predictor_results.columns:
-        # Ensure predictor values are numeric
-        predictor_values = analysis_df[clock].astype(float)
+        predictor_values = compute_scores(analysis_df, clock)
 
         # Calculate the C-index directly
         c_index = concordance_index(
@@ -77,6 +109,25 @@ def calculate_c_index(data, predictor_results):
         )
         c_index_values.append(c_index)
 
+        if ci_bootstrap_samples > 0:
+            rng = np.random.default_rng(seed)
+            boot_values = []
+            n = len(analysis_df)
+            for _ in range(ci_bootstrap_samples):
+                idx = rng.integers(0, n, size=n)
+                sample = analysis_df.iloc[idx]
+                if sample["dead"].astype(bool).sum() == 0:
+                    continue
+                boot_values.append(
+                    concordance_index(
+                        event_times=sample["years_until_death"],
+                        predicted_scores=-compute_scores(sample, clock),
+                        event_observed=sample["dead"],
+                    )
+                )
+            ci_lows.append(np.quantile(boot_values, 0.025))
+            ci_highs.append(np.quantile(boot_values, 0.975))
+
     # Create a DataFrame with the results
     results_df = pd.DataFrame(
         {
@@ -84,6 +135,9 @@ def calculate_c_index(data, predictor_results):
             "C_index": c_index_values,
         }
     )
+    if ci_bootstrap_samples > 0:
+        results_df["CI95_low"] = ci_lows
+        results_df["CI95_high"] = ci_highs
 
     return results_df
 
diff --git a/biolearn/test/test_mortality_c_index.py b/biolearn/test/test_mortality_c_index.py
new file mode 100644
index 0000000..ac48450
--- /dev/null
+++ b/biolearn/test/test_mortality_c_index.py
@@ -0,0 +1,86 @@
+import numpy as np
+import pandas as pd
+from lifelines.utils import concordance_index
+
+from biolearn.mortality import calculate_c_index
+
+
+class _StubData:
+    def __init__(self, metadata):
+        self.metadata = metadata
+
+
+def _make_data(n=200, seed=7):
+    rng = np.random.default_rng(seed)
+    ids = [f"s{i}" for i in range(n)]
+    age = rng.uniform(50, 90, size=n)
+    # mortality driven by age plus noise
+    years_until_death = np.clip(
+        (100 - age) / 4 + rng.normal(0, 2, size=n), 0.1, None
+    )
+    dead = rng.uniform(size=n) < 0.6
+    metadata = pd.DataFrame(
+        {"age": age, "dead": dead.astype(int), "years_until_death": years_until_death},
+        index=ids,
+    )
+    predictors = pd.DataFrame(
+        {
+            # a clock that is essentially chronological age (plus assay noise):
+            # predictive of mortality, but adds nothing beyond age. Small noise
+            # keeps the age-regression residuals well-defined; an EXACT copy
+            # leaves only float-epsilon residuals that are still monotone in age.
+            "AgeCopy": age + rng.normal(0, 0.3, size=n),
+            # a clock with genuine age-independent signal
+            "Informative": (100 - years_until_death * 4) + rng.normal(0, 1, size=n),
+        },
+        index=ids,
+    )
+    return _StubData(metadata), predictors
+
+
+def test_default_matches_direct_concordance_and_columns():
+    data, predictors = _make_data()
+    result = calculate_c_index(data, predictors)
+    assert list(result.columns) == ["Clock", "C_index"]
+    expected = concordance_index(
+        event_times=data.metadata["years_until_death"],
+        predicted_scores=-predictors["AgeCopy"].astype(float),
+        event_observed=data.metadata["dead"],
+    )
+    got = result.loc[result["Clock"] == "AgeCopy", "C_index"].iloc[0]
+    assert abs(got - expected) < 1e-12
+
+
+def test_bootstrap_ci_brackets_estimate_and_is_deterministic():
+    data, predictors = _make_data()
+    result = calculate_c_index(data, predictors, ci_bootstrap_samples=200)
+    assert {"CI95_low", "CI95_high"} <= set(result.columns)
+    for _, row in result.iterrows():
+        assert row["CI95_low"] < row["C_index"] < row["CI95_high"]
+    again = calculate_c_index(data, predictors, ci_bootstrap_samples=200)
+    pd.testing.assert_frame_equal(result, again)
+
+
+def test_ci_narrows_with_more_subjects():
+    small_data, small_pred = _make_data(n=60, seed=3)
+    large_data, large_pred = _make_data(n=600, seed=3)
+    small = calculate_c_index(small_data, small_pred, ci_bootstrap_samples=200)
+    large = calculate_c_index(large_data, large_pred, ci_bootstrap_samples=200)
+    width = lambda df: (df["CI95_high"] - df["CI95_low"]).iloc[0]
+    assert width(large) < width(small)
+
+
+def test_age_adjustment_removes_age_only_signal():
+    data, predictors = _make_data(n=400, seed=11)
+    unadjusted = calculate_c_index(data, predictors)
+    adjusted = calculate_c_index(data, predictors, adjust_for_age=True)
+
+    def value(df, clock):
+        return df.loc[df["Clock"] == clock, "C_index"].iloc[0]
+
+    # the age-copy clock predicts mortality before adjustment...
+    assert value(unadjusted, "AgeCopy") > 0.6
+    # ...and collapses to chance once age is regressed out
+    assert abs(value(adjusted, "AgeCopy") - 0.5) < 0.05
+    # while the genuinely informative clock retains signal beyond age
+    assert value(adjusted, "Informative") > 0.55