Clinical Survival Analysis with Python

Analyze lung cancer data with Kaplan-Meier curves and Cox models.

Introduction

30 Second Summary

When a clinical trial tests a new cancer treatment, the most important question is how long patients survive. Answering that rigorously requires specialized statistical methods that handle incomplete data from patients who drop out of studies or are still alive when data collection ends.

In this project, you will build a complete survival analysis pipeline in Python using real lung cancer clinical trial data. You will produce publication-quality survival curves, run statistical tests comparing patient groups, and fit a multivariate regression model with validated assumptions.

What You'll Build

You'll have a Jupyter notebook that takes raw patient data from a real lung cancer trial and produces survival curves, statistical tests, and individualized survival predictions you can walk through in a portfolio review or job interview.

By the end of this project, you'll have:

  • Kaplan-Meier survival curves with shaded confidence intervals, showing exactly how survival probability declines over time for different patient groups.
  • A log-rank test result that definitively answers whether male and female patients have statistically different survival outcomes.
  • A fitted Cox Proportional Hazards model with hazard ratios telling you which clinical factors increase or decrease risk of death, plus individualized patient survival predictions.
  • Secret Mission: Compare your Cox model against a Weibull Accelerated Failure Time model using concordance index scoring.

Are there any prerequisites?

You should be comfortable with basic Python syntax and familiar with pandas DataFrames. A conceptual understanding of survival analysis basics (censoring, hazard functions) is helpful but not strictly required.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Your Environment

You have lifelines installed and your notebook open. Before you can start analyzing patient data, you need to import all the libraries you'll use throughout this project.

This first cell acts as the foundation for everything that follows. It pulls in pandas for data manipulation, matplotlib for plotting, and the specific lifelines classes for survival modeling and statistical testing.

In this step, get ready to:
  • Import the survival analysis and data science libraries into your notebook.
  • Verify all imports run without errors.
Write your imports cell

Every Jupyter analysis notebook starts with a dedicated imports cell. This keeps all your dependencies in one place and makes the notebook easy to restart and re-run from scratch.

  • In your survival_analysis.ipynb notebook, click into the first empty cell.
  • Type the following import code into the cell:
# Cell 1: Imports
import pandas as pd
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter, CoxPHFitter, WeibullAFTFitter
from lifelines.datasets import load_lung
from lifelines.statistics import logrank_test
from lifelines.plotting import add_at_risk_counts

What does this code do?

  • pandas gives you DataFrames to store and manipulate the patient data.
  • matplotlib.pyplot handles all the survival curve plotting.
  • KaplanMeierFitter estimates survival probabilities without assuming a distribution shape.
  • CoxPHFitter fits a multivariate regression model to identify which covariates predict survival.
  • WeibullAFTFitter is a parametric alternative you'll use in the Secret Mission.
  • load_lung loads the built-in NCCTG lung cancer dataset directly from lifelines.
  • logrank_test compares survival distributions between two groups statistically.
  • add_at_risk_counts adds patient counts below your survival curves for publication-quality figures.
  • Run the cell by pressing Shift+Enter.

A successful import cell produces no output. You should see the cell execution counter increment (showing [1] to the left of the cell) and no error messages.

Seeing a ModuleNotFoundError?

  • Make sure your Jupyter kernel is using the same Python environment where you installed lifelines. You can check by running import sys; print(sys.executable) in a cell and comparing it to your terminal's which python output.
  • If the paths don't match, restart Jupyter from the same terminal session where you ran pip install.

Still stuck? Help me fix a ModuleNotFoundError for lifelines in Jupyter Notebook.

✔️ Awesome, I've got everything!

Great. Double check you see [1] next to the cell and no error output, then move on.

ⓧ I'd like to double check the full code

# Cell 1: Imports
import pandas as pd
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter, CoxPHFitter, WeibullAFTFitter
from lifelines.datasets import load_lung
from lifelines.statistics import logrank_test
from lifelines.plotting import add_at_risk_counts

Your notebook is loaded with every tool you need. Next up, you'll pull in the real lung cancer dataset and build your first survival curve.

Load Data and Build Your First Survival Curve

Before you can build any survival model, you need to understand what clinical survival data actually looks like. The NCCTG Lung Cancer dataset contains real patient records from a clinical trial, and it needs some cleanup before it's ready for analysis.

