mb-cdar 0.3.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mb_cdar-0.3.1/LICENSE +21 -0
- mb_cdar-0.3.1/PKG-INFO +350 -0
- mb_cdar-0.3.1/README.md +324 -0
- mb_cdar-0.3.1/pyproject.toml +130 -0
- mb_cdar-0.3.1/setup.cfg +4 -0
- mb_cdar-0.3.1/src/mb_cdar/__init__.py +3 -0
- mb_cdar-0.3.1/src/mb_cdar/__main__.py +6 -0
- mb_cdar-0.3.1/src/mb_cdar/_archive_v2/backbone/mb_cdar_backbone.py +379 -0
- mb_cdar-0.3.1/src/mb_cdar/_archive_v2/data/augmentation.py +326 -0
- mb_cdar-0.3.1/src/mb_cdar/_archive_v2/diffusion/schedules.py +43 -0
- mb_cdar-0.3.1/src/mb_cdar/_archive_v2/loss/beta_diversity.py +96 -0
- mb_cdar-0.3.1/src/mb_cdar/_archive_v2/loss/mmd_loss.py +60 -0
- mb_cdar-0.3.1/src/mb_cdar/backbone/__init__.py +11 -0
- mb_cdar-0.3.1/src/mb_cdar/backbone/discriminator.py +120 -0
- mb_cdar-0.3.1/src/mb_cdar/backbone/unet.py +288 -0
- mb_cdar-0.3.1/src/mb_cdar/backbone/vae.py +124 -0
- mb_cdar-0.3.1/src/mb_cdar/cli.py +396 -0
- mb_cdar-0.3.1/src/mb_cdar/cond/__init__.py +10 -0
- mb_cdar-0.3.1/src/mb_cdar/data/__init__.py +44 -0
- mb_cdar-0.3.1/src/mb_cdar/data/composition.py +433 -0
- mb_cdar-0.3.1/src/mb_cdar/data/dataset.py +247 -0
- mb_cdar-0.3.1/src/mb_cdar/data/image_coding.py +68 -0
- mb_cdar-0.3.1/src/mb_cdar/data/normalizer.py +165 -0
- mb_cdar-0.3.1/src/mb_cdar/data/phylogeny.py +241 -0
- mb_cdar-0.3.1/src/mb_cdar/diffusion/__init__.py +11 -0
- mb_cdar-0.3.1/src/mb_cdar/diffusion/ddpm_scheduler.py +67 -0
- mb_cdar-0.3.1/src/mb_cdar/loss/__init__.py +8 -0
- mb_cdar-0.3.1/src/mb_cdar/loss/ecology_loss.py +104 -0
- mb_cdar-0.3.1/src/mb_cdar/reg/__init__.py +8 -0
- mb_cdar-0.3.1/src/mb_cdar/train.py +195 -0
- mb_cdar-0.3.1/src/mb_cdar/training/__init__.py +7 -0
- mb_cdar-0.3.1/src/mb_cdar/training/sampler.py +404 -0
- mb_cdar-0.3.1/src/mb_cdar/training/tail_calib.py +110 -0
- mb_cdar-0.3.1/src/mb_cdar/training/trainer.py +796 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/PKG-INFO +350 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/SOURCES.txt +49 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/dependency_links.txt +1 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/entry_points.txt +2 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/requires.txt +18 -0
- mb_cdar-0.3.1/src/mb_cdar.egg-info/top_level.txt +1 -0
- mb_cdar-0.3.1/tests/test_ablation_switches.py +157 -0
- mb_cdar-0.3.1/tests/test_backbone_cfg.py +184 -0
- mb_cdar-0.3.1/tests/test_beta_eval_fix.py +269 -0
- mb_cdar-0.3.1/tests/test_end_to_end_normalization.py +45 -0
- mb_cdar-0.3.1/tests/test_eval_refactor.py +131 -0
- mb_cdar-0.3.1/tests/test_generate_cond.py +76 -0
- mb_cdar-0.3.1/tests/test_model_io.py +100 -0
- mb_cdar-0.3.1/tests/test_normalizer.py +107 -0
- mb_cdar-0.3.1/tests/test_sampler_denormalize.py +88 -0
- mb_cdar-0.3.1/tests/test_trainer_normalizer.py +78 -0
- mb_cdar-0.3.1/tests/test_v31_arch.py +107 -0
mb_cdar-0.3.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 MB-CDAR PengXiangCao
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
mb_cdar-0.3.1/PKG-INFO
ADDED
|
@@ -0,0 +1,350 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: mb-cdar
|
|
3
|
+
Version: 0.3.1
|
|
4
|
+
Summary: MB-CDAR: Microbiome data generation via MB-DDPM (2D UNet + image_coding + discrete DDPM) with a VAE conditional encoder (channel-concatenated latent + CFG)
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: numpy>=1.24
|
|
10
|
+
Requires-Dist: torch>=2.0
|
|
11
|
+
Requires-Dist: einops>=0.6
|
|
12
|
+
Requires-Dist: biopython>=1.80
|
|
13
|
+
Requires-Dist: pandas>=2.0
|
|
14
|
+
Requires-Dist: scipy>=1.10
|
|
15
|
+
Provides-Extra: test
|
|
16
|
+
Requires-Dist: pytest>=7.0; extra == "test"
|
|
17
|
+
Requires-Dist: pytest-cov; extra == "test"
|
|
18
|
+
Requires-Dist: matplotlib>=3.6; extra == "test"
|
|
19
|
+
Requires-Dist: seaborn>=0.12; extra == "test"
|
|
20
|
+
Requires-Dist: scikit-learn>=1.2; extra == "test"
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: mb-cdar[test]; extra == "dev"
|
|
23
|
+
Requires-Dist: ruff; extra == "dev"
|
|
24
|
+
Requires-Dist: mypy<1.12,>=1.11; extra == "dev"
|
|
25
|
+
Dynamic: license-file
|
|
26
|
+
|
|
27
|
+
<div align="center">
|
|
28
|
+
|
|
29
|
+
# MB-CDAR
|
|
30
|
+
|
|
31
|
+
**Microbiome Conditional Diffusion with Adversarial Regularization** — a conditional denoising diffusion model with adversarial regularization for high-fidelity, zero-inflated, compositional microbiome data
|
|
32
|
+
|
|
33
|
+
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.
|
|
34
|
+
|
|
35
|
+
[](https://www.python.org/)
|
|
36
|
+
[](https://pytorch.org/)
|
|
37
|
+
[](pyproject.toml)
|
|
38
|
+
[](https://pypi.org/project/mb-cdar/)
|
|
39
|
+
[](#license)
|
|
40
|
+
[](tests/)
|
|
41
|
+
|
|
42
|
+
</div>
|
|
43
|
+
|
|
44
|
+
---
|
|
45
|
+
|
|
46
|
+
## Overview
|
|
47
|
+
|
|
48
|
+
**MB-CDAR** generates synthetic 16S microbiome abundance profiles that faithfully reproduce the
|
|
49
|
+
distributional, compositional, and phylogenetic structure of a real cohort. Microbiome OTU tables are
|
|
50
|
+
notoriously hard to simulate: they are **highly sparse** (~82% zeros), **compositional** (rows sum to 1),
|
|
51
|
+
and exhibit **heavy-tailed** rank-abundance structure. Classical simulators handle these properties with
|
|
52
|
+
hand-tuned parametric models; MB-CDAR instead learns them end-to-end with a **denoising diffusion
|
|
53
|
+
probabilistic model (DDPM)** operating in a 2D *image-coded* representation, conditioned on a **VAE
|
|
54
|
+
community-shape latent** injected by channel-concatenation with **classifier-free guidance (CFG)**.
|
|
55
|
+
|
|
56
|
+
The production architecture (**v3.1**) extends the diffusion backbone with a spectral-normalized
|
|
57
|
+
patch discriminator, a marginal-correlation auxiliary loss, and a group-specific zero-inflation
|
|
58
|
+
frequency-matching post-processor — while keeping the UNet's primary objective a pure noise-prediction MSE.
|
|
59
|
+
|
|
60
|
+
<p align="center">
|
|
61
|
+
<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"/>
|
|
62
|
+
<br/>
|
|
63
|
+
<em>Reverse diffusion from pure Gaussian noise. Each frame shows the model's predicted final sample
|
|
64
|
+
(<span>x̂</span><sub>0</sub>) at a denoising step <em>t</em>, conditioned on one real sample's VAE latent.
|
|
65
|
+
Left: the image-coded 40×32 representation; right: the emerging rank-abundance curve vs. the real sample.</em>
|
|
66
|
+
</p>
|
|
67
|
+
|
|
68
|
+
---
|
|
69
|
+
|
|
70
|
+
## Key results
|
|
71
|
+
|
|
72
|
+
Benchmarked on the real **Gordon cohort** (281 samples × 1033 taxa; 220 case / 61 ctrl) against two
|
|
73
|
+
established statistical simulators (**MIDASim**, **SparseDOSSA 2**) and two pooled generative controls
|
|
74
|
+
(**MB-DDPM**, **MB-GAN**), MB-CDAR (v3.1) attains **the highest per-taxon correlation fidelity of all
|
|
75
|
+
compared simulators**, is **the only method whose Shannon-diversity distribution is statistically
|
|
76
|
+
indistinguishable from the real cohort in every dataset × group** (MWU *p* > 0.05 throughout), and keeps
|
|
77
|
+
the Bray-Curtis / UniFrac distance-distribution discrepancies at or near the best of the field.
|
|
78
|
+
All numbers below are from the 2026-09-20 full re-run under the final evaluation convention
|
|
79
|
+
(permutation KS *p*-values, scipy-aligned tie handling, degenerate-taxon backfill for the R baselines).
|
|
80
|
+
|
|
81
|
+
| 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 |
|
|
82
|
+
| ------------------------------------------- | -------------- | -------------- | ------------ | ------------ | ----------------- | ----------------- | ------------ | ------------ | ------------ | ------------ |
|
|
83
|
+
| **Spearman ρ² (per-taxon mean)** | **0.9342** | **0.9939** | 0.8745 | 0.6544 | 0.8371 | 0.6442 | 0.7933 | 0.3978 | 0.5087 | 0.4065 |
|
|
84
|
+
| **Pearson r** | **0.997** | **0.999** | 0.991 | 0.639 | 0.989 | 0.711 | 0.995 | 0.968 | 0.980 | 0.958 |
|
|
85
|
+
| **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) |
|
|
86
|
+
| **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 |
|
|
87
|
+
| **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 |
|
|
88
|
+
| **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) |
|
|
89
|
+
| **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) |
|
|
90
|
+
| **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) |
|
|
91
|
+
|
|
92
|
+
Three-dataset summary (MB-CDAR, n = 1000 synthetic samples per group; full tables in
|
|
93
|
+
`script_test/results*/summary.txt`):
|
|
94
|
+
|
|
95
|
+
| Dataset | MB-CDAR Spearman ρ² case / ctrl | MB-DDPM ρ² case / ctrl | MB-GAN ρ² case / ctrl | Paired Mantel BC / WUF | Shannon MWU case / ctrl |
|
|
96
|
+
| -------------------------------- | -------------------------------- | ---------------------- | --------------------- | ---------------------- | ----------------------- |
|
|
97
|
+
| OB / Gordon (281 × 1033) | 0.934 / 0.994 | 0.793 / 0.398 | 0.509 / 0.407 | 0.999 / ≈1.000 | ns / ns |
|
|
98
|
+
| IBD / HMP2 biopsy (178 × 959) | 0.993 / 0.999 | 0.712 / 0.427 | 0.416 / 0.314 | 0.999 / ≈1.000 | ns / ns |
|
|
99
|
+
| HIV / Noguera-Julian (240 × 822) | 0.941 / 0.995 | 0.893 / 0.514 | 0.646 / 0.521 | ≈1.000 / ≈1.000 | ns / ns |
|
|
100
|
+
|
|
101
|
+
**Ablation evidence** (9 arms on the OB cohort, `Ablation Experiment/`; table + cross-arm figures):
|
|
102
|
+
the Diffusion-GAN discriminator materially improves the unweighted-UniFrac distance distribution
|
|
103
|
+
(KS D 0.096 / 0.121 → 0.072 / 0.107, case / ctrl) at unchanged per-taxon fidelity; group-specific
|
|
104
|
+
zero-inflation frequency matching is indispensable (disabling it inflates the UWUF dispersion ratio
|
|
105
|
+
to 1.27 / 1.16); the VAE condition is essential (unconditional sampling collapses Spearman ρ² to
|
|
106
|
+
0.731 / 0.642); CFG strength w ∈ {0, 1, 2, 3} has only a marginal effect (default w = 2).
|
|
107
|
+
|
|
108
|
+
**Conditional geometry.** For MB-CDAR alone, each synthetic sample is conditioned on the latent of a
|
|
109
|
+
specific real sample, enabling a 1:1 **paired Mantel test** on within-group distance matrices:
|
|
110
|
+
**BC *r* = 0.999, weighted-UniFrac *r* ≈ 1.000** (*p* = 0.001) — direct evidence that conditional
|
|
111
|
+
generation preserves pairwise sample-to-sample geometry.
|
|
112
|
+
|
|
113
|
+
<p align="center">
|
|
114
|
+
<img src="script_test/results/figures/correlation_tests_summary.png" width="760" alt="Correlation-test square matrix across simulators"/>
|
|
115
|
+
<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>
|
|
116
|
+
</p>
|
|
117
|
+
|
|
118
|
+
<p align="center">
|
|
119
|
+
<img src="script_test/results/figures/rank_abundance.png" width="480" alt="rank abundance"/>
|
|
120
|
+
<img src="script_test/results/figures/spearman_r2_bars.png" width="480" alt="spearman r2 bars"/>
|
|
121
|
+
<br/><em>Rank-abundance curves (left) and per-taxon Spearman ρ² by group (right).</em>
|
|
122
|
+
</p>
|
|
123
|
+
|
|
124
|
+
---
|
|
125
|
+
|
|
126
|
+
## Architecture
|
|
127
|
+
|
|
128
|
+
```
|
|
129
|
+
raw counts (B, 1033)
|
|
130
|
+
│
|
|
131
|
+
├─ CountNormalizer count → proportion → per-feature z-score [Architecture Decision Record 019]
|
|
132
|
+
│ ├─ VAE encoder → latent (B, 256) ──broadcast──► (B, 256, 40, 32)
|
|
133
|
+
│ └─ image_coding → (B, 1, 40, 32) ──concat──► (B, 257, 40, 32)
|
|
134
|
+
│ │
|
|
135
|
+
│ DDPM forward / UNet ε-prediction / MSE
|
|
136
|
+
│ + [v3.1] PatchDiscriminator + MarginalCorrLoss
|
|
137
|
+
│ │
|
|
138
|
+
│ DDPM reverse (T=1000, CFG extrapolation)
|
|
139
|
+
│ │
|
|
140
|
+
│ CountNormalizer.inverse_transform (clamp ≥ 0, row-normalize)
|
|
141
|
+
│ + group-specific zero-inflation frequency matching [Architecture Decision Record 026]
|
|
142
|
+
│ │
|
|
143
|
+
└──────────────► synthetic proportions (rows sum to 1, non-negative)
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
**Design principles** (each grounded in a recorded architecture decision, see `.agents/memory/decisions.md`):
|
|
147
|
+
|
|
148
|
+
- **Normalized proportion space** — the model trains on per-feature *z-scored proportions*, not raw counts
|
|
149
|
+
nor ILR/softmax. This matches the DDPM scheduler's `N(0,1)` assumption while remaining a linear per-feature
|
|
150
|
+
rescale (not a compositional transform).
|
|
151
|
+
- **Image coding** — 1033 features are reshaped into a 40×32 2D image so 2D convolutions capture local
|
|
152
|
+
taxon co-occurrence; the transform is losslessly invertible.
|
|
153
|
+
- **VAE channel-concatenation conditioning + CFG** — the 256-d community latent is broadcast over the spatial
|
|
154
|
+
plane and concatenated as extra channels (input 257 ch). Training uses 10% CFG dropout; inference applies
|
|
155
|
+
extrapolation `pred = pred_uncond + (1+w)·(pred_cond − pred_uncond)` with `w = cfg_scale` (default 2.0).
|
|
156
|
+
- **Discrete DDPM schedule** — standard Ho et al. β = linspace(1e-4, 0.02, T=1000); **not** altered.
|
|
157
|
+
- **Two-stage training** — Stage 1 VAE pretrain (10k epochs) + Stage 2 UNet DDPM (150k epochs, EMA 0.9999).
|
|
158
|
+
- **v3.1 adversarial & marginal terms** — spectral-norm patch discriminator (random-*t* noised inputs,
|
|
159
|
+
Diffusion-GAN style); the BC/Aitchison pairwise-distance losses are auto-zeroed (they traded per-taxon
|
|
160
|
+
fidelity for β-geometry — a v3 regression culprit); `MarginalCorrLoss` = 1 − Pearson correlation between
|
|
161
|
+
batch marginal and the fixed real per-taxon marginal, `α̂`-weighted.
|
|
162
|
+
- **Group-specific zero-inflation matching** — post-sampling, per taxon the smallest surplus non-zeros are
|
|
163
|
+
zeroed so the batch non-zero rate matches the observed group frequency. This removes the dense tiny-positive
|
|
164
|
+
tail without degrading (it *raises*) Spearman ρ².
|
|
165
|
+
|
|
166
|
+
---
|
|
167
|
+
|
|
168
|
+
## Installation
|
|
169
|
+
|
|
170
|
+
### From PyPI (recommended for users)
|
|
171
|
+
|
|
172
|
+
In a fresh Python 3.10+ environment:
|
|
173
|
+
|
|
174
|
+
```bash
|
|
175
|
+
pip install mb-cdar
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
This installs the `mb_cdar` package into the environment's `site-packages` and places the
|
|
179
|
+
`mb-cdar` command on the PATH (`Scripts\mb-cdar.exe` on Windows, `bin/mb-cdar` on Linux/macOS).
|
|
180
|
+
Runtime dependencies (numpy, torch, einops, biopython, pandas, scipy) are installed automatically.
|
|
181
|
+
|
|
182
|
+
### From source (for development)
|
|
183
|
+
|
|
184
|
+
```bash
|
|
185
|
+
git clone https://github.com/MKSW1/MB-CDAR.git
|
|
186
|
+
cd MB-CDAR
|
|
187
|
+
python -m venv .venv
|
|
188
|
+
.venv\Scripts\activate # Windows (cmd / PowerShell)
|
|
189
|
+
source .venv/bin/activate # Linux / macOS
|
|
190
|
+
|
|
191
|
+
pip install -e . # core (numpy, torch, biopython, pandas, scipy)
|
|
192
|
+
pip install -e ".[test]" # + pytest for the test suite
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
GPU builds: see `requirements_CUDA.txt` (NVIDIA) or `requirements_XPU.txt` (Intel GPU).
|
|
196
|
+
|
|
197
|
+
---
|
|
198
|
+
|
|
199
|
+
## Quick start
|
|
200
|
+
|
|
201
|
+
The `mb-cdar` CLI has two subcommands: `train` and `sample`.
|
|
202
|
+
|
|
203
|
+
**Train** (two-stage VAE pretrain + UNet DDPM, v3.1):
|
|
204
|
+
|
|
205
|
+
```bash
|
|
206
|
+
mb-cdar train \
|
|
207
|
+
--otu_path data/otu.tab.csv \
|
|
208
|
+
--arch v3.1 \
|
|
209
|
+
--vae_epochs 10000 --unet_epochs 150000 \
|
|
210
|
+
--cfg_scale 2.0 --device cuda
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
`python -m mb_cdar train ...` is equivalent. The `train` subcommand forwards every flag
|
|
214
|
+
verbatim to the training module, so all options in `mb-cdar train --help` are available.
|
|
215
|
+
|
|
216
|
+
**Sample** synthetic profiles from a trained checkpoint (1000 case samples conditioned on the
|
|
217
|
+
real group; the group's own presence-frequency prior is applied so taxa never observed in the
|
|
218
|
+
group stay exact zeros):
|
|
219
|
+
|
|
220
|
+
```bash
|
|
221
|
+
mb-cdar sample \
|
|
222
|
+
--ckpt output_v31/mb_cdar_unet_best.pt \
|
|
223
|
+
--data-dir data --group case --case-states OW,OB \
|
|
224
|
+
--n 1000 --out case_samples.csv
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
For a pure prior draw with no VAE condition: `mb-cdar sample --ckpt <ckpt> --unconditional --out prior.csv`.
|
|
228
|
+
|
|
229
|
+
**Resume a crashed long run** (restores weights + epoch counter + global step + discriminator state,
|
|
230
|
+
skips the VAE stage) — mutually exclusive with `--resume_ckpt` (weights-only warm start):
|
|
231
|
+
|
|
232
|
+
```bash
|
|
233
|
+
mb-cdar train --otu_path data/otu.tab.csv --arch v3.1 \
|
|
234
|
+
--resume output_v31/mb_cdar_unet_epoch_145000.pt \
|
|
235
|
+
--unet_epochs 150000 --device cuda
|
|
236
|
+
```
|
|
237
|
+
|
|
238
|
+
**Evaluate** against the R baselines (generates case/ctrl samples, computes Python metrics + R figures):
|
|
239
|
+
|
|
240
|
+
```bash
|
|
241
|
+
python script_test/run_evaluation.py \
|
|
242
|
+
--n-case 1000 --n-ctrl 1000 \
|
|
243
|
+
--device cuda --cfg-scale 2.0 \
|
|
244
|
+
--ckpt output_v31/mb_cdar_unet_best.pt
|
|
245
|
+
```
|
|
246
|
+
|
|
247
|
+
Fast smoke test (no training, no R):
|
|
248
|
+
|
|
249
|
+
```bash
|
|
250
|
+
python script_test/run_evaluation.py --fast --skip-r \
|
|
251
|
+
--ckpt output_v31/mb_cdar_unet_best.pt
|
|
252
|
+
```
|
|
253
|
+
|
|
254
|
+
**Run the test suite:**
|
|
255
|
+
|
|
256
|
+
```bash
|
|
257
|
+
pytest tests/ -v # 60 tests
|
|
258
|
+
```
|
|
259
|
+
|
|
260
|
+
---
|
|
261
|
+
|
|
262
|
+
## Repository layout
|
|
263
|
+
|
|
264
|
+
```
|
|
265
|
+
MB-CDAR/
|
|
266
|
+
├── src/mb_cdar/
|
|
267
|
+
│ ├── backbone/ # UNet (2D DDPM) + VAE + PatchDiscriminator
|
|
268
|
+
│ ├── data/ # normalizer, image_coding, composition/phylogeny (eval-only)
|
|
269
|
+
│ ├── diffusion/ # DDPMScheduler (discrete DDPM)
|
|
270
|
+
│ ├── loss/ # ecology_loss (MarginalCorrLoss; BC/Aitchison auto-zeroed)
|
|
271
|
+
│ ├── training/ # Trainer (two-stage, --arch, --resume) + Sampler + tail_calib
|
|
272
|
+
│ ├── cli.py # `mb-cdar` CLI (train/sample subcommands)
|
|
273
|
+
│ ├── __main__.py # `python -m mb_cdar`
|
|
274
|
+
│ └── train.py # training entry point
|
|
275
|
+
├── script_test/ # evaluation pipeline (vs. MIDASim / SparseDOSSA2
|
|
276
|
+
│ └── baselines/ # pooled MB-DDPM / MB-GAN control wrappers (Architecture Decision Record 035)
|
|
277
|
+
├── output_v31/ # v3.1 production checkpoint, Git LFS (--ckpt target; Architecture Decision Record 031)
|
|
278
|
+
├── Ablation Experiment/ # ablation protocol, cross-platform drivers, ablation comparison figures, per-dataset data (Architecture Decision Record 033/036)
|
|
279
|
+
├── tests/ # unit + end-to-end tests (60)
|
|
280
|
+
├── docs/ # architecture, training/testing, dataset & literature notes
|
|
281
|
+
├── .agents/ # agent-framework-agnostic project memory & conventions
|
|
282
|
+
│ ├── memory/ # project_brief, active_context, progress, decisions (Architecture Decision Records)
|
|
283
|
+
│ └── skills/ # project-core-conventions (architecture bottom lines, SOPs)
|
|
284
|
+
├── assets/ # README figures & the reverse-diffusion animation
|
|
285
|
+
└── data/ # Gordon OTU table + metadata + phylogenetic tree
|
|
286
|
+
```
|
|
287
|
+
|
|
288
|
+
---
|
|
289
|
+
|
|
290
|
+
## Reproducibility notes
|
|
291
|
+
|
|
292
|
+
- **Checkpoint contract** — UNet checkpoints embed `normalizer_state`, `cfg_scale`, `arch`,
|
|
293
|
+
`tail_calib_state`, `taxon_freq_state`, and resume keys (`global_step`, `best_loss`,
|
|
294
|
+
`best_loss_gan`, `disc_state_dict`, `epoch`). Checkpoints lacking `normalizer_state` or with
|
|
295
|
+
`data_scale != "zscore_proportion"` are explicitly rejected by the evaluation script.
|
|
296
|
+
- **β-diversity testing** — cross-matrix Mantel (real-real vs. synth-synth with no sample correspondence)
|
|
297
|
+
was removed as statistically invalid (Architecture Decision Record 027); the toolkit now uses distribution-level KS + dispersion
|
|
298
|
+
ratio on within-group pairwise distances (all simulators) and a paired Mantel for MB-CDAR only.
|
|
299
|
+
KS p-values come from a 999-permutation label-shuffle test because within-group pairwise
|
|
300
|
+
distances are not independent — asymptotic KS p-values are invalid there (Architecture Decision Record 032).
|
|
301
|
+
- **Long-training stability** — Windows 150k-epoch runs follow strict memory rules
|
|
302
|
+
(numpy RNG, no tqdm, periodic `gc.collect`/`empty_cache`, atomic checkpoint saves); see
|
|
303
|
+
`.agents/skills/project-core-conventions/SKILL.md`.
|
|
304
|
+
- **Three-group evaluation** — every evaluation now reports case / ctrl / the pooled
|
|
305
|
+
all slot (case+ctrl concatenated) for all simulators, and the pooled-trained
|
|
306
|
+
MB-DDPM / MB-GAN control models join the comparison via `--controls-dir`
|
|
307
|
+
(their CSVs are produced by `script_test/baselines/` wrappers; Architecture Decision Record 035).
|
|
308
|
+
- **Ablation figures** — the cross-arm comparison figures (arm x (group x metric)
|
|
309
|
+
matrix, training-level component-attribution ladder, inference-level sweeps) are
|
|
310
|
+
rendered by `Ablation Experiment/scripts/plot_ablation_figures.py` from the
|
|
311
|
+
ablation table or the per-arm results, in the same style as the main evaluation
|
|
312
|
+
figures (Architecture Decision Record 036).
|
|
313
|
+
- **Figure style** — the nine evaluation figures share one centralised style constant
|
|
314
|
+
(sans-serif, white background, no gridlines, thin black frames, inward ticks):
|
|
315
|
+
row=group × column=simulator identity and alpha-diversity grids, a corrplot-style
|
|
316
|
+
correlation-test square matrix (square area and colour = signed value / row max), marginal-density PCoA
|
|
317
|
+
biplot and RdBu_r heatmaps (Architecture Decision Record 035).
|
|
318
|
+
|
|
319
|
+
---
|
|
320
|
+
|
|
321
|
+
## References
|
|
322
|
+
|
|
323
|
+
1. **MB-DDPM backbone** — Huo, B. et al. *MB-DDPM: microbiome simulation based on the denoising
|
|
324
|
+
diffusion probabilistic model.* **Information Fusion** 127, 103880 (2026).
|
|
325
|
+
doi:[10.1016/j.inffus.2025.103880](https://doi.org/10.1016/j.inffus.2025.103880)
|
|
326
|
+
2. **DDPM** — Ho, J., Jain, A. & Abbeel, P. *Denoising diffusion probabilistic models.* NeurIPS (2020).
|
|
327
|
+
3. **Classifier-free guidance** — Ho, J. & Salimans, T. *Classifier-free diffusion guidance.* (2022).
|
|
328
|
+
4. **MIDASim** — He, M., Zhao, N. & Satten, G. A. *MIDASim: a fast and simple simulator for realistic
|
|
329
|
+
microbiome data.* **Microbiome** (2024). doi:[10.1186/s40168-024-01822-z](https://doi.org/10.1186/s40168-024-01822-z)
|
|
330
|
+
5. **SparseDOSSA 2** — Huttenhower Lab. *SparseDOSSA 2: a computational framework for simulating
|
|
331
|
+
microbiome data.* (see also Ma, S. et al., **PLOS Comput. Biol.** 2021,
|
|
332
|
+
doi:[10.1371/journal.pcbi.1008913](https://doi.org/10.1371/journal.pcbi.1008913)).
|
|
333
|
+
6. **Real cohort** — the Gordon-lab obesity 16S cohort used by Ref. 1 (281 samples × 1033 OTUs;
|
|
334
|
+
220 overweight/obese case vs. 61 healthy ctrl), bundled under [`data/`](data/).
|
|
335
|
+
|
|
336
|
+
> Project-specific design decisions, regression post-mortems (v3 → v3.1), and evaluation-statistics
|
|
337
|
+
> corrections are recorded as Architecture Decision Records in
|
|
338
|
+
> [`.agents/memory/decisions.md`](.agents/memory/decisions.md).
|
|
339
|
+
|
|
340
|
+
---
|
|
341
|
+
|
|
342
|
+
## License
|
|
343
|
+
|
|
344
|
+
Released under the **MIT License** — see [`pyproject.toml`](pyproject.toml). The bundled Gordon dataset
|
|
345
|
+
under [`data/`](data/) is provided for research reproducibility; please cite the original study (Ref. 6)
|
|
346
|
+
when reusing it.
|
|
347
|
+
|
|
348
|
+
## Citation
|
|
349
|
+
|
|
350
|
+
If you use MB-CDAR in your research, please cite the MB-DDPM backbone (Ref. 1) and this repository.
|