Coxphfitter

ehrapy/tools/_sa.py

@@ -6,7 +6,7 @@
 import pandas as pd
 import statsmodels.api as sm
 import statsmodels.formula.api as smf
-from lifelines import KaplanMeierFitter
+from lifelines import CoxPHFitter, KaplanMeierFitter
 from lifelines.statistics import StatisticalResult, logrank_test
 from scipy import stats
 
@@ -262,3 +262,15 @@ def anova_glm(result_1: GLMResultsWrapper, result_2: GLMResultsWrapper, formula_
     }
     dataframe = pd.DataFrame(data=table)
     return dataframe
+
+
+def cph(ad: AnnData, duration_col: str, event_col: str, entry_col: str = None) -> KaplanMeierFitter:
+    """Fit the Cox’s proportional hazard for the survival function.
+
+    See https://lifelines.readthedocs.io/en/latest/fitters/regression/CoxPHFitter.html
+    """
+    df = ad.to_df()
+    df = df[[duration_col, event_col, entry_col]]
+    cph = CoxPHFitter()
+    cph.fit(df, duration_col, event_col, entry_col=entry_col)
+    return cph

tests/tools/test_sa.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 import statsmodels
-from lifelines import KaplanMeierFitter
+from lifelines import CoxPHFitter, KaplanMeierFitter
 
 import ehrapy as ep
 
@@ -75,3 +75,12 @@ def test_anova_glm(self):
         assert dataframe.shape == (2, 6)
         assert dataframe.iloc[1, 4] == 2
         assert pytest.approx(dataframe.iloc[1, 5], 0.1) == 0.103185
+
+    def test_cph(self):
+        adata = ep.dt.mimic_2(encoded=False)
+        adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
+        cph = ep.tl.cph(adata, "mort_day_censored", "censor_flg")
+
+        assert isinstance(cph, CoxPHFitter)
+        assert len(cph.durations) == 1776
+        assert sum(cph.event_observed) == 497