In this step, you'll load the dataset, prepare it for the lifelines library, and produce the most fundamental visualization in survival analysis: the Kaplan-Meier survival curve.

In this step, get ready to:
  • Load and clean the NCCTG lung cancer dataset for survival analysis.
  • Fit and plot an overall Kaplan-Meier survival curve with median survival annotation.
  • Compare male vs. female survival curves on the same plot.
Load and explore the NCCTG lung cancer dataset

The NCCTG Lung Cancer dataset contains 288 patients with advanced lung cancer from a North Central Cancer Treatment Group trial. Each row is a patient, with columns for survival time, event status, and clinical measurements.

  • In your survival_analysis.ipynb notebook, create a new cell below Cell 1 and paste the following code:
# Cell 2: Load and prepare data
lung = load_lung()
print(f"Dataset shape: {lung.shape}")
print(lung.head())
print(f"\nColumn names: {lung.columns.tolist()}")
print(f"\nMissing values:\n{lung.isnull().sum()}")

# Recode status: original has 1=censored, 2=dead -> convert to 0=censored, 1=dead
lung['status'] = lung['status'] - 1

# Rename columns to remove dots
lung.columns = lung.columns.str.replace('.', '_', regex=False)

# Drop rows with missing values in key columns
lung_clean = lung[['time', 'status', 'age', 'sex', 'ph_ecog', 'ph_karno', 'pat_karno', 'wt_loss']].dropna()
print(f"\nCleaned dataset shape: {lung_clean.shape}")

What does this code do?

The original dataset codes status as 1=censored and 2=dead. The lifelines library expects 0=censored and 1=event, so subtracting 1 converts to the correct convention.

Column names like ph.ecog contain dots that cause issues with CoxPHFitter later. Renaming them to use underscores (ph_ecog) avoids that problem.

Finally, dropping rows with NaN values in the key analysis columns gives you a clean DataFrame to work with.

  • Run the cell by pressing Shift+Enter.
  • You should see the dataset shape (228, 10) printed, followed by a preview of the first 5 rows, column names, missing value counts, and the cleaned shape.

Seeing a NameError?

If you see NameError: name 'load_lung' is not defined, make sure Cell 1 (your imports) was executed first. Switch back to Cell 1 and run it with Shift+Enter, then re-run Cell 2.

help me fix my NameError in Jupyter

Fit and plot the overall Kaplan-Meier survival curve

A Kaplan-Meier curve shows the probability of survival over time as a step function. Each step down represents a patient death, and the shaded region shows the confidence interval.

  • Create a new cell below Cell 2 and paste this code to fit the Kaplan-Meier estimator and plot the survival curve:
# Cell 3: Overall Kaplan-Meier curve
kmf = KaplanMeierFitter()
kmf.fit(lung_clean['time'], event_observed=lung_clean['status'], label='Overall Survival')

fig, ax = plt.subplots(figsize=(10, 6))
kmf.plot_survival_function(ax=ax)
ax.set_title('Kaplan-Meier Survival Curve - NCCTG Lung Cancer')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
ax.annotate(f'Median survival: {kmf.median_survival_time_:.0f} days',
            xy=(kmf.median_survival_time_, 0.5), xytext=(kmf.median_survival_time_ + 50, 0.6),
            arrowprops=dict(arrowstyle='->', color='gray'), fontsize=10)
plt.tight_layout()
plt.show()

print(f"Median survival time: {kmf.median_survival_time_} days")

What does this code do?

  • kmf.fit() takes the duration column (time) and the event indicator (status) to estimate the survival function non-parametrically.
  • The dashed horizontal line at y=0.5 marks where 50% of patients have experienced the event. The annotation labels the exact median survival time.
  • median_survival_time_ is the time at which the survival probability drops below 0.5.

Before you run this cell, what shape do you expect the survival curve to take? Will it be a smooth line or something else?

  • Run the cell with Shift+Enter.
  • You should see a step-function survival curve with shaded confidence intervals, starting at 1.0 and declining over approximately 1000 days. A gray arrow points to the median survival time on the curve.

Plot not showing?

