Align ticks in at-risk count table #747

christopherahern · 2019-06-18T19:31:06Z

Currently, add_at_risk_counts generates a set of xticks that get added to a separate axes object below the original axes. In some cases the ticks can be mis-aligned due to differences in how the xlim and xticks objects from the initial axes object being passed in are used to generate the new axes object.

The documentation has a good example of this, which can be reproduced with the code below.

from lifelines.datasets import load_waltons
from lifelines import KaplanMeierFitter
from lifelines.plotting import add_at_risk_counts
import matplotlib.pyplot as plt

waltons = load_waltons()
T = waltons['T']
E = waltons['E']

ix = waltons['group'] == 'control'
kmf_control = KaplanMeierFitter()
kmf_control.fit(waltons.loc[ix]['T'], waltons.loc[ix]['E'], label='control')
kmf_exp = KaplanMeierFitter()
kmf_exp.fit(waltons.loc[~ix]['T'], waltons.loc[~ix]['E'], label='exp')

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)

plt.show()

There is one extra tick in the at-risk counts than in the plot. This can be fixed by re-ordering how the xlim and xticks objects are created.

This MR does three things:

Re-orders the creation of the axes attributes so that there are the same number of plot ticks and counts, and they line up
Adds padding to the stratum label on the first tick so that labels are all right-aligned to the same point
Allows any kwargs to be passed in to add_at_risk_counts that can be passed into the set_xticklabels method (e.g. fontsize)

Before

After

Update lifelines fork

…labels, explicitly pop kwargs to allow tick labels configuration

…ython 3.5

CamDavidsonPilon · 2019-06-18T20:20:03Z

This is pretty darn exciting! I'll review and test it locally now

CamDavidsonPilon · 2019-06-18T20:48:56Z

Great work, I appreciate the detail PR post as well!

In diagnosing this, I've also turned up a few other bugs as well, so thanks for turning my attention to this!

CamDavidsonPilon · 2019-06-18T20:51:40Z

Actually, I did find one bug. If you change the KaplanMeierFitter to WeibullFitter, the bottom labels no longer show up. Example:

from lifelines.datasets import load_waltons
from lifelines import KaplanMeierFitter, WeibullFitter
from lifelines.plotting import add_at_risk_counts
import matplotlib.pyplot as plt

waltons = load_waltons()
T = waltons['T']
E = waltons['E']

ix = waltons['group'] == 'control'
kmf_control = WeibullFitter()
kmf_control.fit(waltons.loc[ix]['T'], waltons.loc[ix]['E'], label='control')
kmf_exp = WeibullFitter()
kmf_exp.fit(waltons.loc[~ix]['T'], waltons.loc[~ix]['E'], label='exp')

ax = plt.subplot(111)
ax = kmf_control.plot_survival_function(ax=ax)
ax = kmf_exp.plot_survival_function(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)

plt.show()

master

branch

CamDavidsonPilon · 2019-06-18T21:01:26Z

From what I can tell, it has to do with the ax.get_xlim() being different between the two.

Edit: the problem is that we write the labels on the first tick, but this tick location is assigned to 0, and if the lower-bound of xlim is greater than 0, that tick is never shown. A solution is to skip ticks that will never be shown, and then on the first shown tick ( greater than lower-bound), write the labels. Ex code that works:

    # Add population size at times
    ticklabels = []
    has_written_labels = False

    for tick in ax2.get_xticks():
        lbl = ""

        if tick < ax2.get_xlim()[0]:
            ticklabels.append(lbl)
            continue

        # Get counts at tick
        counts = [f.durations[f.durations >= tick].shape[0] for f in fitters]
        # Create tick label
        for l, c in zip(labels, counts):
            # First shown tick is prepended with the label
            if not has_written_labels and l is not None:
                # Get length of largest count
                max_length = len(str(max(counts)))
                if is_latex_enabled():
                    s = "\n{}\\quad".format(l) + "{{:>{}d}}".format(max_length)
                else:
                    s = "\n{}   ".format(l) + "{{:>{}d}}".format(max_length)

            else:
                s = "\n{}"
            lbl += s.format(c)

        has_written_labels = True
        ticklabels.append(lbl.strip())

christopherahern · 2019-06-19T03:01:36Z

Hey, thanks for taking a look so quickly!

The Weibull fitter example is really useful. I think a good way to cover both cases is to take the xticks from the plot axis and filter them down to those that fall within the xlim of the plot axis. I think that will work with any subset of times plotted (e.g. automatically or via a loc that gets passed in). I'll test that out and make sure it covers Kaplan-Meier and Weibull.

Any idea why the xlim for the Weibull fitter starts at 6? The same limit is generated by plot and plot_survival_function, so there's something in _plot_estimate that's generating that. If the Weibull plot should start at zero, then that would be where it's happening.

CamDavidsonPilon · 2019-06-19T12:14:20Z

Any idea why the xlim for the Weibull fitter starts at 6?

The xlim automatically set by mpl after we provide mpl the data to plot. In all cases, the x-variable is the timeline attribute (which can be user provided in the call to fit). In the KaplanMeier case, the timeline attribute starts at 0, and for parametric models, the timeline starts at the minimum observation.

CamDavidsonPilon · 2019-06-19T12:14:52Z

I think a good way to cover both cases is to take the xticks from the plot axis and filter them down to those that fall within the xlim of the plot axis. I think that will work with any subset of times plotted (e.g. automatically or via a loc that gets passed in).

