Add unit test for yli.kaplanmeier

This commit is contained in:
RunasSudo 2023-04-22 01:18:02 +10:00
parent aff4287ccc
commit 3d30045832
Signed by: RunasSudo
GPG Key ID: 7234E476BF21C61A
2 changed files with 70 additions and 3 deletions

58
tests/test_kaplanmeier.py Normal file
View File

@ -0,0 +1,58 @@
# scipy-yli: Helpful SciPy utilities and recipes
# Copyright © 2022–2023 Lee Yingtong Li (RunasSudo)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from pytest import approx
import pandas as pd
import yli
def test_kaplanmeier_simple():
"""Compare yli.kaplanmeier for simple example"""
df = pd.DataFrame({
'SurvTime': [2, 4, 6, 8],
'Status': [True, True, True, True]
})
xpoints, ypoints, _, _ = yli.survival.calc_survfunc_kaplanmeier(df['SurvTime'], df['Status'], False)
assert xpoints[0] == 0
assert ypoints[0] == 1
assert xpoints[1] == 2
assert ypoints[1] == 1
assert xpoints[2] == 2
assert ypoints[2] == 0.75
assert xpoints[3] == 4
assert ypoints[3] == 0.75
assert xpoints[4] == 4
assert ypoints[4] == approx(0.5)
assert xpoints[5] == 6
assert ypoints[5] == approx(0.5)
assert xpoints[6] == 6
assert ypoints[6] == approx(0.25)
assert xpoints[7] == 8
assert ypoints[7] == approx(0.25)
assert xpoints[8] == 8
assert ypoints[8] == 0

View File

@ -87,6 +87,16 @@ def kaplanmeier(df, time, status, by=None, *, ci=True, transform_x=None, transfo
return fig, ax
def plot_survfunc_kaplanmeier(ax, time, status, ci, transform_x=None, transform_y=None):
xpoints, ypoints, ypoints0, ypoints1 = calc_survfunc_kaplanmeier(time, status, ci, transform_x, transform_y)
handle = ax.plot(xpoints, ypoints)[0]
if ci:
ax.fill_between(xpoints, ypoints0, ypoints1, alpha=0.3, label='_')
return handle
def calc_survfunc_kaplanmeier(time, status, ci, transform_x=None, transform_y=None):
# Estimate the survival function
sf = sm.SurvfuncRight(time, status)
@ -94,7 +104,6 @@ def plot_survfunc_kaplanmeier(ax, time, status, ci, transform_x=None, transform_
# np.concatenate(...) to force starting drawing from time 0, survival 100%
xpoints = np.concatenate([[0], sf.surv_times]).repeat(2)[1:]
ypoints = np.concatenate([[1], sf.surv_prob]).repeat(2)[:-1]
handle = ax.plot(xpoints, ypoints)[0]
if transform_x:
xpoints = transform_x(xpoints)
@ -116,9 +125,9 @@ def plot_survfunc_kaplanmeier(ax, time, status, ci, transform_x=None, transform_
ypoints0 = transform_y(ypoints0)
ypoints1 = transform_y(ypoints1)
ax.fill_between(xpoints, ypoints0, ypoints1, alpha=0.3, label='_')
return xpoints, ypoints, ypoints0, ypoints1
return handle
return xpoints, ypoints, None, None
def turnbull(df, time_left, time_right, by=None, *, step_loc=0.5, transform_x=None, transform_y=None, nan_policy='warn'):
"""