From 541e505c9e7c59624f755a3463ad881897ffe611 Mon Sep 17 00:00:00 2001
From: Carl Buchholz <32228189+aGuyLearning@users.noreply.github.com>
Date: Fri, 13 Dec 2024 09:36:27 +0100
Subject: [PATCH] renamed function and updated documentation to mention, that
 it is a lifelines object

---
 ehrapy/plot/__init__.py           |  2 +-
 ehrapy/plot/_survival_analysis.py | 19 ++++++++++++-------
 tests/plot/test_catplot.py        |  2 +-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/ehrapy/plot/__init__.py b/ehrapy/plot/__init__.py
index de0b75f7..4a57b84e 100644
--- a/ehrapy/plot/__init__.py
+++ b/ehrapy/plot/__init__.py
@@ -2,6 +2,6 @@
 from ehrapy.plot._colormaps import *  # noqa: F403
 from ehrapy.plot._missingno_pl_api import *  # noqa: F403
 from ehrapy.plot._scanpy_pl_api import *  # noqa: F403
-from ehrapy.plot._survival_analysis import coxph_forestplot, kaplan_meier, ols
+from ehrapy.plot._survival_analysis import cox_ph_forestplot, kaplan_meier, ols
 from ehrapy.plot.causal_inference._dowhy import causal_effect
 from ehrapy.plot.feature_ranking._feature_importances import rank_features_supervised
diff --git a/ehrapy/plot/_survival_analysis.py b/ehrapy/plot/_survival_analysis.py
index 9302d119..8a05ef6c 100644
--- a/ehrapy/plot/_survival_analysis.py
+++ b/ehrapy/plot/_survival_analysis.py
@@ -301,8 +301,8 @@ def kaplan_meier(
         return None
 
 
-def coxph_forestplot(
-    coxph: CoxPHFitter,
+def cox_ph_forestplot(
+    cox_ph: CoxPHFitter,
     labels: list[str] | None = None,
     fig_size: tuple = (10, 10),
     t_adjuster: float = 0.1,
@@ -313,11 +313,12 @@ def coxph_forestplot(
     text_size: int = 12,
     color: str = "k",
 ):
-    """Plots a forest plot of the Cox Proportional Hazard model.
+    """Generates a forest plot to visualize the coefficients and confidence intervals of a Cox Proportional Hazards model. 
+    The method requires a fitted CoxPHFitter object from the lifelines library.
     Inspired by `zepid.graphics.EffectMeasurePlot <https://readthedocs.org>`_ (zEpid Package, https://pypi.org/project/zepid/).
 
     Args:
-        coxph: Fitted CoxPHFitter object.
+        coxph: Fitted CoxPHFitter object from the lifelines library.
         labels: List of labels for each coefficient, default uses the index of the coxph.summary.
         fig_size: Width, height in inches.
         t_adjuster: Adjust the table to the right.
@@ -332,13 +333,17 @@ def coxph_forestplot(
         >>> import ehrapy as ep
         >>> adata = ep.dt.mimic_2(encoded=False)
         >>> adata_subset = adata[:, ["mort_day_censored", "censor_flg", "gender_num", "afib_flg", "day_icu_intime_num"]]
-        >>> coxph = ep.tl.coxph(adata_subset, event_col="censor_flg", duration_col="mort_day_censored")
-        >>> ep.pl.coxph_forestplot(coxph)
+        >>> coxph = ep.tl.cox_ph(adata_subset, event_col="censor_flg", duration_col="mort_day_censored")
+        >>> ep.pl.cox_ph_forestplot(coxph)
 
         .. image:: /_static/docstring_previews/coxph_forestplot.png
 
     """
-    data = coxph.summary
+    # check that the coxph object is fitted
+    if not cox_ph._fitted:
+        raise ValueError("The CoxPHFitter object must be fitted")
+    
+    data = cox_ph.summary
     auc_col = "coef"
 
     if labels is None:
diff --git a/tests/plot/test_catplot.py b/tests/plot/test_catplot.py
index 986de59f..90b2b5f3 100644
--- a/tests/plot/test_catplot.py
+++ b/tests/plot/test_catplot.py
@@ -19,7 +19,7 @@ def test_catplot_vanilla(adata_mini, check_same_image):
 def test_coxph_forestplot(mimic_2, check_same_image):
     adata_subset = mimic_2[:, ["mort_day_censored", "censor_flg", "gender_num", "afib_flg", "day_icu_intime_num"]]
     coxph = ep.tl.cox_ph(adata_subset, duration_col="mort_day_censored", event_col="censor_flg")
-    fig, ax = ep.pl.coxph_forestplot(coxph, fig_size=(12, 3), t_adjuster=0.15, marker="o", size=2, text_size=14)
+    fig, ax = ep.pl.cox_ph_forestplot(coxph, fig_size=(12, 3), t_adjuster=0.15, marker="o", size=2, text_size=14)
 
     check_same_image(
         fig=fig,