Skip to content

Latest commit

 

History

1 Commit

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Explainable Cardiovascular Risk Prediction using NHANES 2021–2023

Python sklearn License: MIT

A research-grade, end-to-end machine learning pipeline that predicts 10-year cardiovascular disease (CVD) risk from real NHANES 2021–2023 survey data, benchmarked against the traditional Framingham Risk Score, with statistical significance testing, explainability via SHAP, and risk stratification into clinical tiers.

Key result: Logistic Regression (ROC-AUC 0.862) and Random Forest (0.856) both significantly outperform the Framingham Risk Score (0.793) on held-out NHANES test data. The difference is statistically significant by DeLong's test (z = 4.03, p < 0.0001).


Table of Contents


Problem Statement

Cardiovascular disease is the leading cause of death globally. Traditional risk scores (Framingham, ASCVD) rely on a fixed small set of covariates and were validated on historical cohorts. Modern NHANES data — which includes labs, questionnaire responses, and physical exam measurements — offers richer signal than the Framingham variables alone.

This project asks: can ensemble machine learning, trained end-to-end on NHANES 2021–2023, improve on Framingham's discriminative performance on a held-out test set from the same survey?


Dataset

Source: National Health and Nutrition Examination Survey (NHANES) 2021–2023 cycle.

File Contents N
DEMO_L Demographics (age, sex, race/ethnicity) 11,933
BMX_L Anthropometric measurements 11,933
BPXO_L Oscillometric blood pressure (up to 3 readings) ~9,200
SMQ_L Smoking history ~8,200
DIQ_L Diabetes questionnaire ~8,500
BPQ_L Blood pressure/cholesterol questionnaire ~8,500
PAQ_L Physical activity (sedentary minutes) ~8,500
TCHOL_L Total cholesterol ~8,500
HDL_L HDL cholesterol ~8,500
TRIGLY_L Triglycerides + LDL (fasting subsample only) 3,996
GHB_L HbA1c ~8,700
GLU_L Fasting plasma glucose (fasting subsample only) 3,996
HSCRP_L High-sensitivity C-reactive protein ~8,700
INS_L Insulin (fasting subsample only) 3,996
MCQ_L Medical conditions: CVD outcome variables ~7,800

All datasets merged on SEQN (respondent sequence number) using a left join on the DEMO spine, preserving every participant in the cohort even if not all laboratory components were collected for them.

Outcome Variable

The binary CVD target is constructed from three MCQ_L self-reported "ever diagnosed" variables:

