durations_to_survival.py

"""
How to convert a duration table to a survival table, then compute the Kaplan-Meier and
Nelson-Aalen estimates. Accompanies the article at
https://crosstab.io/articles/durations-to-survivals.
"""

import altair as alt
import numpy as np
import pandas as pd


## Load data
durations = pd.read_parquet("data/retailrocket_durations.parquet")
print(durations.head())


## Round days up for a meaningful grouping
durations["duration_days"] = np.ceil(durations["duration_days"]).astype(int)

grp = durations.groupby("duration_days")
survival = pd.DataFrame(
    {"num_obs": grp.size(), "events": grp["endpoint_observed"].sum()}
)
print(survival.head())


# Get the total number of observations that had happened *before* each row. The
# complement of that is the number still at-risk of the event *in* each row.
num_subjects = len(durations)
prior_count = survival["num_obs"].cumsum().shift(1, fill_value=0)
survival.insert(0, "at_risk", num_subjects - prior_count)
print(survival.head())


# Keep only the rows with at least one event, and insert an initial row of 0's.
survival = survival.loc[survival["events"] > 0]


# The number of censored includes units censored at any duration between an event
# duration and the next event duration. This has to be backed out of the number at-risk:
# if the number at-risk falls from 15 to 11, but only one event was observed, then the
# number of censored in that interval must be 3.
survival["censored"] = (
    survival["at_risk"]
    - survival["at_risk"].shift(-1, fill_value=0)
    - survival["events"]
)

print(survival)


## Kaplan-Meier survival function estimate
inverse_hazard = 1 - survival["events"] / survival["at_risk"]
survival["survival_proba"] = inverse_hazard.cumprod()
print(survival)

survival["conversion_pct"] = 100 * (1 - survival["survival_proba"])


fig = (
    alt.Chart(survival.reset_index())
    .mark_line(interpolate="step-after")
    .encode(
        x=alt.X("duration_days", axis=alt.Axis(title="Duration (days)")),
        y=alt.Y(
            "survival_proba",
            axis=alt.Axis(title="Survival probability"),
            scale=alt.Scale(zero=False),
        ),
    )
)

fig.save("rocketretail_survival.svg")

fig = (
    alt.Chart(survival.reset_index())
    .mark_line(interpolate="step-after")
    .encode(
        x=alt.X("duration_days", axis=alt.Axis(title="Duration (days)")),
        y=alt.Y("conversion_pct", axis=alt.Axis(title="Conversion rate (%)")),
    )
)

fig.save("rocketretail_conversion.svg")


## Aalen-Nelson cumulative hazard estimate
survival["cumulative_hazard"] = (survival["events"] / survival["at_risk"]).cumsum()

fig = (
    alt.Chart(survival.reset_index())
    .mark_line(interpolate="step-after")
    .encode(
        x=alt.X("duration_days", axis=alt.Axis(title="Duration (days)")),
        y=alt.Y("cumulative_hazard", axis=alt.Axis(title="Cumulative hazard")),
    )
)

fig.save("rocketretail_hazard.svg")