import pytest
import numpy as np
from numpy.testing import assert_equal, assert_allclose
from scipy import stats
from scipy.stats import _survival
def _kaplan_meier_reference(times, censored):
# This is a very straightforward implementation of the Kaplan-Meier
# estimator that does almost everything differently from the implementation
# in stats.ecdf.
# Begin by sorting the raw data. Note that the order of death and loss
# at a given time matters: death happens first. See [2] page 461:
# "These conventions may be paraphrased by saying that deaths recorded as
# of an age t are treated as if they occurred slightly before t, and losses
# recorded as of an age t are treated as occurring slightly after t."
# We implement this by sorting the data first by time, then by `censored`,
# (which is 0 when there is a death and 1 when there is only a loss).
dtype = [('time', float), ('censored', int)]
data = np.array([(t, d) for t, d in zip(times, censored)], dtype=dtype)
data = np.sort(data, order=('time', 'censored'))
times = data['time']
died = np.logical_not(data['censored'])
m = times.size
n = np.arange(m, 0, -1) # number at risk
sf = np.cumprod((n - died) / n)
# Find the indices of the *last* occurrence of unique times. The
# corresponding entries of `times` and `sf` are what we want.
_, indices = np.unique(times[::-1], return_index=True)
ref_times = times[-indices - 1]
ref_sf = sf[-indices - 1]
return ref_times, ref_sf
class TestSurvival:
@staticmethod
def get_random_sample(rng, n_unique):
# generate random sample
unique_times = rng.random(n_unique)
# convert to `np.int32` to resolve `np.repeat` failure in 32-bit CI
repeats = rng.integers(1, 4, n_unique).astype(np.int32)
times = rng.permuted(np.repeat(unique_times, repeats))
censored = rng.random(size=times.size) > rng.random()
sample = stats.CensoredData.right_censored(times, censored)
return sample, times, censored
def test_input_validation(self):
message = '`sample` must be a one-dimensional sequence.'
with pytest.raises(ValueError, match=message):
stats.ecdf([[1]])
with pytest.raises(ValueError, match=message):
stats.ecdf(1)
message = '`sample` must not contain nan'
with pytest.raises(ValueError, match=message):
stats.ecdf([np.nan])
message = 'Currently, only uncensored and right-censored data...'
with pytest.raises(NotImplementedError, match=message):
stats.ecdf(stats.CensoredData.left_censored([1], censored=[True]))
message = 'method` must be one of...'
res = stats.ecdf([1, 2, 3])
with pytest.raises(ValueError, match=message):
res.cdf.confidence_interval(method='ekki-ekki')
with pytest.raises(ValueError, match=message):
res.sf.confidence_interval(method='shrubbery')
message = 'confidence_level` must be a scalar between 0 and 1'
with pytest.raises(ValueError, match=message):
res.cdf.confidence_interval(-1)
with pytest.raises(ValueError, match=message):
res.sf.confidence_interval([0.5, 0.6])
message = 'The confidence interval is undefined at some observations.'
with pytest.warns(RuntimeWarning, match=message):
ci = res.cdf.confidence_interval()
message = 'Confidence interval bounds do not implement...'
with pytest.raises(NotImplementedError, match=message):
ci.low.confidence_interval()
with pytest.raises(NotImplementedError, match=message):
ci.high.confidence_interval()
def test_edge_cases(self):
res = stats.ecdf([])
assert_equal(res.cdf.quantiles, [])
assert_equal(res.cdf.probabilities, [])
res = stats.ecdf([1])
assert_equal(res.cdf.quantiles, [1])
assert_equal(res.cdf.probabilities, [1])
def test_unique(self):
# Example with unique observations; `stats.ecdf` ref. [1] page 80
sample = [6.23, 5.58, 7.06, 6.42, 5.20]
res = stats.ecdf(sample)
ref_x = np.sort(np.unique(sample))
ref_cdf = np.arange(1, 6) / 5
ref_sf = 1 - ref_cdf
assert_equal(res.cdf.quantiles, ref_x)
assert_equal(res.cdf.probabilities, ref_cdf)
assert_equal(res.sf.quantiles, ref_x)
assert_equal(res.sf.probabilities, ref_sf)
def test_nonunique(self):
# Example with non-unique observations; `stats.ecdf` ref. [1] page 82
sample = [0, 2, 1, 2, 3, 4]
res = stats.ecdf(sample)
ref_x = np.sort(np.unique(sample))
ref_cdf = np.array([1/6, 2/6, 4/6, 5/6, 1])
ref_sf = 1 - ref_cdf
assert_equal(res.cdf.quantiles, ref_x)
assert_equal(res.cdf.probabilities, ref_cdf)
assert_equal(res.sf.quantiles, ref_x)
assert_equal(res.sf.probabilities, ref_sf)
def test_evaluate_methods(self):
# Test CDF and SF `evaluate` methods
rng = np.random.default_rng(1162729143302572461)
sample, _, _ = self.get_random_sample(rng, 15)
res = stats.ecdf(sample)
x = res.cdf.quantiles
xr = x + np.diff(x, append=x[-1]+1)/2 # right shifted points
assert_equal(res.cdf.evaluate(x), res.cdf.probabilities)
assert_equal(res.cdf.evaluate(xr), res.cdf.probabilities)
assert_equal(res.cdf.evaluate(x[0]-1), 0) # CDF starts at 0
assert_equal(res.cdf.evaluate([-np.inf, np.inf]), [0, 1])
assert_equal(res.sf.evaluate(x), res.sf.probabilities)
assert_equal(res.sf.evaluate(xr), res.sf.probabilities)
assert_equal(res.sf.evaluate(x[0]-1), 1) # SF starts at 1
assert_equal(res.sf.evaluate([-np.inf, np.inf]), [1, 0])
# ref. [1] page 91
t1 = [37, 43, 47, 56, 60, 62, 71, 77, 80, 81] # times
d1 = [0, 0, 1, 1, 0, 0, 0, 1, 1, 1] # 1 means deaths (not censored)
r1 = [1, 1, 0.875, 0.75, 0.75, 0.75, 0.75, 0.5, 0.25, 0] # reference SF
# https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html
t2 = [8, 12, 26, 14, 21, 27, 8, 32, 20, 40]
d2 = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0]
r2 = [0.9, 0.788, 0.675, 0.675, 0.54, 0.405, 0.27, 0.27, 0.27]
t3 = [33, 28, 41, 48, 48, 25, 37, 48, 25, 43]
d3 = [1, 1, 1, 0, 0, 0, 0, 0, 0, 0]
r3 = [1, 0.875, 0.75, 0.75, 0.6, 0.6, 0.6]
# https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/bs704_survival4.html
t4 = [24, 3, 11, 19, 24, 13, 14, 2, 18, 17,
24, 21, 12, 1, 10, 23, 6, 5, 9, 17]
d4 = [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1]
r4 = [0.95, 0.95, 0.897, 0.844, 0.844, 0.844, 0.844, 0.844, 0.844,
0.844, 0.76, 0.676, 0.676, 0.676, 0.676, 0.507, 0.507]
# https://www.real-statistics.com/survival-analysis/kaplan-meier-procedure/confidence-interval-for-the-survival-function/
t5 = [3, 5, 8, 10, 5, 5, 8, 12, 15, 14, 2, 11, 10, 9, 12, 5, 8, 11]
d5 = [1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1]
r5 = [0.944, 0.889, 0.722, 0.542, 0.542, 0.542, 0.361, 0.181, 0.181, 0.181]
@pytest.mark.parametrize("case", [(t1, d1, r1), (t2, d2, r2), (t3, d3, r3),
(t4, d4, r4), (t5, d5, r5)])
def test_right_censored_against_examples(self, case):
# test `ecdf` against other implementations on example problems
times, died, ref = case
sample = stats.CensoredData.right_censored(times, np.logical_not(died))
res = stats.ecdf(sample)
assert_allclose(res.sf.probabilities, ref, atol=1e-3)
assert_equal(res.sf.quantiles, np.sort(np.unique(times)))
# test reference implementation against other implementations
res = _kaplan_meier_reference(times, np.logical_not(died))
assert_equal(res[0], np.sort(np.unique(times)))
assert_allclose(res[1], ref, atol=1e-3)
@pytest.mark.parametrize('seed', [182746786639392128, 737379171436494115,
576033618403180168, 308115465002673650])
def test_right_censored_against_reference_implementation(self, seed):
# test `ecdf` against reference implementation on random problems
rng = np.random.default_rng(seed)
n_unique = rng.integers(10, 100)
sample, times, censored = self.get_random_sample(rng, n_unique)
res = stats.ecdf(sample)
ref = _kaplan_meier_reference(times, censored)
assert_allclose(res.sf.quantiles, ref[0])
assert_allclose(res.sf.probabilities, ref[1])
# If all observations are uncensored, the KM estimate should match
# the usual estimate for uncensored data
sample = stats.CensoredData(uncensored=times)
res = _survival._ecdf_right_censored(sample) # force Kaplan-Meier
ref = stats.ecdf(times)
assert_equal(res[0], ref.sf.quantiles)
assert_allclose(res[1], ref.cdf.probabilities, rtol=1e-14)
assert_allclose(res[2], ref.sf.probabilities, rtol=1e-14)
def test_right_censored_ci(self):
# test "greenwood" confidence interval against example 4 (URL above).
times, died = self.t4, self.d4
sample = stats.CensoredData.right_censored(times, np.logical_not(died))
res = stats.ecdf(sample)
ref_allowance = [0.096, 0.096, 0.135, 0.162, 0.162, 0.162, 0.162,
0.162, 0.162, 0.162, 0.214, 0.246, 0.246, 0.246,
0.246, 0.341, 0.341]
sf_ci = res.sf.confidence_interval()
cdf_ci = res.cdf.confidence_interval()
allowance = res.sf.probabilities - sf_ci.low.probabilities
assert_allclose(allowance, ref_allowance, atol=1e-3)
assert_allclose(sf_ci.low.probabilities,
np.clip(res.sf.probabilities - allowance, 0, 1))
assert_allclose(sf_ci.high.probabilities,
np.clip(res.sf.probabilities + allowance, 0, 1))
assert_allclose(cdf_ci.low.probabilities,
np.clip(res.cdf.probabilities - allowance, 0, 1))
assert_allclose(cdf_ci.high.probabilities,
np.clip(res.cdf.probabilities + allowance, 0, 1))
# test "log-log" confidence interval against Mathematica
# e = {24, 3, 11, 19, 24, 13, 14, 2, 18, 17, 24, 21, 12, 1, 10, 23, 6, 5,
# 9, 17}
# ci = {1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0}
# R = EventData[e, ci]
# S = SurvivalModelFit[R]
# S["PointwiseIntervals", ConfidenceLevel->0.95,
# ConfidenceTransform->"LogLog"]
ref_low = [0.694743, 0.694743, 0.647529, 0.591142, 0.591142, 0.591142,
0.591142, 0.591142, 0.591142, 0.591142, 0.464605, 0.370359,
0.370359, 0.370359, 0.370359, 0.160489, 0.160489]
ref_high = [0.992802, 0.992802, 0.973299, 0.947073, 0.947073, 0.947073,
0.947073, 0.947073, 0.947073, 0.947073, 0.906422, 0.856521,
0.856521, 0.856521, 0.856521, 0.776724, 0.776724]
sf_ci = res.sf.confidence_interval(method='log-log')
assert_allclose(sf_ci.low.probabilities, ref_low, atol=1e-6)
assert_allclose(sf_ci.high.probabilities, ref_high, atol=1e-6)
def test_right_censored_ci_example_5(self):
# test "exponential greenwood" confidence interval against example 5
times, died = self.t5, self.d5
sample = stats.CensoredData.right_censored(times, np.logical_not(died))
res = stats.ecdf(sample)
lower = np.array([0.66639, 0.624174, 0.456179, 0.287822, 0.287822,
0.287822, 0.128489, 0.030957, 0.030957, 0.030957])
upper = np.array([0.991983, 0.970995, 0.87378, 0.739467, 0.739467,
0.739467, 0.603133, 0.430365, 0.430365, 0.430365])
sf_ci = res.sf.confidence_interval(method='log-log')
cdf_ci = res.cdf.confidence_interval(method='log-log')
assert_allclose(sf_ci.low.probabilities, lower, atol=1e-5)
assert_allclose(sf_ci.high.probabilities, upper, atol=1e-5)
assert_allclose(cdf_ci.low.probabilities, 1-upper, atol=1e-5)
assert_allclose(cdf_ci.high.probabilities, 1-lower, atol=1e-5)
# Test against R's `survival` library `survfit` function, 90%CI
# library(survival)
# options(digits=16)
# time = c(3, 5, 8, 10, 5, 5, 8, 12, 15, 14, 2, 11, 10, 9, 12, 5, 8, 11)
# status = c(1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 1)
# res = survfit(Surv(time, status)
# ~1, conf.type = "log-log", conf.int = 0.90)
# res$time; res$lower; res$upper
low = [0.74366748406861172, 0.68582332289196246, 0.50596835651480121,
0.32913131413336727, 0.32913131413336727, 0.32913131413336727,
0.15986912028781664, 0.04499539918147757, 0.04499539918147757,
0.04499539918147757]
high = [0.9890291867238429, 0.9638835422144144, 0.8560366823086629,
0.7130167643978450, 0.7130167643978450, 0.7130167643978450,
0.5678602982997164, 0.3887616766886558, 0.3887616766886558,
0.3887616766886558]
sf_ci = res.sf.confidence_interval(method='log-log',
confidence_level=0.9)
assert_allclose(sf_ci.low.probabilities, low)
assert_allclose(sf_ci.high.probabilities, high)
# And with conf.type = "plain"
low = [0.8556383113628162, 0.7670478794850761, 0.5485720663578469,
0.3441515412527123, 0.3441515412527123, 0.3441515412527123,
0.1449184105424544, 0., 0., 0.]
high = [1., 1., 0.8958723780865975, 0.7391817920806210,
0.7391817920806210, 0.7391817920806210, 0.5773038116797676,
0.3642270254596720, 0.3642270254596720, 0.3642270254596720]
sf_ci = res.sf.confidence_interval(confidence_level=0.9)
assert_allclose(sf_ci.low.probabilities, low)
assert_allclose(sf_ci.high.probabilities, high)
def test_right_censored_ci_nans(self):
# test `ecdf` confidence interval on a problem that results in NaNs
times, died = self.t1, self.d1
sample = stats.CensoredData.right_censored(times, np.logical_not(died))
res = stats.ecdf(sample)
# Reference values generated with Matlab
# format long
# t = [37 43 47 56 60 62 71 77 80 81];
# d = [0 0 1 1 0 0 0 1 1 1];
# censored = ~d1;
# [f, x, flo, fup] = ecdf(t, 'Censoring', censored, 'Alpha', 0.05);
x = [37, 47, 56, 77, 80, 81]
flo = [np.nan, 0, 0, 0.052701464070711, 0.337611126231790, np.nan]
fup = [np.nan, 0.35417230377, 0.5500569798, 0.9472985359, 1.0, np.nan]
i = np.searchsorted(res.cdf.quantiles, x)
message = "The confidence interval is undefined at some observations"
with pytest.warns(RuntimeWarning, match=message):
ci = res.cdf.confidence_interval()
# Matlab gives NaN as the first element of the CIs. Mathematica agrees,
# but R's survfit does not. It makes some sense, but it's not what the
# formula gives, so skip that element.
assert_allclose(ci.low.probabilities[i][1:], flo[1:])
assert_allclose(ci.high.probabilities[i][1:], fup[1:])
# [f, x, flo, fup] = ecdf(t, 'Censoring', censored, 'Function',
# 'survivor', 'Alpha', 0.05);
flo = [np.nan, 0.64582769623, 0.449943020228, 0.05270146407, 0, np.nan]
fup = [np.nan, 1.0, 1.0, 0.947298535929289, 0.662388873768210, np.nan]
i = np.searchsorted(res.cdf.quantiles, x)
with pytest.warns(RuntimeWarning, match=message):
ci = res.sf.confidence_interval()
assert_allclose(ci.low.probabilities[i][1:], flo[1:])
assert_allclose(ci.high.probabilities[i][1:], fup[1:])
# With the same data, R's `survival` library `survfit` function
# doesn't produce the leading NaN
# library(survival)
# options(digits=16)
# time = c(37, 43, 47, 56, 60, 62, 71, 77, 80, 81)
# status = c(0, 0, 1, 1, 0, 0, 0, 1, 1, 1)
# res = survfit(Surv(time, status)
# ~1, conf.type = "plain", conf.int = 0.95)
# res$time
# res$lower
# res$upper
low = [1., 1., 0.64582769623233816, 0.44994302022779326,
0.44994302022779326, 0.44994302022779326, 0.44994302022779326,
0.05270146407071086, 0., np.nan]
high = [1., 1., 1., 1., 1., 1., 1., 0.9472985359292891,
0.6623888737682101, np.nan]
assert_allclose(ci.low.probabilities, low)
assert_allclose(ci.high.probabilities, high)
# It does with conf.type="log-log", as do we
with pytest.warns(RuntimeWarning, match=message):
ci = res.sf.confidence_interval(method='log-log')
low = [np.nan, np.nan, 0.38700001403202522, 0.31480711370551911,
0.31480711370551911, 0.31480711370551911, 0.31480711370551911,
0.08048821148507734, 0.01049958986680601, np.nan]
high = [np.nan, np.nan, 0.9813929658789660, 0.9308983170906275,
0.9308983170906275, 0.9308983170906275, 0.9308983170906275,
0.8263946341076415, 0.6558775085110887, np.nan]
assert_allclose(ci.low.probabilities, low)
assert_allclose(ci.high.probabilities, high)
def test_right_censored_against_uncensored(self):
rng = np.random.default_rng(7463952748044886637)
sample = rng.integers(10, 100, size=1000)
censored = np.zeros_like(sample)
censored[np.argmax(sample)] = True
res = stats.ecdf(sample)
ref = stats.ecdf(stats.CensoredData.right_censored(sample, censored))
assert_equal(res.sf.quantiles, ref.sf.quantiles)
assert_equal(res.sf._n, ref.sf._n)
assert_equal(res.sf._d[:-1], ref.sf._d[:-1]) # difference @ [-1]
assert_allclose(res.sf._sf[:-1], ref.sf._sf[:-1], rtol=1e-14)
def test_plot_iv(self):
rng = np.random.default_rng(1769658657308472721)
n_unique = rng.integers(10, 100)
sample, _, _ = self.get_random_sample(rng, n_unique)
res = stats.ecdf(sample)
try:
import matplotlib.pyplot as plt # noqa: F401
res.sf.plot() # no other errors occur
except (ModuleNotFoundError, ImportError):
# Avoid trying to call MPL with numpy 2.0-dev, because that fails
# too often due to ABI mismatches and is hard to avoid. This test
# will work fine again once MPL has done a 2.0-compatible release.
if not np.__version__.startswith('2.0.0.dev0'):
message = r"matplotlib must be installed to use method `plot`."
with pytest.raises(ModuleNotFoundError, match=message):
res.sf.plot()
class TestLogRank:
@pytest.mark.parametrize(
"x, y, statistic, pvalue",
# Results validate with R
# library(survival)
# options(digits=16)
#
# futime_1 <- c(8, 12, 26, 14, 21, 27, 8, 32, 20, 40)
# fustat_1 <- c(1, 1, 1, 1, 1, 1, 0, 0, 0, 0)
# rx_1 <- c(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
#
# futime_2 <- c(33, 28, 41, 48, 48, 25, 37, 48, 25, 43)
# fustat_2 <- c(1, 1, 1, 0, 0, 0, 0, 0, 0, 0)
# rx_2 <- c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1)
#
# futime <- c(futime_1, futime_2)
# fustat <- c(fustat_1, fustat_2)
# rx <- c(rx_1, rx_2)
#
# survdiff(formula = Surv(futime, fustat) ~ rx)
#
# Also check against another library which handle alternatives
# library(nph)
# logrank.test(futime, fustat, rx, alternative = "two.sided")
# res["test"]
[(
# https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html
# uncensored, censored
[[8, 12, 26, 14, 21, 27], [8, 32, 20, 40]],
[[33, 28, 41], [48, 48, 25, 37, 48, 25, 43]],
# chi2, ["two-sided", "less", "greater"]
6.91598157449,
[0.008542873404, 0.9957285632979385, 0.004271436702061537]
),
(
# https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_survival/BS704_Survival5.html
[[19, 6, 5, 4], [20, 19, 17, 14]],
[[16, 21, 7], [21, 15, 18, 18, 5]],
0.835004855038,
[0.3608293039, 0.8195853480676912, 0.1804146519323088]
),
(
# Bland, Altman, "The logrank test", BMJ, 2004
# https://www.bmj.com/content/328/7447/1073.short
[[6, 13, 21, 30, 37, 38, 49, 50, 63, 79, 86, 98, 202, 219],
[31, 47, 80, 82, 82, 149]],
[[10, 10, 12, 13, 14, 15, 16, 17, 18, 20, 24, 24, 25, 28, 30,
33, 35, 37, 40, 40, 46, 48, 76, 81, 82, 91, 112, 181],
[34, 40, 70]],
7.49659416854,
[0.006181578637, 0.003090789318730882, 0.9969092106812691]
)]
)
def test_log_rank(self, x, y, statistic, pvalue):
x = stats.CensoredData(uncensored=x[0], right=x[1])
y = stats.CensoredData(uncensored=y[0], right=y[1])
for i, alternative in enumerate(["two-sided", "less", "greater"]):
res = stats.logrank(x=x, y=y, alternative=alternative)
# we return z and use the normal distribution while other framework
# return z**2. The p-value are directly comparable, but we have to
# square the statistic
assert_allclose(res.statistic**2, statistic, atol=1e-10)
assert_allclose(res.pvalue, pvalue[i], atol=1e-10)
def test_raises(self):
sample = stats.CensoredData([1, 2])
msg = r"`y` must be"
with pytest.raises(ValueError, match=msg):
stats.logrank(x=sample, y=[[1, 2]])
msg = r"`x` must be"
with pytest.raises(ValueError, match=msg):
stats.logrank(x=[[1, 2]], y=sample)