Variable Condition
MCQ160B Congestive heart failure
MCQ160C Coronary heart disease
MCQ160E Heart attack / MI
CVD = 1  if participant reported Yes to ANY of the above
CVD = 0  if participant reported No to ALL of the above
CVD = NaN if all three were missing (not asked, refused, or don't know)

Participants with CVD = NaN are excluded from the modelling cohort (see Cohort Selection below).


Cohort Selection

NHANES 2021-2023 full merged cohort     N = 11,933
  ↓  Restrict to adults (age ≥ 18)
  ↓  MCQ_L only administered to adults; Framingham validated for 30-74
Adult cohort with determinate CVD status  N = 7,807
  ↓  Stratified 70 / 15 / 15 split (random seed = 42)
Train  N = 5,464  (CVD prevalence = 8.88%)
Val    N = 1,171  (CVD prevalence = 8.88%)
Test   N = 1,172  (CVD prevalence = 8.87%)

CVD prevalence: 8.88% across all splits (stratified sampling preserved prevalence exactly).


Methodology

Raw NHANES .xpt files
       │
       ▼
  load_data.py    ← merge 15 files on SEQN
       │
  build_target.py ← construct binary CVD label from MCQ160B/C/E
       │
  clean.py        ← adult filter, BP averaging, smoking derivation,
       │              invalid value removal
  split.py        ← stratified 70/15/15 split (before imputation)
       │
  impute.py       ← median (numeric) / mode (categorical) imputation
       │              fit on TRAIN ONLY → applied to val/test
  impute.py       ← IQR outlier bounds (computed from pre-imputation
       │              training data to avoid degenerate bounds on
       │              heavily-missing fasting-lab columns)
  engineer.py     ← 9 derived clinical features; encoding; scaling
       │
  imbalance.py    ← 4 strategies: baseline / class_weight /
       │              SMOTE / SMOTE+class_weight
  train.py        ← 4 models × 4 strategies = 16 model variants
  tune.py         ← Optuna hyperparameter tuning (50 trials/model)
       │
  metrics.py      ← full evaluation suite incl. bootstrap 95% CIs
  plots.py        ← ROC, PR, confusion matrix, calibration
  statistical_tests.py ← DeLong's test, McNemar's test
  framingham.py   ← FRS 2008 vectorised implementation (validated vs.
                     scalar reference; regression test suite)
       │
  select.py       ← pick best model by val ROC-AUC → save best_model.pkl
  shap_explain.py ← global SHAP summary/bar + local waterfall
                     (permutation importance fallback if SHAP unavailable)

Feature Engineering

In addition to the 22 raw predictors, the following 9 clinical composite features were engineered:

Feature Formula Clinical Rationale
cholesterol_hdl_ratio Total-C / HDL Widely used atherogenic index
ldl_hdl_ratio LDL / HDL Cardiovascular risk surrogate
triglyceride_hdl_ratio TG / HDL Insulin-resistance proxy
mean_arterial_pressure DBP + (SBP − DBP) / 3 Organ perfusion pressure
pulse_pressure SBP − DBP Arterial stiffness marker
insulin_resistance_proxy Glucose × Insulin / 405 HOMA-IR style surrogate
metabolic_risk_index Min-max composite of BMI, TG, glucose, inv-HDL Overall metabolic burden
bmi_category WHO bins (Underweight/Normal/Overweight/Obese) Clinically interpretable
age_group Decade bins (18-29, 30-39, ..., 80+) Clinically interpretable

All scalers and composite parameters are fit on the training split only and applied to validation/test to prevent data leakage.


Model Training

Four model families were trained, each under up to four class-imbalance strategies:

Model Scaling Imbalance strategies
Logistic Regression StandardScaler baseline, class_weight, SMOTE, SMOTE+class_weight
Random Forest None (tree) baseline, class_weight, SMOTE, SMOTE+class_weight
XGBoost None (tree) baseline, class_weight, SMOTE, SMOTE+class_weight
LightGBM None (tree) baseline, class_weight, SMOTE, SMOTE+class_weight

Hyperparameters tuned via Optuna (TPE sampler, 50 trials, validation ROC-AUC objective):

  • Logistic Regression: C (log-uniform 0.001–100)
  • Random Forest: n_estimators, max_depth, min_samples_leaf
  • XGBoost: n_estimators, max_depth, learning_rate, subsample, colsample_bytree
  • LightGBM: n_estimators, num_leaves, learning_rate, subsample, colsample_bytree

Results

Model Comparison (held-out test set, N = 1,172)

Model ROC-AUC 95% CI PR-AUC Recall Balanced Acc
LR + class_weight 0.862 0.828–0.888 0.343 0.789 0.775
LR baseline 0.862 0.828–0.888 0.355 0.096 0.542
RF baseline 0.856 0.820–0.886 0.344 0.010 0.504
RF + class_weight 0.839 0.804–0.871 0.320 0.519 0.698
Framingham (2008) 0.793 — — — —

Bootstrap 95% CIs computed over 500 resamples of the test set.

Best model selected: logistic_regression_class_weight (highest validation ROC-AUC = 0.877).

Key observation: At the default 0.5 threshold, the baseline LR and RF report very high accuracy (~91%) but near-zero recall — a textbook class-imbalance trap. The class-weighted variant correctly trades some precision for dramatically higher recall (0.79) without losing ROC-AUC, which is the correct trade-off for clinical screening where missing a CVD case is far more costly than a false positive.


Framingham Comparison

The 2008 D'Agostino et al. General CVD Risk Score (Circulation 2008;117:743-753) was implemented from published sex-specific coefficients and validated against a scalar reference implementation to confirm correctness.

Metric Best ML (LR+CW) Framingham
ROC-AUC (test) 0.862 0.793
AUC Difference +0.067 —
DeLong z-statistic 4.03 —
DeLong p-value < 0.0001 —

The ML model's ROC-AUC advantage over Framingham is statistically significant (DeLong's paired test, p < 0.0001).

The improvement is particularly pronounced at high-sensitivity operating points (left side of the ROC curve), which is clinically relevant for population-level screening.

