univi 0.3.4__py3-none-any.whl

univi/objectives.py

@@ -0,0 +1,46 @@
+# univi/objectives.py
+
+from __future__ import annotations
+from typing import Dict
+import torch
+
+
+def kl_diag_gaussians(
+    mu_q: torch.Tensor,
+    logvar_q: torch.Tensor,
+    mu_p: torch.Tensor,
+    logvar_p: torch.Tensor,
+) -> torch.Tensor:
+    """
+    KL(q||p) for diagonal Gaussians, summed over latent dim, per sample.
+    """
+    var_q = torch.exp(logvar_q)
+    var_p = torch.exp(logvar_p)
+    kl = (
+        logvar_p
+        - logvar_q
+        + (var_q + (mu_q - mu_p) ** 2) / var_p
+        - 1.0
+    )
+    return 0.5 * kl.sum(dim=-1)
+
+
+def symmetric_alignment_loss(
+    mu_per_mod: Dict[str, torch.Tensor],
+) -> torch.Tensor:
+    """
+    Simple symmetric cross-modal alignment: mean pairwise L2 distance
+    between latent means across modalities.
+    """
+    names = list(mu_per_mod.keys())
+    if len(names) < 2:
+        return torch.zeros(mu_per_mod[names[0]].size(0), device=mu_per_mod[names[0]].device)
+
+    losses = []
+    for i in range(len(names)):
+        for j in range(i + 1, len(names)):
+            mu_i = mu_per_mod[names[i]]
+            mu_j = mu_per_mod[names[j]]
+            losses.append(((mu_i - mu_j) ** 2).sum(dim=-1))
+    stacked = torch.stack(losses, dim=0)
+    return stacked.mean(dim=0)

univi/pipeline.py

