From ae1a4cb5b883f4c49cac4b8f4a9768ba74a0cf7f Mon Sep 17 00:00:00 2001 From: RunasSudo Date: Sat, 22 Apr 2023 00:43:33 +1000 Subject: [PATCH] Do not attempt to draw legend on survival curves if only one curve --- yli/survival.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/yli/survival.py b/yli/survival.py index 9b98827..9f68e54 100644 --- a/yli/survival.py +++ b/yli/survival.py @@ -80,7 +80,9 @@ def kaplanmeier(df, time, status, by=None, *, ci=True, transform_x=None, transfo ax.set_ylabel('Survival probability ({:.0%} CI)'.format(1-config.alpha) if ci else 'Survival probability') ax.set_xlim(left=0) ax.set_ylim(0, 1) - ax.legend() + + if by is not None: + ax.legend() return fig, ax @@ -174,7 +176,9 @@ def turnbull(df, time_left, time_right, by=None, *, step_loc=0.5, transform_x=No ax.set_ylabel('Survival probability') ax.set_xlim(left=0) ax.set_ylim(0, 1) - ax.legend() + + if by is not None: + ax.legend() return fig, ax