From d566c80be6a19e41b4025dfdb443fa8c42950278 Mon Sep 17 00:00:00 2001
From: Carl Buchholz <32228189+aGuyLearning@users.noreply.github.com>
Date: Thu, 16 Jan 2025 13:05:54 +0100
Subject: [PATCH] Fix return values in kaplan_meier and update documentation
 for cox_ph_forestplot

---
 ehrapy/plot/_survival_analysis.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/ehrapy/plot/_survival_analysis.py b/ehrapy/plot/_survival_analysis.py
index b02caeb6..948523fc 100644
--- a/ehrapy/plot/_survival_analysis.py
+++ b/ehrapy/plot/_survival_analysis.py
@@ -336,7 +336,7 @@ def kaplan_meier(
         ax_table.spines["left"].set_visible(False)
 
     if not show:
-        return ax
+        return fig, ax
 
     else:
         return None
@@ -361,13 +361,13 @@ def cox_ph_forestplot(
     """Generates a forest plot to visualize the coefficients and confidence intervals of a Cox Proportional Hazards model.
 
     The `adata` object must first be populated using the :func:`~ehrapy.tools.cox_ph` function. This function stores the summary table of the `CoxPHFitter` in the `.uns` attribute of `adata`.
-    The summary table is created when the model is fitted using the :func:`ehrapy.tl.cox_ph` function.
+    The summary table is created when the model is fitted using the :func:`~ehrapy.tl.cox_ph` function.
     For more information on the `CoxPHFitter`, see the `Lifelines documentation <https://lifelines.readthedocs.io/en/latest/fitters/regression/CoxPHFitter.html>`_.
 
     Inspired by `zepid.graphics.EffectMeasurePlot <https://readthedocs.org>`_ (zEpid Package, https://pypi.org/project/zepid/).
 
     Args:
-        adata: :class:`~anndata.AnnData` object containing the summary table from the CoxPHFitter. This is stored in the `.uns` attribute, after fitting the model using :func:`~ehrapy.tl.cox_ph`.
+        adata: :class:`~anndata.AnnData` object containing the summary table from the CoxPHFitter. This is stored in the `.uns` attribute, after fitting the model using :func:`~ehrapy.tools.cox_ph`.
         uns_key: Key in `.uns` where :func:`~ehrapy.tools.cox_ph` function stored the summary table. See argument `uns_key` in :func:`~ehrapy.tools.cox_ph`.
         labels: List of labels for each coefficient, default uses the index of the summary ta
         fig_size: Width, height in inches.