If you see text output but no plot, make sure you are running inside a Jupyter notebook (not a plain .py script). Matplotlib renders inline automatically in Jupyter.

help me display matplotlib plots in Jupyter

Compare survival curves by sex

A single overall curve hides differences between patient subgroups. Splitting the data by sex (1=male, 2=female in this dataset) lets you visually compare whether one group survives longer.

  • Create a new cell below Cell 3 and paste this code to plot sex-stratified survival curves:
# Cell 4: Kaplan-Meier by sex
male = lung_clean[lung_clean['sex'] == 1]
female = lung_clean[lung_clean['sex'] == 2]

kmf_male = KaplanMeierFitter()
kmf_male.fit(male['time'], event_observed=male['status'], label='Male')

kmf_female = KaplanMeierFitter()
kmf_female.fit(female['time'], event_observed=female['status'], label='Female')

fig, ax = plt.subplots(figsize=(10, 6))
kmf_male.plot_survival_function(ax=ax)
kmf_female.plot_survival_function(ax=ax)
ax.set_title('Kaplan-Meier Survival Curves by Sex')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
add_at_risk_counts(kmf_male, kmf_female, ax=ax)
plt.tight_layout()
plt.show()

What does this code do?

  • Two separate KaplanMeierFitter instances are fitted, one for each sex group, so they produce independent survival estimates.
  • Plotting both on the same axes (ax) overlays the curves for direct visual comparison.
  • add_at_risk_counts() adds a table below the x-axis showing how many patients are still being followed at each time point. This is standard in publication-quality survival plots.
  • Run the cell with Shift+Enter.
  • You should see two survival curves (Male and Female) on the same plot, with at-risk count tables below. The female curve sits above the male curve, suggesting better survival for females.

Can you trust the visual difference?

The curves clearly diverge, but looks can be deceiving. With only ~228 patients, that gap might be random chance. You cannot yet say the difference is statistically significant. That is exactly what the log-rank test in the next step will answer.

Seeing overlapping or garbled text in at-risk counts?

This can happen if your figure size is too small. The code uses figsize=(10, 6) which should be large enough. If the labels still overlap, try increasing to figsize=(12, 7).

help me fix overlapping at-risk count labels

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all four cells.

ⓧ I'd like to double check the full code

Here is the complete code for all cells in your notebook so far. Scroll through each cell and confirm your code matches.

Cell 1 (Imports):

# Cell 1: Imports
import pandas as pd
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter, CoxPHFitter, WeibullAFTFitter
from lifelines.datasets import load_lung
from lifelines.statistics import logrank_test
from lifelines.plotting import add_at_risk_counts

Cell 2 (Load and prepare data):

# Cell 2: Load and prepare data
lung = load_lung()
print(f"Dataset shape: {lung.shape}")
print(lung.head())
print(f"\nColumn names: {lung.columns.tolist()}")
print(f"\nMissing values:\n{lung.isnull().sum()}")

# Recode status: original has 1=censored, 2=dead -> convert to 0=censored, 1=dead
lung['status'] = lung['status'] - 1

# Rename columns to remove dots
lung.columns = lung.columns.str.replace('.', '_', regex=False)

# Drop rows with missing values in key columns
lung_clean = lung[['time', 'status', 'age', 'sex', 'ph_ecog', 'ph_karno', 'pat_karno', 'wt_loss']].dropna()
print(f"\nCleaned dataset shape: {lung_clean.shape}")

Cell 3 (Overall Kaplan-Meier curve):

# Cell 3: Overall Kaplan-Meier curve
kmf = KaplanMeierFitter()
kmf.fit(lung_clean['time'], event_observed=lung_clean['status'], label='Overall Survival')

fig, ax = plt.subplots(figsize=(10, 6))
kmf.plot_survival_function(ax=ax)
ax.set_title('Kaplan-Meier Survival Curve - NCCTG Lung Cancer')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
ax.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
ax.annotate(f'Median survival: {kmf.median_survival_time_:.0f} days',
            xy=(kmf.median_survival_time_, 0.5), xytext=(kmf.median_survival_time_ + 50, 0.6),
            arrowprops=dict(arrowstyle='->', color='gray'), fontsize=10)
plt.tight_layout()
plt.show()

print(f"Median survival time: {kmf.median_survival_time_} days")

