removing axes from errorbar_kwargs if provided (#708)

* removing axes from errorbar_kwargs if provided * removing pyplot calls if axes is present
CamDavidsonPilon · Apr 13, 2019 · 225962e · 225962e
1 parent 059b49b
commit 225962e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/lifelines/fitters/cox_time_varying_fitter.py b/lifelines/fitters/cox_time_varying_fitter.py
@@ -723,7 +723,7 @@ def plot(self, columns=None, **errorbar_kwargs):
         """
         from matplotlib import pyplot as plt
 
-        ax = errorbar_kwargs.get("ax", None) or plt.figure().add_subplot(111)
+        ax = errorbar_kwargs.pop("ax", None) or plt.figure().add_subplot(111)
 
         errorbar_kwargs.setdefault("c", "k")
         errorbar_kwargs.setdefault("fmt", "s")
@@ -750,8 +750,9 @@ def plot(self, columns=None, **errorbar_kwargs):
 
         tick_labels = [columns[i] for i in order]
 
-        plt.yticks(yaxis_locations, tick_labels)
-        plt.xlabel("log(HR) (%g%% CI)" % ((1 - self.alpha) * 100))
+        ax.set_yticks(yaxis_locations)
+        ax.set_yticklabels(tick_labels)
+        ax.set_xlabel("log(HR) (%g%% CI)" % ((1 - self.alpha) * 100))
 
         return ax