CamDavidsonPilon / lifelines

Survival analysis in Python
lifelines.readthedocs.org
MIT License
2.37k stars 560 forks source link

color settings in rmst_plot #1258

Open StephanHolgerD opened 3 years ago

StephanHolgerD commented 3 years ago

Hello, is it possible to specify colors for the rmst_plot ?

CamDavidsonPilon commented 3 years ago

Hi there, you should be able to pass in color=... into the function. Is that not working for you?

StephanHolgerD commented 3 years ago

I got this error

<ipython-input-7-37ca617b0f7d> in PlotFigure(df_noBev, df_Bev, title, pdfName)
     61 
     62 
---> 63     rmst_plot(kmf_nobev, t=time_limit, ax=ax3, color='blues')
     64     #rmst_plot(kmf_nobev, ax=ax3)
     65 

/software/miniconda3/envs/lifelines/lib/python3.9/site-packages/lifelines/plotting.py in rmst_plot(model, model2, t, ax, text_position, **plot_kwargs)
    170     rmst = restricted_mean_survival_time(model, t=t)
    171     c = ax._get_lines.get_next_color()
--> 172     model.plot_survival_function(ax=ax, color=c, ci_show=False, **plot_kwargs)
    173 
    174     if text_position is None:

TypeError: lifelines.fitters.kaplan_meier_fitter.KaplanMeierFitter.plot_survival_function() got multiple values for keyword argument 'color'

I fixed it with editing the source code a little but was not sure if I made a mistake in using the rms_plot function thats why I opened this issue

My fix in plotting.py, added color1 and color2, so it also works with 2 two models

def rmst_plot(model, model2=None, t=np.inf, ax=None, text_position=None,color1=None,color2=None, **plot_kwargs):
    """
    This functions plots the survival function of the model plus it's area-under-the-curve (AUC) up
    until the point ``t``. The AUC is known as the restricted mean survival time (RMST).

    To compare the difference between two models' survival curves, you can supply an
    additional model in ``model2``.

    Parameters
    -----------
    model: lifelines.UnivariateFitter
    model2: lifelines.UnivariateFitter, optional
        used to compute the delta RMST of two models
    t: float
        the upper bound of the expectation
    ax: axis
    text_position: tuple
        move the text position of the RMST.

    Examples
    ---------
    .. code:: python

        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')
        kmf_con = KaplanMeierFitter().fit(T[~ix], E[~ix], label='control')

        ax = plt.subplot(311)
        rmst_plot(kmf_exp, t=time_limit, ax=ax)

        ax = plt.subplot(312)
        rmst_plot(kmf_con, t=time_limit, ax=ax)

        ax = plt.subplot(313)
        rmst_plot(kmf_exp, model2=kmf_con, t=time_limit, ax=ax)

    """
    from lifelines.utils import restricted_mean_survival_time
    from matplotlib import pyplot as plt

    if ax is None:
        ax = plt.gca()

    rmst = restricted_mean_survival_time(model, t=t)
    if color1==None:
        c = ax._get_lines.get_next_color()
    else:
        c=color1
    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:
        if color2==None:
            c2 = ax._get_lines.get_next_color()
        else:
            c2=color2
CamDavidsonPilon commented 3 years ago

Hm, looks related to #1231. Can you try using the latest version of lifelines, v0.25.11? If you are using that, I think your solution is what should end up going into lifelines.

StephanHolgerD commented 3 years ago

I will check the version, If it is the most recent I would fork and include the fix

StephanHolgerD commented 3 years ago

I check the version, it is 0.25.11.

I forked, changed the plotting func and pushed, you can used it if you want.

CamDavidsonPilon commented 3 years ago

Great! We'll incorporate it into the next release