m3-sc 0.3.0__py3-none-any.whl

m3/__init__.py

@@ -0,0 +1,29 @@
+"""m3 -- multimodal, multi-batch, condition-aware single-cell model."""
+
+from m3.io import read_matrix, from_anndata, read_h5ad, read_h5
+from m3._concat import concat
+from m3 import datasets
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "read_matrix", "from_anndata", "read_h5ad", "read_h5", "concat",
+    "M3", "Attribution", "M3CapabilityError",
+    "datasets",
+]
+
+# The model and its readouts need torch; import them lazily so that `import m3`
+# and the torch-free readers (io) and bundled datasets stay usable without a
+# torch install. Accessing m3.M3 (etc.) triggers the torch-backed import.
+_LAZY = {"M3", "Attribution", "M3CapabilityError"}
+
+
+def __getattr__(name):
+    if name in _LAZY:
+        from m3 import _model
+        return getattr(_model, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(__all__)

m3/_bridge.py

@@ -0,0 +1,77 @@
+"""Bridge: marshal a :class:`m3.Dataset` to the engine's file-based inputs.
+
+The engine (``m3._engine.run_M3``) consumes per-modality lists of
+paper-format HDF5 files (one element per batch, ``None`` where a batch lacks a
+modality) plus a per-batch metadata CSV list. This module writes those temp
+files from an in-memory ``Dataset`` so the proven engine path runs unchanged.
+
+Paper HDF5 layout (what ``m3._engine.util.read_h5_data`` expects):
+``matrix/data`` is genes x cells (transpose of our cells x genes), and
+``matrix/features`` holds the feature names.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import h5py
+import numpy as np
+
+_MODALITY_ORDER = ("rna", "adt", "atac")  # maps to engine modality1/2/3_path
+
+
+def _write_paper_h5(path: str, counts_cells_by_features, feature_names) -> None:
+    dense = np.asarray(counts_cells_by_features.todense()
+                       if hasattr(counts_cells_by_features, "todense")
+                       else counts_cells_by_features, dtype=np.float32)
+    # Store cells x features (the Liu/Stephenson layout). The engine's
+    # read_h5_data + load_data_from_list double-transpose nets back to cells x features.
+    with h5py.File(path, "w") as f:
+        g = f.create_group("matrix")
+        g.create_dataset("data", data=dense)
+        g.create_dataset("features", data=np.array([str(x) for x in feature_names], dtype="S"))
+
+
+def marshal(dataset, tmpdir: str | None = None) -> dict:
+    """Write a Dataset to temp paper-format h5 + metadata CSVs.
+
+    Returns a dict with:
+      modality_paths: {"rna": [path|None per batch], "adt": [...], "atac": [...]}
+      metadata_paths: [csv path per batch]
+      batches: ordered batch labels (row order of the engine's concatenation)
+      tmpdir: the directory holding the temp files (caller may clean up)
+    """
+    if tmpdir is None:
+        tmpdir = tempfile.mkdtemp(prefix="m3_marshal_")
+    os.makedirs(tmpdir, exist_ok=True)
+
+    obs = dataset.obs
+    batches = list(dataset.batches)
+    batch_rows = {c: (obs["batch"].to_numpy() == c) for c in batches}
+
+    modality_paths: dict[str, list] = {m: [] for m in _MODALITY_ORDER}
+    metadata_paths: list[str] = []
+
+    for ci, c in enumerate(batches):
+        rows = batch_rows[c]
+        # metadata CSV for this batch (engine reads with pandas; needs the role columns)
+        meta_path = os.path.join(tmpdir, f"metadata_{ci}.csv")
+        obs.loc[rows].reset_index(drop=True).to_csv(meta_path, index=False)
+        metadata_paths.append(meta_path)
+
+        for m in _MODALITY_ORDER:
+            if m in dataset.modalities and bool(dataset.present[m][rows].any()):
+                mat = dataset.modalities[m][rows]
+                path = os.path.join(tmpdir, f"{m}_{ci}.h5")
+                _write_paper_h5(path, mat, dataset.var[m])
+                modality_paths[m].append(path)
+            else:
+                modality_paths[m].append(None)
+
+    return {
+        "modality_paths": modality_paths,
+        "metadata_paths": metadata_paths,
+        "batches": batches,
+        "tmpdir": tmpdir,
+    }

m3/_concat.py

@@ -0,0 +1,72 @@
+"""Multi-batch assembly: stack single-batch Datasets into one."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+from m3._dataset import Dataset
+
+
+def _intersect_ordered(indices: list[pd.Index]) -> pd.Index:
+    common = set(indices[0])
+    for idx in indices[1:]:
+        common &= set(idx)
+    # stable order: follow the first batch's ordering
+    return pd.Index([g for g in indices[0] if g in common])
+
+
+def concat(datasets: list[Dataset]) -> Dataset:
+    """Combine single-batch Datasets into one multi-batch Dataset.
+
+    - features per modality are harmonized by NAME intersection (first batch's order)
+    - a batch missing a whole modality is zero-filled with present[modality]=False
+    - batch labels must be unique across inputs
+    """
+    if not datasets:
+        raise ValueError("concat requires at least one Dataset.")
+    if len(datasets) == 1:
+        return datasets[0]
+
+    labels = [c for d in datasets for c in d.batches]
+    if len(labels) != len(set(labels)):
+        raise ValueError(f"batch labels must be unique across inputs; got {labels}.")
+
+    all_mods = sorted({m for d in datasets for m in d.modality_names})
+    n_total = sum(d.n_cells for d in datasets)
+
+    out_mods: dict[str, sp.csr_matrix] = {}
+    out_var: dict[str, pd.Index] = {}
+    out_present: dict[str, np.ndarray] = {}
+
+    for mod in all_mods:
+        carriers = [d for d in datasets if mod in d.modalities]
+        common = _intersect_ordered([d.var[mod] for d in carriers])
+        if len(common) == 0:
+            raise ValueError(f"modality '{mod}': empty feature intersection across batches.")
+        if any(len(common) < len(d.var[mod]) for d in carriers):
+            warnings.warn(
+                f"modality '{mod}': intersection kept {len(common)} features "
+                f"(batches had {[len(d.var[mod]) for d in carriers]}).",
+                stacklevel=2,
+            )
+        blocks = []
+        present = []
+        for d in datasets:
+            if mod in d.modalities:
+                col_pos = [d.var[mod].get_loc(g) for g in common]
+                blocks.append(d.modalities[mod][:, col_pos])
+                present.append(np.ones(d.n_cells, dtype=bool))
+            else:
+                blocks.append(sp.csr_matrix((d.n_cells, len(common)), dtype=np.float32))
+                present.append(np.zeros(d.n_cells, dtype=bool))
+        out_mods[mod] = sp.vstack(blocks).tocsr()
+        out_var[mod] = common
+        out_present[mod] = np.concatenate(present)
+
+    obs = pd.concat([d.obs for d in datasets], ignore_index=True)
+    assert obs.shape[0] == n_total
+    return Dataset(modalities=out_mods, obs=obs, var=out_var, present=out_present)

m3/_dataset.py

@@ -0,0 +1,63 @@
+"""Canonical m3 input container."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+
+
+@dataclass
+class Dataset:
+    """Multimodal, multi-batch raw-count container.
+
+    modalities: {name -> csr_matrix [n_cells, n_features]} raw counts
+    obs:        DataFrame [n_cells, ...]; must contain a 'batch' column
+    var:        {name -> pd.Index of feature names}, aligned to each modality's columns
+    present:    {name -> bool ndarray [n_cells]} per-cell modality presence
+    """
+
+    modalities: dict[str, sp.csr_matrix]
+    obs: pd.DataFrame
+    var: dict[str, pd.Index]
+    present: dict[str, np.ndarray]
+
+    def __post_init__(self) -> None:
+        if not self.modalities:
+            raise ValueError("Dataset requires at least one modality.")
+        n = self.obs.shape[0]
+        if "batch" not in self.obs.columns:
+            raise ValueError("obs must contain a 'batch' column.")
+        for name, mat in self.modalities.items():
+            if mat.shape[0] != n:
+                raise ValueError(
+                    f"modality '{name}' has {mat.shape[0]} rows but obs has {n} (n_cells)."
+                )
+            if name not in self.var or len(self.var[name]) != mat.shape[1]:
+                raise ValueError(
+                    f"var['{name}'] length must equal modality '{name}' column count."
+                )
+            if pd.Index(self.var[name]).has_duplicates:
+                raise ValueError(
+                    f"var['{name}'] has duplicate feature names; names must be unique."
+                )
+            if name not in self.present or self.present[name].shape[0] != n:
+                raise ValueError(f"present['{name}'] must be a length-{n} boolean array.")
+
+    @property
+    def n_cells(self) -> int:
+        return self.obs.shape[0]
+
+    @property
+    def modality_names(self) -> list[str]:
+        return list(self.modalities.keys())
+
+    @property
+    def batches(self) -> list[str]:
+        return list(pd.unique(self.obs["batch"]))
+
+    def __repr__(self) -> str:
+        mods = ", ".join(f"{k}:{v.shape[1]}" for k, v in self.modalities.items())
+        return f"Dataset(n_cells={self.n_cells}, batches={self.batches}, modalities=[{mods}])"

m3/_engine/__init__.py

@@ -0,0 +1,4 @@
+"""Not part of the public API. The user-facing :class:`m3.M3` orchestrates these
+modules. Modules keep their original relative imports (``from .util import ...``)
+and run unchanged; only their location moved into the ``m3._engine`` package.
+"""

m3/_engine/feature_selection.py

@@ -0,0 +1,300 @@
+"""
+End-to-end multi-resolution feature attribution for M3.
+
+Connects the cell-level generator and patient-level corrector into one
+differentiable forward pass, then runs Integrated Gradients from patient
+disease prediction all the way back to raw gene expression.
+
+Two levels of attribution:
+1. Input-level IG: gene -> disease prediction (for gene and cell importance)
+2. Patient vector-level IG: patient vector -> disease prediction (for cell type and donor importance)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+from captum.attr import IntegratedGradients
+from .util import poe
+
+
+class EndToEndWrapper(nn.Module):
+
+    def __init__(self, generator, corrector,
+                 batch_labels, mask_poe,
+                 donor_codes, celltype_codes,
+                 n_celltypes, cond_index=0):
+        super().__init__()
+        self.generator = generator
+        self.corrector = corrector
+
+        self.register_buffer('batch_labels', batch_labels.long())
+        self.register_buffer('mask_poe', mask_poe.float())
+        self.register_buffer('donor_codes', donor_codes.long())
+        self.register_buffer('celltype_codes', celltype_codes.long())
+
+        self.n_celltypes = n_celltypes
+        self.cond_index = cond_index
+        self.n_conds = len(generator.condition_dim)
+        self.batch_classify_dim = generator.batch_classify_dim
+        self.n_cells = batch_labels.shape[0]
+
+    def _get_z_slices(self, mu):
+        z_dim = mu.shape[1]
+        n_conds = self.n_conds
+        emb_dim = z_dim - 2 - 2 * n_conds
+
+        z_embedding = mu[:, :emb_dim]
+        z_conds = []
+        start = emb_dim
+        for _ in range(n_conds):
+            z_conds.append(mu[:, start:start + 2])
+            start += 2
+
+        return z_embedding, z_conds
+
+    def _differentiable_aggregate(self, z, donor_codes, celltype_codes,
+                                   n_celltypes, min_cells=5):
+        feat_dim = z.shape[1]
+        unique_donors = donor_codes.unique()
+        n_donors = unique_donors.shape[0]
+
+        patient_vectors = torch.zeros(n_donors, n_celltypes * feat_dim,
+                                       device=z.device, dtype=z.dtype)
+
+        for d_idx, d in enumerate(unique_donors):
+            for c in range(n_celltypes):
+                mask = (donor_codes == d) & (celltype_codes == c)
+                if mask.sum() >= min_cells:
+                    patient_vectors[d_idx, c * feat_dim:(c + 1) * feat_dim] = z[mask].mean(dim=0)
+
+        return patient_vectors, unique_donors
+
+    def _core_forward(self, x):
+        mu_list, logvar_list = self.generator.encoder(x)
+        b_onehot = F.one_hot(self.batch_labels,
+                             num_classes=self.batch_classify_dim).float()
+        batch_mu, batch_logvar = self.generator.batch_encoder(b_onehot)
+
+        mu, logvar = poe(mu_list + [batch_mu],
+                         logvar_list + [batch_logvar],
+                         self.mask_poe)
+
+        z_embedding, z_conds = self._get_z_slices(mu)
+        z_for_agg = torch.cat([z_embedding, z_conds[self.cond_index]], dim=1)
+
+        patient_vectors, unique_donors = self._differentiable_aggregate(
+            z_for_agg, self.donor_codes, self.celltype_codes,
+            self.n_celltypes
+        )
+
+        z_corrected = self.corrector.correct(patient_vectors)
+        logits = self.corrector.status_head(z_corrected)
+        return logits, patient_vectors, unique_donors
+
+    def forward(self, x):
+        logits, _, _ = self._core_forward(x)
+        return logits
+
+    def forward_flat(self, x_flat):
+        x = x_flat.view(self.n_cells, -1)
+        logits = self.forward(x)
+        return logits.sum(dim=0, keepdim=True)
+
+    def forward_full(self, x):
+        logits, patient_vectors, unique_donors = self._core_forward(x)
+
+        return {
+            'logits': logits,
+            'patient_vectors': patient_vectors,
+            'unique_donors': unique_donors,
+        }
+
+
+def run_captum_ig(wrapper, x, target_class=1, n_steps=50, baseline=None):
+    n_cells, n_genes = x.shape
+    x_flat = x.detach().view(1, -1)
+    if baseline is None:
+        baseline_flat = torch.zeros_like(x_flat)
+    else:
+        baseline_flat = baseline.detach().reshape(1, -1)
+
+    ig = IntegratedGradients(wrapper.forward_flat)
+
+    attr_flat = ig.attribute(
+        x_flat,
+        baselines=baseline_flat,
+        target=target_class,
+        n_steps=n_steps,
+        internal_batch_size=1,
+    )
+
+    attribution = attr_flat.view(n_cells, n_genes)
+    return attribution
+
+
+def run_attribution(generator, corrector,
+                    ref_data, ref_b, ref_mask_poe,
+                    ref_metadata, donor_name, cty_name,
+                    n_celltypes, target_class=1,
+                    cond_index=0, n_steps=50, device='cpu',
+                    baseline=None, condition_col=None, reference_labels=None,
+                    skip_patient_level=False):
+    """
+    Run end-to-end attribution on reference data using Captum IG.
+
+    Two levels of IG:
+    1. Input-level: gene expression -> disease prediction (for gene & cell importance)
+    2. Patient vector-level: patient vector -> disease prediction (for cell type & donor importance)
+
+    Args:
+        condition_col: str, column name for condition in ref_metadata (e.g. 'cond_group').
+        reference_labels: list of str, labels considered as reference/healthy 
+                         (e.g. ['HC'] for COVID, ['Healthy'] for Stephenson).
+        If condition_col and reference_labels are not provided, uses zero baseline 
+        for patient vector level IG.
+    """
+
+    generator = generator.to(device).eval()
+    corrector = corrector.to(device).eval()
+
+    ref_data = ref_data.to(device).float()
+    ref_b = ref_b.to(device)
+    ref_mask_poe = ref_mask_poe.to(device)
+
+    cty_categories = pd.Categorical(ref_metadata[cty_name])
+    celltype_codes = torch.tensor(cty_categories.codes, dtype=torch.long, device=device)
+    celltype_names = list(cty_categories.categories)
+
+    donor_categories = pd.Categorical(ref_metadata[donor_name])
+    donor_codes = torch.tensor(donor_categories.codes, dtype=torch.long, device=device)
+    donor_names = list(donor_categories.categories)
+
+    # build wrapper
+    wrapper = EndToEndWrapper(
+        generator, corrector,
+        batch_labels=ref_b,
+        mask_poe=ref_mask_poe,
+        donor_codes=donor_codes,
+        celltype_codes=celltype_codes,
+        n_celltypes=n_celltypes,
+        cond_index=cond_index
+    ).to(device).eval()
+
+    for p in wrapper.parameters():
+        p.requires_grad_(False)
+
+    # --- Patient-level info (no grad needed); skipped entirely when skip_patient_level ---
+    if not skip_patient_level:
+        print("Computing patient-level predictions...")
+        with torch.no_grad():
+            full_out = wrapper.forward_full(ref_data)
+            logits = full_out['logits']
+            patient_vectors = full_out['patient_vectors']
+            unique_donors = full_out['unique_donors']
+            disease_prob = F.softmax(logits, dim=1)[:, target_class]
+
+    # =============================================
+    # Level 1: Input-level IG (gene & cell importance)
+    # =============================================
+    print(f"Running Input-level Captum IG with {n_steps} steps...")
+    if baseline is not None:
+        baseline = baseline.to(device).float()
+    attribution = run_captum_ig(
+        wrapper, ref_data, target_class=target_class, n_steps=n_steps,
+        baseline=baseline
+    )
+    print("Done.")
+
+    # Gene importance and cell importance from input-level attribution
+    attr_abs = attribution.abs()
+    n_genes = attribution.shape[1]
+    K = len(celltype_names)
+
+    gene_celltype_matrix = torch.zeros(K, n_genes)
+    for c in range(K):
+        mask = (celltype_codes == c)
+        if mask.sum() > 0:
+            gene_celltype_matrix[c] = attribution[mask].mean(dim=0)
+
+    gene_importance = attr_abs.mean(dim=0)          # absolute: gene involvement strength
+    cell_importance = attribution.sum(dim=1)        # SIGNED: net per-cell push (+ -> target_class/Severe, - -> HC)
+
+    # When patient-level attribution is skipped, derive cell-type importance from the
+    # cell-level attribution (mean |cell push| per cell type) and return early -- no
+    # patient-vector IG, no per-donor outputs are computed at all.
+    if skip_patient_level:
+        celltype_importance = torch.zeros(K, device=cell_importance.device)
+        for c in range(K):
+            mask = (celltype_codes == c)
+            if mask.any():
+                celltype_importance[c] = cell_importance[mask].abs().mean()
+        return {
+            'attribution': attribution.detach().cpu(),
+            'gene_celltype_matrix': gene_celltype_matrix.detach().cpu(),
+            'gene_importance': gene_importance.detach().cpu(),
+            'cell_importance': cell_importance.detach().cpu(),
+            'celltype_importance': celltype_importance.detach().cpu(),
+            'celltype_names': celltype_names,
+        }
+
+    # =============================================
+    # Level 2: Patient vector-level IG (cell type & donor importance)
+    # =============================================
+    print(f"Running Patient vector-level Captum IG with {n_steps} steps...")
+
+    def corrector_forward(pv):
+        z = corrector.correct(pv)
+        logits = corrector.status_head(z)
+        return logits
+
+    # Healthy baseline for patient vector
+    if condition_col is not None and reference_labels is not None:
+        donor_cond = ref_metadata.groupby(donor_name)[condition_col].first()
+        healthy_donor_idx = [i for i, d in enumerate(donor_names)
+                             if donor_cond.get(d, '') in reference_labels]
+        if len(healthy_donor_idx) > 0:
+            pv_baseline = patient_vectors[healthy_donor_idx].mean(dim=0, keepdim=True) \
+                .expand_as(patient_vectors).contiguous()
+        else:
+            pv_baseline = torch.zeros_like(patient_vectors)
+    else:
+        pv_baseline = torch.zeros_like(patient_vectors)
+
+
+    ig_patient = IntegratedGradients(corrector_forward)
+    patient_level_attr = ig_patient.attribute(
+        patient_vectors.detach().to(device),
+        baselines=pv_baseline.detach().to(device),
+        target=target_class,
+        n_steps=n_steps
+    )
+    print("Done.")
+
+    # Cell type importance from patient vector attribution
+    feat_dim = patient_vectors.shape[1] // n_celltypes
+    patient_attr_reshaped = patient_level_attr.view(-1, n_celltypes, feat_dim)
+    celltype_importance = patient_attr_reshaped.abs().mean(dim=0).sum(dim=1)  # [K]
+
+    # Donor importance from patient vector attribution
+    donor_attribution = patient_level_attr.abs().sum(dim=1)  # [n_donors]
+
+    return {
+        # Input-level attribution (for gene & cell analysis)
+        'attribution': attribution.detach().cpu(),
+        'gene_celltype_matrix': gene_celltype_matrix.detach().cpu(),
+        'gene_importance': gene_importance.detach().cpu(),
+        'cell_importance': cell_importance.detach().cpu(),
+
+        # Patient vector-level attribution (for cell type & donor analysis)
+        'patient_level_attr': patient_level_attr.detach().cpu(),
+        'celltype_importance': celltype_importance.detach().cpu(),
+        'donor_attribution': donor_attribution.detach().cpu(),
+
+        # Patient info
+        'donor_disease_prob': disease_prob.detach().cpu(),
+        'donor_names': donor_names,
+        'patient_vectors': patient_vectors.detach().cpu(),
+        'celltype_names': celltype_names,
+    }