Cell 4 (Kaplan-Meier by sex):

# Cell 4: Kaplan-Meier by sex
male = lung_clean[lung_clean['sex'] == 1]
female = lung_clean[lung_clean['sex'] == 2]

kmf_male = KaplanMeierFitter()
kmf_male.fit(male['time'], event_observed=male['status'], label='Male')

kmf_female = KaplanMeierFitter()
kmf_female.fit(female['time'], event_observed=female['status'], label='Female')

fig, ax = plt.subplots(figsize=(10, 6))
kmf_male.plot_survival_function(ax=ax)
kmf_female.plot_survival_function(ax=ax)
ax.set_title('Kaplan-Meier Survival Curves by Sex')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
add_at_risk_counts(kmf_male, kmf_female, ax=ax)
plt.tight_layout()
plt.show()

You now have two powerful visualizations: an overall survival curve and a sex-stratified comparison. Next up, you'll run a statistical test to determine whether that visual gap between male and female survival is real or just noise.

Run the Log-Rank Test

Your sex-stratified Kaplan-Meier curves from the last step show a visible gap between male and female survival. But looking at a chart and knowing the difference is real are two different things.

The log-rank test is the standard statistical test for comparing survival distributions between two groups. It tells you whether observed differences are statistically significant or just noise from sample variation.

In this step, get ready to:
  • Run the log-rank test to statistically compare male and female survival.
  • Interpret the p-value and test statistic to determine significance.
  • Recognize the limitation of the log-rank test for multivariate analysis.
Run the log-rank test on your sex-stratified groups

The log-rank test compares the entire survival experience of two groups, not just a single time point. It produces a test statistic and a p-value that tells you the probability of observing this large a difference by chance alone.

  • In your notebook, create a new cell below Cell 4 and add the following code to run the log-rank test and print an interpretation:
# Cell 5: Log-rank test
results = logrank_test(
    male['time'], female['time'],
    event_observed_A=male['status'],
    event_observed_B=female['status']
)
results.print_summary()
print(f"\nTest statistic: {results.test_statistic:.4f}")
print(f"P-value: {results.p_value:.4f}")
if results.p_value < 0.05:
    print("Result: Statistically significant difference in survival between sexes.")
else:
    print("Result: No statistically significant difference in survival between sexes.")

What does this code do?

logrank_test() takes the duration and event arrays for both groups and compares their survival distributions. It uses the male and female subsets you already created in Cell 4.

results.print_summary() outputs a formatted summary table. The if/else block at the bottom translates the raw p-value into a human-readable conclusion using the standard 0.05 threshold.

Before you run this cell, think: based on the visual gap you saw in the KM curves, do you expect the difference to be statistically significant or not?

  • Run the cell by pressing Shift+Enter.

You should see a summary table followed by the test statistic, p-value, and a printed interpretation. If the p-value is below 0.05, the difference in survival between male and female patients is statistically significant.

Seeing a NameError?

If you see NameError: name 'male' is not defined, make sure you ran Cell 4 first. The male and female DataFrames are created there.

If you see NameError: name 'logrank_test' is not defined, go back and run Cell 1 which imports logrank_test from lifelines.statistics.

help me debug my log-rank test cell

Recognize the limitation of the log-rank test

Your log-rank test confirmed that male and female patients have statistically different survival outcomes. That is useful, but it only answers one narrow question about one variable at a time.

Why isn't the log-rank test enough?

The log-rank test compares two groups on a single variable. It cannot tell you which of the many clinical factors (age, ECOG score, Karnofsky score, weight loss) simultaneously influence survival.

It also cannot quantify how much each factor increases or decreases risk. To answer those questions, you need a regression model that handles multiple covariates at once. That is exactly what the Cox Proportional Hazards model does.

✔️ Awesome, I've got everything!

Great. Double check that you have saved your notebook with Cell 5 included.

ⓧ I'd like to double check the full code

Here is the complete Cell 5 code. Make sure yours matches exactly:

# Cell 5: Log-rank test
results = logrank_test(
    male['time'], female['time'],
    event_observed_A=male['status'],
    event_observed_B=female['status']
)
results.print_summary()
print(f"\nTest statistic: {results.test_statistic:.4f}")
print(f"P-value: {results.p_value:.4f}")
if results.p_value < 0.05:
    print("Result: Statistically significant difference in survival between sexes.")
