Improve survival analysis interface

ehrapy/plot/_survival_analysis.py

@@ -186,22 +186,20 @@ def kmf(
         # So we need to flip `censor_fl` when pass `censor_fl` to KaplanMeierFitter
 
         >>> adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
-        >>> kmf = ep.tl.kmf(adata[:, ["mort_day_censored"]].X, adata[:, ["censor_flg"]].X)
+        >>> kmf = ep.tl.kmf(adata, "mort_day_censored", "censor_flg")
         >>> ep.pl.kmf(
         ...     [kmf], color=["r"], xlim=[0, 700], ylim=[0, 1], xlabel="Days", ylabel="Proportion Survived", show=True
         ... )
 
         .. image:: /_static/docstring_previews/kmf_plot_1.png
 
-        >>> T = adata[:, ["mort_day_censored"]].X
-        >>> E = adata[:, ["censor_flg"]].X
         >>> groups = adata[:, ["service_unit"]].X
-        >>> ix1 = groups == "FICU"
-        >>> ix2 = groups == "MICU"
-        >>> ix3 = groups == "SICU"
-        >>> kmf_1 = ep.tl.kmf(T[ix1], E[ix1], label="FICU")
-        >>> kmf_2 = ep.tl.kmf(T[ix2], E[ix2], label="MICU")
-        >>> kmf_3 = ep.tl.kmf(T[ix3], E[ix3], label="SICU")
+        >>> adata_ficu = adata[groups == "FICU"]
+        >>> adata_micu = adata[groups == "MICU"]
+        >>> adata_sicu = adata[groups == "SICU"]
+        >>> kmf_1 = ep.tl.kmf(adata_ficu, "mort_day_censored", "censor_flg", label="FICU")
+        >>> kmf_2 = ep.tl.kmf(adata_micu, "mort_day_censored", "censor_flg", label="MICU")
+        >>> kmf_3 = ep.tl.kmf(adata_sicu, "mort_day_censored", "censor_flg", label="SICU")
         >>> ep.pl.kmf([kmf_1, kmf_2, kmf_3], ci_show=[False,False,False], color=['k','r', 'g'],
         >>>           xlim=[0, 750], ylim=[0, 1], xlabel="Days", ylabel="Proportion Survived")
 

ehrapy/tools/_sa.py

@@ -116,15 +116,10 @@ def glm(
 
 
 def kmf(
-    durations: Iterable,
-    event_observed: Iterable | None = None,
-    timeline: Iterable = None,
-    entry: Iterable | None = None,
-    label: str | None = None,
-    alpha: float | None = None,
-    ci_labels: tuple[str, str] = None,
-    weights: Iterable | None = None,
-    censoring: Literal["right", "left"] = None,
+    adata: AnnData,
+    duration_col: str,
+    event_col: str | None = None,
+    **kwargs,
 ) -> KaplanMeierFitter:
     """Fit the Kaplan-Meier estimate for the survival function.
 
@@ -156,24 +151,38 @@ def kmf(
         >>> adata = ep.dt.mimic_2(encoded=False)
         >>> # Flip 'censor_fl' because 0 = death and 1 = censored
         >>> adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
-        >>> kmf = ep.tl.kmf(adata[:, ["mort_day_censored"]].X, adata[:, ["censor_flg"]].X)
+        >>> kmf = ep.tl.kmf(adata, "mort_day_censored", "censor_flg", label="Mortality")
     """
+
     kmf = KaplanMeierFitter()
-    if censoring == "None" or "right":
-        kmf.fit(
-            durations=durations,
-            event_observed=event_observed,
+    df = anndata_to_df(adata)
+    T = df[duration_col]
+    E = df[event_col]
+
+    # unpack kwargs
+    timeline = kwargs.get("timeline", None)
+    entry = kwargs.get("entry", None)
+    label = kwargs.get("label", None)
+    alpha = kwargs.get("alpha", None)
+    ci_labels = kwargs.get("ci_labels", None)
+    weights = kwargs.get("weights", None)
+    censoring = kwargs.get("censoring", "right")
+
+    if censoring == "left":
+        kmf.fit_left_censoring(
+            durations=T,
+            event_observed=E,
             timeline=timeline,
             entry=entry,
             label=label,
             alpha=alpha,
             ci_labels=ci_labels,
             weights=weights,
         )
-    elif censoring == "left":
-        kmf.fit_left_censoring(
-            durations=durations,
-            event_observed=event_observed,
+    else:
+        kmf.fit(
+            durations=T,
+            event_observed=E,
             timeline=timeline,
             entry=entry,
             label=label,
@@ -376,7 +385,9 @@ def log_logistic_aft(adata: AnnData, duration_col: str, event_col: str, entry_co
     )
 
 
-def _univariate_model(adata: AnnData, duration_col: str, event_col: str, model_class, accept_zero_duration=True):
+def _univariate_model(
+    adata: AnnData, duration_col: str, event_col: str, model_class, accept_zero_duration=True, **kwargs
+):
     """Convenience function for univariate models."""
     df = anndata_to_df(adata)
 
@@ -385,13 +396,32 @@ def _univariate_model(adata: AnnData, duration_col: str, event_col: str, model_c
     T = df[duration_col]
     E = df[event_col]
 
+    # unpack kwargs
+    timeline = kwargs.get("timeline", None)
+    entry = kwargs.get("entry", None)
+    label = kwargs.get("label", None)
+    alpha = kwargs.get("alpha", None)
+    ci_labels = kwargs.get("ci_labels", None)
+    weights = kwargs.get("weights", None)
+    fit_options = kwargs.get("fit_options", None)
+
     model = model_class()
-    model.fit(T, event_observed=E)
+    model.fit(
+        T,
+        event_observed=E,
+        timeline=timeline,
+        entry=entry,
+        label=label,
+        alpha=alpha,
+        ci_labels=ci_labels,
+        weights=weights,
+        fit_options=fit_options,
+    )
 
     return model
 
 
-def nelson_aalen(adata: AnnData, duration_col: str, event_col: str) -> NelsonAalenFitter:
+def nelson_aalen(adata: AnnData, duration_col: str, event_col: str, **kwargs) -> NelsonAalenFitter:
     """Employ the Nelson-Aalen estimator to estimate the cumulative hazard function from censored survival data
 
     The Nelson-Aalen estimator is a non-parametric method used in survival analysis to estimate the cumulative hazard function.
@@ -415,10 +445,10 @@ def nelson_aalen(adata: AnnData, duration_col: str, event_col: str) -> NelsonAal
         >>> adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
         >>> naf = ep.tl.nelson_aalen(adata, "mort_day_censored", "censor_flg")
     """
-    return _univariate_model(adata, duration_col, event_col, NelsonAalenFitter)
+    return _univariate_model(adata, duration_col, event_col, NelsonAalenFitter, True, **kwargs)
 
 
-def weibull(adata: AnnData, duration_col: str, event_col: str) -> WeibullFitter:
+def weibull(adata: AnnData, duration_col: str, event_col: str, **kwargs) -> WeibullFitter:
     """Employ the Weibull model in univariate survival analysis to understand event occurrence dynamics.
 
     In contrast to the non-parametric Nelson-Aalen estimator, the Weibull model employs a parametric approach with shape and scale parameters,
@@ -445,4 +475,4 @@ def weibull(adata: AnnData, duration_col: str, event_col: str) -> WeibullFitter:
         >>> adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
         >>> wf = ep.tl.weibull(adata, "mort_day_censored", "censor_flg")
     """
-    return _univariate_model(adata, duration_col, event_col, WeibullFitter, accept_zero_duration=False)
+    return _univariate_model(adata, duration_col, event_col, WeibullFitter, accept_zero_duration=False, **kwargs)

tests/tools/test_sa.py

@@ -97,7 +97,7 @@ def _sa_func_test(self, sa_function, sa_class, mimic_2_sa):
 
     def test_kmf(self, mimic_2_sa):
         adata, _, _ = mimic_2_sa
-        kmf = ep.tl.kmf(adata[:, ["mort_day_censored"]].X, adata[:, ["censor_flg"]].X)
+        kmf = ep.tl.kmf(adata, "mort_day_censored", "censor_flg")
         self._sa_function_assert(kmf, KaplanMeierFitter)
 
     def test_cox_ph(self, mimic_2_sa):