diff --git a/scripts/create_datasets/test_resources.sh b/scripts/create_datasets/test_resources.sh index 5e084c3..c0c46cb 100755 --- a/scripts/create_datasets/test_resources.sh +++ b/scripts/create_datasets/test_resources.sh @@ -138,6 +138,27 @@ for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/s fi fi + # scbutterfly only does multiome (GEX<->ATAC) + if [[ "$name" == bmmc_multiome/* ]]; then + echo "pre-train scbutterfly on $name" + if up_to_date $DATASET_DIR/$name/models/scbutterfly/ $STATE; then + echo " already up to date, skipping" + else + rm -rf $DATASET_DIR/$name/models/scbutterfly/ + mkdir -p $DATASET_DIR/$name/models/scbutterfly/ + viash run src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml -- \ + --input_train_mod1 $DATASET_DIR/$name/train_mod1.h5ad \ + --input_train_mod2 $DATASET_DIR/$name/train_mod2.h5ad \ + --input_test_mod1 $DATASET_DIR/$name/test_mod1.h5ad \ + --rna_pretrain_epoch 2 \ + --atac_pretrain_epoch 2 \ + --translator_epoch 2 \ + --patience 5 \ + --n_top_genes 500 \ + --output $DATASET_DIR/$name/models/scbutterfly + fi + fi + 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 0000000..4ddf855 --- /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 0000000..e33802e --- /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 0000000..be01fbe --- /dev/null +++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml @@ -0,0 +1,27 @@ +__merge__: /src/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 0000000..ceb9e91 --- /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 0000000..595b4e3 --- /dev/null +++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml @@ -0,0 +1,33 @@ +__merge__: /src/api/comp_method_predict.yaml +name: scbutterfly_predict +info: + 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 +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: + - 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 + - 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] diff --git a/src/methods/scbutterfly/scbutterfly_predict/script.py b/src/methods/scbutterfly/scbutterfly_predict/script.py new file mode 100644 index 0000000..06a639c --- /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/*.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 0000000..87a917c --- /dev/null +++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml @@ -0,0 +1,62 @@ +__merge__: ../../../api/comp_method_train.yaml +name: scbutterfly_train +info: + 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 +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 + 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: + - 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 + - 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] diff --git a/src/methods/scbutterfly/scbutterfly_train/script.py b/src/methods/scbutterfly/scbutterfly_train/script.py new file mode 100644 index 0000000..ed8dbed --- /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 /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 /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 089a0f5..3cda3cc 100644 --- a/src/workflows/run_benchmark/config.vsh.yaml +++ b/src/workflows/run_benchmark/config.vsh.yaml @@ -76,6 +76,7 @@ dependencies: - name: methods/simple_mlp - name: methods/babel - name: methods/senkin_tmp + - 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 86679fa..92a80c0 100644 --- a/src/workflows/run_benchmark/main.nf +++ b/src/workflows/run_benchmark/main.nf @@ -20,7 +20,8 @@ methods = [ novel, simple_mlp, babel, - senkin_tmp + senkin_tmp, + scbutterfly ] // construct list of metrics