Return the RMST estimator's sampling variance (Greenwood, K&M eq. 4.4.4) from restricted_mean_survival_time for Kaplan-Meier fits

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)

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