mb-cdar 0.3.1__py3-none-any.whl

mb_cdar/__init__.py

@@ -0,0 +1,3 @@
+"""MB-CDAR: Microbiome Conditional Diffusion with Adversarial Regularization."""
+
+__version__ = "0.3.1"

mb_cdar/__main__.py

@@ -0,0 +1,6 @@
+"""Enable ``python -m mb_cdar`` as an alias for the ``mb-cdar`` CLI."""
+
+from mb_cdar.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mb_cdar/_archive_v2/backbone/mb_cdar_backbone.py

@@ -0,0 +1,379 @@
+"""Continuous-only DiT backbone for MB-CDAR.
+
+A Diffusion Transformer (DiT) that operates in ILR (Isometric Log-Ratio)
+space for continuous compositional diffusion.  Features:
+
+1. 2-layer MLP numerical tokenisation for the continuous ILR channel.
+2. Hand-written Multi-Head Attention with fused SDPA (memory-efficient at
+   D=1033 — avoids materialising the (B, nhead, D, D) score tensor).
+3. AdaLN-Zero conditioning from the time embedding (and an optional external
+   condition embedding, e.g. simple categorical metadata).
+4. Classifier-Free Guidance (CFG) condition dropout during training (10%
+   default) — the backbone learns both conditional and unconditional
+   trajectories.
+5. Absolute per-taxon identity injection via a learnable ``taxon_id_embed``
+   that breaks permutation-equivariance over taxa (the precondition for
+   non-trivial Spearman R² on species co-occurrence).
+
+The previous V1 architecture (PhyloBias attention bias, phylogenetic spectral
+identity, CtrTab Laplace noise injection, CADD cross-attention, discrete
+RADD track) has been removed.  The phylogenetic tree is now used only for
+UniFrac evaluation, not for model conditioning.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from mb_cdar.data.composition import helmert_contrast_matrix
+
+
+class SinusoidalEmbedding(nn.Module):
+    """Sinusoidal time embedding module."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+    def forward(self, t: Tensor) -> Tensor:
+        device = t.device
+        half_dim = self.d_model // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device, dtype=t.dtype) * -emb)
+        emb = t.unsqueeze(-1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class MBCDARBackbone(nn.Module):
+    """Continuous-only DiT architecture operating in ILR space.
+
+    Parameters
+    ----------
+    D : int
+        Number of OTUs (taxa).
+    d_model : int
+        DiT hidden dimension.  Default 256.
+    nhead : int
+        Number of attention heads.  Default 4.
+    num_layers : int
+        Number of DiT blocks.  Default 8.
+    dropout : float
+        Dropout rate.  Default 0.1.
+    id_dim : int
+        Dimensionality of the learnable per-taxon identity residual.  Default
+        24.  Set to 0 to disable (not recommended — permutation-equivariance
+        returns and Spearman R² collapses to ~0).
+    cond_dim : Optional[int]
+        Dimensionality of an external condition embedding (e.g. simple
+        categorical metadata projected to a fixed vector).  When provided,
+        ``cond_proj`` maps it to ``d_model`` and adds it to the time
+        condition.  Default None (time-only conditioning).
+    """
+
+    def __init__(
+        self,
+        D: int = 1033,
+        d_model: int = 256,
+        nhead: int = 4,
+        num_layers: int = 8,
+        dropout: float = 0.1,
+        id_dim: int = 24,
+        cond_dim: int | None = None,
+    ):
+        super().__init__()
+        self.D = D
+        self.K = D - 1  # ILR dimension
+        self.d_model = d_model
+        self.id_dim = id_dim
+        self.cond_dim = cond_dim
+
+        # Token embedding for Taxon space (shared across taxa — symmetric)
+        self.otu_embed = nn.Sequential(
+            nn.Linear(1, d_model),
+            nn.SiLU(),
+            nn.Linear(d_model, d_model),
+            nn.LayerNorm(d_model),
+        )
+
+        # === Absolute per-taxon identity (learnable) =========================
+        # The shared ``otu_embed`` above is permutation-equivariant over taxa.
+        # We break that symmetry by adding a *unique* learnable identity vector
+        # per taxon, projected to d_model and added to the input embedding
+        # (broadcast over the batch).  This is the necessary precondition for
+        # learning species-specific co-occurrence structure (Spearman R² > 0).
+        #
+        # The previous V1 architecture combined this with a *fixed*
+        # phylogenetic spectral prior; that has been removed (the tree is now
+        # used only for evaluation, not model conditioning).  Identity comes
+        # entirely from the learnable code below.
+        if id_dim > 0:
+            self.taxon_id_embed = nn.Parameter(torch.randn(D, id_dim) * 0.02)
+            self.taxon_id_proj = nn.Linear(id_dim, d_model)
+        # ======================================================================
+
+        # Helmert contrast matrix cache for projections
+        self._V_np = helmert_contrast_matrix(D)
+        self._V_cache: dict[tuple[torch.device, torch.dtype], Tensor] = {}
+
+        # Time embedding
+        self.t_embed = SinusoidalEmbedding(d_model)
+        self.t_mlp = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.SiLU(),
+            nn.Linear(d_model, d_model),
+        )
+
+        # External condition projection (e.g. simple categorical metadata).
+        # Projects an external condition embedding (cond_dim) to d_model so it
+        # can be added to the time-conditioning vector c.  Zero-init the output
+        # layer so the condition starts as a no-op (DiT-style gradual onset).
+        if cond_dim is not None:
+            self.cond_proj = nn.Sequential(
+                nn.Linear(cond_dim, d_model),
+                nn.SiLU(),
+                nn.Linear(d_model, d_model),
+            )
+            nn.init.zeros_(self.cond_proj[-1].weight)
+            nn.init.zeros_(self.cond_proj[-1].bias)
+
+        # DiT blocks
+        self.blocks = nn.ModuleList([
+            DiTBlock(d_model, nhead, dropout=dropout)
+            for _ in range(num_layers)
+        ])
+
+        # Output — final AdaLN-Zero modulation applied just before the output
+        # LayerNorm.  SiLU + Linear(d → 2d) producing (shift, scale); zero-init
+        # so the network starts as identity → final_norm only.
+        self.final_adaln = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(d_model, 2 * d_model),
+        )
+        nn.init.zeros_(self.final_adaln[1].weight)
+        nn.init.zeros_(self.final_adaln[1].bias)
+
+        self.final_norm = nn.LayerNorm(d_model)
+        self.out_proj = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.SiLU(),
+            nn.Linear(d_model, 1),
+        )
+
+    def _get_V(self, ref: Tensor) -> Tensor:
+        key = (ref.device, ref.dtype)
+        if key not in self._V_cache:
+            self._V_cache[key] = torch.from_numpy(self._V_np).to(
+                device=ref.device, dtype=ref.dtype
+            )
+        return self._V_cache[key]
+
+    def _taxon_identity(self) -> Tensor | None:
+        """Absolute per-taxon identity embedding, shape (D, d_model) or None."""
+        if self.id_dim > 0:
+            return self.taxon_id_proj(self.taxon_id_embed)  # (D, d_model)
+        return None
+
+    def forward(
+        self,
+        x_ilr: Tensor,
+        t: Tensor,
+        cond_emb: Tensor | None = None,
+        cond_dropout_prob: float = 0.1,
+    ) -> Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x_ilr : (B, D-1) Tensor
+            ILR-transformed compositions.
+        t : (B,) Tensor
+            Diffusion time step.
+        cond_emb : Optional[Tensor]
+            External condition embedding, shape (B, cond_dim).  When provided
+            it is projected to ``d_model`` and added to the time-conditioning
+            vector ``c``.  Default None (time-only conditioning).
+        cond_dropout_prob : float
+            Probability of dropping the condition (CFG training).  Set to 0
+            to disable (e.g. at inference).  Default 0.1.
+
+        Returns
+        -------
+        noise_pred : (B, D-1) Tensor
+        """
+        # 1. Project ILR to taxon space
+        V = self._get_V(x_ilr)
+        x_taxon = torch.matmul(x_ilr, V.t())  # (B, D)
+
+        # 2. Embed taxon coordinates + inject absolute per-taxon identity.
+        #    The identity (D, d_model) breaks permutation-equivariance so the
+        #    network can learn species-specific structure (Spearman R² fix).
+        x = self.otu_embed(x_taxon.unsqueeze(-1))  # (B, D, d_model)
+        ident = self._taxon_identity()
+        if ident is not None:
+            x = x + ident.unsqueeze(0)  # broadcast (1, D, d_model) → (B, D, d_model)
+
+        # 3. Time embedding (the condition c)
+        c = self.t_mlp(self.t_embed(t))  # (B, d_model)
+
+        # External condition (e.g. simple categorical metadata) fused into c.
+        # cond_emb is projected to d_model and added to the time condition.
+        cond_signal = None
+        if cond_emb is not None and hasattr(self, 'cond_proj'):
+            cond_signal = self.cond_proj(cond_emb)
+
+        # CFG condition dropout: zero out the condition projection (NOT the
+        # time embedding!) on a random subset of the batch during training so
+        # the network learns both conditional and unconditional trajectories.
+        # The unconditional trajectory uses c = t_emb only (no condition
+        # signal), matching what is passed at inference time when
+        # cond_emb=None.
+        if self.training and cond_dropout_prob > 0 and cond_signal is not None:
+            drop_mask = torch.rand(c.shape[0], device=c.device) < cond_dropout_prob
+            cond_signal = torch.where(
+                drop_mask.unsqueeze(-1),
+                torch.zeros_like(cond_signal),
+                cond_signal,
+            )
+
+        if cond_signal is not None:
+            c = c + cond_signal
+
+        # 4. DiT blocks (no phylo attention bias — removed in V1 cleanup)
+        for block in self.blocks:
+            x = block(x, c)  # (B, D, d_model)
+
+        # 5. Final AdaLN-Zero modulation + output projection back to ILR
+        #    noise/score.  c is (B, d_model) and x is (B, D, d_model), so the
+        #    shift/scale are unsqueezed to (B, 1, d_model) for broadcasting.
+        final_shift, final_scale = self.final_adaln(c).chunk(2, dim=-1)
+        x = self.final_norm(x) * (1 + final_scale.unsqueeze(1)) + final_shift.unsqueeze(1)
+        score_taxon = self.out_proj(x).squeeze(-1)  # (B, D)
+        noise_pred = torch.matmul(score_taxon, V)  # (B, K)
+
+        return noise_pred
+
+
+# =====================================================================
+# DiTBlock — DiT block with adaLN modulation
+# =====================================================================
+
+
+class DiTBlock(nn.Module):
+    """DiT block with adaLN modulation.
+
+    Components:
+      norm1, MultiheadAttention(d_model, nhead, dropout)
+      norm2, MLP(d → 4d → d) with GELU + Dropout
+      adaln: SiLU + Linear(d → 6d) 生成 6 个调制参数
+
+    Parameters
+    ----------
+    d_model : int
+        Model hidden dimension.
+    nhead : int
+        Number of attention heads.
+    dropout : float
+        Dropout rate.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 256,
+        nhead: int = 4,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+        assert d_model % nhead == 0
+
+        # Self-attention QKV (attention itself uses fused SDPA in forward)
+        self.W_qkv = nn.Linear(d_model, 3 * d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.attn_dropout_p = dropout  # applied inside SDPA during training
+        self.resid_dropout = nn.Dropout(dropout)
+
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # Feedforward MLP
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model, 4 * d_model),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(4 * d_model, d_model),
+            nn.Dropout(dropout),
+        )
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # adaLN: SiLU + Linear(d → 6d) 生成 6 个调制参数 (DiT-style).
+        # SiLU gate before the projection matches the DiT architecture; the
+        # Linear is zero-initialised so the network starts as identity (all
+        # shifts/scales/gates = 0 → modulation is a no-op at init).
+        self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(d_model, 6 * d_model))
+        nn.init.zeros_(self.adaln[1].weight)
+        nn.init.zeros_(self.adaln[1].bias)
+
+    def forward(self, x: Tensor, c: Tensor) -> Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : (B, K, d_model) Tensor
+            Input sequence.
+        c : (B, d_model) Tensor
+            Condition vector (time embedding, optionally with cond signal).
+
+        Returns
+        -------
+        (B, K, d_model) Tensor
+        """
+        B, K, C = x.shape
+
+        # --- adaLN modulation parameters ---
+        # 6 params: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp
+        params = self.adaln(c)  # (B, 6*d_model)
+        params = params.unsqueeze(1)  # (B, 1, 6*d_model)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+            params.chunk(6, dim=-1)  # each (B, 1, d_model)
+
+        # --- Self-Attention ---
+        # Modulate input
+        h = self.norm1(x)
+        h = h * (1 + scale_msa) + shift_msa
+
+        # QKV projection
+        qkv = self.W_qkv(h)  # (B, K, 3*C)
+        qkv = qkv.reshape(B, K, 3, self.nhead, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, nhead, K, head_dim)
+        q, k, v = qkv.unbind(0)
+
+        # Fused scaled dot-product attention.  SDPA computes
+        # softmax(QKᵀ/√d)·V without ever materialising the (B, nhead, K, K)
+        # score tensor — essential at K=1033 (the hand-written path allocated
+        # ~0.5 GB/layer here).  SDPA's default scale is 1/√head_dim.
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.attn_dropout_p if self.training else 0.0,
+        )  # (B, nhead, K, head_dim)
+        attn_out = attn_out.transpose(1, 2).reshape(B, K, C)
+        attn_out = self.resid_dropout(self.out_proj(attn_out))
+
+        # Gate + residual
+        x = x + gate_msa * attn_out
+
+        # --- MLP ---
+        h = self.norm2(x)
+        h = h * (1 + scale_mlp) + shift_mlp
+        mlp_out = self.mlp(h)
+
+        # Gate + residual
+        x = x + gate_mlp * mlp_out
+
+        return x