more rmst

docs/Examples.rst

@@ -132,45 +132,34 @@ This is a good metric for comparing two survival curves, as their difference rep
 
     from lifelines.utils import restricted_mean_survival_time
     from lifelines.datasets import load_waltons
+    from lifelines.plotting import rmst_plot
 
     df = load_waltons()
     ix = df['group'] == 'miR-137'
     T, E = df['T'], df['E']
 
+    time_limit = 50
+
     kmf_exp = KaplanMeierFitter().fit(T[ix], E[ix], label='exp')
+    rmst_exp = restricted_mean_survival_time(kmf_exp, t=time_limit)
+
     kmf_con = KaplanMeierFitter().fit(T[~ix], E[~ix], label='control')
+    rmst_con = restricted_mean_survival_time(kmf_con, t=time_limit)
+
 
-    limit = 50
 
-    rmst_exp = restricted_mean_survival_time(kmf_exp, t=limit)
-    rmst_con = restricted_mean_survival_time(kmf_con, t=limit)
 
     ax = plt.subplot(311)
-    kmf_exp.plot(ax=ax, c="#0C7BDC", ci_show=False)
-    sf_exp_at_limit = kmf_exp.predict(np.append(kmf_exp.timeline, limit)).sort_index().loc[:limit]
-    ax.fill_between(sf_exp_at_limit.index, sf_exp_at_limit.values, step='post', color="#0C7BDC", alpha=0.20)
-    ax.axvline(limit, ls='--', c='k')
-    ax.text(10, 0.3, "%.3f" % rmst_exp)
-    ax.set_xlim(0, 65)
+    rmst_plot(kmf_exp, t=time_limit, ax=ax)
 
 
     ax = plt.subplot(312)
-    kmf_con.plot(ax=ax, c="#FFC20A", ci_show=False)
-    sf_con_at_limit = kmf_con.predict(np.append(kmf_con.timeline, limit)).sort_index().loc[:limit]
-    ax.fill_between(sf_con_at_limit.index, sf_con_at_limit.values, step='post', color="#FFC20A", alpha=0.20)
-    ax.axvline(limit, ls='--', c='k')
-    ax.text(10, 0.4, "%.3f" % rmst_con)
-    ax.set_xlim(0, 65)
+    rmst_plot(kmf_con, t=time_limit, ax=ax)
 
 
     ax = plt.subplot(313)
-    kmf_con.plot(ax=ax, c="#FFC20A", ci_show=False)
-    kmf_exp.plot(ax=ax, c="#0C7BDC", ci_show=False)
-    timeline = np.unique(T.tolist() + [limit])
-    ax.axvline(limit, ls='--', c='k')
-    ax.fill_between(timeline[timeline<=limit], kmf_con.predict(timeline).loc[:limit], kmf_exp.predict(timeline).loc[:limit], step="post", color='k', alpha=0.10)
-    ax.text(34, 0.4, "%.3f" % (rmst_con - rmst_exp))
-    ax.set_xlim(0, 65)
+    rmst_plot(kmf_exp, model2=kmf_con, t=time_limit, ax=ax)
+
 
 
 .. image:: images/rmst_example.png

lifelines/plotting.py

@@ -78,7 +78,7 @@ def cdf_plot(model, timeline=None, ax=None, **plot_kwargs):
     return ax
 
 
-def rmst_plot(model, t=np.inf, model2=None, ax=None, text_position=None):
+def rmst_plot(model, model2=None, t=np.inf, ax=None, text_position=None, **plot_kwargs):
     """
 
     Parameters
@@ -100,15 +100,15 @@ def rmst_plot(model, t=np.inf, model2=None, ax=None, text_position=None):
 
     rmst = restricted_mean_survival_time(model, t=t)
     c = ax._get_lines.get_next_color()
-    model.plot_survival_function(ax=ax, color=c, ci_show=False)
+    model.plot_survival_function(ax=ax, color=c, ci_show=False, **plot_kwargs)
 
     if text_position is None:
         text_position = (np.percentile(model.timeline, 10), 0.15)
 
     if model2 is not None:
         c2 = ax._get_lines.get_next_color()
         rmst2 = restricted_mean_survival_time(model2, t=t)
-        model2.plot_survival_function(ax=ax, color=c2, ci_show=False)
+        model2.plot_survival_function(ax=ax, color=c2, ci_show=False, **plot_kwargs)
         timeline = np.unique(model.timeline.tolist() + model2.timeline.tolist() + [t])
         predict1 = model.predict(timeline).loc[:t]
         predict2 = model2.predict(timeline).loc[:t]