CogniGuard-ML: A Leakage-Resistant Explainable Automated Machine Learning Framework for Early Cognitive Function Decline Risk Prediction Using Population Health Data

Overview

CogniGuard-ML is a research-oriented machine learning framework for early cognitive function decline risk prediction using population health data. The repository implements a leakage-resistant pipeline that combines data auditing, harmonization, feature engineering, target validation, model selection, LightGBM optimization, probability calibration, threshold optimization, temporal validation, subgroup robustness analysis, explainability, and reproducibility reporting.

The implementation is organized as a set of runnable analysis stages so that each part of the workflow can be executed and audited independently while preserving the research methodology already encoded in the repository.

Key Features

• Dataset auditing and variable profiling for BRFSS-style input data
• Target validation across multiple years
• Dataset merging and feature inventory generation
• Leakage detection and leakage-resistance auditing
• Cross-validated AutoML model selection using development-only folds
• Leakage-resistant LightGBM hyperparameter optimization with Optuna
• Development-only probability calibration analysis
• Threshold optimization under a minimum recall constraint
• Locked temporal validation using a held-out temporal cohort
• TreeSHAP-based explainability analysis and feature mapping
• Subgroup robustness and uncertainty analysis
• Reproducibility manifests and experiment audit outputs

Repository Structure

• main.py – performs initial dataset verification, target analysis, dataset merging, feature inventory generation, and leakage reporting.
• preprocessing_runner.py – creates an ML-ready dataset from the cleaned and selected feature set.
• automl_runner.py – ranks development-only cross-validated candidate models and selects the best model.
• optimization_runner.py – tunes a LightGBM classifier with Optuna while preserving the held-out test set.
• calibration_runner.py – evaluates probability calibration using development-only out-of-fold predictions.
• threshold_optimization_runner.py – searches for a classification threshold that satisfies a minimum recall requirement.
• temporal_validation_runner.py – evaluates the locked temporal model and threshold on BRFSS 2024.
• shap_runner.py – generates TreeSHAP explainability outputs and feature mapping artifacts.
• subgroup_audit_runner.py – evaluates model performance across clinically relevant subgroups.
• figure_generation_runner.py – produces the publication-style figures stored in the results directory.
• reproducibility_runner.py – writes reproducibility and experiment audit artifacts.
• src/ – contains the core implementation modules for auditing, harmonization, preprocessing, feature selection, leakage analysis, optimization, training, validation, and explainability.
• results/ – stores generated reports, predictions, figures, explainability outputs, subgroup analysis artifacts, and reproducibility manifests.
• Dataset/ – intended for the raw BRFSS XPT files used by the workflow.
• paper/ – reserved for manuscript-related materials.

Installation

1. Clone the repository.
2. Create a Python environment (Python 3.14.4 was used in the current environment).
3. Install the required dependencies:

pip install -r requirements.txt

Usage

The workflow is organized as a set of runner scripts that should be executed from the repository root.

Basic entry point

python main.py

This script performs the initial dataset inspection, target assessment, multi-year target validation, dataset merging, feature inventory creation, and leakage audit setup.

Recommended analysis flow

Run the runner scripts in the order that matches the workflow stages:

python preprocessing_runner.py
python baseline_runner.py
python cv_baseline_runner.py
python automl_runner.py
python optimization_runner.py
python calibration_runner.py
python threshold_optimization_runner.py
python temporal_preprocessing_runner.py
python temporal_optimization_runner.py
python temporal_cv_baseline_runner.py
python temporal_validation_runner.py
python shap_runner.py
python subgroup_audit_runner.py
python subgroup_uncertainty_runner.py
python figure_generation_runner.py
python final_report_runner.py
python paper_table_runner.py
python reproducibility_runner.py

The exact ordering may vary depending on the local execution plan, but the repository is designed so that each runner writes its outputs into the results/ directory.

Results

The repository produces a structured set of outputs under results/, including:

• results/figures/ – publication-ready figures such as the framework overview, calibration diagrams, ROC comparisons, and subgroup analysis visuals.
• results/reports/ – model summaries, thresholds, optimization logs, baseline comparisons, and other experiment tables.
• results/explainability/ – TreeSHAP importance outputs, feature mappings, and analysis metadata.
• results/subgroup/ – subgroup performance, uncertainty, and bootstrap confidence interval outputs.
• results/reproducibility/ – environment details, package versions, artifact hashes, and reproducibility status files.
• results/final_report/ – curated tables for manuscript-style reporting.
• results/predictions/ – model predictions generated during development and temporal validation.

Dataset

The raw BRFSS datasets are intentionally excluded from this repository because of GitHub file size limitations. To reproduce the workflow locally, place the downloaded datasets in the Dataset/ directory using the expected file names:

• Dataset/2021.XPT
• Dataset/2022.XPT
• Dataset/2023.XPT
• Dataset/2024.XPT

Requirements

• Python: 3.14.4 was used in the current environment.
• Dependencies are listed in requirements.txt and can be installed with:

pip install -r requirements.txt

Citation

If you use this repository or its workflow in a publication, please cite the project as follows:

@software{sadik2026cogniguardml,
  title = {CogniGuard-ML: A Leakage-Resistant Explainable Automated Machine Learning Framework for Early Cognitive Function Decline Risk Prediction Using Population Health Data},
  author = {Sadik, SK Fahim},
  year = {2026},
  url = {https://github.com/your-username/CogniGuard-ML}
}

License

This project is licensed under the MIT License. See the LICENSE file for details.

Acknowledgements

This repository was prepared as a publication-oriented research artifact. The work builds on the open-source Python machine learning ecosystem and the publicly available BRFSS data source used in the analysis workflow.