fusemap 0.0.2__tar.gz → 1.1.0__tar.gz

{fusemap-0.0.2 → fusemap-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusemap
-Version: 0.0.2
+Version: 1.1.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: scanpy==1.9.3
@@ -71,7 +71,8 @@ seaborn
 
 
 ## Installation and Tutorial
-- Read the tutorial [here](https://fusemap.readthedocs.io/en/latest/) .
+- Read the FuseMap tutorial [here](https://fusemap.readthedocs.io/en/latest/).
+- FuseMap-Agent set up tutorial [here](https://github.com/wanglab-broad/FuseMap/tree/main/agent_setup).
 
 
 ## Citation

{fusemap-0.0.2 → fusemap-1.1.0}/README.md

@@ -40,7 +40,8 @@ seaborn
 
 
 ## Installation and Tutorial
-- Read the tutorial [here](https://fusemap.readthedocs.io/en/latest/) .
+- Read the FuseMap tutorial [here](https://fusemap.readthedocs.io/en/latest/).
+- FuseMap-Agent set up tutorial [here](https://github.com/wanglab-broad/FuseMap/tree/main/agent_setup).
 
 
 ## Citation

fusemap-1.1.0/fusemap/__init__.py

@@ -0,0 +1,13 @@
+from fusemap.training.integrate import *
+from fusemap.training.map import *
+from fusemap.logger import *
+from fusemap.utils import *
+from fusemap.config import *
+from fusemap.models.network import *
+from fusemap.data.loaders import *
+from fusemap.models.losses import *
+from fusemap.training.train_model import *
+from fusemap.data.graph import *
+from fusemap.api import *
+from fusemap import metrics
+from fusemap import pl

fusemap-1.1.0/fusemap/api.py

@@ -0,0 +1,285 @@
+"""One-call convenience API.
+
+These wrappers make the common workflows single function calls::
+
+    import fusemap
+
+    fusemap.integrate("./data", "./output")
+    fusemap.map_to_reference("./new_data", "./output_map", "./pretrained_model")
+    fusemap.deconvolve_beads("./output", "./data",
+                             bead_files="slideseq", sig_ref="starmap")
+
+Every argument that ``main.py`` accepts on the command line is available as a
+keyword argument; everything else uses the validated defaults.
+"""
+
+import logging
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import scanpy as sc
+
+from fusemap.logger import setup_logging
+from fusemap.training.integrate import spatial_integrate
+from fusemap.training.map import spatial_map
+from fusemap.utils import seed_all
+
+__all__ = ["integrate", "map_to_reference", "deconvolve_beads", "transfer_labels", "read_input_folder"]
+
+
+def read_input_folder(input_data_folder_path):
+    """Read every ``.h5ad`` in a folder and normalize spatial coordinates.
+
+    Coordinates are taken from ``obs['x']/obs['y']``, falling back to
+    ``obs['col']/obs['row']`` and then ``obsm['spatial']``.
+
+    Returns
+    -------
+    list of :class:`anndata.AnnData`, one per section.
+    """
+    folder = Path(input_data_folder_path)
+    file_names = sorted(str(p) for p in folder.iterdir() if p.is_file() and p.suffix == ".h5ad")
+    if not file_names:
+        raise ValueError(f"No .h5ad files found in {folder}")
+
+    X_input = []
+    for ind, file_name_i in enumerate(file_names):
+        X = sc.read_h5ad(file_name_i)
+        if "x" not in X.obs.columns:
+            if "col" in X.obs.columns and "row" in X.obs.columns:
+                X.obs["x"] = X.obs["col"]
+                X.obs["y"] = X.obs["row"]
+            elif "spatial" in X.obsm:
+                X.obs["x"] = X.obsm["spatial"][:, 0]
+                X.obs["y"] = X.obsm["spatial"][:, 1]
+            else:
+                raise ValueError(
+                    f"{file_name_i}: provide spatial coordinates in obs['x']/obs['y'], "
+                    "obs['col']/obs['row'], or obsm['spatial']"
+                )
+        X.obs["name"] = f"section{ind}"
+        X.obs["file_name"] = os.path.basename(file_name_i)
+        X_input.append(X)
+    return X_input
+
+
+def _make_args(output_save_dir, keep_celltype, keep_tissueregion,
+               use_llm_gene_embedding, pretrain_model_path):
+    Path(output_save_dir).mkdir(parents=True, exist_ok=True)
+    return SimpleNamespace(
+        output_save_dir=str(output_save_dir),
+        keep_celltype=keep_celltype,
+        keep_tissueregion=keep_tissueregion,
+        use_llm_gene_embedding=use_llm_gene_embedding,
+        pretrain_model_path=str(pretrain_model_path),
+    )
+
+
+def integrate(input_data_folder_path, output_save_dir,
+              keep_celltype="", keep_tissueregion="",
+              use_llm_gene_embedding="false",
+              bead_files=None, sig_ref=None):
+    """Integrate all spatial sections in a folder into shared embeddings.
+
+    Equivalent to ``python main.py --mode integrate``. Training resumes
+    automatically from ``snapshot.pt`` if the same command is re-run.
+
+    Parameters
+    ----------
+    input_data_folder_path
+        Folder with one ``.h5ad`` file per section.
+    output_save_dir
+        Output directory; embeddings are written here as
+        ``ad_celltype_embedding.h5ad``, ``ad_tissueregion_embedding.h5ad``,
+        and ``ad_gene_embedding.h5ad``.
+    keep_celltype, keep_tissueregion
+        Optional ``obs`` column names with existing labels to carry into the
+        output embedding files.
+    bead_files, sig_ref
+        Declare bead/spot-resolution datasets (Slide-seq, Visium ...) and the
+        single-cell signature reference(s), as comma-separated file-name
+        substrings. When declared, Stage-B deconvolution runs AUTOMATICALLY
+        after training (recommended default for any bead platform), and the
+        DECOMPOSED embeddings become the canonical outputs
+        (``ad_celltype_embedding.h5ad`` / ``ad_tissueregion_embedding.h5ad``);
+        the pre-decomposition versions are kept as ``ad_*_embedding_nodeconv.h5ad``.
+        Roles are always user-declared, never auto-detected.
+    """
+    seed_all(0)
+    args = _make_args(output_save_dir, keep_celltype, keep_tissueregion,
+                      use_llm_gene_embedding, "")
+    setup_logging(args.output_save_dir)
+    logging.info("Arguments: %s", vars(args))
+
+    X_input = read_input_folder(input_data_folder_path)
+    kneighbor = ["delaunay"] * len(X_input)
+    input_identity = ["ST"] * len(X_input)
+    spatial_integrate(X_input, args, kneighbor, input_identity)
+
+    if bead_files:
+        if not sig_ref:
+            raise ValueError("bead_files declared without sig_ref (single-cell signature reference)")
+        logging.info("[FuseMap] bead datasets declared - running Stage-B deconvolution")
+        deconvolve_beads(output_save_dir, input_data_folder_path, bead_files, sig_ref)
+
+
+def map_to_reference(input_data_folder_path, output_save_dir, pretrain_model_path,
+                     keep_celltype="", keep_tissueregion="",
+                     use_llm_gene_embedding="false"):
+    """Map each section in a folder onto a pretrained FuseMap model.
+
+    Equivalent to ``python main.py --mode map``. Results are written to one
+    subdirectory of ``output_save_dir`` per input file.
+    """
+    import copy as _copy
+
+    seed_all(0)
+    X_input = read_input_folder(input_data_folder_path)
+    for X in X_input:
+        args_i = _make_args(
+            os.path.join(str(output_save_dir), X.obs["file_name"].unique()[0]),
+            keep_celltype, keep_tissueregion,
+            use_llm_gene_embedding, pretrain_model_path,
+        )
+        setup_logging(args_i.output_save_dir)
+        spatial_map([X], args_i, ["delaunay"], ["ST"])
+
+
+def deconvolve_beads(output_save_dir, input_data_folder_path,
+                     bead_files, sig_ref,
+                     entropy_weight=None, signature_mode=None):
+    """Deconvolve bead/spot-resolution datasets after an integration run.
+
+    Each bead is decomposed into a mixture over cell archetypes learned from
+    the single-cell sections (Stage-B); its cell and tissue embeddings are
+    rebuilt from the mixture. The decomposed embeddings become the canonical
+    ``ad_celltype_embedding.h5ad`` / ``ad_tissueregion_embedding.h5ad`` (bead
+    rows replaced; single-cell rows unchanged). The pre-decomposition versions
+    are kept as ``ad_*_embedding_nodeconv.h5ad``; ``ad_*_embedding_stageB.h5ad``
+    (identical to canonical) and ``stageB_pi.npz`` (per-bead mixtures) are also
+    written into ``output_save_dir``.
+
+    Parameters
+    ----------
+    output_save_dir
+        Output directory of the finished :func:`integrate` run.
+    input_data_folder_path
+        The data folder that was integrated.
+    bead_files
+        Comma-separated file-name substrings, each matching exactly one input
+        file: the bead-resolution dataset(s) to deconvolve.
+    sig_ref
+        Comma-separated substrings selecting the single-cell reference
+        dataset(s) used to build archetype signatures.
+    entropy_weight
+        Entropy regularization of the mixture weights (default ``5e-4``).
+    signature_mode
+        ``"empirical"`` (default, recommended) or ``"decode"``.
+    """
+    import runpy
+
+    os.environ["STAGEB_OUT_DIR"] = str(output_save_dir)
+    os.environ["STAGEB_DATA_DIR"] = str(input_data_folder_path)
+    os.environ["FUSEMAP_BEAD_FILES"] = bead_files
+    os.environ["FUSEMAP_SIG_REF"] = sig_ref
+    if entropy_weight is not None:
+        os.environ["STAGEB_ENT_W"] = str(entropy_weight)
+    if signature_mode is not None:
+        os.environ["STAGEB_SIG"] = str(signature_mode)
+
+    script = Path(__file__).resolve().parent / "postprocess" / "stage_b_script.py"
+    runpy.run_path(str(script), run_name="__main__")
+
+    # Decomposed embeddings become the CANONICAL outputs (default for any run
+    # with declared bead datasets). The pre-decomposition versions are kept as
+    # *_nodeconv.h5ad; *_stageB.h5ad names are kept for compatibility.
+    import shutil as _sh
+    for lvl in ("celltype", "tissueregion"):
+        canon = os.path.join(str(output_save_dir), f"ad_{lvl}_embedding.h5ad")
+        stage = os.path.join(str(output_save_dir), f"ad_{lvl}_embedding_stageB.h5ad")
+        keep = os.path.join(str(output_save_dir), f"ad_{lvl}_embedding_nodeconv.h5ad")
+        if os.path.exists(stage) and os.path.exists(canon):
+            if not os.path.exists(keep):
+                os.rename(canon, keep)
+            _sh.copyfile(stage, canon)
+    logging.info("[FuseMap] Stage-B embeddings promoted to canonical outputs "
+                 "(pre-decomposition kept as *_nodeconv.h5ad)")
+
+def transfer_labels(adata, label_key, batch_size=256, epochs=200,
+                    unlabeled_values=("nan", "Unannotated", ""), device=None):
+    """Transfer labels between datasets through the shared embedding — one call.
+
+    Trains a small classifier on all cells that carry a label in
+    ``obs[label_key]`` (the reference) and predicts that label for every cell
+    in ``adata``, adding two columns to ``obs``:
+
+    - ``transfer_<label_key>`` — the transferred label
+    - ``transfer_<label_key>_uncertainty`` — 1 - prediction confidence
+
+    Works on any embedding AnnData produced by :func:`integrate` or
+    :func:`map_to_reference` (latent in ``X``). Cells whose label is missing
+    or in ``unlabeled_values`` are excluded from training.
+
+    Returns
+    -------
+    dict with ``labels`` (np.ndarray), ``uncertainty`` (np.ndarray) and
+    ``test_accuracy`` (balanced per-class accuracy on a held-out 10% of the reference).
+    """
+    import numpy as np
+    import torch
+    from sklearn import preprocessing
+    from torch import nn
+    from torch.utils.data import DataLoader, TensorDataset, random_split
+
+    from fusemap.models.network import NNTransfer
+    from fusemap.utils import NNTransferPredictWithUncertainty, NNTransferTrain
+
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    X = np.asarray(adata.X, dtype=np.float32)
+    lab = adata.obs[label_key].astype(str).values
+    is_ref = ~np.isin(lab, list(unlabeled_values)) & (lab != "nan")
+    if is_ref.sum() < 100:
+        raise ValueError(f"only {is_ref.sum()} labeled reference cells for '{label_key}'")
+
+    le = preprocessing.LabelEncoder()
+    y_ref = le.fit_transform(lab[is_ref]).astype(int)
+
+    ds = TensorDataset(torch.tensor(X[is_ref]), torch.tensor(y_ref).long())
+    n_train = int(0.8 * len(ds))
+    n_val = int(0.1 * len(ds))
+    train_ds, val_ds, test_ds = random_split(
+        ds, [n_train, n_val, len(ds) - n_train - n_val])
+    loaders = dict(
+        train=DataLoader(train_ds, batch_size=batch_size, shuffle=True),
+        val=DataLoader(val_ds, batch_size=batch_size, shuffle=False),
+        test=DataLoader(test_ds, batch_size=batch_size, shuffle=False),
+    )
+
+    import sklearn.utils
+
+    model = NNTransfer(input_dim=X.shape[1], output_dim=len(le.classes_)).to(device)
+    class_weight = torch.Tensor(sklearn.utils.class_weight.compute_class_weight(
+        class_weight="balanced", classes=np.unique(y_ref), y=y_ref))
+    criterion = nn.CrossEntropyLoss(weight=class_weight.to(device))
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    NNTransferTrain(model, criterion, optimizer, loaders["train"], loaders["val"],
+                    device, epochs=epochs)
+
+    test_pred, _ = NNTransferPredictWithUncertainty(model, loaders["test"], device)
+    test_true = np.array([y for _, y in test_ds])
+    test_pred = np.array(test_pred)
+    # balanced (per-class mean) accuracy: plain accuracy hides rare classes
+    accs = [float((test_pred[test_true == c] == c).mean())
+            for c in np.unique(test_true) if (test_true == c).sum() > 0]
+    test_acc = float(np.mean(accs))
+
+    all_loader = DataLoader(TensorDataset(torch.tensor(X)), batch_size=batch_size, shuffle=False)
+    pred, unc = NNTransferPredictWithUncertainty(model, all_loader, device)
+    labels = le.inverse_transform(np.array(pred))
+
+    adata.obs[f"transfer_{label_key}"] = labels
+    adata.obs[f"transfer_{label_key}_uncertainty"] = np.array(unc)
+    return {"labels": labels, "uncertainty": np.array(unc), "test_accuracy": test_acc}

fusemap-1.1.0/fusemap/config.py

@@ -0,0 +1,170 @@
+import argparse
+import os
+from enum import Enum
+
+
+def parse_input_args():
+    parser = argparse.ArgumentParser(description="FuseMap")
+
+    parser.add_argument(
+        "--input_data_folder_path",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--output_save_dir",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--keep_celltype",
+        type=str,
+        default="",
+    )
+    parser.add_argument(
+        "--keep_tissueregion",
+        type=str,
+        default="",
+    )
+    parser.add_argument(
+        "--use_llm_gene_embedding",
+        default='false',
+    )
+    parser.add_argument(
+        "--pretrain_model_path",
+        default="",
+    )
+    parser.add_argument(
+        "--bead_files",
+        type=str,
+        default="",
+        help="comma-separated substrings declaring bead/spot-resolution datasets; Stage-B runs automatically after integration",
+    )
+    parser.add_argument(
+        "--sig_ref",
+        type=str,
+        default="",
+        help="comma-separated substrings selecting single-cell signature reference dataset(s) for Stage-B",
+    )
+    args = parser.parse_args()
+    return args
+
+
+class FlagConfig:
+    lambda_disc_single = 1
+    align_anneal = 1e10
+
+
+class ModelType(Enum):
+    pca_dim = 50
+    hidden_dim = 512
+    latent_dim = 64
+    dropout_rate = 0.2
+    n_epochs = 16
+    batch_size = 64
+    learning_rate = 0.001
+    optim_kw = "RMSprop"
+    use_input = "norm"
+    lambda_ae_single = 1
+    lambda_disc_spatial = 1
+    lambda_ae_spatial = 1
+    align_noise_coef = 1.5
+    lr_patience_pretrain = 2
+    lr_factor_pretrain = 0.5
+    lr_limit_pretrain = 0.00001
+    patience_limit_final = 5
+    lr_patience_final = 3
+    lr_factor_final = 0.5
+    lr_limit_final = 0.00001
+    patience_limit_pretrain = 3
+    EPS = 1e-10
+    DIS_LAMDA = 2
+    TRAIN_WITHOUT_EVAL = 10
+    USE_REFERENCE_PCT = 0.02
+    verbose = False
+    use_key='final'
+
+
+class AnchorConfig:
+    """Cross-sample MNN anchor alignment (STAligner-style), loss-level only.
+
+    This is a plain class (NOT a ``ModelType`` Enum member; the Enum has
+    value-aliasing quirks). Setting ``anchor_lambda = 0`` disables every anchor
+    code path and reproduces the original FuseMap behavior exactly.
+    """
+
+    # default 0.3 = validated optimal (2026-09-01); override via env var
+    anchor_lambda = float(os.environ.get("FUSEMAP_ANCHOR_LAMBDA", "0.3"))
+    anchor_refresh_epochs = 2    # recompute MNN anchors every N epochs
+    anchor_k = 15                # kNN used per direction for MNN
+    anchor_start_epoch = int(os.environ.get("FUSEMAP_ANCHOR_START", "2"))
+    # quality filters (2026-09-01): drop low-similarity MNN pairs and pairs not
+    # stable across consecutive refreshes; disable anchors during pretrain by
+    # default (immature latent locks in wrong anchors).
+    anchor_sim_threshold = float(os.environ.get("FUSEMAP_ANCHOR_SIM", "0.5"))
+    # comma-separated file-name substrings marking QUERY datasets (e.g. bead
+    # data): they get pulled toward others, but nothing is pulled toward them.
+    anchor_query_files = os.environ.get("FUSEMAP_ANCHOR_QUERY", "")
+    anchor_query_atlases = set()  # resolved at runtime from input file names
+    anchor_stable_only = os.environ.get("FUSEMAP_ANCHOR_STABLE", "1") == "1"
+    anchor_in_pretrain = os.environ.get("FUSEMAP_ANCHOR_PRETRAIN", "0") == "1"
+
+    # Within-dataset structure-preservation triplet loss (opt-in).
+    # struct_lambda = 0 keeps the feature fully OFF and reproduces current
+    # behavior EXACTLY (no extra RNG draws, no extra compute anywhere).
+    # When > 0, spatial_integrate precomputes per-atlas expression-space kNN
+    # (stored at runtime as ``AnchorConfig.struct_knn``) and
+    # ``fusemap.loss.compute_struct_loss`` adds a triplet term in the final
+    # training phase that keeps each cell's expression neighbors closer in
+    # latent space than random same-dataset cells.
+    struct_lambda = float(os.environ.get("FUSEMAP_STRUCT_LAMBDA", "0"))
+    struct_k = int(os.environ.get("FUSEMAP_STRUCT_K", "10"))
+    struct_margin = float(os.environ.get("FUSEMAP_STRUCT_MARGIN", "1.0"))
+    # Semi-hard negative mining (default ON when struct loss is enabled):
+    # uniformly random negatives are already far apart in the compact latent,
+    # so the repulsion term is inert and the triplet degenerates into extra
+    # compression (measured: type purity 0.578 -> 0.502). Instead mine the
+    # hardest VALID negative among struct_neg_candidates random candidates.
+    # Set FUSEMAP_STRUCT_HARDNEG=0 to reproduce the legacy random-negative
+    # path bit-for-bit (same number and order of RNG draws).
+    struct_hardneg = os.environ.get("FUSEMAP_STRUCT_HARDNEG", "1") == "1"
+    struct_neg_candidates = int(os.environ.get("FUSEMAP_STRUCT_NEGC", "16"))
+
+
+class AnchorState:
+    """Runtime container for anchor alignment state.
+
+    Populated by ``fusemap.train_model.refresh_anchors`` and consumed by the AE
+    loss (``fusemap.loss.compute_anchor_loss``). Kept dependency-free (plain
+    attributes) so it can live in the config module.
+
+    Attributes
+    ----------
+    single_cache : list of torch.Tensor or None
+        Per-atlas full-dataset single latent (z_mean), detached, [n_obs_i, d].
+    spatial_cache : list of torch.Tensor or None
+        Per-atlas full-dataset spatial latent (z_spatial), detached.
+    partner : dict
+        (i, j) -> int64 tensor of length n_obs_i giving, for each cell in atlas
+        i, the index of its MNN partner cell in atlas j (or -1 if none).
+    n_pairs : int
+        Total number of MNN anchor pairs found at the last refresh.
+    mean_dist : float
+        Mean cross-atlas single-latent distance of anchored pairs at refresh.
+    has_anchors : bool
+        Whether any anchors are currently stored.
+    """
+
+    def __init__(self):
+        self.prev_pairs = {}
+        self.single_cache = None
+        self.spatial_cache = None
+        self.partner = {}
+        self.n_pairs = 0
+        self.mean_dist = float("nan")
+        self.has_anchors = False