@@ -0,0 +1,194 @@
+# univi/pipeline.py
+from __future__ import annotations
+
+from dataclasses import asdict
+from typing import Any, Dict, Optional, Tuple
+
+import os
+import numpy as np
+import torch
+import anndata as ad
+from anndata import AnnData
+
+from .config import UniVIConfig, ModalityConfig, TrainingConfig
+from .models import UniVIMultiModalVAE
+from .data import align_paired_obs_names, infer_input_dim, _get_matrix
+from .utils.io import load_config, load_checkpoint
+from .utils.seed import set_seed
+
+
+def load_anndata_dict(cfg: Dict[str, Any], *, data_root: Optional[str] = None) -> Dict[str, AnnData]:
+    adata_dict: Dict[str, AnnData] = {}
+    for m in cfg["data"]["modalities"]:
+        name = m["name"]
+        path = m["h5ad_path"]
+        if data_root is not None and not os.path.isabs(path):
+            path = os.path.join(data_root, path)
+        adata_dict[name] = ad.read_h5ad(path)
+    return adata_dict
+
+
+def build_univi_from_config(
+    cfg: Dict[str, Any],
+    adata_dict: Dict[str, AnnData],
+) -> Tuple[UniVIMultiModalVAE, UniVIConfig, Dict[str, Optional[str]], Dict[str, str]]:
+    """Build UniVI model + selectors from a loaded config and loaded AnnData dict."""
+    mcfgs = cfg["data"]["modalities"]
+    model_cfg = cfg.get("model", {})
+
+    # selectors
+    layer_by = {m["name"]: m.get("layer", None) for m in mcfgs}
+    xkey_by = {m["name"]: m.get("X_key", "X") for m in mcfgs}
+
+    modalities = []
+    for m in mcfgs:
+        name = m["name"]
+        input_dim = infer_input_dim(adata_dict[name], layer=layer_by[name], X_key=xkey_by[name])
+
+        hidden_default = model_cfg.get("hidden_dims_default", [256, 128])
+        enc = m.get("encoder_hidden", m.get("hidden_dims", hidden_default))
+        dec = m.get("decoder_hidden", m.get("decoder_hidden", list(enc)[::-1]))
+        modalities.append(
+            ModalityConfig(
+                name=name,
+                input_dim=int(input_dim),
+                encoder_hidden=list(enc),
+                decoder_hidden=list(dec),
+                likelihood=m.get("likelihood", "gaussian"),
+            )
+        )
+
+    univi_cfg = UniVIConfig(
+        latent_dim=int(model_cfg.get("latent_dim", 32)),
+        modalities=modalities,
+        beta=float(model_cfg.get("beta", 1.0)),
+        gamma=float(model_cfg.get("gamma", 1.0)),
+        encoder_dropout=float(model_cfg.get("encoder_dropout", model_cfg.get("dropout", 0.0))),
+        decoder_dropout=float(model_cfg.get("decoder_dropout", model_cfg.get("dropout", 0.0))),
+        encoder_batchnorm=bool(model_cfg.get("encoder_batchnorm", model_cfg.get("batchnorm", True))),
+        decoder_batchnorm=bool(model_cfg.get("decoder_batchnorm", False)),
+        kl_anneal_start=int(model_cfg.get("kl_anneal_start", 0)),
+        kl_anneal_end=int(model_cfg.get("kl_anneal_end", 0)),
+        align_anneal_start=int(model_cfg.get("align_anneal_start", 0)),
+        align_anneal_end=int(model_cfg.get("align_anneal_end", 0)),
+    )
+
+    model = UniVIMultiModalVAE(
+        univi_cfg,
+        loss_mode=model_cfg.get("loss_mode", "v2"),
+        v1_recon=model_cfg.get("v1_recon", "cross"),
+        v1_recon_mix=float(model_cfg.get("v1_recon_mix", 0.0)),
+        normalize_v1_terms=bool(model_cfg.get("normalize_v1_terms", True)),
+    )
+    return model, univi_cfg, layer_by, xkey_by
+
+
+@torch.no_grad()
+def encode_latents_single_modality(
+    model: UniVIMultiModalVAE,
+    adata: AnnData,
+    modality: str,
+    *,
+    layer: Optional[str] = None,
+    X_key: str = "X",
+    batch_size: int = 512,
+    device: str = "cpu",
+) -> np.ndarray:
+    """Encode *one modality only* (for unimodal / independent dataset experiments)."""
+    model.eval()
+    model.to(device)
+
+    X = _get_matrix(adata, layer=layer, X_key=X_key)
+    # materialize sparse batches safely
+    n = X.shape[0]
+    zs = []
+    for start in range(0, n, batch_size):
+        end = min(n, start + batch_size)
+        xb = X[start:end]
+        if hasattr(xb, "A"):
+            xb = xb.A
+        xb = torch.as_tensor(np.asarray(xb), dtype=torch.float32, device=device)
+        mu_dict, logvar_dict = model.encode_modalities({modality: xb})
+        mu_z, logvar_z = model.mixture_of_experts(mu_dict, logvar_dict)
+        zs.append(mu_z.detach().cpu().numpy())
+    return np.vstack(zs)
+
+
+@torch.no_grad()
+def encode_latents_paired(
+    model: UniVIMultiModalVAE,
+    adata_dict: Dict[str, AnnData],
+    *,
+    layer_by: Optional[Dict[str, Optional[str]]] = None,
+    xkey_by: Optional[Dict[str, str]] = None,
+    batch_size: int = 512,
+    device: str = "cpu",
+    fused: bool = True,
+) -> Dict[str, np.ndarray]:
+    """Encode paired cells for each modality and optionally the fused MoE latent.
+
+    Returns dict with keys for each modality (mu of that encoder) and optionally 'fused'.
+    """
+    model.eval()
+    model.to(device)
+
+    names = list(adata_dict.keys())
+    n = adata_dict[names[0]].n_obs
+    # require paired order
+    for nm in names[1:]:
+        if not np.array_equal(adata_dict[nm].obs_names.values, adata_dict[names[0]].obs_names.values):
+            raise ValueError(f"obs_names mismatch between {names[0]} and {nm}")
+
+    out = {nm: [] for nm in names}
+    if fused:
+        out["fused"] = []
+
+    for start in range(0, n, batch_size):
+        end = min(n, start + batch_size)
+        x_dict = {}
+        for nm, adata in adata_dict.items():
+            layer = None if layer_by is None else layer_by.get(nm, None)
+            xkey  = "X" if xkey_by is None else xkey_by.get(nm, "X")
+            X = _get_matrix(adata, layer=layer, X_key=xkey)[start:end]
+            if hasattr(X, "A"):
+                X = X.A
+            x_dict[nm] = torch.as_tensor(np.asarray(X), dtype=torch.float32, device=device)
+
+        mu_dict, logvar_dict = model.encode_modalities(x_dict)
+        # per-modality mus
+        for nm in names:
+            out[nm].append(mu_dict[nm].detach().cpu().numpy())
+        if fused:
+            mu_z, logvar_z = model.mixture_of_experts(mu_dict, logvar_dict)
+            out["fused"].append(mu_z.detach().cpu().numpy())
+
+    for k in list(out.keys()):
+        out[k] = np.vstack(out[k])
+    return out
+
+
+def load_model_and_data(
+    config_path: str,
+    *,
+    checkpoint_path: Optional[str] = None,
+    data_root: Optional[str] = None,
+    device: str = "cpu",
+    align_obs: bool = True,
+) -> Tuple[Dict[str, Any], Dict[str, AnnData], UniVIMultiModalVAE, Dict[str, Optional[str]], Dict[str, str]]:
+    cfg = load_config(config_path)
+    seed = int(cfg.get("training", {}).get("seed", 0))
+    set_seed(seed, deterministic=False)
+
+    adata_dict = load_anndata_dict(cfg, data_root=data_root)
+    if align_obs:
+        adata_dict = align_paired_obs_names(adata_dict)
+
+    model, univi_cfg, layer_by, xkey_by = build_univi_from_config(cfg, adata_dict)
+    if checkpoint_path is not None:
+        ck = load_checkpoint(checkpoint_path)
+        # common patterns
+        state = ck.get("model_state", ck.get("state_dict", ck))
+        model.load_state_dict(state, strict=False)
+
+    model.to(device)
+    return cfg, adata_dict, model, layer_by, xkey_by

univi/plotting.py

@@ -0,0 +1,126 @@
+# univi/plotting.py
+
+from __future__ import annotations
+from typing import Dict, Optional, Sequence, List
+
+import numpy as np
+import matplotlib.pyplot as plt
+import scanpy as sc
+import seaborn as sns
+from anndata import AnnData
+import anndata as ad
+
+
+def set_style(font_scale: float = 1.25, dpi: int = 150) -> None:
+    """Readable, manuscript-friendly plotting defaults."""
+    import matplotlib as mpl
+
+    base = 10.0 * float(font_scale)
+    mpl.rcParams.update({
+        "figure.dpi": int(dpi),
+        "savefig.dpi": 300,
+        "font.size": base,
+        "axes.titlesize": base * 1.2,
+        "axes.labelsize": base * 1.1,
+        "xtick.labelsize": base * 0.95,
+        "ytick.labelsize": base * 0.95,
+        "legend.fontsize": base * 0.95,
+        "pdf.fonttype": 42,
+        "ps.fonttype": 42,
+    })
+    sc.settings.set_figure_params(dpi=int(dpi), dpi_save=300, frameon=False)
+
+
+def umap_single_adata(
+    adata_obj: AnnData,
+    obsm_key: str = "X_univi",
+    color: Optional[Sequence[str]] = None,
+    savepath: Optional[str] = None,
+    n_neighbors: int = 30,
+    random_state: int = 0,
+) -> None:
+    if obsm_key not in adata_obj.obsm:
+        raise KeyError("Missing obsm[%r]. Available: %s" % (obsm_key, list(adata_obj.obsm.keys())))
+
+    # Compute neighbors/umap if missing
+    if "neighbors" not in adata_obj.uns:
+        sc.pp.neighbors(adata_obj, use_rep=obsm_key, n_neighbors=int(n_neighbors))
+    if "X_umap" not in adata_obj.obsm:
+        sc.tl.umap(adata_obj, random_state=int(random_state))
+
+    if color is None:
+        color = []
+    sc.pl.umap(adata_obj, color=list(color), show=False)
+
+    if savepath is not None:
+        plt.savefig(savepath, dpi=300, bbox_inches="tight")
+    plt.close()
+
+
+def umap_by_modality(
+    adata_dict: Dict[str, AnnData],
+    obsm_key: str = "X_univi",
+    color: str = "celltype",
+    savepath: Optional[str] = None,
+    n_neighbors: int = 30,
+    random_state: int = 0,
+) -> None:
+    """
+    Concatenate adatas; expects each input adata already has the same obsm_key embedding
+    (or you should add it before calling this).
+    """
+    annotated: List[AnnData] = []
+    for mod, a in adata_dict.items():
+        aa = a.copy()
+        aa.obs["univi_modality"] = str(mod)
+        annotated.append(aa)
+
+    combined = annotated[0].concatenate(
+        *annotated[1:],
+        batch_key="univi_source",
+        batch_categories=list(adata_dict.keys()),
+        index_unique="-",
+        join="outer",
+    )
+
+    # Carry embeddings if needed (common pitfall: concatenate drops obsm keys)
+    if obsm_key not in combined.obsm:
+        # try to rebuild from parts
+        try:
+            Zs = [adata_dict[m].obsm[obsm_key] for m in adata_dict.keys()]
+            combined.obsm[obsm_key] = np.vstack(Zs)
+        except Exception:
+            raise KeyError("combined is missing obsm[%r] after concatenation; add it manually." % obsm_key)
+
+    umap_single_adata(
+        combined,
+        obsm_key=obsm_key,
+        color=[color, "univi_modality"],
+        savepath=savepath,
+        n_neighbors=n_neighbors,
+        random_state=random_state,
+    )
+
+
+def plot_confusion_matrix(
+    cm: np.ndarray,
+    labels: np.ndarray,
+    title: str = "Label transfer (source \u2192 target)",
+    savepath: Optional[str] = None,
+) -> None:
+    plt.figure(figsize=(6, 5))
+    sns.heatmap(
+        cm,
+        annot=False,
+        xticklabels=labels,
+        yticklabels=labels,
+        cmap="viridis",
+    )
+    plt.xlabel("Predicted")
+    plt.ylabel("True")
+    plt.title(title)
+    plt.tight_layout()
+    if savepath is not None:
+        plt.savefig(savepath, dpi=300, bbox_inches="tight")
+    plt.close()
+