Survival Analysis Guide: Predicting "When" Instead of "If"

Imagine you are a doctor running a clinical trial. You know if a patient recovered, but that’s only half the story. Did they recover in 3 days or 3 months?

Or imagine you’re a data scientist at a subscription company. Knowing if a customer will churn is useful, but knowing when they will churn is a goldmine. It tells you exactly when to intervene.

This is the domain of Survival Analysis.

While standard classification answers "Will the event happen?" and regression answers "How much?", Survival Analysis answers "How long until the event happens?"

In this guide, we will master the unique mathematics of time-to-event data, understand why standard regression fails miserably here, and build professional-grade models using Python.

Why Standard Regression Fails (The Censoring Problem)

You might ask: "Why can't I just use Linear Regression? Time is just a number, right?"

If you try to predict days_to_event using Linear Regression, you will fail because of Censoring.

In any time-based study, there are always subjects who haven't experienced the event by the time the study ends.

• Patient A recovered on Day 12. (We know the exact time).
• Patient B was still sick when the study ended on Day 30.

For Patient B, we don't know the true time-to-recovery. It could be Day 31 or Day 100. We only know it is > 30. This is called Right Censoring.

The Two Traps of Ignoring Censoring

1. The Drop Trap: If you delete censored patients (Patient B), you throw away the "survivors," biasing your model to predict artificially short times.
2. The Plug-In Trap: If you treat Patient B's time as "30", you are lying to your model. You are saying they recovered on Day 30, when they actually didn't.

Survival Analysis is explicitly designed to handle this "partial information." It learns from Patient B that "the event takes at least 30 days," without assuming it happened on day 30.

The Two Pillars: Survival and Hazard Functions

To understand survival analysis, you need to understand two fundamental functions that describe time.

1. The Survival Function $S(t)$

The Survival Function tells you the probability that the event has not yet happened by time $t$.

$S(t) = P(T > t)$

• $T$: The actual time of the event.
• $t$: A specific point in time.

At the start ($t=0$), $S(t) = 1$ (100% survival). As time goes on, $S(t)$ drops.

2. The Hazard Function $h(t)$

The Hazard Function is trickier. It is not a probability. It is an instantaneous rate.

$h(t) = \lim_{\Delta t \to 0} \frac{P(t \le T < t + \Delta t | T \ge t)}{\Delta t}$

Hands-On: The Kaplan-Meier Estimator

The Kaplan-Meier (KM) estimator is the most popular way to estimate the Survival Function $S(t)$ without making strong assumptions about the data's shape. It calculates probabilities step-by-step every time an event occurs.

Let's use our clinical trial dataset. In this context:

• Event: Recovery (1 = Recovered, 0 = Censored/Still Sick).
• Time: Days until recovery.
• Goal: Compare how fast patients recover on different drugs.

Step 1: Loading the Data

We'll use the lifelines library, the industry standard for survival analysis in Python.

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter

# Load the dataset
url = "https://letsdatascience.com/datasets/playground/lds_stats_probability.csv"
df = pd.read_csv(url)

# Select relevant columns for survival analysis
# days_to_event: Time (T)
# event_occurred: Event (E) - 1 if recovered, 0 if censored
survival_df = df[['treatment_group', 'days_to_event', 'event_occurred']].copy()

print(survival_df.head())

Step 2: Plotting the Survival Curve

We will plot the curve for the Placebo group to see their recovery trajectory.

kmf = KaplanMeierFitter()

# Filter for Placebo group
placebo_df = survival_df[survival_df['treatment_group'] == 'Placebo']

# Fit the model: fit(Time, Event)
kmf.fit(durations=placebo_df['days_to_event'], 
        event_observed=placebo_df['event_occurred'], 
        label='Placebo')

# Plot the survival function
plt.figure(figsize=(10, 6))
kmf.plot_survival_function(linewidth=2)
plt.title('Survival Function for Placebo Group (Time to Recovery)')
plt.xlabel('Days')
plt.ylabel('Probability of Still Being Sick S(t)')
plt.grid(alpha=0.3)
plt.show()

# Probability of still being sick at Day 30?
p_30 = kmf.predict(30)
print(f"Probability of being sick after 30 days (Placebo): {p_30:.2%}")

Expected Output: The plot shows a staircase decreasing over time.

Probability of being sick after 30 days (Placebo): 43.55%

Step 3: Comparing Multiple Groups

Now, let's compare Placebo vs. Drug_B. If Drug_B works, patients should recover faster, meaning the "Still Sick" curve ($S(t)$) should drop more steeply.

groups = ['Placebo', 'Drug_B']
plt.figure(figsize=(10, 6))

for group in groups:
    subset = survival_df[survival_df['treatment_group'] == group]
    
    kmf = KaplanMeierFitter()
    kmf.fit(subset['days_to_event'], subset['event_occurred'], label=group)
    kmf.plot_survival_function(linewidth=2)

plt.title('Comparison: Placebo vs Drug B Recovery Times')
plt.xlabel('Days')
plt.ylabel('Probability of Still Being Sick')
plt.grid(alpha=0.3)
plt.show()

Visually, you should see the Drug_B curve dropping significantly faster than the Placebo curve. This confirms that Drug_B patients are "exiting" the sick state earlier.

The Log-Rank Test: Is the Difference Real?

The curves look different, but is the difference statistically significant? We use the Log-Rank Test.

The Log-Rank test compares the observed number of events at each time point against what we would expect if the two groups were identical (the Null Hypothesis).

from lifelines.statistics import logrank_test

