From 88509c71a3727456225319a657f1114ee1d5e70d Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Sat, 25 Feb 2023 17:23:20 +1100
Subject: [PATCH] Implement yli.logrank

---
 yli/__init__.py  |  2 +-
 yli/sig_tests.py |  2 +-
 yli/survival.py  | 14 ++++++++++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/yli/__init__.py b/yli/__init__.py
index 7c88db0..e9e4c66 100644
--- a/yli/__init__.py
+++ b/yli/__init__.py
@@ -22,7 +22,7 @@ from .graphs import init_fonts
 from .io import pickle_read_compressed, pickle_read_encrypted, pickle_write_compressed, pickle_write_encrypted
 from .regress import OrdinalLogit, PenalisedLogit, logit_then_regress, regress, vif
 from .sig_tests import anova_oneway, auto_univariable, chi2, mannwhitney, pearsonr, spearman, ttest_ind
-from .survival import kaplanmeier
+from .survival import kaplanmeier, logrank
 from .utils import as_ordinal
 
 def reload_me():
diff --git a/yli/sig_tests.py b/yli/sig_tests.py
index 6365f92..59e7b6f 100644
--- a/yli/sig_tests.py
+++ b/yli/sig_tests.py
@@ -513,7 +513,7 @@ class ChiSquaredResult:
 		:rtype: str
 		"""
 		
-		return 'χ²({}) = {:.2f}; p {}'.format(self.ct, self.dof, self.statistic, fmt_p(self.pvalue, PValueStyle.RELATION))
+		return 'χ²({}) = {:.2f}; p {}'.format(self.dof, self.statistic, fmt_p(self.pvalue, PValueStyle.RELATION))
 
 class PearsonChiSquaredResult(ChiSquaredResult):
 	"""
diff --git a/yli/survival.py b/yli/survival.py
index 4e05792..7c02830 100644
--- a/yli/survival.py
+++ b/yli/survival.py
@@ -18,6 +18,7 @@ from scipy import stats
 import statsmodels.api as sm
 
 from .config import config
+from .sig_tests import ChiSquaredResult
 from .utils import check_nan
 
 def kaplanmeier(df, time, status, by=None, ci=True, nan_policy='warn'):
@@ -102,3 +103,16 @@ def plot_survfunc(ax, time, status, ci):
 		ax.fill_between(xpoints, ypoints0, ypoints1, alpha=0.3, label='_')
 	
 	return handle
+
+def logrank(df, time, status, by, nan_policy='warn'):
+	# TODO: Documentation
+	
+	# Check for/clean NaNs
+	df = check_nan(df[[time, status, by]], nan_policy)
+	
+	if df[time].dtype == '<m8[ns]':
+		df[time] = df[time].dt.total_seconds()
+	
+	statistic, pvalue = sm.duration.survdiff(df[time], df[status], df[by])
+	
+	return ChiSquaredResult(statistic=statistic, dof=1, pvalue=pvalue)