Metadata-Version: 2.4
Name: mb-cdar
Version: 0.3.2
Summary: MB-CDAR: Microbiome data generation via MB-DDPM (2D UNet + image_coding + discrete DDPM) with a VAE conditional encoder (channel-concatenated latent + CFG)
License-Expression: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.24
Requires-Dist: torch>=2.0
Requires-Dist: einops>=0.6
Requires-Dist: biopython>=1.80
Requires-Dist: pandas>=2.0
Requires-Dist: scipy>=1.10
Provides-Extra: eval
Requires-Dist: matplotlib>=3.6; extra == "eval"
Requires-Dist: seaborn>=0.12; extra == "eval"
Requires-Dist: scikit-learn>=1.2; extra == "eval"
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Requires-Dist: mb-cdar[eval]; extra == "test"
Provides-Extra: dev
Requires-Dist: mb-cdar[test]; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: mypy<1.12,>=1.11; extra == "dev"
Dynamic: license-file

<div align="center">

# MB-CDAR

**Microbiome Conditional Diffusion with Adversarial Regularization** — a conditional denoising diffusion model with adversarial regularization for high-fidelity, zero-inflated, compositional microbiome data

A conditional generative framework for high-fidelity, zero-inflated, compositional microbiome data — built on a 2D-UNet denoising diffusion backbone with a VAE conditional encoder.

