def test_plot_km(): if pdf_output: from matplotlib.backends.backend_pdf import PdfPages pdf = PdfPages("test_survfunc.pdf") else: pdf = None sr1 = SurvfuncRight(ti1, st1) sr2 = SurvfuncRight(ti2, st2) fig = plot_survfunc(sr1) close_or_save(pdf, fig) fig = plot_survfunc(sr2) close_or_save(pdf, fig) fig = plot_survfunc([sr1, sr2]) close_or_save(pdf, fig) # Plot the SAS BMT data gb = bmt.groupby("Group") sv = [] for g in gb: s0 = SurvfuncRight(g[1]["T"], g[1]["Status"], title=g[0]) sv.append(s0) fig = plot_survfunc(sv) ax = fig.get_axes()[0] ax.set_position([0.1, 0.1, 0.64, 0.8]) ha, lb = ax.get_legend_handles_labels() fig.legend([ha[k] for k in (0, 2, 4)], [lb[k] for k in (0, 2, 4)], 'center right') close_or_save(pdf, fig) # Simultaneous CB for BMT data ii = bmt.Group == "ALL" sf = SurvfuncRight(bmt.loc[ii, "T"], bmt.loc[ii, "Status"]) fig = sf.plot() ax = fig.get_axes()[0] ax.set_position([0.1, 0.1, 0.64, 0.8]) ha, lb = ax.get_legend_handles_labels() lcb, ucb = sf.simultaneous_cb(transform="log") plt.fill_between(sf.surv_times, lcb, ucb, color="lightgrey") lcb, ucb = sf.simultaneous_cb(transform="arcsin") plt.plot(sf.surv_times, lcb, color="darkgrey") plt.plot(sf.surv_times, ucb, color="darkgrey") plt.plot(sf.surv_times, sf.surv_prob - 2*sf.surv_prob_se, color="red") plt.plot(sf.surv_times, sf.surv_prob + 2*sf.surv_prob_se, color="red") plt.xlim(100, 600) close_or_save(pdf, fig) if pdf_output: pdf.close()
def test_simultaneous_cb(): # The exact numbers here are regression tests, but they are close # to page 103 of Klein and Moeschberger. df = bmt.loc[bmt["Group"] == "ALL", :] sf = SurvfuncRight(df["T"], df["Status"]) lcb1, ucb1 = sf.simultaneous_cb(transform="log") lcb2, ucb2 = sf.simultaneous_cb(transform="arcsin") ti = sf.surv_times.tolist() ix = [ti.index(x) for x in (110, 122, 129, 172)] assert_allclose(lcb1[ix], np.r_[0.43590582, 0.42115592, 0.4035897, 0.38785927]) assert_allclose(ucb1[ix], np.r_[0.93491636, 0.89776803, 0.87922239, 0.85894181]) assert_allclose(lcb2[ix], np.r_[0.52115708, 0.48079378, 0.45595321, 0.43341115]) assert_allclose(ucb2[ix], np.r_[0.96465636, 0.92745068, 0.90885428, 0.88796708])