Update grid specification in kaplan_meier function and rerendered tes…

…t images
theislab · Jan 21, 2025 · db91335 · db91335
1 parent ce89e61
commit db91335
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 7 deletions.
diff --git a/ehrapy/plot/_survival_analysis.py b/ehrapy/plot/_survival_analysis.py
@@ -269,7 +269,7 @@ def kaplan_meier(
         color = [None] * len(kmfs)
 
     fig = plt.figure(constrained_layout=True, figsize=figsize)
-    spec = fig.add_gridspec(2, 1)
+    spec = fig.add_gridspec(2, 1) if display_survival_statistics else fig.add_gridspec(1, 1)
     ax = plt.subplot(spec[0, 0])
 
     for i, kmf in enumerate(kmfs):
@@ -301,7 +301,6 @@ def kaplan_meier(
     if title:
         ax.set_title(title)
 
-    # Prepare data for the table
     if display_survival_statistics:
         xticks = [x for x in ax.get_xticks() if x >= 0]
         xticks_space = xticks[1] - xticks[0]

diff --git a/tests/_scripts/kaplain_meier_create_expected_plots.ipynb b/tests/_scripts/kaplain_meier_create_expected_plots.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -27,7 +27,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -58,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -86,7 +86,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [

diff --git a/tests/plot/_images/kaplan_meier_expected.png b/tests/plot/_images/kaplan_meier_expected.png