[![Python](https://img.shields.io/badge/Python-3.10%2B-3776AB?logo=python&logoColor=white)](https://www.python.org/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.0%2B-EE4C2C?logo=pytorch&logoColor=white)](https://pytorch.org/)
[![Version](https://img.shields.io/badge/version-0.3.1-blue)](pyproject.toml)
[![PyPI](https://img.shields.io/pypi/v/mb-cdar?logo=pypi&logoColor=white)](https://pypi.org/project/mb-cdar/)
[![License](https://img.shields.io/badge/license-MIT-green)](#license)
[![Tests](https://img.shields.io/badge/tests-60%20passed-brightgreen)](tests/)

</div>

---

## Overview

**MB-CDAR** generates synthetic 16S microbiome abundance profiles that faithfully reproduce the
distributional, compositional, and phylogenetic structure of a real cohort. Microbiome OTU tables are
notoriously hard to simulate: they are **highly sparse** (~82% zeros), **compositional** (rows sum to 1),
and exhibit **heavy-tailed** rank-abundance structure. Classical simulators handle these properties with
hand-tuned parametric models; MB-CDAR instead learns them end-to-end with a **denoising diffusion
probabilistic model (DDPM)** operating in a 2D *image-coded* representation, conditioned on a **VAE
community-shape latent** injected by channel-concatenation with **classifier-free guidance (CFG)**.

The production architecture (**v3.1**) extends the diffusion backbone with a spectral-normalized
patch discriminator, a marginal-correlation auxiliary loss, and a group-specific zero-inflation
frequency-matching post-processor — while keeping the UNet's primary objective a pure noise-prediction MSE.

<p align="center">
  <img src="assets/diffusion_demo.gif" width="720" alt="Reverse diffusion: the model's predicted sample sharpens as the noise is removed step by step"/>
  <br/>
  <em>Reverse diffusion from pure Gaussian noise. Each frame shows the model's predicted final sample
  (<span>x̂</span><sub>0</sub>) at a denoising step <em>t</em>, conditioned on one real sample's VAE latent.
  Left: the image-coded 40×32 representation; right: the emerging rank-abundance curve vs. the real sample.</em>
</p>

---

## Key results

Benchmarked on the real **Gordon cohort** (281 samples × 1033 taxa; 220 case / 61 ctrl) against two
established statistical simulators (**MIDASim**, **SparseDOSSA 2**) and two pooled generative controls
(**MB-DDPM**, **MB-GAN**), MB-CDAR (v3.1) attains **the highest per-taxon correlation fidelity of all
compared simulators**, is **the only method whose Shannon-diversity distribution is statistically
indistinguishable from the real cohort in every dataset × group** (MWU *p* > 0.05 throughout), and keeps
the Bray-Curtis / UniFrac distance-distribution discrepancies at or near the best of the field.
All numbers below are from the 2026-09-20 full re-run under the final evaluation convention
(permutation KS *p*-values, scipy-aligned tie handling, degenerate-taxon backfill for the R baselines).

| Metric                                      | MB-CDAR case  | MB-CDAR ctrl  | MIDASim case | MIDASim ctrl | SparseDOSSA2 case | SparseDOSSA2 ctrl | MB-DDPM case | MB-DDPM ctrl | MB-GAN case  | MB-GAN ctrl  |
| ------------------------------------------- | -------------- | -------------- | ------------ | ------------ | ----------------- | ----------------- | ------------ | ------------ | ------------ | ------------ |
| **Spearman ρ² (per-taxon mean)**            | **0.9342**     | **0.9939**     | 0.8745       | 0.6544       | 0.8371            | 0.6442            | 0.7933       | 0.3978       | 0.5087       | 0.4065       |
| **Pearson r**                               | **0.997**      | **0.999**      | 0.991        | 0.639        | 0.989             | 0.711             | 0.995        | 0.968        | 0.980        | 0.958        |
| **Shannon (real 4.048 / 4.034; MWU)**       | **4.020** (ns) | **4.002** (ns) | 3.959 (ns)   | 2.512 (sig)  | 3.738 (sig)       | 2.405 (sig)       | 4.068 (ns)   | 4.081 (ns)   | 4.081 (ns)   | 4.051 (ns)   |
| **Sparsity (real 0.880 / 0.888)**           | 0.882          | 0.890          | 0.880        | 0.931        | 0.888             | 0.933             | 0.886        | 0.885        | 0.906        | 0.907        |
| **Mean max abundance (real 0.152 / 0.151)** | **0.1516**     | **0.1507**     | 0.1732       | 0.4746       | 0.2127            | 0.4846            | 0.1570       | 0.1576       | 0.2325       | 0.2353       |
| **BC KS-D (↓)**                             | 0.049 (ns)     | **0.035** (ns) | 0.038 (ns)   | 0.573 (sig)  | 0.226 (sig)       | 0.633 (sig)       | 0.023 (ns)   | 0.024 (ns)   | 0.343 (sig)  | 0.293 (sig)  |
| **Weighted-UniFrac KS-D (↓)**               | **0.014** (ns) | 0.048 (ns)     | 0.197 (sig)  | 0.551 (sig)  | 0.226 (sig)       | 0.534 (sig)       | 0.031 (ns)   | 0.042 (ns)   | 0.107 (sig)  | 0.089 (sig)  |
| **Unweighted-UniFrac KS-D (↓)**             | 0.071 (ns)     | 0.106 (ns)     | 0.098 (sig)  | 0.104 (sig)  | 0.542 (sig)       | 0.401 (sig)       | 0.472 (sig)  | 0.345 (sig)  | 1.000 (sig)  | 1.000 (sig)  |

Three-dataset summary (MB-CDAR, n = 1000 synthetic samples per group; full tables in
`script_test/results*/summary.txt`):

| Dataset                          | MB-CDAR Spearman ρ² case / ctrl | MB-DDPM ρ² case / ctrl | MB-GAN ρ² case / ctrl | Paired Mantel BC / WUF | Shannon MWU case / ctrl |
| -------------------------------- | -------------------------------- | ---------------------- | --------------------- | ---------------------- | ----------------------- |
| OB / Gordon (281 × 1033)         | 0.934 / 0.994                    | 0.793 / 0.398          | 0.509 / 0.407         | 0.999 / ≈1.000         | ns / ns                 |
| IBD / HMP2 biopsy (178 × 959)    | 0.993 / 0.999                    | 0.712 / 0.427          | 0.416 / 0.314         | 0.999 / ≈1.000         | ns / ns                 |
| HIV / Noguera-Julian (240 × 822) | 0.941 / 0.995                    | 0.893 / 0.514          | 0.646 / 0.521         | ≈1.000 / ≈1.000        | ns / ns                 |

**Ablation evidence** (9 arms on the OB cohort, `Ablation Experiment/`; table + cross-arm figures):
the Diffusion-GAN discriminator materially improves the unweighted-UniFrac distance distribution
(KS D 0.096 / 0.121 → 0.072 / 0.107, case / ctrl) at unchanged per-taxon fidelity; group-specific
zero-inflation frequency matching is indispensable (disabling it inflates the UWUF dispersion ratio
to 1.27 / 1.16); the VAE condition is essential (unconditional sampling collapses Spearman ρ² to
0.731 / 0.642); CFG strength w ∈ {0, 1, 2, 3} has only a marginal effect (default w = 2).

**Conditional geometry.** For MB-CDAR alone, each synthetic sample is conditioned on the latent of a
specific real sample, enabling a 1:1 **paired Mantel test** on within-group distance matrices:
**BC *r* = 0.999, weighted-UniFrac *r* ≈ 1.000** (*p* = 0.001) — direct evidence that conditional
generation preserves pairwise sample-to-sample geometry.

<p align="center">
  <img src="script_test/results/figures/correlation_tests_summary.png" width="760" alt="Correlation-test square matrix across simulators"/>
  <br/><em>Per-group fidelity vs. real data across simulators: a corrplot-style square matrix of the correlation tests (square area and colour = signed value / row max, stars = significance) for case / ctrl / pooled all.</em>
</p>

<p align="center">
  <img src="script_test/results/figures/rank_abundance.png" width="480" alt="rank abundance"/>
  <img src="script_test/results/figures/spearman_r2_bars.png" width="480" alt="spearman r2 bars"/>
  <br/><em>Rank-abundance curves (left) and per-taxon Spearman ρ² by group (right).</em>
</p>

---

## Architecture

```
raw counts (B, 1033)
   │
   ├─ CountNormalizer   count → proportion → per-feature z-score       [Architecture Decision Record 019]
   │     ├─ VAE encoder → latent (B, 256) ──broadcast──► (B, 256, 40, 32)
   │     └─ image_coding → (B, 1, 40, 32) ──concat──► (B, 257, 40, 32)
   │                                        │
   │                        DDPM forward / UNet ε-prediction / MSE
   │                        + [v3.1] PatchDiscriminator + MarginalCorrLoss
   │                                        │
   │                        DDPM reverse (T=1000, CFG extrapolation)
   │                                        │
   │     CountNormalizer.inverse_transform (clamp ≥ 0, row-normalize)
   │     + group-specific zero-inflation frequency matching            [Architecture Decision Record 026]
   │                                        │
   └──────────────► synthetic proportions (rows sum to 1, non-negative)
```

**Design principles** (each grounded in a recorded architecture decision, see `.agents/memory/decisions.md`):

- **Normalized proportion space** — the model trains on per-feature *z-scored proportions*, not raw counts
  nor ILR/softmax. This matches the DDPM scheduler's `N(0,1)` assumption while remaining a linear per-feature
  rescale (not a compositional transform).
- **Image coding** — 1033 features are reshaped into a 40×32 2D image so 2D convolutions capture local
  taxon co-occurrence; the transform is losslessly invertible.
- **VAE channel-concatenation conditioning + CFG** — the 256-d community latent is broadcast over the spatial
  plane and concatenated as extra channels (input 257 ch). Training uses 10% CFG dropout; inference applies
  extrapolation `pred = pred_uncond + (1+w)·(pred_cond − pred_uncond)` with `w = cfg_scale` (default 2.0).
- **Discrete DDPM schedule** — standard Ho et al. β = linspace(1e-4, 0.02, T=1000); **not** altered.
- **Two-stage training** — Stage 1 VAE pretrain (10k epochs) + Stage 2 UNet DDPM (150k epochs, EMA 0.9999).
- **v3.1 adversarial & marginal terms** — spectral-norm patch discriminator (random-*t* noised inputs,
  Diffusion-GAN style); the BC/Aitchison pairwise-distance losses are auto-zeroed (they traded per-taxon
  fidelity for β-geometry — a v3 regression culprit); `MarginalCorrLoss` = 1 − Pearson correlation between
  batch marginal and the fixed real per-taxon marginal, `α̂`-weighted.
- **Group-specific zero-inflation matching** — post-sampling, per taxon the smallest surplus non-zeros are
  zeroed so the batch non-zero rate matches the observed group frequency. This removes the dense tiny-positive
  tail without degrading (it *raises*) Spearman ρ².

---

## Installation

### From PyPI (recommended for users)

In a fresh Python 3.10+ environment:

```bash
pip install mb-cdar
```

This installs the `mb_cdar` package into the environment's `site-packages` and places the
`mb-cdar` command on the PATH (`Scripts\mb-cdar.exe` on Windows, `bin/mb-cdar` on Linux/macOS).
Runtime dependencies (numpy, torch, einops, biopython, pandas, scipy) are installed automatically.

The evaluation pipeline (metrics + nine figures) needs an optional plotting stack. Install the
`eval` extra when you plan to run `mb-cdar eval`:

```bash
pip install "mb-cdar[eval]"      # + matplotlib, seaborn, scikit-learn
```

The computation helpers import without this extra; the plotting libraries are loaded lazily only
when figures are actually rendered.

### From source (for development)

```bash
git clone https://github.com/MKSW1/MB-CDAR.git
cd MB-CDAR
python -m venv .venv
.venv\Scripts\activate          # Windows (cmd / PowerShell)
source .venv/bin/activate        # Linux / macOS

pip install -e .                # core (numpy, torch, biopython, pandas, scipy)
pip install -e ".[eval]"        # + matplotlib/seaborn/scikit-learn for `mb-cdar eval`
pip install -e ".[test]"        # + pytest for the test suite (pulls in the eval extra)
```

GPU builds: see `requirements_CUDA.txt` (NVIDIA) or `requirements_XPU.txt` (Intel GPU).

---

## Quick start

The `mb-cdar` CLI has three subcommands: `train`, `sample`, and `eval`.

**Train** (two-stage VAE pretrain + UNet DDPM, v3.1):

```bash
mb-cdar train \
  --otu_path data/otu.tab.csv \
  --arch v3.1 \
  --vae_epochs 10000 --unet_epochs 150000 \
  --cfg_scale 2.0 --device cuda
```

`python -m mb_cdar train ...` is equivalent. The `train` subcommand forwards every flag
verbatim to the training module, so all options in `mb-cdar train --help` are available.

**Sample** synthetic profiles from a trained checkpoint (1000 case samples conditioned on the
real group; the group's own presence-frequency prior is applied so taxa never observed in the
group stay exact zeros):

```bash
mb-cdar sample \
  --ckpt output_v31/mb_cdar_unet_best.pt \
  --data-dir data --group case --case-states OW,OB \
  --n 1000 --out case_samples.csv
```

For a pure prior draw with no VAE condition: `mb-cdar sample --ckpt <ckpt> --unconditional --out prior.csv`.

**Resume a crashed long run** (restores weights + epoch counter + global step + discriminator state,
skips the VAE stage) — mutually exclusive with `--resume_ckpt` (weights-only warm start):

```bash
mb-cdar train --otu_path data/otu.tab.csv --arch v3.1 \
  --resume output_v31/mb_cdar_unet_epoch_145000.pt \
  --unet_epochs 150000 --device cuda
```

**Evaluate** against the R baselines (generates case/ctrl samples, computes Python metrics and
renders the nine figures). The evaluation engine ships inside the package and is exposed through
the `eval` subcommand — every evaluation flag is forwarded verbatim:

```bash
mb-cdar eval \
  --n-case 1000 --n-ctrl 1000 \
  --device cuda --cfg-scale 2.0 \
  --ckpt output_v31/mb_cdar_unet_best.pt
```

Fast smoke test (no training, no R):

```bash
mb-cdar eval --fast --skip-r \
  --ckpt output_v31/mb_cdar_unet_best.pt
```

`python script_test/run_evaluation.py ...` is an equivalent compatibility shim that forwards to
the same packaged engine. When the package is installed outside a source checkout, no dataset is
bundled: pass `--data-dir` (or explicit `--otu-path/--meta-path/--tree-path`) and outputs default
to a CWD-relative `mb_cdar_eval_results/` folder.

**Run the test suite:**

```bash
pytest tests/ -v        # 60 tests
```

---

## Repository layout

```
MB-CDAR/
├── src/mb_cdar/
│   ├── backbone/        # UNet (2D DDPM) + VAE + PatchDiscriminator
│   ├── data/            # normalizer, image_coding, composition/phylogeny (eval-only)
│   ├── diffusion/       # DDPMScheduler (discrete DDPM)
│   ├── evaluation/      # evaluation engine + nine-figure plotting (`mb-cdar eval`),
│   │                    #   Rdata->CSV R helper shipped as package data
│   ├── loss/            # ecology_loss (MarginalCorrLoss; BC/Aitchison auto-zeroed)
│   ├── training/        # Trainer (two-stage, --arch, --resume) + Sampler + tail_calib
│   ├── cli.py           # `mb-cdar` CLI (train/sample/eval subcommands)
│   ├── __main__.py      # `python -m mb_cdar`
│   └── train.py         # training entry point
├── script_test/         # evaluation shim + R baseline generators (vs. MIDASim / SparseDOSSA2
│   └── baselines/       # pooled MB-DDPM / MB-GAN control wrappers (Architecture Decision Record 035)
├── output_v31/          # v3.1 production checkpoint (--ckpt target; .pt files are 0-byte
│                       #   placeholders in git, real weights stay local-only via skip-worktree)
├── Ablation Experiment/ # ablation protocol, cross-platform drivers, ablation comparison figures, per-dataset data (Architecture Decision Record 033/036)
├── tests/               # unit + end-to-end tests (60)
├── docs/                # architecture, training/testing, dataset & literature notes
├── .agents/              # agent-framework-agnostic project memory & conventions
│   ├── memory/          # project_brief, active_context, progress, decisions (Architecture Decision Records)
│   └── skills/          # project-core-conventions (architecture bottom lines, SOPs)
├── assets/              # README figures & the reverse-diffusion animation
└── data/                # Gordon OTU table + metadata + phylogenetic tree
```

---

## Reproducibility notes

- **Checkpoint contract** — UNet checkpoints embed `normalizer_state`, `cfg_scale`, `arch`,
  `tail_calib_state`, `taxon_freq_state`, and resume keys (`global_step`, `best_loss`,
  `best_loss_gan`, `disc_state_dict`, `epoch`). Checkpoints lacking `normalizer_state` or with
  `data_scale != "zscore_proportion"` are explicitly rejected by the evaluation script.
- **β-diversity testing** — cross-matrix Mantel (real-real vs. synth-synth with no sample correspondence)
  was removed as statistically invalid (Architecture Decision Record 027); the toolkit now uses distribution-level KS + dispersion
  ratio on within-group pairwise distances (all simulators) and a paired Mantel for MB-CDAR only.
  KS p-values come from a 999-permutation label-shuffle test because within-group pairwise
  distances are not independent — asymptotic KS p-values are invalid there (Architecture Decision Record 032).
- **Long-training stability** — Windows 150k-epoch runs follow strict memory rules
  (numpy RNG, no tqdm, periodic `gc.collect`/`empty_cache`, atomic checkpoint saves); see
  `.agents/skills/project-core-conventions/SKILL.md`.
- **Three-group evaluation** — every evaluation now reports case / ctrl / the pooled
  all slot (case+ctrl concatenated) for all simulators, and the pooled-trained
  MB-DDPM / MB-GAN control models join the comparison via `--controls-dir`
  (their CSVs are produced by `script_test/baselines/` wrappers; Architecture Decision Record 035).
- **Ablation figures** — the cross-arm comparison figures (arm x (group x metric)
  matrix, training-level component-attribution ladder, inference-level sweeps) are
  rendered by `Ablation Experiment/scripts/plot_ablation_figures.py` from the
  ablation table or the per-arm results, in the same style as the main evaluation
  figures (Architecture Decision Record 036).
- **Figure style** — the nine evaluation figures share one centralised style constant
  (sans-serif, white background, no gridlines, thin black frames, inward ticks):
  row=group × column=simulator identity and alpha-diversity grids, a corrplot-style
  correlation-test square matrix (square area and colour = signed value / row max), marginal-density PCoA
  biplot and RdBu_r heatmaps (Architecture Decision Record 035).

---

## References

1. **MB-DDPM backbone** — Huo, B. et al. *MB-DDPM: microbiome simulation based on the denoising
   diffusion probabilistic model.* **Information Fusion** 127, 103880 (2026).
   doi:[10.1016/j.inffus.2025.103880](https://doi.org/10.1016/j.inffus.2025.103880)
2. **DDPM** — Ho, J., Jain, A. & Abbeel, P. *Denoising diffusion probabilistic models.* NeurIPS (2020).
3. **Classifier-free guidance** — Ho, J. & Salimans, T. *Classifier-free diffusion guidance.* (2022).
4. **MIDASim** — He, M., Zhao, N. & Satten, G. A. *MIDASim: a fast and simple simulator for realistic
   microbiome data.* **Microbiome** (2024). doi:[10.1186/s40168-024-01822-z](https://doi.org/10.1186/s40168-024-01822-z)
5. **SparseDOSSA 2** — Huttenhower Lab. *SparseDOSSA 2: a computational framework for simulating
   microbiome data.* (see also Ma, S. et al., **PLOS Comput. Biol.** 2021,
   doi:[10.1371/journal.pcbi.1008913](https://doi.org/10.1371/journal.pcbi.1008913)).
6. **Real cohort** — the Gordon-lab obesity 16S cohort used by Ref. 1 (281 samples × 1033 OTUs;
   220 overweight/obese case vs. 61 healthy ctrl), bundled under [`data/`](data/).

> Project-specific design decisions, regression post-mortems (v3 → v3.1), and evaluation-statistics
> corrections are recorded as Architecture Decision Records in
> [`.agents/memory/decisions.md`](.agents/memory/decisions.md).

---

## License

Released under the **MIT License** — see [`pyproject.toml`](pyproject.toml). The bundled Gordon dataset
under [`data/`](data/) is provided for research reproducibility; please cite the original study (Ref. 6)
when reusing it.

## Citation

If you use MB-CDAR in your research, please cite the MB-DDPM backbone (Ref. 1) and this repository.