else:
    print("Result: No statistically significant difference in survival between sexes.")

You now have statistical proof that survival differs by sex. But sex is just one factor among many. In the next step, you will fit a Cox Proportional Hazards model that assesses all your clinical covariates simultaneously and quantifies each one's effect as a hazard ratio.

Fit the Cox Proportional Hazards Model

Your log-rank test confirmed that male and female patients have statistically different survival outcomes. But it only compares two groups on a single variable.

What if age, ECOG performance score, Karnofsky rating, and weight loss all influence survival at the same time? You need a model that handles multiple covariates simultaneously and quantifies exactly how much each one shifts survival risk. That's what the Cox Proportional Hazards model does.

In this step, get ready to:
  • Fit a Cox Proportional Hazards model on multiple clinical covariates.
  • Visualize hazard ratios in a forest-plot style chart.
  • Test the proportional hazards assumption using Schoenfeld residuals.
Fit the Cox model to your clinical data

Your lung_clean DataFrame already has columns renamed (no dots) and missing values dropped from Cell 2. CoxPHFitter can fit directly on this data, using every column except the duration and event columns as covariates.

  • Create a new cell in your notebook (Cell 6) and add the following code to fit the Cox model and display its regression table:
# Cell 6: Cox Proportional Hazards Model
cph = CoxPHFitter()
cph.fit(lung_clean, duration_col='time', event_col='status')
cph.print_summary()

What does this code do?

CoxPHFitter() creates a semi-parametric regression model. Unlike a fully parametric model, it makes no assumptions about the shape of the baseline hazard function.

The fit() call estimates a coefficient for each covariate (age, sex, ph_ecog, ph_karno, pat_karno, wt_loss) that tells you how that variable shifts the hazard of death. print_summary() outputs the full regression table with hazard ratios.

Before you run this cell, what do you think the summary table will show you about each covariate?

  • Run Cell 6.

You see a regression table with columns for each covariate's coefficient, exp(coef) (the hazard ratio), confidence intervals, and p-values.

How do I read hazard ratios?

The exp(coef) column is the hazard ratio. A value greater than 1 means that covariate increases the risk of death. A value less than 1 means it is protective.

For example, if ph_ecog has a hazard ratio of 1.6, that means each one-unit increase in ECOG score is associated with a 60% increase in hazard of death. Check the p-value column to see which covariates are statistically significant.

Seeing a column name error?

If you see an error about column names containing dots, make sure Cell 2 ran successfully. The line lung.columns = lung.columns.str.replace('.', '_', regex=False) must execute before fitting the Cox model.

Help me debug my CoxPHFitter column name error.

Visualize the hazard ratios

A table of numbers is hard to interpret at a glance. A forest plot shows each covariate's hazard ratio as a point with confidence interval bars, making it immediately obvious which factors increase risk and which are protective.

  • Create Cell 7 and add this code to generate a forest-plot style visualization:
# Cell 7: Plot hazard ratios
fig, ax = plt.subplots(figsize=(8, 5))
cph.plot(ax=ax, hazard_ratios=True)
ax.set_title('Cox PH Model - Hazard Ratios with 95% CI')
plt.tight_layout()
plt.show()

What does this plot show?

The hazard_ratios=True parameter plots hazard ratios instead of log-hazard coefficients. Each covariate appears as a horizontal point-and-whisker. Points to the right of 1.0 increase hazard. Points to the left of 1.0 are protective.

If a covariate's confidence interval crosses 1.0, it is not statistically significant at the 95% level.

  • Run Cell 7.

You see a forest plot with each covariate's hazard ratio and 95% confidence interval displayed as horizontal bars.

Plot is empty or shows no data?

Make sure Cell 6 ran without errors first. The cph.plot() method requires a fitted model. Re-run Cell 6, then Cell 7.

Help me fix my Cox model forest plot.

Test the proportional hazards assumption

The Cox model has a critical assumption: each covariate's effect on hazard must remain constant over time. If a covariate's effect changes as time passes, the hazard ratio you read from the summary is unreliable.