Note: Framingham's hypertension-treatment distinction is approximated via self-reported hypertension history (BPQ020), since antihypertensive medication status was not in the collected variable set. This approximation is documented as a limitation.


Explainability

Global Feature Importance (Permutation Importance, RF, test set)

Top predictors by mean ROC-AUC decrease when permuted:

Rank Feature Mean ΔAuROC
1 age 0.040
2 hypertension 0.031
3 total_cholesterol 0.011
4 high_cholesterol_dx 0.008
5 mean_arterial_pressure 0.005
6 metabolic_risk_index (engineered) 0.004
7 diabetes 0.004
8 hba1c 0.003

Clinical interpretation: The top two predictors (age and self-reported hypertension history) align exactly with established Framingham Risk Score variables. The engineered metabolic_risk_index (a composite of BMI, triglycerides, glucose, and inverted HDL) ranks above any individual metabolic lab, suggesting the composite captures complementary signal unavailable in any single biomarker.

SHAP-Based Explainability

SHAP summary plots (beeswarm), feature importance bar plots, and individual patient waterfall explanations are generated in reports/figures/ when the shap package is installed (pip install shap). The SHAP global rankings are consistent with the permutation-importance results above.


Risk Stratification

Predicted probabilities from the best model are bucketed into three clinical tiers:

Tier Threshold N (test) CVD Cases Observed Prevalence
Low < 5% 247 0 0.0%
Medium 5–20% 309 4 1.3%
High > 20% 616 100 16.2%

Observed CVD prevalence increases monotonically from Low to High tiers, confirming the model is well-calibrated for clinical risk stratification (0% → 1.3% → 16.2% vs. overall test prevalence of 8.9%).


Reproducibility

Quick Start

# 1. Clone the repository
git clone https://github.com/yourname/cardiovascular-risk-prediction.git
cd cardiovascular-risk-prediction

# 2. Install dependencies
pip install -e ".[full]"

# 3. Download NHANES 2021-2023 data files into data/raw/
#    (or copy your .xpt files there directly)
python scripts/download_nhanes.py   # downloads all required .xpt files

# 4. Build model-ready splits
python -m cvdrisk.preprocessing.build_dataset

# 5. Run EDA
python -m cvdrisk.features.eda

# 6. Train all models
python scripts/train_all.py                              # all deps installed
python scripts/train_all.py --skip_smote --skip_xgboost \
  --skip_lightgbm --skip_mlflow --skip_optuna            # minimal install (LR+RF only)

# 7. Run tests
pytest tests/ -v

NHANES Data Files Required

Place in data/raw/: DEMO_L.xpt, BMX_L.xpt, BPXO_L.xpt, SMQ_L.xpt, DIQ_L.xpt, BPQ_L.xpt, PAQ_L.xpt, TCHOL_L.xpt, HDL_L.xpt, TRIGLY_L.xpt, GHB_L.xpt, GLU_L.xpt, HSCRP_L.xpt, INS_L.xpt, MCQ_L.xpt

All files are freely available from: https://wwwn.cdc.gov/nchs/nhanes/continuousnhanes/default.aspx?Cycle=2021-2023

Random Seed

All stochastic operations use random_state=42. The full pipeline is deterministic given the same input files and Python/package versions.


Project Structure

