RavshanjonEminov/microbiome-predictor
GI Cancer Microbiome Predictor
A research-prototype ML pipeline + Streamlit app that predicts tissue of origin among five gastrointestinal (GI) cancer types — HNSC, ESCA, STAD, COAD, READ — from microbiome abundance data, optionally combined with basic clinical metadata (age, sex, stage).
This is a research prototype, not a clinical decision-making tool. Predictions are not validated for diagnostic use.
Project structure
project/
├── app/ # Streamlit application
│ ├── main.py # Homepage — streamlit run app/main.py
│ ├── pages/
│ │ ├── 1_Data_Overview.py # Cohort sizes, class balance, clinical coverage
│ │ ├── 2_Model_Performance.py # ROC/PR/calibration curves, model comparison
│ │ ├── 3_Patient_Prediction.py # Clinical form + microbiome upload -> prediction
│ │ ├── 4_Explainability.py # SHAP global/local explanations, biomarkers
│ │ └── 5_About_Data.py # Datasets, methodology, limitations
│ ├── components/ # Reusable UI pieces
│ │ ├── charts.py # Plotly chart builders
│ │ ├── tables.py # Styled dataframe helpers
│ │ ├── forms.py # Patient input widgets
│ │ └── layout.py # Header / caveat banner / sidebar
│ └── utils/ # App-layer business logic
│ ├── io.py # Load trained artifacts / reports
│ ├── preprocessing.py # Align uploads to model feature space
│ ├── prediction.py # Run inference, format results
│ ├── explainability.py # SHAP wrapper for the app
│ └── validation.py # Input validation
├── config.py # ALL paths, hyperparameters, label maps live here
├── src/ # Core pipeline (importable + runnable as scripts)
│ ├── data_ingestion.py # Load real cohort files, or synthetic fallback
│ ├── data_cleaning.py # Prevalence filter, CLR transform, label harmonization
│ ├── feature_engineering.py # Microbiome-only / clinical-only / combined features
│ ├── train.py # Classical ensemble + optional Set-Transformer
│ ├── evaluate.py # Full metric suite + bootstrap CIs
│ └── interpret.py # SHAP global/local explanations, biomarkers
├── data/
│ ├── raw/ # Drop real cohort files here (see Datasets below)
│ └── processed/ # Cohort manifest, processed Parquet/CSV outputs
├── models/ # Saved fitted models, preprocessor, label encoders
├── reports/ # Metrics JSON, confusion matrices, biomarker CSVs
├── figures/ # Saved plots (if exported outside the app)
├── tests/ # pytest suite (65 tests across 5 files)
├── requirements.txt
└── README.mdSetup
pip install -r requirements.txtPyTorch is optional — the classical ensemble and the full app work without it. If it's missing, train.py logs a message and skips the Set-Transformer.
Usage
1. Add data (optional — runs on synthetic data otherwise)
Drop TCMA-format files into data/raw/:
- An abundance/relative-abundance table (
.csv/.tsv, samples × taxa or taxa × samples — auto-detected and transposed if needed). - A clinical/sample metadata sheet with at least a cancer-type column.
Get TCMA from the Duke Research Data Repository (Arguijo-Mendoza et al., 2022, DOI: 10.7924/R4BK1J35S).
If `data/raw/` is empty, every script and the app run on a clearly-labeled synthetic cohort so the whole pipeline is testable immediately. Every metric computed on synthetic data is flagged with an explicit warning — never silently presented as a real result.
2. Train
python src/train.py # full run: Optuna tuning + transformer
python src/train.py --no-tune # skip hyperparameter search (faster)
python src/train.py --no-tune --no-transformer # fastest smoke testTrains Logistic Regression, Random Forest, XGBoost, LightGBM, and a soft-voting ensemble on both the microbiome_only and combined (microbiome + clinical) feature sets, with SMOTE applied strictly inside CV folds. Saves fitted models, the fitted CLR preprocessor, and the held-out test split to models/.
3. Evaluate
python src/evaluate.pyWrites per-model metrics (macro-AUROC/PR-AUC with bootstrap 95% CIs, per-class sensitivity/specificity, calibration, adjacency-pair confusion) to reports/.
4. Interpret
python src/interpret.pyComputes SHAP global feature importance for the microbiome-only Random Forest and writes reports/top_biomarkers.csv, flagging taxa that match canonical literature CRC biomarkers.
5. Run the app
streamlit run app/main.py6. Run tests
pytest tests/ -vDatasets
Primary dataset citation
Dohlman, A. B., Iliev, I. D., Arguijo-Mendoza, D., Lipkin, S. M., Ding, S., Dressman, H., Shen, X., & Gao, M. (2020). Data from: The Cancer Microbiome Atlas (TCMA): A pan-cancer comparative analysis to distinguish organ-associated microbiota from contaminants [Dataset]. Duke Research Data Repository. https://doi.org/10.7924/R4RN36833
Please cite the above work if you use or build on results derived from the TCMA dataset. Data use is subject to the terms of the Duke Research Data Repository.
Methodology summary
- Preprocessing: prevalence filter (≥10% of samples) → top-200 most-variable taxa → centered log-ratio (CLR) transform. Clinical features are included only above an 80% non-null coverage threshold.
- No leakage: all imputation, scaling, one-hot encoding, and SMOTE resampling happen inside
sklearn/imblearnpipelines, fit only on training folds. - Models: Logistic Regression, Random Forest, XGBoost, LightGBM, soft-voting ensemble (Optuna-tuned, Platt-calibrated); optional Set-Transformer (ISAB + PMA) treating taxa as an unordered token set.
- Metrics: macro-AUROC (primary), PR-AUC, F1, per-class sensitivity/specificity, calibration, confusion matrix, decision threshold sweep, bootstrap 95% confidence intervals on every headline number.
- Interpretability: SHAP (
TreeExplainerfor tree models,LinearExplainerfor logistic regression) for global importance and per-patient local explanations; results cross-checked against canonical CRC biomarkers (Fusobacterium nucleatum, Parvimonas micra, Peptostreptococcus stomatis).
Known limitations
- Tumor-vs-adjacent-normal detection is near-chance (~0.57 AUROC). TCMA "normal" samples are tumor-adjacent tissue from the same patients, not independent healthy controls.
- COAD/READ and ESCA/STAD confusion is a documented biological ceiling from shared microbial ecology across anatomically adjacent GI segments — reported explicitly rather than tuned away.
- Batch effects make naive multi-cohort pooling unreliable; leave-one-cohort-out (LODO) validation is the planned mitigation once a second cohort is integrated.
- Synthetic-data runs are smoke tests only and are never presented as performance estimates (see the in-app warning banners).
- Not validated on an independent clinical cohort; not for diagnostic use.
License / data use
TCMA and other third-party datasets carry their own usage terms — consult the source repository before redistribution.
