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).
- Problem Statement
- Dataset
- Cohort Selection
- Methodology
- Feature Engineering
- Model Training
- Results
- Framingham Comparison
- Explainability
- Risk Stratification
- Reproducibility
- Project Structure
- Limitations & Future Work
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?
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.
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).
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).
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)
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.
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
| 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.
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.
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 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.
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%).
# 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/ -vPlace 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
All stochastic operations use random_state=42. The full pipeline is deterministic given the same input files and Python/package versions.
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
-
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.
-
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.
-
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.
-
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.
-
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.
- 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.
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}
}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.