univi 0.3.4__py3-none-any.whl

univi/hyperparam_optimization/common.py

@@ -0,0 +1,339 @@
+# univi/hyperparam_optimization/common.py
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Dict, Any, List, Optional, Iterable, Tuple
+
+import json
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+
+from univi import UniVIMultiModalVAE, ModalityConfig, UniVIConfig, TrainingConfig
+from univi.data import MultiModalDataset
+from univi.trainer import UniVITrainer
+from univi import evaluation as univi_eval
+
+
+@dataclass
+class SearchResult:
+    config_id: int
+    hparams: Dict[str, Any]
+    best_val_loss: float
+    metrics: Dict[str, float]
+    runtime_min: float
+
+
+def iter_hparam_configs(
+    space_dict: Dict[str, List[Any]],
+    max_configs: int,
+    seed: int = 0,
+) -> Iterable[Dict[str, Any]]:
+    """
+    Random sampler over a hyperparameter space dict.
+    Each config independently samples a value for each key.
+    """
+    rng = np.random.default_rng(seed)
+    keys = list(space_dict.keys())
+    for _ in range(max_configs):
+        hp = {}
+        for k in keys:
+            options = space_dict[k]
+            idx = rng.integers(len(options))
+            hp[k] = options[idx]
+        yield hp
+
+
+def build_modality_configs(
+    arch_config_per_mod: Dict[str, Dict[str, Any]],
+    likelihood_per_mod: Dict[str, str],
+    input_dims: Dict[str, int],
+) -> List[ModalityConfig]:
+    """
+    arch_config_per_mod: e.g. {"rna": {"enc": [...], "dec": [...]}, ...}
+    likelihood_per_mod:  e.g. {"rna": "nb", "atac": "gaussian"}
+    input_dims:          e.g. {"rna": rna.n_vars, "atac": atac.n_vars}
+    """
+    mod_cfgs: List[ModalityConfig] = []
+    for mod_name, arch_cfg in arch_config_per_mod.items():
+        enc = arch_cfg["enc"]
+        dec = arch_cfg["dec"]
+        lik = likelihood_per_mod.get(mod_name, "gaussian")
+        in_dim = int(input_dims[mod_name])
+        mod_cfgs.append(
+            ModalityConfig(
+                name=mod_name,
+                input_dim=in_dim,
+                encoder_hidden=list(enc),
+                decoder_hidden=list(dec),
+                likelihood=lik,
+            )
+        )
+    return mod_cfgs
+
+
+def make_dataloaders(
+    adata_train: Dict[str, "AnnData"],
+    adata_val: Dict[str, "AnnData"],
+    layer: Optional[str],
+    X_key: str,
+    batch_size: int,
+    num_workers: int,
+    device: Optional[str] = None,
+) -> Tuple[DataLoader, DataLoader]:
+    """
+    Thin wrapper to build MultiModalDataset and DataLoader for train/val.
+    """
+    device_obj = None
+    if device is not None and device != "cpu":
+        device_obj = torch.device(device)
+
+    train_ds = MultiModalDataset(
+        adata_dict=adata_train,
+        layer=layer,
+        X_key=X_key,
+        paired=True,
+        device=device_obj,
+    )
+    val_ds = MultiModalDataset(
+        adata_dict=adata_val,
+        layer=layer,
+        X_key=X_key,
+        paired=True,
+        device=device_obj,
+    )
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+    )
+
+    return train_loader, val_loader
+
+
+def train_single_config(
+    config_id: int,
+    hparams: Dict[str, Any],
+    mod_arch_space: Dict[str, List[Dict[str, Any]]],
+    modalities: List[str],
+    input_dims: Dict[str, int],
+    likelihood_per_mod: Dict[str, List[str]],
+    adata_train: Dict[str, "AnnData"],
+    adata_val: Dict[str, "AnnData"],
+    base_train_cfg: TrainingConfig,
+    layer: Optional[str],
+    X_key: str,
+    celltype_key: Optional[str],
+    device: str = "cuda",
+    multimodal_eval: bool = True,
+) -> SearchResult:
+    """
+    Train one UniVI model for a given hyperparam config and return metrics.
+
+    multimodal_eval:
+        - True: compute FOSCTTM, label transfer, modality mixing where possible
+        - False: only val_loss (for unimodal).
+    """
+    start = time.time()
+
+    # ----- pick architectures -----
+    arch_cfg_per_mod: Dict[str, Dict[str, Any]] = {}
+    like_cfg_per_mod: Dict[str, str] = {}
+    for mod in modalities:
+        arch_list = mod_arch_space[mod]
+        arch_choice = hparams[f"{mod}_arch"]
+        if isinstance(arch_choice, dict) and "name" in arch_choice:
+            arch_cfg = arch_choice
+        else:
+            # if user passed just the dict, keep as-is
+            arch_cfg = arch_choice
+        arch_cfg_per_mod[mod] = arch_cfg
+
+        like_options = likelihood_per_mod.get(mod, ["gaussian"])
+        like_choice = hparams.get(f"{mod}_likelihood", like_options[0])
+        like_cfg_per_mod[mod] = like_choice
+
+    # ----- build modality configs -----
+    mod_cfgs = build_modality_configs(
+        arch_config_per_mod=arch_cfg_per_mod,
+        likelihood_per_mod=like_cfg_per_mod,
+        input_dims=input_dims,
+    )
+
+    # ----- UniVI config -----
+    univi_cfg = UniVIConfig(
+        latent_dim=hparams["latent_dim"],
+        modalities=mod_cfgs,
+        beta=hparams["beta"],
+        gamma=hparams["gamma"],
+        encoder_dropout=hparams["encoder_dropout"],
+        decoder_dropout=0.0,
+        encoder_batchnorm=True,
+        decoder_batchnorm=hparams["decoder_batchnorm"],
+        kl_anneal_start=0,
+        kl_anneal_end=0,
+        align_anneal_start=0,
+        align_anneal_end=0,
+    )
+
+    model = UniVIMultiModalVAE(univi_cfg)
+
+    # ----- dataloaders -----
+    train_loader, val_loader = make_dataloaders(
+        adata_train=adata_train,
+        adata_val=adata_val,
+        layer=layer,
+        X_key=X_key,
+        batch_size=base_train_cfg.batch_size,
+        num_workers=base_train_cfg.num_workers,
+        device=device,
+    )
+
+    # ----- training config -----
+    train_cfg = TrainingConfig(
+        n_epochs=base_train_cfg.n_epochs,
+        batch_size=base_train_cfg.batch_size,
+        lr=hparams["lr"],
+        weight_decay=hparams["weight_decay"],
+        device=device,
+        log_every=base_train_cfg.log_every,
+        grad_clip=base_train_cfg.grad_clip,
+        num_workers=base_train_cfg.num_workers,
+        seed=base_train_cfg.seed,
+        early_stopping=base_train_cfg.early_stopping,
+        patience=base_train_cfg.patience,
+        min_delta=base_train_cfg.min_delta,
+    )
+
+    trainer = UniVITrainer(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        train_cfg=train_cfg,
+        device=device,
+    )
+
+    print("=" * 80)
+    print(f"[Config {config_id}] Hyperparameters:")
+    print(json.dumps(hparams, indent=2))
+    print("=" * 80, flush=True)
+
+    history = trainer.fit()
+    best_val_loss = trainer.best_val_loss
+
+    metrics: Dict[str, float] = {"best_val_loss": float(best_val_loss)}
+
+    # ----- evaluation metrics -----
+    if multimodal_eval and len(modalities) >= 2:
+        # encode all modalities for val set
+        Z_val: Dict[str, np.ndarray] = {}
+        for mod in modalities:
+            # use same layer used in training
+            Z = trainer.encode_modality(
+                adata_val[mod],
+                modality=mod,
+                layer=layer,
+                X_key=X_key,
+                batch_size=1024,
+            )
+            Z_val[mod] = Z
+            print(f"  Encoded modality {mod} into latent shape {Z.shape}")
+
+        # pairwise FOSCTTM
+        mods = modalities
+        fos_vals = []
+        for i in range(len(mods)):
+            for j in range(i + 1, len(mods)):
+                m1, m2 = mods[i], mods[j]
+                fos = univi_eval.compute_foscttm(Z_val[m1], Z_val[m2])
+                key = f"foscttm_{m1}_vs_{m2}"
+                metrics[key] = float(fos)
+                fos_vals.append(fos)
+                print(f"  FOSCTTM ({m1} vs {m2}): {fos:.4f}")
+
+        if fos_vals:
+            metrics["foscttm_mean"] = float(np.mean(fos_vals))
+
+        # label transfer: use first modality as "reference" if celltype_key is given
+        if celltype_key is not None and celltype_key in adata_val[modalities[0]].obs:
+            ref_mod = modalities[0]
+            labels_ref = adata_val[ref_mod].obs[celltype_key].astype(str).values
+
+            for tgt_mod in modalities[1:]:
+                labels_tgt = (
+                    adata_val[tgt_mod].obs[celltype_key].astype(str).values
+                    if celltype_key in adata_val[tgt_mod].obs
+                    else None
+                )
+                _, acc, _ = univi_eval.label_transfer_knn(
+                    Z_source=Z_val[ref_mod],
+                    labels_source=labels_ref,
+                    Z_target=Z_val[tgt_mod],
+                    labels_target=labels_tgt,
+                    k=15,
+                )
+                if acc is not None:
+                    key = f"label_acc_{tgt_mod}_from_{ref_mod}"
+                    metrics[key] = float(acc)
+                    print(f"  Label transfer ({ref_mod}→{tgt_mod}) accuracy: {acc:.3f}")
+
+        # modality mixing
+        Z_joint = np.concatenate(list(Z_val.values()), axis=0)
+        modality_labels = np.concatenate(
+            [[m] * Z_val[m].shape[0] for m in modalities]
+        )
+        mix_score = univi_eval.compute_modality_mixing(Z_joint, modality_labels)
+        metrics["modality_mixing"] = float(mix_score)
+        print(f"  Modality mixing score: {mix_score:.4f}")
+
+        # composite score: simple weighted combo
+        if "foscttm_mean" in metrics:
+            comp = best_val_loss + 1000.0 * metrics["foscttm_mean"]
+        else:
+            comp = best_val_loss
+        metrics["composite_score"] = float(comp)
+        print(f"  Composite score: {comp:.3f}")
+    else:
+        # unimodal: composite == val loss
+        metrics["composite_score"] = float(best_val_loss)
+
+    runtime_min = (time.time() - start) / 60.0
+    print(f"[Config {config_id}] Done in {runtime_min:.2f} min")
+    print(
+        f"  best_val_loss         = {best_val_loss:.3f}\n"
+        f"  composite_score       = {metrics['composite_score']:.3f}"
+    )
+
+    return SearchResult(
+        config_id=config_id,
+        hparams=hparams,
+        best_val_loss=float(best_val_loss),
+        metrics=metrics,
+        runtime_min=runtime_min,
+    )
+
+
+def results_to_dataframe(results: List[SearchResult]) -> pd.DataFrame:
+    rows = []
+    for r in results:
+        row = {
+            "config_id": r.config_id,
+            "runtime_min": r.runtime_min,
+        }
+        row.update(r.hparams)
+        row.update(r.metrics)
+        rows.append(row)
+    df = pd.DataFrame(rows)
+    df = df.sort_values("composite_score", ascending=True).reset_index(drop=True)
+    return df