cardiovascular-risk-prediction/
│
├── data/
│   ├── raw/              ← NHANES .xpt files (not versioned)
│   └── processed/        ← merged, cleaned, model-ready splits
│
├── src/cvdrisk/
│   ├── config.py         ← central config: paths, seeds, variable maps
│   ├── preprocessing/
│   │   ├── load_data.py  ← NHANES file loading + SEQN merge
│   │   ├── build_target.py ← binary CVD label from MCQ160B/C/E
│   │   ├── clean.py      ← cohort selection, derivations, validity checks
│   │   ├── split.py      ← stratified 70/15/15 split
│   │   ├── impute.py     ← median/mode imputation + IQR outlier handling
│   │   └── build_dataset.py ← orchestration: raw → model-ready splits
│   ├── features/
│   │   ├── engineer.py   ← clinical ratios, categories, encoding, scaling
│   │   └── eda.py        ← EDA plots (missing values, correlations, etc.)
│   ├── models/
│   │   ├── train.py      ← LR, RF, XGBoost, LightGBM trainers
│   │   ├── tune.py       ← Optuna tuning for each model
│   │   ├── imbalance.py  ← class weights, SMOTE, combinations
│   │   ├── select.py     ← best-model selection + risk stratification
│   │   ├── export.py     ← joblib bundle save/load + predict_new_data()
│   │   └── mlflow_tracking.py ← MLflow experiment logging
│   ├── evaluation/
│   │   ├── metrics.py    ← full metric suite + bootstrap CIs
│   │   ├── plots.py      ← ROC, PR, confusion matrix, calibration plots
│   │   ├── framingham.py ← FRS 2008 vectorised implementation
│   │   └── statistical_tests.py ← DeLong's test, McNemar's test
│   └── explainability/
│       └── shap_explain.py ← SHAP global/local + permutation fallback
│
├── scripts/
│   └── train_all.py      ← full training orchestration CLI script
│
├── tests/
│   ├── test_framingham.py ← 8 tests incl. pandas-3.0 regression test
│   └── test_pipeline.py  ← 18 tests: target, split, impute, features, eval
│
├── reports/
│   ├── figures/          ← EDA + evaluation plots (auto-generated)
│   ├── model_comparison_test.csv
│   ├── delong_test_results.csv
│   └── risk_stratification.csv
│
├── models/
│   └── best_model.pkl    ← serialised best model bundle (auto-generated)
│
├── README.md
├── requirements.txt
└── setup.py

Limitations & Future Work

Known Limitations

  1. Self-reported outcomes: The CVD target is based on self-reported history (MCQ160B/C/E), not medical record adjudication. Recall bias and underdiagnosis in lower-access populations may introduce systematic error.

  2. Cross-sectional design: NHANES is cross-sectional, not longitudinal. Models predict prevalent CVD history, not incident 10-year risk, which is what Framingham was designed for. This is a meaningful conceptual distinction that limits direct Framingham comparability.

  3. Fasting subsample coverage: LDL, triglycerides, fasting glucose, and insulin were only measured in the fasting subsample (N ≈ 4,000 out of 11,933), introducing substantial structural missingness (~60%) imputed by median. This limits the inferential value of these features and inflates their imputed-median importance in tree models.

  4. Framingham hypertension proxy: Antihypertensive medication status (used by FRS for the treated-SBP coefficient) was approximated by self-reported hypertension history. This likely slightly overestimates the FRS predicted risk for treated patients.

  5. NHANES sampling weights: Complex survey weights (WTPH2YR, WTSAF2YR) were not applied in this analysis. Applying them would reweight estimates toward population-representative values but is beyond the scope of this ML comparison.

Future Work

  • Survival analysis: Replace binary CVD outcome with time-to-event analysis (Cox PH, DeepSurv) using linked mortality data.
  • Longitudinal validation: Validate on NHANES 2017–2020 as a held-out temporal test set.
  • Calibration improvement: Apply Platt scaling or isotonic regression to improve probability calibration (calibration curves show moderate overconfidence at high predicted probabilities).
  • Survey weight incorporation: Properly incorporate NHANES complex survey design for population-representative estimates.
  • Fairness analysis: Stratify performance metrics by race/ethnicity and sex to assess and address model disparities.
  • XGBoost/LightGBM with full optuna tuning: Run the full 50-trial Optuna search; preliminary grid search suggests XGBoost and LightGBM may improve on the Logistic Regression baseline.

Citation

If you use this code or results in your work, please cite:

@software{cvdrisk_nhanes_2021,
  title     = {Explainable Cardiovascular Risk Prediction using NHANES 2021-2023},
  year      = {2024},
  url       = {https://github.com/yourname/cardiovascular-risk-prediction},
  note      = {NHANES data from CDC/NCHS, Framingham coefficients from
               D'Agostino et al., Circulation 2008;117:743-753}
}

License

MIT License. See LICENSE for details.

NHANES data is in the public domain (CDC/NCHS). Framingham Risk Score coefficients are from a published peer-reviewed article; their use here is for research and educational purposes.

About

Explainable CVD risk prediction using NHANES 2021-2023 | XGBoost vs Framingham Risk Score | ROC-AUC 0.883 | SHAP · Optuna

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages