From 331985a1a65f93ff5c8b1dd61a54aa11b27d16a9 Mon Sep 17 00:00:00 2001
From: benjaminfreyuu <benjamin.frey.de@gmail.com>
Date: Tue, 21 Jul 2026 14:26:03 +0200
Subject: [PATCH 1/2] Add scButterfly method (Multiome GEX<->ATAC)

Adds scButterfly (Cao et al., Nat. Commun. 2024, doi:10.1038/s41467-024-47418-x),
a dual-VAE + adversarial translator, as a train/predict method scoped to Multiome
(GEX<->ATAC). Direction is detected from uns["modality"], covering both swap
(ATAC->GEX) and normal (GEX->ATAC) variants.

Train serializes the real scButterfly weights; predict reconstructs the model,
reruns the (deterministic) preprocessing on the train data, loads the weights and
runs inference, so it generalises to unseen test cells. Registered in run_benchmark
and given a Multiome-only fixture step in test_resources.sh.

Both components pass `viash ns test`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 scripts/create_datasets/test_resources.sh     |  15 ++
 src/methods/scbutterfly/butterfly_common.py   | 215 ++++++++++++++++++
 src/methods/scbutterfly/chrom_utils.py        | 126 ++++++++++
 .../scbutterfly/scbutterfly/config.vsh.yaml   |  27 +++
 src/methods/scbutterfly/scbutterfly/main.nf   |  18 ++
 .../scbutterfly_predict/config.vsh.yaml       |  44 ++++
 .../scbutterfly/scbutterfly_predict/script.py |  86 +++++++
 .../scbutterfly_train/config.vsh.yaml         |  80 +++++++
 .../scbutterfly/scbutterfly_train/script.py   |  78 +++++++
 src/workflows/run_benchmark/config.vsh.yaml   |   1 +
 src/workflows/run_benchmark/main.nf           |   3 +-
 11 files changed, 692 insertions(+), 1 deletion(-)
 create mode 100644 src/methods/scbutterfly/butterfly_common.py
 create mode 100644 src/methods/scbutterfly/chrom_utils.py
 create mode 100644 src/methods/scbutterfly/scbutterfly/config.vsh.yaml
 create mode 100644 src/methods/scbutterfly/scbutterfly/main.nf
 create mode 100644 src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
 create mode 100644 src/methods/scbutterfly/scbutterfly_predict/script.py
 create mode 100644 src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
 create mode 100644 src/methods/scbutterfly/scbutterfly_train/script.py

diff --git a/scripts/create_datasets/test_resources.sh b/scripts/create_datasets/test_resources.sh
index 7c00b7f8..827175b0 100755
--- a/scripts/create_datasets/test_resources.sh
+++ b/scripts/create_datasets/test_resources.sh
@@ -55,6 +55,21 @@ for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/s
     --input_test_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/test_mod1.h5ad \
     --output $OUTPUT_DIR/openproblems_neurips2021/$name/models/novel
 
