Predicting Patient Readmission with XGBoost and Claims Data

Introduction

Hospital readmission within 30 days is a critical quality metric and cost driver in healthcare. In this tutorial, we will use structured claims data to train an XGBoost model that predicts the probability of a patient being readmitted. We will also use SHAP values to explain the model's predictions to clinicians.

Architecture Overview

graph LR
    Syn(Synthea Claims Data) --> FE(Feature Engineering)
    FE --> Join(Join Encounters & Patients)
    Join --> Matrix(XGBoost DMatrix)
    Matrix --> Model(XGBClassifier)
    Model --> Pred(30-Day Risk Score)
    Model --> SHAP(SHAP Explainability)
        

Prerequisites

• Python 3.8+
• Pandas, Scikit-learn, XGBoost, SHAP
• A synthetic claims dataset (e.g., Synthea)

Step 1: Data Preparation and Feature Engineering

Claims data requires significant feature engineering. We need to extract demographics, calculate the length of stay (LOS), and aggregate past utilization counts to build a predictive feature set.

import pandas as pd
import numpy as np

# Load synthetic encounters and patients
encounters = pd.read_csv('encounters.csv')
patients = pd.read_csv('patients.csv')

# Calculate Length of Stay (LOS)
encounters['START'] = pd.to_datetime(encounters['START'])
encounters['STOP'] = pd.to_datetime(encounters['STOP'])
encounters['LENGTH_OF_STAY'] = (encounters['STOP'] - encounters['START']).dt.days

# Join with Patients table to get demographics
df = pd.merge(encounters, patients, left_on='PATIENT', right_on='Id')

# Calculate Age at Encounter
df['BIRTHDATE'] = pd.to_datetime(df['BIRTHDATE'])
df['AGE'] = (df['START'] - df['BIRTHDATE']).dt.days // 365

# Create a target variable: Readmitted within 30 days (1=Yes, 0=No)
# We sort by patient and date to find the next admission
df = df.sort_values(['PATIENT', 'START'])
df['NEXT_ADMISSION'] = df.groupby('PATIENT')['START'].shift(-1)
df['DAYS_TO_READMISSION'] = (df['NEXT_ADMISSION'] - df['STOP']).dt.days
df['READMITTED_30D'] = np.where((df['DAYS_TO_READMISSION'] <= 30) & (df['DAYS_TO_READMISSION'] >= 0), 1, 0)

# Select predictive features
features = ['AGE', 'LENGTH_OF_STAY'] # In a real scenario, add comorbidity flags here
X = df[features]
y = df['READMITTED_30D']

Step 2: Training the XGBoost Model

XGBoost handles tabular healthcare data exceptionally well, including missing values natively.

import xgboost as xgb
from sklearn.model_selection import train_test_split

# Split into Train / Test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

model = xgb.XGBClassifier(
    n_estimators=200,
    max_depth=4,
    learning_rate=0.05,
    objective='binary:logistic',
    scale_pos_weight=5, # Handle class imbalance (readmissions are rare)
    use_label_encoder=False,
    eval_metric='logloss'
)

model.fit(X_train, y_train)

Step 3: Evaluation and Feature Importance

We evaluate the model using the ROC-AUC score.

from sklearn.metrics import roc_auc_score

preds = model.predict_proba(X_test)[:, 1]
auc = roc_auc_score(y_test, preds)
print(f"ROC-AUC: {auc:.3f}")

Step 4: Clinical Explainability with SHAP

Clinicians need to know why a patient is high risk. SHAP (SHapley Additive exPlanations) values provide individual-level feature attributions.

import shap
import matplotlib.pyplot as plt

# Create the explainer
explainer = shap.TreeExplainer(model)
shap_values = explainer.shap_values(X_test)

# Plot summary plot for the entire cohort
shap.summary_plot(shap_values, X_test)

# Plot a force plot for a single high-risk patient
# shap.initjs()
# shap.force_plot(explainer.expected_value, shap_values[0,:], X_test.iloc[0,:])

Conclusion

Using XGBoost on structured claims data provides a fast, interpretable baseline for readmission risk scoring in population health management. By layering SHAP on top, we build trust with clinical end-users.