univi/hyperparam_optimization/run_adt_hparam_search.py

@@ -0,0 +1,109 @@
+# univi/hyperparam_optimization/run_adt_hparam_search.py
+
+from __future__ import annotations
+from typing import List, Any, Optional
+
+from anndata import AnnData
+
+from univi.config import TrainingConfig
+from .common import (
+    iter_hparam_configs,
+    train_single_config,
+    results_to_dataframe,
+)
+
+
+def run_adt_hparam_search(
+    adt_train: AnnData,
+    adt_val: AnnData,
+    device: str = "cuda",
+    layer: Optional[str] = "counts",  # raw ADT counts; for CLR/log1p use appropriate layer
+    X_key: str = "X",
+    max_configs: int = 50,
+    seed: int = 0,
+    base_train_cfg: Optional[TrainingConfig] = None,
+):
+    """
+    Hyperparameter search for *unimodal* ADT.
+    """
+
+    adata_train = {"adt": adt_train}
+    adata_val = {"adt": adt_val}
+    modalities = ["adt"]
+
+    if base_train_cfg is None:
+        base_train_cfg = TrainingConfig(
+            n_epochs=80,
+            batch_size=256,
+            lr=1e-3,
+            weight_decay=1e-5,
+            device=device,
+            log_every=5,
+            grad_clip=5.0,
+            num_workers=0,
+            seed=42,
+            early_stopping=True,
+            patience=15,
+            min_delta=0.0,
+        )
+
+    adt_arch_options = [
+        {"name": "adt_small2", "enc": [128, 64], "dec": [64, 128]},
+        {"name": "adt_med2", "enc": [256, 128], "dec": [128, 256]},
+    ]
+
+    mod_arch_space = {"adt": adt_arch_options}
+    likelihood_per_mod = {
+        "adt": ["nb", "zinb", "gaussian"],
+    }
+
+    search_space = {
+        "latent_dim":          [10, 20, 32, 40, 50, 64],
+        "beta":                [0.0, 1.0, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 100.0, 160.0, 200.0, 300.0],
+        "gamma":               [0.0],
+        "lr":                  [1e-3, 5e-4],
+        "weight_decay":        [1e-4, 1e-5],
+        "encoder_dropout":     [0.0, 0.1],
+        "decoder_batchnorm":   [False, True],
+        "adt_arch":            adt_arch_options,
+        "adt_likelihood":      likelihood_per_mod["adt"],
+    }
+
+    input_dims = {"adt": adt_train.n_vars}
+
+    results: List[Any] = []
+    best_score = float("inf")
+    best_cfg = None
+    best_model = None
+
+    for cfg_id, hp in enumerate(
+        iter_hparam_configs(search_space, max_configs=max_configs, seed=seed),
+        start=1,
+    ):
+        res = train_single_config(
+            config_id=cfg_id,
+            hparams=hp,
+            mod_arch_space=mod_arch_space,
+            modalities=modalities,
+            input_dims=input_dims,
+            likelihood_per_mod={"adt": likelihood_per_mod["adt"]},
+            adata_train=adata_train,
+            adata_val=adata_val,
+            base_train_cfg=base_train_cfg,
+            layer=layer,
+            X_key=X_key,
+            celltype_key=None,
+            device=device,
+            multimodal_eval=False,
+        )
+        results.append(res)
+
+        score = res.metrics["composite_score"]
+        if score < best_score:
+            best_score = score
+            best_model = res
+            best_cfg = hp
+            print(f"--> New best ADT-only config (id={cfg_id}) with score={score:.3f}")
+
+    df = results_to_dataframe(results)
+    return df, best_model, best_cfg