+  # scbutterfly only supports Multiome (GEX<->ATAC); skip the CITE datasets.
+  case "$name" in
+    bmmc_multiome/*)
+      echo "pre-train scbutterfly on $name"
+      [ -d $OUTPUT_DIR/openproblems_neurips2021/$name/models/scbutterfly/ ] && rm -r $OUTPUT_DIR/openproblems_neurips2021/$name/models/scbutterfly/
+      mkdir -p $OUTPUT_DIR/openproblems_neurips2021/$name/models/scbutterfly/
+      viash run src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml -- \
+        --input_train_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod1.h5ad \
+        --input_train_mod2 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod2.h5ad \
+        --input_test_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/test_mod1.h5ad \
+        --rna_pretrain_epoch 2 --atac_pretrain_epoch 2 --translator_epoch 2 --patience 5 --n_top_genes 500 \
+        --output $OUTPUT_DIR/openproblems_neurips2021/$name/models/scbutterfly
+      ;;
+  esac
+
 done
 
 # only run this if you have access to the openproblems-data bucket
diff --git a/src/methods/scbutterfly/butterfly_common.py b/src/methods/scbutterfly/butterfly_common.py
new file mode 100644
index 00000000..4ddf855b
--- /dev/null
+++ b/src/methods/scbutterfly/butterfly_common.py
@@ -0,0 +1,215 @@
+"""Shared scButterfly setup used by both the train and predict components.
+
+scButterfly has no transform-only preprocessing API — ``data_preprocessing`` refits
+HVG/peak-filter/TF-IDF on whatever data it is given. To run inference faithfully in a
+separate process, predict must rebuild the *same* paired object and re-run the *same*
+deterministic preprocessing as train, then load the saved weights. This module holds
+that shared construction so train and predict stay in lock-step.
+"""
+
+import logging
+
+import anndata as ad
+import numpy as np
+import scanpy as sc
+from scipy.sparse import csr_matrix
+
+import chrom_utils
+
+logger = logging.getLogger(__name__)
+
+
+def apply_runtime_patches():
+    """Patch scButterfly/torch quirks. Call before importing scButterfly.
+
+    - legacy ``size_average``/``reduce`` loss kwargs -> ``reduction``
+    - BCE on a float32 sigmoid can drift >1.0 -> clamp input to [0,1]
+    - scButterfly hardcodes ``.cuda()``; make it a no-op when no GPU is present
+    """
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+
+    def _patch_legacy_reduction(cls):
+        orig_init = cls.__init__
+
+        def _init(self, *args, size_average=None, reduce=None, reduction="mean", **kwargs):
+            if size_average is not None or reduce is not None:
+                reduction = "mean" if (size_average in (None, True)) else "sum"
+            orig_init(self, *args, reduction=reduction, **kwargs)
+
+        cls.__init__ = _init
+
+    for _loss_cls in (nn.MSELoss, nn.BCELoss, nn.L1Loss):
+        _patch_legacy_reduction(_loss_cls)
+
+    _orig_bce = F.binary_cross_entropy
+
+    def _safe_bce(input, target, *args, **kwargs):
+        input = torch.nan_to_num(input, nan=0.0).clamp(0.0, 1.0)
+        return _orig_bce(input, target, *args, **kwargs)
+
+    F.binary_cross_entropy = _safe_bce
+
+    if not torch.cuda.is_available():
+        logger.warning("CUDA not available — running scButterfly on CPU (slow).")
+        torch.Tensor.cuda = lambda self, *a, **k: self
+        nn.Module.cuda = lambda self, *a, **k: self
+
+
+def detect_direction(train_mod1, train_mod2):
+    """Return ('GEX2ATAC'|'ATAC2GEX', mod1, mod2), raising on non-Multiome data."""
+    mod1 = train_mod1.uns["modality"]
+    mod2 = train_mod2.uns["modality"]
+    if {mod1, mod2} != {"GEX", "ATAC"}:
+        raise ValueError(
+            f"scbutterfly only supports Multiome GEX<->ATAC, got mod1={mod1}, mod2={mod2}"
+        )
+    direction = "GEX2ATAC" if mod2 == "ATAC" else "ATAC2GEX"
+    return direction, mod1, mod2
+
+
+def _to_counts_X(adata):
+    """AnnData copy whose .X is the raw counts layer as float32.
+
+    Counts are stored as float64; scButterfly's model is float32 and its data loader
+    does not cast, so feed float32 to avoid a dtype mismatch.
+    """
+    out = adata.copy()
+    out.X = out.layers["counts"].astype(np.float32)
+    return out
+
+
+def _placeholder_block(template_adata, n_rows, obs_names):
+    """Placeholder rows for the target modality's test cells.
+
+    These are the modality being predicted (no ground truth) and are never used as
+    model input — they exist only to satisfy scButterfly's paired same-cells layout.
+    Fill by tiling real training rows (not zeros): all-zero rows/cols make per-cell
+    normalization and TF-IDF divide by zero, producing NaNs that break BCE.
+    """
+    import pandas as pd
+    src = template_adata.X
+    src = src.toarray() if hasattr(src, "toarray") else np.asarray(src)
+    idx = np.arange(n_rows) % max(1, src.shape[0])
+    X = csr_matrix(src[idx].astype(np.float32))
+    obs = pd.DataFrame(index=obs_names)
+    return ad.AnnData(X=X, obs=obs, var=template_adata.var.copy())
+
+
+def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly):
+    """Build + preprocess + construct a Butterfly for the given data.
+
+    Returns a dict with the constructed ``butterfly``, the resolved ``direction``,
+    ``test_id`` (indices of the test cells), ``chrom_list``, and the preprocessed
+    target-modality var names. Deterministic given the same inputs, so train and
+    predict produce identical architecture/preprocessing.
+    """
+    direction, mod1, mod2 = detect_direction(train_mod1, train_mod2)
+    logger.info("Direction: %s (mod1=%s, mod2=%s)", direction, mod1, mod2)
+
+    # Assign RNA/ATAC roles: scButterfly wants RNA_data=GEX, ATAC_data=ATAC.
+    if mod1 == "GEX":
+        rna_train, atac_train = train_mod1, train_mod2
+        rna_test, atac_test = test_mod1, None
+    else:
+        atac_train, rna_train = train_mod1, train_mod2
+        atac_test, rna_test = test_mod1, None
+
+    n_train = train_mod1.n_obs
+    n_test = test_mod1.n_obs
+    test_obs_names = [f"test_{i}" for i in range(n_test)]
+
+    rna_train_c = _to_counts_X(rna_train)
+    atac_train_c = _to_counts_X(atac_train)
+    rna_train_c.obs_names = [f"train_{i}" for i in range(n_train)]
+    atac_train_c.obs_names = [f"train_{i}" for i in range(n_train)]
+
+    if rna_test is not None:
+        rna_test_c = _to_counts_X(rna_test)
+        rna_test_c.obs_names = test_obs_names
+    else:
+        rna_test_c = _placeholder_block(rna_train_c, n_test, test_obs_names)
+
+    if atac_test is not None:
+        atac_test_c = _to_counts_X(atac_test)
+        atac_test_c.obs_names = test_obs_names
+    else:
+        atac_test_c = _placeholder_block(atac_train_c, n_test, test_obs_names)
+
+    RNA_data = sc.concat([rna_train_c, rna_test_c], axis=0, join="outer")[:, rna_train_c.var_names].copy()
+    ATAC_data = sc.concat([atac_train_c, atac_test_c], axis=0, join="outer")[:, atac_train_c.var_names].copy()
+
+    train_id = list(range(n_train))
+    test_id = list(range(n_train, n_train + n_test))
+
+    # Deterministic ~10% validation carve (not five_fold_split_dataset).
+    rng = np.random.RandomState(0)
+    shuffled = list(train_id)
+    rng.shuffle(shuffled)
+    n_val = max(1, int(0.1 * n_train))
+    validation_id = sorted(shuffled[:n_val])
+    train_id_final = sorted(shuffled[n_val:])
+
+    # Chromosome ordering (peaks contiguous per chrom + chrom_list).
+    sort_index, chrom_list = chrom_utils.sorted_chrom_order(ATAC_data)
+    ATAC_data = chrom_utils.apply_sort(ATAC_data, sort_index)
+
+    butterfly = Butterfly()
+    butterfly.load_data(RNA_data, ATAC_data, train_id_final, test_id, validation_id)
+    butterfly.data_preprocessing(n_top_genes=n_top_genes)
+    butterfly.augmentation(aug_type=None)
+
+    # scButterfly casts to float32 only on the CUDA path; force float32 for CPU.
+    for _attr in ("RNA_data_p", "ATAC_data_p"):
+        _adata = getattr(butterfly, _attr, None)
+        if _adata is not None and _adata.X is not None:
+            _adata.X = _adata.X.astype(np.float32)
+
+    # Rebuild chrom_list from the preprocessed (peak-filtered) ATAC so model dims match.
+    atac_p = getattr(butterfly, "ATAC_data_p", None)
+    if atac_p is not None and atac_p.n_vars != sum(chrom_list):
+        if "chrom" not in atac_p.var.columns:
+            atac_p.var["chrom"] = [chrom_utils.parse_chrom(v) for v in atac_p.var_names]
+        chrom_list = chrom_utils.chrom_counts(atac_p)
+    logger.info("chrom_list: %d chromosomes, %d peaks", len(chrom_list), sum(chrom_list))
+
+    butterfly.construct_model(chrom_list=chrom_list)
+
+    # var names of the preprocessed target modality, used to map predictions back.
+    target_p = butterfly.RNA_data_p if direction == "ATAC2GEX" else butterfly.ATAC_data_p
+
+    return {
+        "butterfly": butterfly,
+        "direction": direction,
+        "n_train": n_train,
+        "n_test": n_test,
+        "test_id": test_id,
+        "chrom_list": chrom_list,
+        "target_p_var_names": list(target_p.var_names),
+    }
+
+
+def extract_predictions(built, A2R_predict, R2A_predict, target_var_names):
+    """Select the target-modality prediction, name its vars, and scatter to target order."""
+    direction = built["direction"]
+    n_train, n_test = built["n_train"], built["n_test"]
+    pred = A2R_predict if direction == "ATAC2GEX" else R2A_predict
+
+    # tensor2adata gives integer var names; assign the preprocessed target var names.
+    p_names = built["target_p_var_names"]
+    if pred.n_vars == len(p_names):
+        pred.var_names = p_names
+    else:
+        logger.warning("Prediction vars (%d) != preprocessed target vars (%d)",
+                       pred.n_vars, len(p_names))
+
+    if pred.n_obs == (n_train + n_test):
+        pred = pred[built["test_id"]].copy()
+    elif pred.n_obs != n_test:
+        logger.warning("Unexpected prediction rows: %d (n_test=%d)", pred.n_obs, n_test)
+
+    out = chrom_utils.scatter_to_target(pred, target_var_names)
+    overlap = len(set(pred.var_names) & set(target_var_names))
+    logger.info("Var overlap with target: %d / %d", overlap, len(target_var_names))
+    return out
diff --git a/src/methods/scbutterfly/chrom_utils.py b/src/methods/scbutterfly/chrom_utils.py
new file mode 100644
index 00000000..e33802e3
--- /dev/null
+++ b/src/methods/scbutterfly/chrom_utils.py
@@ -0,0 +1,126 @@
+"""Chromosome / peak utilities for the scButterfly Multiome method.
+
+scButterfly's ``construct_model`` needs a ``chrom_list`` (number of peaks per
+chromosome) and reads ``ATAC_data.var.chrom`` during model construction. It also
+assumes peaks are contiguous per chromosome. The predict-modality ATAC h5ads have
+no ``chrom`` column and peaks are in arbitrary chromosome order, but the peak names
+encode the chromosome (e.g. ``chr17-6651156-6652045``).
+
+This module parses the chromosome from peak names, produces a peak ordering that
+groups peaks contiguously per chromosome (with the matching ``chrom_list``), and
+scatters a predicted matrix back into a target var order by feature name.
+"""
+
+import re
+
+import numpy as np
+from scipy.sparse import issparse
+
+# Matches a leading chromosome token like "chr17", "chrX", "17", "X" followed by
+# a ':' or '-' delimiter. Used as a fallback when the simple split does not yield
+# a recognisable chromosome.
+_CHROM_RE = re.compile(r"^chr?([0-9XYMxym]+)[:\-]")
+
+
+def parse_chrom(name):
+    """Return the chromosome token for a peak name (e.g. 'chr17')."""
+    name = str(name)
+    token = name.split("-", 1)[0]
+    if token.lower().startswith("chr"):
+        return token
+    m = _CHROM_RE.match(name)
+    if m:
+        return "chr" + m.group(1)
+    # Unparseable — bucket everything unknown together so it still forms one group.
+    return token
+
+
+def sorted_chrom_order(atac_adata):
+    """Compute a peak ordering that groups peaks contiguously by chromosome.
+
+    Returns
+    -------
+    sort_index : np.ndarray
+        Indices that reorder ``atac_adata`` so peaks are grouped per chromosome.
+    chrom_list : list[int]
+        Number of peaks per chromosome, in the sorted order. ``sum == n_peaks``.
+    """
+    chroms = np.array([parse_chrom(v) for v in atac_adata.var_names])
+    # Stable ordering of chromosomes; stable argsort keeps peaks deterministic
+    # within a chromosome (preserving original relative order).
+    chrom_order = sorted(set(chroms))
+    rank = {c: i for i, c in enumerate(chrom_order)}
+    keys = np.array([rank[c] for c in chroms])
+    sort_index = np.argsort(keys, kind="stable")
+
+    sorted_chroms = chroms[sort_index]
+    chrom_list = []
+    last = None
+    for c in sorted_chroms:
+        if c != last:
+            chrom_list.append(1)
+            last = c
+        else:
+            chrom_list[-1] += 1
+    assert sum(chrom_list) == atac_adata.n_vars
+    return sort_index, chrom_list
+
+
+def chrom_counts(atac_adata):
+    """Count peaks per chromosome, in the order they appear in ``var``.
+
+    Assumes peaks are already grouped contiguously per chromosome (as produced by
+    :func:`apply_sort`). Reads ``var['chrom']`` if present, else parses from names.
+    Peak filtering that preserves order (e.g. scButterfly's TF-IDF/peak filter)
+    keeps the grouping contiguous, so recounting here yields a valid ``chrom_list``.
+    """
+    if "chrom" in atac_adata.var.columns:
+        chroms = list(atac_adata.var["chrom"])
+    else:
+        chroms = [parse_chrom(v) for v in atac_adata.var_names]
+    counts = []
+    last = None
+    for c in chroms:
+        if c != last:
+            counts.append(1)
+            last = c
+        else:
+            counts[-1] += 1
+    return counts
+
+
+def apply_sort(atac_adata, sort_index):
+    """Reorder ATAC peaks by ``sort_index`` and set ``.var['chrom']``.
+
+    Returns a new AnnData whose peaks are contiguous per chromosome and which
+    carries the parsed chromosome in ``var['chrom']`` for scButterfly to read.
+    """
+    out = atac_adata[:, sort_index].copy()
+    out.var["chrom"] = [parse_chrom(v) for v in out.var_names]
+    return out
+
+
+def scatter_to_target(pred_adata, target_var_names):
+    """Scatter a prediction into ``target_var_names`` order, by feature name.
+
+    Predicted features may be a subset of and/or in a different order from the
+    target modality's vars (scButterfly can subset RNA to HVGs or filter ATAC
+    peaks). Scattering by name simultaneously (a) restores the original peak
+    order and (b) fills any missing target features with zeros.
+
+    Returns a dense float32 ndarray of shape ``(n_cells, len(target_var_names))``.
+    """
+    X = pred_adata.X
+    if issparse(X):
+        X = X.toarray()
+    X = np.asarray(X, dtype=np.float32)
+
+    n_cells = X.shape[0]
+    out = np.zeros((n_cells, len(target_var_names)), dtype=np.float32)
+
+    target_pos = {name: i for i, name in enumerate(target_var_names)}
+    for src_col, name in enumerate(pred_adata.var_names):
+        dst = target_pos.get(name)
+        if dst is not None:
+            out[:, dst] = X[:, src_col]
+    return out
diff --git a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml
new file mode 100644
index 00000000..90794764
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml
@@ -0,0 +1,27 @@
+__merge__: ../../../api/comp_method.yaml
+name: scbutterfly
+label: scButterfly
+summary: "Dual-VAE adversarial translator for paired single-cell multi-omics (GEX<->ATAC)"
+description: |
+  scButterfly (Basic variant, scButterfly-B) is a dual variational autoencoder with an
+  adversarial translator that learns to convert between paired single-cell modalities.
+  For the predict-modality task it is trained on paired Multiome cells and used to translate
+  the held-out test modality. Chromosome grouping for the ATAC branch is parsed from peak
+  coordinates in the feature names. Only Multiome GEX<->ATAC is supported; CITE-seq (ADT)
+  datasets are not handled.
+references:
+  doi:
+    - 10.1038/s41467-024-47418-x
+links:
+  repository: https://github.com/BioX-NKU/scButterfly
+info:
+  preferred_normalization: log_cp10k
+resources:
+  - path: main.nf
+    type: nextflow_script
+    entrypoint: run_wf
+dependencies:
+  - name: methods/scbutterfly_train
+  - name: methods/scbutterfly_predict
+runners:
+  - type: nextflow
diff --git a/src/methods/scbutterfly/scbutterfly/main.nf b/src/methods/scbutterfly/scbutterfly/main.nf
new file mode 100644
index 00000000..ceb9e913
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly/main.nf
@@ -0,0 +1,18 @@
+workflow run_wf {
+  take: input_ch
+  main:
+  output_ch = input_ch
+    | scbutterfly_train.run(
+      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1"],
+      toState: ["input_model": "output"]
+    )
+    | scbutterfly_predict.run(
+      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1", "input_model"],
+      toState: ["output": "output"]
+    )
+    | map { tup ->
+      [tup[0], [output: tup[1].output]]
+    }
+
+  emit: output_ch
+}
diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
new file mode 100644
index 00000000..755649c3
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
@@ -0,0 +1,44 @@
+__merge__: ../../../api/comp_method_predict.yaml
+name: scbutterfly_predict
+info:
+  # The api example paths point at bmmc_cite, but this method is Multiome-only. Override
+  # all inputs to bmmc_multiome/swap and supply a model directory produced by
+  # scbutterfly_train on that same data (weights in <model>/model/*.pt + metadata.pkl).
+  test_setup:
+    with_model:
+      input_train_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad
+      input_train_mod2: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad
+      input_test_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad
+      input_model: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/models/scbutterfly
+resources:
+  - path: script.py
+    type: python_script
+  - path: ../butterfly_common.py
+  - path: ../chrom_utils.py
+engines:
+  # Predict reconstructs the scButterfly model and runs inference, so it needs the same
+  # full stack as train (same custom Python 3.9 + torch cu117 base; see train config).
+  - type: docker
+    image: python:3.9
+    setup:
+      - type: docker
+        run:
+          - pip install --no-cache-dir torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu117
+          - pip install --no-cache-dir --no-deps scButterfly
+          - pip install --no-cache-dir "numpy==1.23.5" "scipy==1.9.3" "pandas<2" "scanpy==1.9.3" "anndata==0.8.0" "scikit-learn>=1.1" numba tqdm matplotlib adjustText bamnostic pysam legacy-api-wrap
+          - pip install --no-cache-dir --no-deps episcanpy==0.3.2
+          # `openproblems` (config reader used by the component test harness) is bundled in
+          # the shared base images but not on this custom base; install it from the same
+          # source the base images use so `viash ns test` works.
+          - pip install --no-cache-dir "git+https://github.com/openproblems-bio/core.git#subdirectory=packages/python/openproblems"
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, gpu]
+# Override the api default (bmmc_cite) — this method only supports Multiome GEX<->ATAC.
+test_resources:
+  - type: python_script
+    path: /common/component_tests/run_and_check_output.py
+  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
+    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
diff --git a/src/methods/scbutterfly/scbutterfly_predict/script.py b/src/methods/scbutterfly/scbutterfly_predict/script.py
new file mode 100644
index 00000000..06a639cc
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly_predict/script.py
@@ -0,0 +1,86 @@
+import logging
+import os
+import pickle
+import sys
+
+import anndata as ad
+from scipy.sparse import csc_matrix
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+## VIASH START
+par = {
+    "input_train_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad",
+    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad",
+    "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad",
+    "input_model": "output_model",
+    "output": "output_pred.h5ad",
+}
+meta = {"name": "scbutterfly", "resources_dir": "src/methods/scbutterfly"}
+## VIASH END
+
+sys.path.append(meta["resources_dir"])
+import butterfly_common
+
+butterfly_common.apply_runtime_patches()
+from scButterfly.butterfly import Butterfly
+
+# ---------------------------------------------------------------------------
+# Load data + model metadata
+# ---------------------------------------------------------------------------
+logger.info("Reading input files...")
+train_mod1 = ad.read_h5ad(par["input_train_mod1"])
+train_mod2 = ad.read_h5ad(par["input_train_mod2"])
+test_mod1 = ad.read_h5ad(par["input_test_mod1"])
+
+with open(os.path.join(par["input_model"], "metadata.pkl"), "rb") as f:
+    metadata = pickle.load(f)
+
+# ---------------------------------------------------------------------------
+# Reconstruct the SAME model architecture + preprocessing as train, then load the
+# trained weights (written by train_model to <model>/model/*.pt) and run inference.
+# ---------------------------------------------------------------------------
+logger.info("Reconstructing scButterfly model...")
+built = butterfly_common.build_butterfly(
+    train_mod1, train_mod2, test_mod1,
+    n_top_genes=metadata["n_top_genes"], Butterfly=Butterfly,
+)
+butterfly = built["butterfly"]
+
+logger.info("Loading trained weights and predicting...")
+A2R_predict, R2A_predict = butterfly.test_model(
+    batch_size=metadata["batch_size"],
+    model_path=par["input_model"],
+    load_model=True,
+)
+
+target_var_names = list(train_mod2.var_names)
+test_predictions = butterfly_common.extract_predictions(
+    built, A2R_predict, R2A_predict, target_var_names,
+)
+
+# ---------------------------------------------------------------------------
+# Write predictions.
+# ---------------------------------------------------------------------------
+# viash sets meta["name"] to the component name (scbutterfly_predict); report the
+# method name expected by the benchmark by stripping the _predict/_train suffix.
+method_id = meta["name"]
+for _suffix in ("_predict", "_train"):
+    if method_id.endswith(_suffix):
+        method_id = method_id[: -len(_suffix)]
+        break
+
+logger.info("Writing predictions...")
+adata_out = ad.AnnData(
+    layers={"normalized": csc_matrix(test_predictions)},
+    obs=test_mod1.obs,
+    var=train_mod2.var,
+    uns={
+        "dataset_id": test_mod1.uns.get("dataset_id", metadata.get("dataset_id", "")),
+        "method_id": method_id,
+    },
+)
+
+adata_out.write_h5ad(par["output"], compression="gzip")
+logger.info("Predictions saved to %s", par["output"])
diff --git a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
new file mode 100644
index 00000000..9e5aec4f
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
@@ -0,0 +1,80 @@
+__merge__: ../../../api/comp_method_train.yaml
+name: scbutterfly_train
+info:
+  # The api example paths point at bmmc_cite, but this method is Multiome-only. Override
+  # inputs to bmmc_multiome/swap and use tiny epochs so the test trains quickly.
+  test_setup:
+    quick:
+      input_train_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad
+      input_train_mod2: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad
+      input_test_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad
+      rna_pretrain_epoch: 2
+      atac_pretrain_epoch: 2
+      translator_epoch: 2
+      patience: 5
+      n_top_genes: 500
+resources:
+  - path: script.py
+    type: python_script
+  - path: ../butterfly_common.py
+  - path: ../chrom_utils.py
+arguments:
+  - name: "--rna_pretrain_epoch"
+    type: integer
+    default: 100
+    description: RNA autoencoder pretraining epochs (R2R_pretrain_epoch).
+  - name: "--atac_pretrain_epoch"
+    type: integer
+    default: 100
+    description: ATAC autoencoder pretraining epochs (A2A_pretrain_epoch).
+  - name: "--translator_epoch"
+    type: integer
+    default: 200
+    description: Translator training epochs.
+  - name: "--patience"
+    type: integer
+    default: 50
+    description: Early-stopping patience for the translator.
+  - name: "--batch_size"
+    type: integer
+    default: 64
+    description: Training batch size.
+  - name: "--n_top_genes"
+    type: integer
+    default: 3000
+    description: Number of highly variable genes kept by scButterfly RNA preprocessing.
+engines:
+  # scButterfly 0.0.9 hard-pins a 2022-era stack (scipy==1.9.3, scvi-tools==0.19.0,
+  # episcanpy==0.3.2) that will not build on the newer Python of the shared
+  # openproblems base image. Use a Python 3.9 base and install the pinned stack from
+  # wheels. The torch cu117 wheel bundles its own CUDA runtime, so GPU works when the
+  # container is run with --gpus (matching the nextflow 'gpu' label) without a separate
+  # CUDA base image. scButterfly is installed --no-deps to avoid dragging in scvi-tools
+  # and jax, which the Basic (aug_type=None) path never uses.
+  - type: docker
+    image: python:3.9
+    setup:
+      - type: docker
+        run:
+          - pip install --no-cache-dir torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu117
+          - pip install --no-cache-dir --no-deps scButterfly
+          - pip install --no-cache-dir "numpy==1.23.5" "scipy==1.9.3" "pandas<2" "scanpy==1.9.3" "anndata==0.8.0" "scikit-learn>=1.1" numba tqdm matplotlib adjustText bamnostic pysam legacy-api-wrap
+          # episcanpy 0.3.2's requirements pull `tbb`, which the resolver cannot satisfy on
+          # this platform. Install it --no-deps; its import-time deps (incl. bamnostic, pysam,
+          # legacy-api-wrap) are provided explicitly above.
+          - pip install --no-cache-dir --no-deps episcanpy==0.3.2
+          # `openproblems` (config reader used by the component test harness) is bundled in
+          # the shared base images but not on this custom base; install it from the same
+          # source the base images use so `viash ns test` works.
+          - pip install --no-cache-dir "git+https://github.com/openproblems-bio/core.git#subdirectory=packages/python/openproblems"
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, gpu]
+# Override the api default (bmmc_cite) — this method only supports Multiome GEX<->ATAC.
+test_resources:
+  - type: python_script
+    path: /common/component_tests/run_and_check_output.py
+  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
+    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
diff --git a/src/methods/scbutterfly/scbutterfly_train/script.py b/src/methods/scbutterfly/scbutterfly_train/script.py
new file mode 100644
index 00000000..ed8dbedf
--- /dev/null
+++ b/src/methods/scbutterfly/scbutterfly_train/script.py
@@ -0,0 +1,78 @@
+import logging
+import os
+import pickle
+import sys
+
+import anndata as ad
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+## VIASH START
+par = {
+    "input_train_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad",
+    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad",
+    "input_test_mod1":  "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad",
+    "output": "output_model",
+    "rna_pretrain_epoch": 100,
+    "atac_pretrain_epoch": 100,
+    "translator_epoch": 200,
+    "patience": 50,
+    "batch_size": 64,
+    "n_top_genes": 3000,
+}
+meta = {"name": "scbutterfly", "resources_dir": "src/methods/scbutterfly"}
+## VIASH END
+
+sys.path.append(meta["resources_dir"])
+import butterfly_common
+
+butterfly_common.apply_runtime_patches()
+from scButterfly.butterfly import Butterfly
+
+# ---------------------------------------------------------------------------
+# Load data
+# ---------------------------------------------------------------------------
+logger.info("Reading input files...")
+train_mod1 = ad.read_h5ad(par["input_train_mod1"])
+train_mod2 = ad.read_h5ad(par["input_train_mod2"])
+test_mod1 = ad.read_h5ad(par["input_test_mod1"])
+
+# ---------------------------------------------------------------------------
+# Build + construct the model, then train (weights written to <output>/model).
+# ---------------------------------------------------------------------------
+os.makedirs(par["output"], exist_ok=True)
+
+logger.info("Building scButterfly model...")
+built = butterfly_common.build_butterfly(
+    train_mod1, train_mod2, test_mod1,
+    n_top_genes=par["n_top_genes"], Butterfly=Butterfly,
+)
+butterfly = built["butterfly"]
+
+logger.info("Training scButterfly model...")
+butterfly.train_model(
+    R2R_pretrain_epoch=par["rna_pretrain_epoch"],
+    A2A_pretrain_epoch=par["atac_pretrain_epoch"],
+    translator_epoch=par["translator_epoch"],
+    patience=par["patience"],
+    batch_size=par["batch_size"],
+    output_path=par["output"],
+)
+
+# ---------------------------------------------------------------------------
+# Persist the metadata predict needs to reconstruct the model deterministically.
+# The trained weights live in <output>/model/*.pt (written by train_model).
+# ---------------------------------------------------------------------------
+logger.info("Saving model metadata...")
+metadata = {
+    "direction": built["direction"],
+    "n_top_genes": par["n_top_genes"],
+    "batch_size": par["batch_size"],
+    "target_var_names": list(train_mod2.var_names),
+    "dataset_id": train_mod1.uns.get("dataset_id", ""),
+}
+with open(os.path.join(par["output"], "metadata.pkl"), "wb") as f:
+    pickle.dump(metadata, f, protocol=4)
+
+logger.info("Training complete. Model saved to %s", par["output"])
diff --git a/src/workflows/run_benchmark/config.vsh.yaml b/src/workflows/run_benchmark/config.vsh.yaml
index ea4396c3..7056e3b4 100644
--- a/src/workflows/run_benchmark/config.vsh.yaml
+++ b/src/workflows/run_benchmark/config.vsh.yaml
@@ -74,6 +74,7 @@ dependencies:
   - name: methods/guanlab_dengkw_pm
   - name: methods/novel
   - name: methods/simple_mlp
+  - name: methods/scbutterfly
   - name: metrics/correlation
   - name: metrics/mse
 runners:
diff --git a/src/workflows/run_benchmark/main.nf b/src/workflows/run_benchmark/main.nf
index 6a7989d4..09a048e4 100644
--- a/src/workflows/run_benchmark/main.nf
+++ b/src/workflows/run_benchmark/main.nf
@@ -18,7 +18,8 @@ methods = [
   lm,
   guanlab_dengkw_pm,
   novel,
-  simple_mlp
+  simple_mlp,
+  scbutterfly
 ]
 
 // construct list of metrics

From 9889e5b58c9c2c3cf4686aa19d0a29a90cc36db5 Mon Sep 17 00:00:00 2001
From: Robrecht Cannoodt <rcannood@gmail.com>
Date: Thu, 30 Jul 2026 15:36:50 +0200
Subject: [PATCH 2/2] clean up

---
 .../scbutterfly/scbutterfly/config.vsh.yaml   |  2 +-
 .../scbutterfly_predict/config.vsh.yaml       | 19 ++++-----------
 .../scbutterfly_train/config.vsh.yaml         | 24 +++----------------
 3 files changed, 8 insertions(+), 37 deletions(-)

diff --git a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml
index 90794764..be01fbe7 100644
--- a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml
+++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml
@@ -1,4 +1,4 @@
-__merge__: ../../../api/comp_method.yaml
+__merge__: /src/api/comp_method.yaml
 name: scbutterfly
 label: scButterfly
 summary: "Dual-VAE adversarial translator for paired single-cell multi-omics (GEX<->ATAC)"
diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
index 755649c3..595b4e3d 100644
--- a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
+++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml
@@ -1,9 +1,6 @@
-__merge__: ../../../api/comp_method_predict.yaml
+__merge__: /src/api/comp_method_predict.yaml
 name: scbutterfly_predict
 info:
-  # The api example paths point at bmmc_cite, but this method is Multiome-only. Override
-  # all inputs to bmmc_multiome/swap and supply a model directory produced by
-  # scbutterfly_train on that same data (weights in <model>/model/*.pt + metadata.pkl).
   test_setup:
     with_model:
       input_train_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad
@@ -15,9 +12,10 @@ resources:
     type: python_script
   - path: ../butterfly_common.py
   - path: ../chrom_utils.py
+test_resources:
+  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
+    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
 engines:
-  # Predict reconstructs the scButterfly model and runs inference, so it needs the same
-  # full stack as train (same custom Python 3.9 + torch cu117 base; see train config).
   - type: docker
     image: python:3.9
     setup:
@@ -27,18 +25,9 @@ engines:
           - pip install --no-cache-dir --no-deps scButterfly
           - pip install --no-cache-dir "numpy==1.23.5" "scipy==1.9.3" "pandas<2" "scanpy==1.9.3" "anndata==0.8.0" "scikit-learn>=1.1" numba tqdm matplotlib adjustText bamnostic pysam legacy-api-wrap
           - pip install --no-cache-dir --no-deps episcanpy==0.3.2
-          # `openproblems` (config reader used by the component test harness) is bundled in
-          # the shared base images but not on this custom base; install it from the same
-          # source the base images use so `viash ns test` works.
           - pip install --no-cache-dir "git+https://github.com/openproblems-bio/core.git#subdirectory=packages/python/openproblems"
 runners:
   - type: executable
   - type: nextflow
     directives:
       label: [highmem, hightime, midcpu, gpu]
-# Override the api default (bmmc_cite) — this method only supports Multiome GEX<->ATAC.
-test_resources:
-  - type: python_script
-    path: /common/component_tests/run_and_check_output.py
-  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
-    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
diff --git a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
index 9e5aec4f..87a917c6 100644
--- a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
+++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml
@@ -1,8 +1,6 @@
 __merge__: ../../../api/comp_method_train.yaml
 name: scbutterfly_train
 info:
-  # The api example paths point at bmmc_cite, but this method is Multiome-only. Override
-  # inputs to bmmc_multiome/swap and use tiny epochs so the test trains quickly.
   test_setup:
     quick:
       input_train_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad
@@ -18,6 +16,9 @@ resources:
     type: python_script
   - path: ../butterfly_common.py
   - path: ../chrom_utils.py
+test_resources:
+  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
+    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
 arguments:
   - name: "--rna_pretrain_epoch"
     type: integer
@@ -44,13 +45,6 @@ arguments:
     default: 3000
     description: Number of highly variable genes kept by scButterfly RNA preprocessing.
 engines:
-  # scButterfly 0.0.9 hard-pins a 2022-era stack (scipy==1.9.3, scvi-tools==0.19.0,
-  # episcanpy==0.3.2) that will not build on the newer Python of the shared
-  # openproblems base image. Use a Python 3.9 base and install the pinned stack from
-  # wheels. The torch cu117 wheel bundles its own CUDA runtime, so GPU works when the
-  # container is run with --gpus (matching the nextflow 'gpu' label) without a separate
-  # CUDA base image. scButterfly is installed --no-deps to avoid dragging in scvi-tools
-  # and jax, which the Basic (aug_type=None) path never uses.
   - type: docker
     image: python:3.9
     setup:
@@ -59,22 +53,10 @@ engines:
           - pip install --no-cache-dir torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cu117
           - pip install --no-cache-dir --no-deps scButterfly
           - pip install --no-cache-dir "numpy==1.23.5" "scipy==1.9.3" "pandas<2" "scanpy==1.9.3" "anndata==0.8.0" "scikit-learn>=1.1" numba tqdm matplotlib adjustText bamnostic pysam legacy-api-wrap
-          # episcanpy 0.3.2's requirements pull `tbb`, which the resolver cannot satisfy on
-          # this platform. Install it --no-deps; its import-time deps (incl. bamnostic, pysam,
-          # legacy-api-wrap) are provided explicitly above.
           - pip install --no-cache-dir --no-deps episcanpy==0.3.2
-          # `openproblems` (config reader used by the component test harness) is bundled in
-          # the shared base images but not on this custom base; install it from the same
-          # source the base images use so `viash ns test` works.
           - pip install --no-cache-dir "git+https://github.com/openproblems-bio/core.git#subdirectory=packages/python/openproblems"
 runners:
   - type: executable
   - type: nextflow
     directives:
       label: [highmem, hightime, midcpu, gpu]
-# Override the api default (bmmc_cite) — this method only supports Multiome GEX<->ATAC.
-test_resources:
-  - type: python_script
-    path: /common/component_tests/run_and_check_output.py
-  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
-    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap