ML AMR Prediction Framework v2 — Alignment-Free WGS Out-of-Core Learning

Abstract

A scalable, multi-organism / multi-antibiotic machine-learning pipeline that predicts Antimicrobial Resistance (AMR) directly from Whole Genome Sequencing (WGS) assemblies. It uses an alignment-free k-mer representation (no reference genome) and out-of-core XGBoost to train over tens of millions of features on commodity hardware, then reverse-translates the most important k-mers into biology via BLAST (CARD + NCBI). The current dataset is Escherichia coli across four antibiotics (ampicillin, cefotaxime, ciprofloxacin, gentamicin); adding a new organism is a registry entry, not a code change.

---

Three core pillars

1. Computer science / informatics

• K-mer extraction with KMC (canonical 21-mers, compressed binary DBs).
• Out-of-core learning: dynamic stratified chunking + compressed .npz sparse matrices; training streams chunks via XGBoost DMatrix, so matrices far larger than RAM are handled on a laptop.

2. Mathematics

• Bayesian HPO with Optuna.
• colsample_bytree search anchored to the 1/√p square-root heuristic (see METHODOLOGY.md).
• Unbiased thresholding: the operating threshold is derived from the training distribution (no test-set leakage); evaluation reports bootstrap 95% confidence intervals.

3. Biology

• Reference-free discovery: top k-mers are BLASTed against CARD (local) and NCBI nt (remote) via Nextflow, then graded into confirmed / candidate / weak confidence tiers.

---

Repository structure

ML_AMR_Prediction_v2/
├── config/
│   ├── config.yaml              # global config (paths, params, organism slug)
│   ├── registry/                # single source of truth
│   │   ├── organisms.yaml        #   organism → data paths + antibiotic set
│   │   └── antibiotics.yaml      #   antibiotic → class mapping
│   └── experiments/{organism}/config_{antibiotic}.yaml  # auto-generated by step 04 (split + best params)
├── scripts/
│   ├── 01..09_*.py              # the numbered pipeline steps
│   ├── 08_blast_pipeline.nf     # Nextflow BLAST workflow (step 08)
│   ├── lib/                     # shared package (no duplicated code)
│   │   ├── config.py            #   load_config + resolve_path (organism-aware)
│   │   ├── registry.py          #   organisms/antibiotics access
│   │   ├── chunking.py          #   get_y_chunk
│   │   ├── io_utils.py          #   run_command (shlex, never shell=True)
│   │   └── run_metadata.py      #   git hash / versions / run_id capture
│   ├── constants.py, utils.py   # thin backward-compat shims -> lib/
│   └── migrate_to_organism_layout.py  # reversible data-layout migration
├── tests/                       # pytest suite (smoke / unit / integration)
├── docs/
│   ├── TECHNICAL_REVIEW.md      # consolidated review + remediation record
│   ├── SCALE_MLOPS_PLAN.md      # multi-organism + KB + MLOps plan
│   └── ROADMAP.md               # thesis roadmap
├── data/                        # raw genomes, metadata, CARD DB, matrices
├── models/  results/  logs/  runs/   # generated (results/logs/runs gitignored)
├── config files: requirements.txt, environment.yml, pytest.ini
└── README.md, QUICKSTART.md, METHODOLOGY.md

Data layout (organism-scoped)

Paths are resolved through scripts/lib/config.py:resolve_path() from the {organism} templates in config.yaml:

• data/raw/{organism}/genomes/ — immutable .fna assemblies
• data/external/{organism}/metadata/amr_phenotypes.csv — phenotype labels
• data/external/blast_db/card_nt/ — CARD BLAST DB (shared)
• data/interim/{organism}/kmc_outputs/ — KMC binary DBs
• data/processed/{organism}/{antibiotic}/matrix/ — .npz chunks, y_*.csv
• models/{organism}/{antibiotic}/ — model + manifest.json
• results/{organism}/{antibiotic}/ — plots + metric CSVs (generated)
• runs/{organism}/{antibiotic}/{run_id}/ — run metadata + metrics (generated)

Generated outputs (results/, logs/, runs/, models, .npz, large features.txt) are not version-controlled — they are reproduced by running the pipeline.

---

Pipeline workflow

Step	Script	Description
01	01_data_validation.py / 01b_…	Validate metadata, class balance, EDA plots
02	02_kmer_extraction.py / 02b_…	KMC k-mer counting + global QC
03	03_matrix_construction.py / 03b_…	Sparse binary k-mer matrices + matrix QC
03u	03u_unitig_matrix.py	Unitig (unitig-caller) presence/absence matrix — same chunk contract, matrix_unitig/ (ROADMAP §0 M12)
04	04_optimization.py	Optuna Bayesian HPO → config_{antibiotic}.yaml + run metadata
05	05_model_training.py	Full-data XGBoost boosting over a streaming QuantileDMatrix → model + manifest.json
06	06_evaluation.py	Metrics, ROC/PR, calibration, bootstrap CIs
07	07_explainability.py	Gain-based top k-mers → CSV + FASTA
08	08_blast_annotation.py / .nf	BLAST top k-mers vs CARD + NCBI (Nextflow)
09	09_biological_summary.py	Confidence-tiered biological report (Markdown)

---

Results (E. coli, unitig features, 3 antibiotics)

Lineage-aware cross-validation ROC-AUC (PopPUNK + StratifiedGroupKFold — the honest, non-inflated estimate; an entire lineage stays on one CV side) alongside the single-split test AUC. The KB currently holds three antibiotics spanning the two AMR mechanism classes (acquired gene vs target-SNP).

Antibiotic	Mechanism (confirmed)	Lineage-CV ROC-AUC	Single-split AUC
Ampicillin	acquired gene — TEM β-lactamase	0.951 ± 0.011	0.924
Cefotaxime	acquired gene — CTX-M / CMY ESBL/AmpC	0.955 ± 0.020	0.969
Ciprofloxacin	target-gene SNP — gyrA S83L + parC S80I	0.950 ± 0.007	0.980

Validation done (see docs/ROADMAP.md / python scripts/kb_report.py): lineage-aware CV (M2), CPSS stability selection + PFER bound (M4), permutation tests (M9), pyseer LMM population-structure-corrected significance (M14), CARD variant-model SNP allele check (M11), genome QC (CheckM2 + QUAST, 97.1% pass, M15), and external validation — AMRFinderPlus/ResFinder concordance + a leakage-free model-vs-tool head-to-head on held-out genomes (M13), where the unitig model matches ResFinder and beats AMRFinderPlus for ciprofloxacin. Remaining: Zenodo DOI (M10) and an optional temporal/geographic hold-out.

---

Installation

Python: 3.10+. External tools (KMC, BLAST+, Nextflow) are not pip-installable — the conda route below installs everything in one step.

Option A — conda (recommended; includes KMC/BLAST/Nextflow)

conda env create -f environment.yml
conda activate amr-prediction

Option B — pip (Python deps only; install tools separately)

python -m venv .venv && source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -r requirements.txt

Then install the external tools (conda recommended: conda env create -f environment.yml installs them all): KMC ≥3.2, BLAST+ ≥2.12, Nextflow ≥22.10. Tools are found on PATH (conda/module); a macOS binary under bin/bin/ is used only as a fallback.

---

Quickstart

# 1) choose target in config/config.yaml (project.organism + project.target_antibiotic)
# 2) acquire data, then run the steps in order
python scripts/00a_download_bvbrc.py --organism ecoli   # BV-BRC AMR data + assemblies
python scripts/00_prepare_metadata.py --organism ecoli  # binary phenotype matrix
python scripts/01_data_validation.py
python scripts/02_kmer_extraction.py
python scripts/02b_global_qc_analysis.py
python scripts/03_matrix_construction.py
python scripts/04_optimization.py
python scripts/05_model_training.py
python scripts/06_evaluation.py
python scripts/07b_feature_stability.py   # 5-seed stability (run before 07)
python scripts/07_explainability.py       # gain top-N ∪ stable set -> CSV + FASTA
python scripts/08_blast_annotation.py     # needs BLAST+ / Nextflow / CARD DB
python scripts/09_biological_summary.py   # tiered report; set ncbi.entrez_email or AMR_ENTREZ_EMAIL
python scripts/10_kmer_background_frequency.py  # resistant-vs-susceptible discriminativeness
python scripts/11_variant_snp_check.py    # CARD variant-model SNP allele check (optional)
python scripts/12_permutation_test.py     # M9: MDA permutation importance (+ BH-FDR)
python scripts/12b_label_permutation_test.py  # M9: label-permutation null (model significance)
python scripts/13_stability_selection.py  # M4: CPSS stability selection (PFER bound) + TreeSHAP
python scripts/13b_stable_annotation.py   # M4: CARD-annotate the stable set (tier + ARO)
python scripts/14_pyseer_lmm.py --mode prep   # M14: pyseer LMM (population-structure-corrected); pyseer CLIs run via amr-tools.sif
python scripts/populate_database.py --antibiotic <ab>   # M8: load everything -> results/{org}/kb/amrk.db (multi-antibiotic, unitig schema 0.4.0)
streamlit run scripts/kb_app.py           # local queryable KB explorer (pip install streamlit pandas)
uvicorn scripts.kb_api:app                # REST API (S8/S9): /api/v1/{kmers,kmers/{seq},overlap,stats,metadata} + /docs

Querying the KB via the REST API (scripts/kb_api.py, FastAPI):

pip install fastapi uvicorn
AMR_KB_DB=results/ecoli/kb/amrk.db uvicorn scripts.kb_api:app   # http://localhost:8000/docs
# GET /api/v1/kmers?antibiotic=cefotaxime&tier=confirmed&stable_only=true
# GET /api/v1/kmers/{sequence}     GET /api/v1/overlap?ab1=ampicillin&ab2=cefotaxime
# GET /api/v1/stats                GET /api/v1/metadata   (FAIR: schema ver, DOI, license)

Feature unit = unitig (compacted de Bruijn graph; ROADMAP §0). The pipeline runs per antibiotic (set project.target_antibiotic in config.yaml); the KB is organism-level and multi-antibiotic. Each run is reproducible — pipeline_runs stamps the git commit, seed, config hash and CARD version.

See QUICKSTART.md for prerequisites, HPC notes and expected outputs.

Testing

A layered suite lets you validate changes in seconds instead of rerunning the multi-day pipeline (see tests/README.md):

pytest                 # fast: smoke + unit (seconds); never touches config
pytest -m integration  # opt-in end-to-end on a tiny synthetic dataset (minutes)

Orchestration & developer workflow

make help                          # list all targets
make dev-install                   # install the QA toolchain (ruff/mypy/pre-commit) + pre-commit hooks
make lint                          # ruff
make test                          # unit + smoke
python scripts/run_pipeline.py --list                         # show the step plan
python scripts/run_pipeline.py --organism ecoli --antibiotic ampicillin   # run the analysis core 01->10

CI (GitHub Actions, .github/workflows/ci.yml) runs ruff + the unit/smoke suite on Python 3.10–3.12.

Reproducibility & data

The repository tracks only code, configuration, docs, and the small CARD homolog BLAST DB. All datasets, k-mer matrices, models, results and run metadata are generated and reproduced by running the pipeline (see QUICKSTART.md). Each run records its git commit, library versions and seed in runs/.../run_metadata.json. Genome data is obtained from BV-BRC; resistance annotations from CARD.

License & citation

The software is released under the MIT License. If you use it, please cite it via CITATION.cff. Contributions are welcome — see CONTRIBUTING.md and CHANGELOG.md.

Data availability & KB versioning (FAIR)

The AMRK-DB knowledge base (results/{organism}/kb/amrk.db) is released under CC-BY-4.0 (distinct from the MIT-licensed code) and versioned with a semantic kb_schema_version (currently 0.4.0, stamped in the kb_metadata table). Every KB record links to a pipeline_runs row carrying the git commit, random seed, config hash and CARD version, so any biomarker is reproducible.

Each release is archived on Zenodo with a persistent DOI (concept DOI resolves to the latest version; each version also gets its own DOI). The DOI is stored in kb_metadata.zenodo_doi and mirrored in CITATION.cff / .zenodo.json.

Zenodo DOI: reserved — added on first deposit (see docs/RELEASE_ZENODO.md).

Raw genome assemblies come from BV-BRC; resistance gene annotations from CARD (v4.0.1). The pipeline regenerates the unitig matrices from the BV-BRC assemblies, so the FAIR deposit ships the KB + evidence tables + run metadata rather than the large raw matrices.