The check_assumptions() method tests each covariate for time-varying effects using scaled Schoenfeld residuals. It transforms time in four different ways and checks whether any covariate's effect significantly correlates with time.

  • Create Cell 8 and add this code to run the assumption test:
# Cell 8: Check proportional hazards assumption
cph.check_assumptions(lung_clean, p_value_threshold=0.05, show_plots=True)

What does this test do?

For each covariate, lifelines tests whether its scaled Schoenfeld residuals correlate with time. If the p-value is below 0.05, the covariate violates the proportional hazards assumption.

The show_plots=True parameter displays residual plots with a LOWESS smoothing line. A flat LOWESS line means the assumption holds. A sloped line indicates a time-varying effect.

Before you run Cell 8, do you expect all six covariates to pass the proportional hazards test? Or do you think some effects might change over time?

  • Run Cell 8.

You see that at least one covariate (likely ph_ecog or ph_karno) fails the test. Its residual plot shows a sloped LOWESS line, and the printed output flags the violation.

This is a real problem. The hazard ratio you read for that covariate in the summary table is not trustworthy because its effect on survival changes over time. You need to fix this before drawing conclusions from your model.

Why is this violation important?

The Cox model estimates a single hazard ratio per covariate, assuming that ratio stays constant. If a covariate's true effect is strong early on but fades later, that single number is misleading. You cannot report it as a reliable finding.

The standard fix is stratification. You will address this in the next step.

No violations reported?

If no violations appear, check that you passed p_value_threshold=0.05. Also confirm your lung_clean DataFrame matches the one from Cell 2 with all key columns present. The exact covariate that fails may vary depending on which rows were dropped during cleaning.

Help me understand my check_assumptions output.

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with Cells 6, 7, and 8 all executed.

ⓧ I'd like to double check the full code

Here are all three cells you added in this step. Compare them against your notebook:

# Cell 6: Cox Proportional Hazards Model
cph = CoxPHFitter()
cph.fit(lung_clean, duration_col='time', event_col='status')
cph.print_summary()

# Cell 7: Plot hazard ratios
fig, ax = plt.subplots(figsize=(8, 5))
cph.plot(ax=ax, hazard_ratios=True)
ax.set_title('Cox PH Model - Hazard Ratios with 95% CI')
plt.tight_layout()
plt.show()

# Cell 8: Check proportional hazards assumption
cph.check_assumptions(lung_clean, p_value_threshold=0.05, show_plots=True)

You now have a multivariate Cox model fitted, its hazard ratios visualized, and a clear diagnosis: at least one covariate breaks the model's core assumption. In the next step, you will fix this violation using stratification and then generate individual patient survival predictions.

Fix the Violation and Generate Predictions

Your Cox model identified which clinical covariates predict survival, but the assumption check revealed a problem. A proportional hazards violation means the offending covariate's effect on survival changes over time, making its hazard ratio unreliable.

In this step, you'll fix that violation using stratification. Stratification lets the model estimate a separate baseline hazard for each level of the offending variable while still pooling coefficient estimates for everything else. Once the model is valid, you'll use it to predict individualized survival curves for specific patients.

In this step, get ready to:
  • Stratify the Cox model on the offending covariate and confirm the violation is resolved.
  • Generate individualized patient survival predictions from the corrected model.
  • Visualize contrasting patient profiles to demonstrate how the model differentiates risk.
Stratify on the offending covariate

The check_assumptions output from the previous step flagged ph_ecog as violating the proportional hazards assumption. Stratification handles this by fitting a separate baseline hazard for each ECOG performance level, so the model no longer forces a single coefficient onto a time-varying effect.

  • Create a new cell in your notebook and add the stratified model code:
# Cell 9: Fix violation with stratification
cph_stratified = CoxPHFitter()
lung_strat = lung_clean.copy()
lung_strat['ph_ecog'] = lung_strat['ph_ecog'].astype(int)
cph_stratified.fit(lung_strat, duration_col='time', event_col='status', strata=['ph_ecog'])
cph_stratified.print_summary()

# Verify assumptions are now satisfied
cph_stratified.check_assumptions(lung_strat, p_value_threshold=0.05)