# Extract data for the two groups
group_placebo = survival_df[survival_df['treatment_group'] == 'Placebo']
group_drug_b = survival_df[survival_df['treatment_group'] == 'Drug_B']

# Run the test
results = logrank_test(
    durations_A=group_placebo['days_to_event'],
    event_observed_A=group_placebo['event_occurred'],
    durations_B=group_drug_b['days_to_event'],
    event_observed_B=group_drug_b['event_occurred']
)

print(f"Log-Rank Test p-value: {results.p_value:.5e}")

Expected Output:

Log-Rank Test p-value: 1.56e-04

Interpretation: The p-value is very small ($< 0.05$). We reject the null hypothesis. There is a statistically significant difference in recovery times between Placebo and Drug B.

Cox Proportional Hazards: Multivariate Modeling

Kaplan-Meier is great for visualizing one variable (Treatment). But what if Age affects recovery? What if Severity matters?

We need a model that can handle multiple variables (covariates) simultaneously. Enter the Cox Proportional Hazards Model.

The Math

$h(t|x) = h_0(t) \cdot \exp(\beta_1 x_1 + \beta_2 x_2 + ... + \beta_k x_k)$

• $h_0(t)$: The Baseline Hazard. The risk over time for an "average" individual (where all $x=0$).
• $\exp(\beta x)$: The Partial Hazard. A multiplier that scales the risk up or down based on covariates.

Hands-On: Implementing Cox Regression

We will model recovery time using treatment_group, age, and gender.

First, we need to prepare the data. The Cox model requires numerical inputs, so we must encode categorical variables.

from lifelines import CoxPHFitter

# 1. Select relevant columns
cols = ['days_to_event', 'event_occurred', 'treatment_group', 'age', 'gender']
cox_data = df[cols].copy()

# 2. Encode categorical variables (One-Hot Encoding)
# drop_first=True to avoid multicollinearity (Dummy Variable Trap)
cox_data = pd.get_dummies(cox_data, columns=['treatment_group', 'gender'], drop_first=True)

print("Columns for modeling:", cox_data.columns.tolist())

# 3. Fit the Cox Model
cph = CoxPHFitter()
cph.fit(cox_data, duration_col='days_to_event', event_col='event_occurred')

# 4. View the Summary
cph.print_summary()

Interpreting the Output (Hazard Ratios)

The most important column in the output is exp(coef), also known as the Hazard Ratio (HR).

$HR = \frac{h(t|x+1)}{h(t|x)}$

Here is how to interpret HR in our context (where Event = Recovery):

• HR = 1: No effect.
• HR > 1: The covariate increases the rate of the event. (Recover faster).
• HR < 1: The covariate decreases the rate of the event. (Recover slower).

Example Interpretations:

1. treatment_group_Drug_B (HR ≈ 1.8): Patients on Drug B recover 1.8x faster (80% higher daily recovery rate) than the baseline (Placebo).
2. age (HR ≈ 0.98): For every extra year of age, the recovery rate drops by ~2%. Older patients recover slower.
3. gender_M (HR ≈ 1.0): If the CI crosses 1, gender has no significant effect on recovery speed.

The Proportional Hazards Assumption

The Cox model makes a strong assumption: Proportional Hazards.

It assumes that the ratio of hazards between two groups is constant over time.

• Correct: Drug B is always 2x better than Placebo (at Day 1, Day 10, and Day 50).
• Violation: Drug B is wonderful for the first 10 days, but after Day 10, it stops working and becomes equal to Placebo.

If the curves cross each other, the assumption is likely violated.

How to Check the Assumption

lifelines makes this easy with a built-in check.

# Check assumption
cph.check_assumptions(cox_data, p_value_threshold=0.05, show_plots=True)

If the assumption is violated, you might need to:

1. Stratify: Run separate models for different groups.
2. Time-Varying Covariates: Use advanced models that allow $\beta$ to change over time (beyond the scope of this intro, but good to know exists).

Conclusion

Survival Analysis unlocks insights that standard regression completely misses. By modeling "time-to-event," you can handle censored data correctly and gain a dynamic view of risk.

Recap of what we covered:

1. Censoring: Why we can't ignore data just because the event hasn't happened yet.
2. Survival Function $S(t)$: The probability of "surviving" (or not recovering) past time $t$.
3. Kaplan-Meier: The staircase plot for visualizing survival curves.
4. Log-Rank Test: The statistical test to confirm if curves are different.
5. Cox Regression: The multivariate model to calculate Hazard Ratios for covariates.

In our dataset, we found that Drug B significantly accelerates recovery compared to the Placebo, and we quantified exactly how much faster using the Hazard Ratio.

Next Steps

• Dive into Hypothesis Testing to understand the p-values in your Log-Rank test.
• Explore A/B Testing Design to learn how to set up the experiments that generate this survival data.

Hands-On Practice

Survival Analysis is essential when we care about 'time-to-event' rather than just binary outcomes. Standard regression fails here because of 'censoring'—we know some patients survived at least until the study ended, but not when they would have eventually recovered. While the lifelines library is the industry standard for this, understanding the underlying mathematics is powerful. In this guide, we will implement the Kaplan-Meier Estimator and the Log-Rank Test from scratch using pandas and scipy to visualize and quantify recovery rates in a clinical trial.

By building the Kaplan-Meier estimator from scratch, we visualized the 'survival' (or in this case, non-recovery) probability over time without relying on black-box libraries. The steep drop in the Drug B curve compared to Placebo indicates faster recovery, and our manual Log-Rank test confirmed this difference is statistically significant. This 'time-to-event' perspective provides much richer actionable insights than simple binary classification.