👍 seems like a good solution

…plot axes object

christopherahern · 2019-06-19T14:33:35Z

Ok, I changed how the xticks for the at-risk counts are being generated. They can now be passed in as a kwarg or they are generated as the visible ticks from the plot axes object.

Here is boilerplate code for the plots below.

from lifelines.datasets import load_waltons
from lifelines import KaplanMeierFitter, WeibullFitter
from lifelines.plotting import add_at_risk_counts
import matplotlib.pyplot as plt

waltons = load_waltons()
T = waltons['T']
E = waltons['E']

ix = waltons['group'] == 'control'

kmf_control = KaplanMeierFitter()
kmf_control.fit(waltons.loc[ix]['T'], waltons.loc[ix]['E'], label='control')
kmf_exp = KaplanMeierFitter()
kmf_exp.fit(waltons.loc[~ix]['T'], waltons.loc[~ix]['E'], label='exp')

wbl_control = WeibullFitter()
wbl_control.fit(waltons.loc[ix]['T'], waltons.loc[ix]['E'], label='control')
wbl_exp = WeibullFitter()
wbl_exp.fit(waltons.loc[~ix]['T'], waltons.loc[~ix]['E'], label='exp')

Default

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)
plt.show()

ax = plt.subplot(111)
ax = wbl_control.plot_survival_function(ax=ax)
ax = wbl_exp.plot_survival_function(ax=ax)
add_at_risk_counts(wbl_exp, wbl_control, ax=ax)
plt.show()

Setting `xlim`

Works when setting xlim outside of the fitter plot methods.

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
ax.set_xlim(0, 30)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)
plt.show()

ax = plt.subplot(111)
ax = wbl_control.plot_survival_function(ax=ax)
ax = wbl_exp.plot_survival_function(ax=ax)
ax.set_xlim(0, 80)
add_at_risk_counts(wbl_exp, wbl_control, ax=ax)
plt.show()

Works when plotting subset of time

Not entirely sure what to expect for the use of loc here for the plots themselves. The at-risk counts do appear to be showing up only at the visible ticks though.

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax, loc=slice(10, 60))
ax = kmf_exp.plot(ax=ax, loc=slice(10, 60))
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)
plt.show()

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax, loc=slice(30, 40))
ax = kmf_exp.plot(ax=ax, loc=slice(30, 40))
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)
plt.show()

Supply `xticks` as kwarg

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax, xticks=[15, 23, 55])
plt.show()

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax, xticks=[17, 30, 60])
plt.show()

Modify at-risk counts text

ax = plt.subplot(111)
ax = kmf_control.plot(ax=ax)
ax = kmf_exp.plot(ax=ax)
add_at_risk_counts(kmf_exp, kmf_control, ax=ax, fontsize=8, color='b', fontstyle='italic')
plt.show()

Let me know if that looks good.

CamDavidsonPilon · 2019-06-19T16:11:31Z

Let me know if that looks good.

looks great! Thanks for this attention to detail!

christopherahern · 2019-06-19T19:24:49Z

Great! Thanks for your help and great work!

christopherahern and others added 8 commits June 18, 2019 10:26

Merge pull request #1 from CamDavidsonPilon/master

92a6d7b

Update lifelines fork

Pad counts based on largest number for the first tick to align group …

e36dda2

…labels, explicitly pop kwargs to allow tick labels configuration

Set xlimit after setting xticks to maintain alignment for at-risk counts

11b5363

Formatting with black

6f88203

Move calculation of longest count into first tick check

f189d93

Use format for compatibility with python 3.5

91e00ff

Remove black inserted comma after kwarg in plot_lifetimes

27b1a62

Remove another trailing black-inserted comma for compatibility with p…

9fcd60f

…ython 3.5

CamDavidsonPilon approved these changes Jun 18, 2019

View reviewed changes

Chris Ahern added 2 commits June 19, 2019 09:09

Set at-risk xticks based on ticks falling within plot xlim

93cd3a6

Allow xticks as kwarg to at-risk counts or set to visible ticks from …

9d7c134

…plot axes object

CamDavidsonPilon merged commit 52fba6e into CamDavidsonPilon:master Jun 19, 2019

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Align ticks in at-risk count table #747

Align ticks in at-risk count table #747

christopherahern commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019 •

edited

christopherahern commented Jun 19, 2019

CamDavidsonPilon commented Jun 19, 2019

CamDavidsonPilon commented Jun 19, 2019

christopherahern commented Jun 19, 2019

CamDavidsonPilon commented Jun 19, 2019

christopherahern commented Jun 19, 2019

Align ticks in at-risk count table #747

Align ticks in at-risk count table #747

Conversation

christopherahern commented Jun 18, 2019

Before

After

CamDavidsonPilon commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019

CamDavidsonPilon commented Jun 18, 2019

master

branch

CamDavidsonPilon commented Jun 18, 2019 • edited

christopherahern commented Jun 19, 2019

CamDavidsonPilon commented Jun 19, 2019

CamDavidsonPilon commented Jun 19, 2019

christopherahern commented Jun 19, 2019

Default

Setting xlim

Works when plotting subset of time

Supply xticks as kwarg

Modify at-risk counts text

CamDavidsonPilon commented Jun 19, 2019

christopherahern commented Jun 19, 2019

CamDavidsonPilon commented Jun 18, 2019 •

edited

Setting `xlim`

Supply `xticks` as kwarg