What does this code do?

  • lung_clean.copy() creates a working copy so the original DataFrame stays unchanged.
  • .astype(int) converts ph_ecog from float to integer, since ECOG scores are discrete levels (0, 1, 2, 3) that lifelines uses as strata labels.
  • strata=['ph_ecog'] tells the model to estimate a separate baseline hazard for each ECOG level, removing ph_ecog from the coefficient table entirely.
  • check_assumptions re-runs the Schoenfeld residuals test to confirm no remaining violations.

Before you run the cell, what do you expect check_assumptions to print now that the offending covariate has been stratified out?

  • Run the cell.

You should see the summary table (now without ph_ecog as a coefficient) followed by the assumption check confirming no violations remain.

Why did ph_ecog disappear from the summary?

When you stratify on a variable, it no longer appears in the coefficient table because the model no longer estimates a single hazard ratio for it. Instead, it gets its own baseline hazard curve at each level.

The remaining covariates (age, sex, ph_karno, pat_karno, wt_loss) still have pooled hazard ratios estimated across all strata.

Still seeing a violation?

  • Confirm that strata=['ph_ecog'] is inside the .fit() call, not passed to the CoxPHFitter() constructor.
  • Make sure you're checking assumptions on lung_strat (the copy with int-converted ph_ecog), not lung_clean.
  • If a different covariate now fails, your dataset's specific subset may have additional violations. For this project, the stratified model should pass.

Why is my stratified model still failing?

Create patient profiles for prediction

With a valid model in hand, you can now predict survival for individual patients based on their clinical characteristics. You'll create two contrasting profiles to see how the model differentiates risk.

  • Create a new cell below and add the patient profile definitions:
# Cell 10: Individual patient survival predictions
patient_healthy = pd.DataFrame({
    'age': [55],
    'sex': [2],
    'ph_karno': [90],
    'pat_karno': [90],
    'wt_loss': [0]
})

patient_poor = pd.DataFrame({
    'age': [72],
    'sex': [1],
    'ph_karno': [60],
    'pat_karno': [50],
    'wt_loss': [15]
})

What do these profiles represent?

  • patient_healthy: A 55-year-old female (sex=2) with high physician-rated Karnofsky score (90), high patient-rated score (90), and no weight loss. This represents a best-case scenario.
  • patient_poor: A 72-year-old male (sex=1) with low physician-rated Karnofsky score (60), very low patient-rated score (50), and significant weight loss (15 lbs). This represents a worst-case scenario.

Notice that ph_ecog is not included in these DataFrames because it was moved to the strata. The model handles it through the baseline hazard instead.

Visualize predicted survival for contrasting patients
  • Add the following plotting and prediction code below the patient profiles in the same cell:
fig, ax = plt.subplots(figsize=(10, 6))
cph_stratified.predict_survival_function(patient_healthy).plot(ax=ax, label='55yo Female, Good Status')
cph_stratified.predict_survival_function(patient_poor).plot(ax=ax, label='72yo Male, Poor Status')
ax.set_title('Predicted Survival Curves - Individual Patients')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
ax.legend()
plt.tight_layout()
plt.show()

print(f"Predicted median survival (healthy profile): {cph_stratified.predict_median(patient_healthy).values[0]:.0f} days")
print(f"Predicted median survival (poor profile): {cph_stratified.predict_median(patient_poor).values[0]:.0f} days")

What does this code do?

  • predict_survival_function() returns a time-indexed DataFrame of survival probabilities for the given patient covariates.
  • predict_median() returns the time at which predicted survival probability drops below 50%.
  • Plotting both on the same axes shows how dramatically different clinical profiles lead to different survival trajectories.

Before you run the cell, which patient do you predict will have the higher survival probability at any given time point?

  • Run the cell.

You should see two distinct survival curves on the same plot. The healthy profile's curve stays higher for longer, while the poor profile's curve drops more steeply. The printed median survival times quantify the gap between these two patients.

Seeing a KeyError or column mismatch?

  • The patient DataFrames must contain columns matching exactly what the stratified model expects: age, sex, ph_karno, pat_karno, wt_loss.
  • Do not include ph_ecog, time, or status in the patient DataFrames.
  • Make sure Cell 9 ran successfully before running Cell 10, since Cell 10 references cph_stratified.

Help me fix my patient prediction error.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook.