univi/hyperparam_optimization/run_atac_hparam_search.py

@@ -0,0 +1,109 @@
+# univi/hyperparam_optimization/run_atac_hparam_search.py
+
+from __future__ import annotations
+from typing import List, Any, Optional
+
+from anndata import AnnData
+
+from univi.config import TrainingConfig
+from .common import (
+    iter_hparam_configs,
+    train_single_config,
+    results_to_dataframe,
+)
+
+
+def run_atac_hparam_search(
+    atac_train: AnnData,
+    atac_val: AnnData,
+    device: str = "cuda",
+    layer: Optional[str] = "counts",  # raw peaks; or change to TF-IDF/LSI etc.
+    X_key: str = "X",
+    max_configs: int = 50,
+    seed: int = 0,
+    base_train_cfg: Optional[TrainingConfig] = None,
+):
+    """
+    Hyperparameter search for *unimodal* ATAC.
+    """
+
+    adata_train = {"atac": atac_train}
+    adata_val = {"atac": atac_val}
+    modalities = ["atac"]
+
+    if base_train_cfg is None:
+        base_train_cfg = TrainingConfig(
+            n_epochs=80,
+            batch_size=256,
+            lr=1e-3,
+            weight_decay=1e-5,
+            device=device,
+            log_every=5,
+            grad_clip=5.0,
+            num_workers=0,
+            seed=42,
+            early_stopping=True,
+            patience=15,
+            min_delta=0.0,
+        )
+
+    atac_arch_options = [
+        {"name": "atac_med2", "enc": [512, 256], "dec": [256, 512]},
+        {"name": "atac_wide2", "enc": [1024, 512], "dec": [512, 1024]},
+    ]
+
+    mod_arch_space = {"atac": atac_arch_options}
+    likelihood_per_mod = {
+        "atac": ["nb", "poisson", "zinb", "gaussian"],
+    }
+
+    search_space = {
+        "latent_dim":       [10, 20, 32, 40, 50, 64, 82, 120, 160, 200],
+        "beta":             [0.0, 1.0, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 100.0, 160.0, 200.0, 300.0],
+        "gamma":               [0.0],
+        "lr":                  [1e-3, 5e-4],
+        "weight_decay":        [1e-4, 1e-5],
+        "encoder_dropout":     [0.0, 0.1],
+        "decoder_batchnorm":   [False, True],
+        "atac_arch":           atac_arch_options,
+        "atac_likelihood":     likelihood_per_mod["atac"],
+    }
+
+    input_dims = {"atac": atac_train.n_vars}
+
+    results: List[Any] = []
+    best_score = float("inf")
+    best_cfg = None
+    best_model = None
+
+    for cfg_id, hp in enumerate(
+        iter_hparam_configs(search_space, max_configs=max_configs, seed=seed),
+        start=1,
+    ):
+        res = train_single_config(
+            config_id=cfg_id,
+            hparams=hp,
+            mod_arch_space=mod_arch_space,
+            modalities=modalities,
+            input_dims=input_dims,
+            likelihood_per_mod={"atac": likelihood_per_mod["atac"]},
+            adata_train=adata_train,
+            adata_val=adata_val,
+            base_train_cfg=base_train_cfg,
+            layer=layer,
+            X_key=X_key,
+            celltype_key=None,
+            device=device,
+            multimodal_eval=False,
+        )
+        results.append(res)
+
+        score = res.metrics["composite_score"]
+        if score < best_score:
+            best_score = score
+            best_model = res
+            best_cfg = hp
+            print(f"--> New best ATAC-only config (id={cfg_id}) with score={score:.3f}")
+
+    df = results_to_dataframe(results)
+    return df, best_model, best_cfg

