From c369a08fe15d6fd8364d1bf71fb8bfab9c81f712 Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Mon, 24 Aug 2026 02:28:54 -0400 Subject: [PATCH] Fix restricted_mean_survival_time return_variance to return the estimator's sampling variance for Kaplan-Meier fits restricted_mean_survival_time(..., return_variance=True) returned E[min(T,t)^2] - E[min(T,t)]^2, the variance of the truncated random variable -- a property of the distribution -- rather than the sampling variance of the RMST estimator, silently inflating standard errors and p-values built from it. For a KaplanMeierFitter fit to right-censored data, return the Greenwood-based sampling variance (Klein & Moeschberger 2003, eq. 4.4.4), matching R's survRM2::rmst2. For parametric fitters and precomputed survival functions, keep the previous value for backwards compatibility but emit a StatisticalWarning that it is not a sampling variance; document both behaviours in the docstring. Closes #1682 --- lifelines/tests/utils/test_utils.py | 58 +++++++++++++++++++++++++ lifelines/utils/__init__.py | 65 ++++++++++++++++++++++++++++- 2 files changed, 122 insertions(+), 1 deletion(-) diff --git a/lifelines/tests/utils/test_utils.py b/lifelines/tests/utils/test_utils.py index f7fe5e2c7..60b4f784b 100644 --- a/lifelines/tests/utils/test_utils.py +++ b/lifelines/tests/utils/test_utils.py @@ -963,6 +963,64 @@ def test_rmst_variance(): assert abs(utils.restricted_mean_survival_time(expf, t=t, return_variance=True)[1] - actual_var) < 0.001 +def test_rmst_variance_of_kaplan_meier_matches_survRM2(): + # issue #1682: return_variance=True on a KaplanMeierFitter must return the sampling + # variance of the RMST estimator (Klein & Moeschberger 2003, eq. 4.4.4), not the + # variance of the truncated random variable. + # Reference values from R's survRM2::rmst2 on the waltons miR-137 group at tau=10: + # RMST = 9.794, SE = 0.123. + df = load_waltons() + ix = df["group"] == "miR-137" + kmf = KaplanMeierFitter().fit(df.loc[ix]["T"], df.loc[ix]["E"]) + + mean, var = utils.restricted_mean_survival_time(kmf, t=10, return_variance=True) + assert abs(mean - 9.794) < 0.001 + assert abs(np.sqrt(var) - 0.123) < 0.0005 + + +def test_rmst_variance_of_kaplan_meier_equals_greenwood_formula_computed_by_hand(): + # Var(RMST) = sum_i [int_{t_i}^{tau} S(u) du]^2 * d_i / (n_i * (n_i - d_i)), + # summed over distinct event times t_i <= tau. Computed here directly from the + # raw data, independent of the library implementation. + T = np.array([1.0, 2.0, 2.0, 3.0, 5.0, 6.0, 8.0, 9.0]) + E = np.array([1, 1, 0, 1, 1, 0, 1, 0]) + tau = 7.0 + + event_times = np.unique(T[(E == 1) & (T <= tau)]) + surv, d_, n_ = [], [], [] + S = 1.0 + for t_i in event_times: + n_i = (T >= t_i).sum() + d_i = ((T == t_i) & (E == 1)).sum() + S *= 1 - d_i / n_i + surv.append(S), d_.append(d_i), n_.append(n_i) + + grid = np.append(event_times, tau) + expected_var = 0.0 + for i in range(len(event_times)): + integral_ti_to_tau = (np.diff(grid[i:]) * np.array(surv[i:])).sum() + expected_var += integral_ti_to_tau**2 * d_[i] / (n_[i] * (n_[i] - d_[i])) + + kmf = KaplanMeierFitter().fit(T, E) + _, var = utils.restricted_mean_survival_time(kmf, t=tau, return_variance=True) + assert abs(var - expected_var) < 1e-10 + + +def test_rmst_variance_of_non_km_models_warns_that_it_is_not_a_sampling_variance(): + # issue #1682: for non-KM inputs the returned "variance" is still the variance of the + # truncated random variable (kept for backwards compatibility), so users must be warned + # that it cannot be used for standard errors or hypothesis tests. + T = np.random.exponential(2, 100) + expf = ExponentialFitter().fit(T) + + with pytest.warns(exceptions.StatisticalWarning, match="sampling variance"): + utils.restricted_mean_survival_time(expf, t=1.0, return_variance=True) + + kmf = KaplanMeierFitter().fit(T) + with pytest.warns(exceptions.StatisticalWarning, match="sampling variance"): + utils.restricted_mean_survival_time(kmf.survival_function_, t=1.0, return_variance=True) + + def test_find_best_parametric_model(): T = np.random.exponential(2, 1000) E = np.ones_like(T) diff --git a/lifelines/utils/__init__.py b/lifelines/utils/__init__.py index ee87e2dcc..2a067006f 100644 --- a/lifelines/utils/__init__.py +++ b/lifelines/utils/__init__.py @@ -20,7 +20,7 @@ import formulaic from lifelines.utils.concordance import concordance_index -from lifelines.exceptions import ConvergenceWarning, ApproximationWarning, ConvergenceError +from lifelines.exceptions import ConvergenceWarning, ApproximationWarning, ConvergenceError, StatisticalWarning __all__ = [ @@ -224,6 +224,19 @@ def restricted_mean_survival_time( This can be a univariate model, or a pandas DataFrame. The former will provide a more accurate estimate however. t: float The upper limit of the integration in the RMST. + return_variance: bool + If True, also return a variance along with the RMST. + + If the input is a ``KaplanMeierFitter`` fit to right-censored data, this is the sampling variance + of the RMST *estimator*, computed with the Greenwood-based formula of Klein & Moeschberger (2003), + eq. 4.4.4 (the same estimator used by R's ``survRM2::rmst2``). Its square root is a standard error, + suitable for confidence intervals and hypothesis tests. + + For any other input (a parametric model, or a precomputed survival function DataFrame), the value + returned is :math:`E[\min(T, t)^2] - E[\min(T, t)]^2`, the variance of the *truncated random + variable* — a property of the fitted distribution, **not** the sampling variance of the estimator. + It is not suitable for standard errors, confidence intervals or hypothesis tests, and a + ``StatisticalWarning`` is emitted. Example -------- @@ -244,17 +257,67 @@ def restricted_mean_survival_time( ------- https://bmcmedresmethodol.biomedcentral.com/articles/10.1186/1471-2288-13-152#Sec27 + Klein, J. P., & Moeschberger, M. L. (2003). Survival Analysis: Techniques for Censored and Truncated + Data (2nd ed.), Section 4.5, eq. 4.4.4. Springer. + """ + import lifelines + t = coalesce(t, np.inf) mean = _expected_value_of_survival_up_to_t(model_or_survival_function, t) if return_variance: + is_right_censored_km = isinstance( + model_or_survival_function, lifelines.KaplanMeierFitter + ) and CensoringType.is_right_censoring(model_or_survival_function) + if is_right_censored_km: + return (mean, _sampling_variance_of_km_rmst(model_or_survival_function, t)) + warnings.warn( + "The variance returned is the variance of the truncated random variable min(T, t) — a property of " + "the fitted distribution — NOT the sampling variance of the RMST estimator. It is not suitable for " + "standard errors, confidence intervals or hypothesis tests. Fit a KaplanMeierFitter to right-censored " + "data and pass the fitter (not its survival function) to get the estimator's sampling variance.", + StatisticalWarning, + ) sq = _expected_value_of_survival_squared_up_to_t(model_or_survival_function, t) return (mean, sq - mean**2) else: return mean +def _sampling_variance_of_km_rmst(model, t: float) -> float: + r""" + Sampling variance of the Kaplan-Meier RMST estimator :math:`\hat{\mu}(t) = \int_0^t \hat{S}(u) du`, + using the Greenwood-based formula of Klein & Moeschberger (2003), eq. 4.4.4: + + .. math:: \widehat{\text{Var}}[\hat{\mu}(t)] = \sum_{i: t_i \le t} \left[ \int_{t_i}^{t} \hat{S}(u) du \right]^2 \frac{d_i}{n_i (n_i - d_i)} + + where the sum is over distinct event times :math:`t_i`, with :math:`d_i` events out of + :math:`n_i` at risk. This matches R's ``survRM2::rmst2``. + """ + event_table = model.event_table + events = event_table.loc[(event_table.index > 0) & (event_table.index <= t) & (event_table["observed"] > 0)] + if events.empty: + return 0.0 + + event_times = events.index.values.astype(float) + d = events["observed"].values.astype(float) + n = events["at_risk"].values.astype(float) + + # S is a step function that only changes at event times, so on [t_i, t_{i+1}) it equals S(t_i). + surv_at_events = np.atleast_1d(np.asarray(model.predict(event_times))) + widths = np.diff(np.append(event_times, t)) + # avoid 0 * inf = nan when t == inf and the survival function reaches 0 at the last event time + with np.errstate(invalid="ignore"): + areas = np.where(surv_at_events == 0.0, 0.0, widths * surv_at_events) + # tail_integrals[i] = integral of S from t_i to t + tail_integrals = np.cumsum(areas[::-1])[::-1] + + # when n_i == d_i (everyone remaining dies), the Greenwood term is conventionally 0 (as in survRM2). + variance_terms = np.where(n > d, d / (n * (n - d)), 0.0) + return float((tail_integrals**2 * variance_terms).sum()) + + def _expected_value_of_survival_up_to_t(model_or_survival_function, t: float = np.inf) -> float: import lifelines