ⓧ I'd like to double check the full code

Here is the complete code for Cell 9:

# Cell 9: Fix violation with stratification
cph_stratified = CoxPHFitter()
lung_strat = lung_clean.copy()
lung_strat['ph_ecog'] = lung_strat['ph_ecog'].astype(int)
cph_stratified.fit(lung_strat, duration_col='time', event_col='status', strata=['ph_ecog'])
cph_stratified.print_summary()

# Verify assumptions are now satisfied
cph_stratified.check_assumptions(lung_strat, p_value_threshold=0.05)

And the complete code for Cell 10:

# Cell 10: Individual patient survival predictions
patient_healthy = pd.DataFrame({
    'age': [55],
    'sex': [2],
    'ph_karno': [90],
    'pat_karno': [90],
    'wt_loss': [0]
})

patient_poor = pd.DataFrame({
    'age': [72],
    'sex': [1],
    'ph_karno': [60],
    'pat_karno': [50],
    'wt_loss': [15]
})

fig, ax = plt.subplots(figsize=(10, 6))
cph_stratified.predict_survival_function(patient_healthy).plot(ax=ax, label='55yo Female, Good Status')
cph_stratified.predict_survival_function(patient_poor).plot(ax=ax, label='72yo Male, Poor Status')
ax.set_title('Predicted Survival Curves - Individual Patients')
ax.set_xlabel('Time (days)')
ax.set_ylabel('Survival Probability')
ax.legend()
plt.tight_layout()
plt.show()

print(f"Predicted median survival (healthy profile): {cph_stratified.predict_median(patient_healthy).values[0]:.0f} days")
print(f"Predicted median survival (poor profile): {cph_stratified.predict_median(patient_poor).values[0]:.0f} days")

Your survival analysis pipeline is complete. You've gone from raw clinical data through non-parametric estimation, statistical testing, multivariate regression, assumption validation, and individualized prediction.

Next up, the Secret Mission challenges you to compare your Cox model against a parametric alternative.

Secret mission

Model Comparison and Concordance

A good data scientist doesn't stop at one model. Evaluate how well your Cox model actually discriminates between patients, and compare it against a fully parametric alternative. You'll compute concordance indices and fit a Weibull Accelerated Failure Time model to see which approach better predicts survival.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • The survival_analysis.ipynb notebook file on your Desktop.
  • The lifelines Python library installed via pip.

Keep everything running

No action needed. Choose this if you're still actively building or plan to extend this notebook with additional survival models.

  • Your survival_analysis.ipynb notebook stays on your Desktop, ready to open in Jupyter at any time.
  • The lifelines library remains installed in your Python environment for future projects.

Pause - I'll come back to this later

Shut down the Jupyter Notebook server to free up memory. Your files stay on disk.

  • Switch to the terminal where Jupyter is running and press Ctrl+C to stop the server.
  • Confirm the shutdown when prompted.
  • Your notebook and all outputs are saved. To resume later, run jupyter notebook from the same directory.

Delete - I don't want to use this again

Remove all project files and uninstall the library to start fresh.

  • Stop the Jupyter server by pressing Ctrl+C in the terminal where it is running.
  • Delete the notebook file by running:
rm ~/Desktop/survival_analysis.ipynb
  • Uninstall the lifelines library by running:
pip uninstall lifelines -y
  • Confirm deletion by running ls ~/Desktop and checking that survival_analysis.ipynb no longer appears.

Nice Work!

Nice Work!

You just built a complete clinical survival analysis pipeline on real lung cancer data. From raw patient records to validated regression models and individualized predictions, your notebook demonstrates the exact workflow that biostatistics and healthcare analytics teams use every day.

You've learned how to:

  • Build publication-ready Kaplan-Meier survival curves with confidence intervals and statistically confirm group differences using the log-rank test.
  • Fit and interpret a multivariate Cox Proportional Hazards model, extracting hazard ratios that quantify each clinical covariate's contribution to survival risk.
  • Diagnose a proportional hazards violation using Schoenfeld residuals, fix it with stratification, and generate individualized patient survival predictions from your validated model.
  • Secret Mission: Compared your Cox PH model against a Weibull AFT model using concordance index scoring and documented a final recommendation.

Ready to quiz yourself?