Survival analysis is a powerful statistical method that helps us understand and predict the timing of events. While traditionally used in medical research, it has broad applications across industries – from customer churn prediction to equipment failure analysis. In this article, we’ll walk through a practical example using Python’s lifelines package to analyze recidivism (repeat offender) data.

Setting Up Your Environment

Before we begin, you’ll need to install the lifelines package. If you haven’t already installed it, you can do so using pip:

pip install lifelines

The lifelines package is a comprehensive Python library for survival analysis that provides easy-to-use implementations of various survival analysis models. It includes the Kaplan-Meier estimator, Cox Proportional Hazard models, and more.

For this tutorial, we’ll need just a few key libraries:

# Import necessary libraries
import lifelines
from lifelines.datasets import load_rossi
from lifelines import KaplanMeierFitter, CoxPHFitter
import matplotlib.pyplot as plt

# Load the dataset
rossi = load_rossi()

Understanding Our Data

We’ll be working with the ‘rossi’ dataset that comes built into lifelines. This dataset tracks whether released prison inmates reoffend (recidivate). The dataset contains several variables including:

• Whether the inmate received financial aid
• Age at time of release
• Previous convictions
• Work experience
• Marital status
• Time until re-arrest (if it occurred)

Basic Survival Curves

Let’s start with the most fundamental visualization in survival analysis – the Kaplan-Meier curve. We’ll first look at the overall survival curve for our population:

# Initialize the KaplanMeierFitter model
kmf = KaplanMeierFitter()

# Fit the model
kmf.fit(rossi['week'], 
        rossi['arrest'],
        label='Overall Population')

# Create the plot
plt.figure(figsize=(10, 6))
kmf.plot()

# Customize the plot
plt.title('Overall Survival Curve')
plt.xlabel('Weeks')
plt.ylabel('Survival Probability')
plt.grid(True)
plt.show()

Image By Author

Cumulative Risk Over Time

While the survival curve shows us the probability of “surviving” (not reoffending), we can also examine how risk accumulates over time using the Nelson-Aalen estimator. This gives us a different perspective on the same phenomenon:

# Initialize two Nelson-Aalen models for cumulative hazard
from lifelines import NelsonAalenFitter
naf_fin_aid = NelsonAalenFitter()
naf_no_aid = NelsonAalenFitter()

# Fit on different groups
naf_fin_aid.fit(rossi[rossi['fin'] == 1]['week'], 
                rossi[rossi['fin'] == 1]['arrest'],
                label='Received Financial Aid')

naf_no_aid.fit(rossi[rossi['fin'] == 0]['week'],
               rossi[rossi['fin'] == 0]['arrest'],
               label='No Financial Aid')

# Create the plot
plt.figure(figsize=(10, 6))

# Plot both curves
naf_fin_aid.plot()
naf_no_aid.plot()

plt.title('Cumulative Hazard of Recidivism by Financial Aid Status')
plt.xlabel('Weeks')
plt.ylabel('Cumulative Hazard')
plt.grid(True)
plt.show()

Image By Author

• The steeper the curve, the faster risk accumulates during that period
• Those without financial aid (orange line) show a steeper increase in cumulative risk
• The gap between the two groups widens over time, suggesting that financial aid’s protective effect may be long-lasting

The cumulative hazard plot complements our survival analysis by showing how risk builds up over time, rather than showing the probability of survival at each point in time.

Understanding Risk Factors

Finally, let’s examine which factors influence recidivism risk using a Cox Proportional Hazards model. This model helps us understand the impact of multiple factors simultaneously:

# Initialize and fit the Cox model
cph = CoxPHFitter()
cph.fit(rossi, duration_col='week', event_col='arrest')

# Get hazard ratios
hazard_ratios = np.exp(cph.params_)

# Create the plot
plt.figure(figsize=(10, 6))

# Create bar colors based on whether hazard ratio > 1
colors = ['red' if x > 1 else 'green' for x in hazard_ratios]

# Create horizontal bar plot
y_pos = np.arange(len(hazard_ratios))
plt.barh(y_pos, hazard_ratios, color=colors)

# Customize the plot
plt.axvline(x=1, color='black', linestyle='--', alpha=0.3)  # Add reference line at 1
plt.yticks(y_pos, hazard_ratios.index)
plt.xlabel('Hazard Ratio')
plt.title('Risk Factors in Recidivism')

# Add value labels on the bars
for i, v in enumerate(hazard_ratios):
    plt.text(v, i, f'{v:.2f}', verticalalignment='center')
plt.tight_layout()
plt.show()

Image By Author

Hazard Ratios:
prio: 1.10
paro: 0.92
mar: 0.65
wexp: 0.86
race: 1.37
age: 0.94
fin: 0.68

The hazard ratios plot reveals several important insights:

• Values greater than 1 (shown in red) indicate increased risk:
  • Race (1.37) shows the highest risk factor
  • Prior convictions (1.10) are also associated with increased risk
• Values less than 1 (shown in green) indicate protective factors:
  • Marriage (0.65) appears to be the strongest protective factor
  • Financial aid (0.68) also significantly reduces risk
  • Work experience (0.86) and age (0.94) also show modest protective effects

The vertical dashed line at 1.0 serves as a reference point – factors to the right of this line increase risk, while factors to the left decrease risk.

Conclusion: Putting It All Together

In this article, we’ve explored three fundamental aspects of survival analysis using Python’s lifelines package:

1. We started with basic survival curves that showed us how the probability of “survival” (not reoffending) changes over time
2. We then examined cumulative hazard curves that helped us understand how risk accumulates differently between groups
3. Finally, we used the Cox Proportional Hazards model to quantify the impact of various risk factors

These analyses suggest several potential interventions for reducing recidivism:

• Expand financial aid programs, which show a significant protective effect
• Provide additional support for those with prior convictions
• Develop employment programs, as work experience appears protective
• Consider age-specific interventions for younger offenders

This example demonstrates how survival analysis can provide actionable insights from time-to-event data. While we used recidivism data, the same techniques can be applied to many other scenarios:

• Customer churn analysis
• Equipment failure prediction
• Employee retention studies
• Medical treatment outcomes

The lifelines package makes it straightforward to implement these analyses in Python, providing both the statistical tools and visualization capabilities needed for comprehensive survival analysis.