univi/hyperparam_optimization/run_citeseq_hparam_search.py

@@ -0,0 +1,137 @@
+# univi/hyperparam_optimization/run_citeseq_hparam_search.py
+
+from __future__ import annotations
+from typing import Dict, Any, List, Optional, Tuple
+
+from anndata import AnnData
+
+from univi.config import TrainingConfig
+from .common import (
+    iter_hparam_configs,
+    train_single_config,
+    results_to_dataframe,
+)
+
+
+def run_citeseq_hparam_search(
+    rna_train: AnnData,
+    adt_train: AnnData,
+    rna_val: AnnData,
+    adt_val: AnnData,
+    celltype_key: Optional[str] = "cell_type",
+    device: str = "cuda",
+    layer: str = "counts",  # raw counts for NB / ZINB
+    X_key: str = "X",
+    max_configs: int = 100,
+    seed: int = 0,
+    base_train_cfg: Optional[TrainingConfig] = None,
+):
+    """
+    Hyperparameter random search for RNA+ADT CITE-seq.
+
+    Assumes:
+    - rna_* and adt_* are paired and share obs_names.
+    - raw counts for RNA and ADT are stored in .layers[layer] (default 'counts').
+    """
+
+    assert rna_train.n_obs == adt_train.n_obs
+    assert rna_val.n_obs == adt_val.n_obs
+
+    adata_train = {"rna": rna_train, "adt": adt_train}
+    adata_val = {"rna": rna_val, "adt": adt_val}
+    modalities = ["rna", "adt"]
+
+    if base_train_cfg is None:
+        base_train_cfg = TrainingConfig(
+            n_epochs=80,
+            batch_size=256,
+            lr=1e-3,
+            weight_decay=1e-5,
+            device=device,
+            log_every=5,
+            grad_clip=5.0,
+            num_workers=0,
+            seed=42,
+            early_stopping=True,
+            patience=15,
+            min_delta=0.0,
+        )
+
+    rna_arch_options = [
+        {"name": "rna_med2", "enc": [512, 256], "dec": [256, 512]},
+        {"name": "rna_wide2", "enc": [1024, 512], "dec": [512, 1024]},
+        {"name": "rna_wide3", "enc": [1024, 512, 256], "dec": [256, 512, 1024]},
+    ]
+    adt_arch_options = [
+        {"name": "adt_small2", "enc": [128, 64], "dec": [64, 128]},
+        {"name": "adt_med2", "enc": [256, 128], "dec": [128, 256]},
+    ]
+
+    mod_arch_space = {
+        "rna": rna_arch_options,
+        "adt": adt_arch_options,
+    }
+
+    likelihood_per_mod = {
+        "rna": ["nb", "zinb"],
+        "adt": ["nb", "zinb", "gaussian"],  # if using CLR/log1p etc.
+    }
+
+    search_space = {
+        "latent_dim":       [10, 20, 32, 40, 50, 64, 82, 120, 160, 200],
+        "beta":             [0.0, 1.0, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 100.0, 160.0, 200.0, 300.0],
+        "gamma":            [0.0, 1.0, 5.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0, 80.0, 100.0, 160.0, 200.0, 300.0],
+        "lr":               [1e-3, 5e-4],
+        "weight_decay":     [1e-4, 1e-5],
+        "encoder_dropout":  [0.0, 0.1],
+        "decoder_batchnorm":[False, True],
+        "rna_arch":         rna_arch_options,
+        "adt_arch":         adt_arch_options,
+        "rna_likelihood":   likelihood_per_mod["rna"],
+        "adt_likelihood":   likelihood_per_mod["adt"],
+    }
+
+    input_dims = {
+        "rna": rna_train.n_vars,
+        "adt": adt_train.n_vars,
+    }
+
+    results: List[Any] = []
+    best_score = float("inf")
+    best_model = None
+    best_cfg = None
+
+    for cfg_id, hp in enumerate(
+        iter_hparam_configs(search_space, max_configs=max_configs, seed=seed),
+        start=1,
+    ):
+        res = train_single_config(
+            config_id=cfg_id,
+            hparams=hp,
+            mod_arch_space=mod_arch_space,
+            modalities=modalities,
+            input_dims=input_dims,
+            likelihood_per_mod={
+                "rna": likelihood_per_mod["rna"],
+                "adt": likelihood_per_mod["adt"],
+            },
+            adata_train=adata_train,
+            adata_val=adata_val,
+            base_train_cfg=base_train_cfg,
+            layer=layer,
+            X_key=X_key,
+            celltype_key=celltype_key,
+            device=device,
+            multimodal_eval=True,
+        )
+        results.append(res)
+
+        score = res.metrics["composite_score"]
+        if score < best_score:
+            best_score = score
+            best_model = res
+            best_cfg = hp
+            print(f"--> New best CITE-seq config (id={cfg_id}) with score={score:.3f}")
+
+    df = results_to_dataframe(results)
+    return df, best_model, best_cfg