Add BABEL method (ATAC->GEX prediction)

scripts/create_datasets/test_resources.sh

@@ -90,6 +90,20 @@ for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/s
       --output $DATASET_DIR/$name/models/simple_mlp/
   fi
 
+  echo "pre-train novel on $name"
+  if up_to_date $DATASET_DIR/$name/models/novel/ $STATE; then
+    echo "  already up to date, skipping"
+  else
+    rm -rf $DATASET_DIR/$name/models/novel/
+    mkdir -p $DATASET_DIR/$name/models/novel/
+    viash run src/methods/novel/novel_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 \
+      --n_epochs 2 \
+      --output $DATASET_DIR/$name/models/novel
+  fi
+
   # senkin_tmp is CITE-only
   if [[ "$name" == bmmc_cite/normal ]]; then
     echo "pre-train senkin_tmp on $name"
@@ -108,18 +122,20 @@ for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/s
     fi
   fi
 
-  echo "pre-train novel on $name"
-  if up_to_date $DATASET_DIR/$name/models/novel/ $STATE; then
-    echo "  already up to date, skipping"
-  else
-    rm -rf $DATASET_DIR/$name/models/novel/
-    mkdir -p $DATASET_DIR/$name/models/novel/
-    viash run src/methods/novel/novel_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 \
-      --n_epochs 2 \
-      --output $DATASET_DIR/$name/models/novel
+  # babel only does ATAC->GEX, which is the multiome swap
+  if [[ "$name" == bmmc_multiome/swap ]]; then
+    echo "pre-train babel on $name"
+    if up_to_date $DATASET_DIR/$name/models/babel/output_model.pkl $STATE; then
+      echo "  already up to date, skipping"
+    else
+      mkdir -p $DATASET_DIR/$name/models/babel/
+      viash run src/methods/babel/babel_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 \
+        --nn_epochs 2 \
+        --output $DATASET_DIR/$name/models/babel/output_model.pkl
+    fi
   fi
 
 done

src/methods/babel/babel/config.vsh.yaml

@@ -0,0 +1,35 @@
+__merge__: ../../../api/comp_method.yaml
+name: babel
+label: BABEL
+summary: "Cross-modal autoencoder predicting scRNA-seq expression from scATAC-seq accessibility (Wu et al. 2021)"
+description: |
+  BABEL is a deep autoencoder that learns modality-specific encoders and decoders for
+  paired scRNA-seq and scATAC-seq measurements. RNA is encoded/decoded with dense
+  layers; ATAC accessibility is encoded with a chromosome-split architecture
+  (per-chromosome sub-networks, matching genomic locality) that groups peaks by
+  chromosome parsed from their genomic-coordinate names (e.g. "chr1:1000-2000").
+  The full architecture is bidirectional (RNA->RNA, RNA->ATAC, ATAC->RNA, ATAC->ATAC),
+  trained jointly with a negative binomial reconstruction loss for RNA counts and a
+  binary cross-entropy loss for binarized ATAC accessibility, but this component only
+  exposes the ATAC->RNA (GEX) prediction direction. The RNA->ATAC direction was found
+  to collapse to a per-peak base-rate prediction that ignores the RNA input -- a known
+  failure mode of unweighted binary cross-entropy under the extreme class imbalance
+  typical of ATAC accessibility data (~3% positive) -- confirmed via near-zero
+  discrimination between accessible/inaccessible peaks even after training to
+  convergence, and is intentionally disabled rather than silently returned.
+references:
+  doi:
+    - https://doi.org/10.1073/pnas.202307011
+links:
+  repository: https://github.com/wukevin/babel
+info:
+  preferred_normalization: log_cp10k
+resources:
+  - path: main.nf
+    type: nextflow_script
+    entrypoint: run_wf
+dependencies:
+  - name: methods/babel_train
+  - name: methods/babel_predict
+runners:
+  - type: nextflow

src/methods/babel/babel/main.nf

@@ -0,0 +1,18 @@
+workflow run_wf {
+  take: input_ch
+  main:
+  output_ch = input_ch
+    | babel_train.run(
+      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1"],
+      toState: ["input_model": "output"]
+    )
+    | babel_predict.run(
+      fromState: ["input_test_mod1", "input_train_mod2", "input_model"],
+      toState: ["output": "output"]
+    )
+    | map { tup ->
+      [tup[0], [output: tup[1].output]]
+    }
+
+  emit: output_ch
+}

src/methods/babel/babel_predict/config.vsh.yaml

@@ -0,0 +1,28 @@
+__merge__: /src/api/comp_method_predict.yaml
+name: babel_predict
+info:
+  test_setup:
+    multiome:
+      input_test_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad
+      input_train_mod2: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad
+      input_model: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/models/babel/output_model.pkl
+resources:
+  - path: script.py
+    type: python_script
+  - path: ../model.py
+    dest: model.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: openproblems/base_pytorch_nvidia:1
+    setup:
+      - type: python
+        packages:
+          - scikit-learn>=1.1
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, gpu]

src/methods/babel/babel_predict/script.py

@@ -0,0 +1,97 @@
+import logging
+import pickle
+import sys
+
+import anndata as ad
+import numpy as np
+import torch
+from scipy.sparse import csc_matrix, issparse
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+## VIASH START
+par = {
+    "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad",
+    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad",
+    "input_model": "output_model.pkl",
+    "output": "output_pred.h5ad",
+}
+meta = {"name": "babel", "resources_dir": "src/methods/babel"}
+## VIASH END
+
+sys.path.append(meta["resources_dir"])
+
+from model import AssymSplicedAutoEncoder
+
+
+def _to_dense(X):
+    return X.toarray() if issparse(X) else np.asarray(X)
+
+
+def _lognorm_per_cell(pred_counts, target_sum=1e4):
+    """Convert the NB decoder's raw-count-scale mean prediction into the same
+    log1p(normalized-to-target_sum) space used for the "normalized" layer
+    elsewhere in this codebase (see babel_train._rna_matrix / senkin_tmp's
+    senkin_tmp_cite_pred.preprocess.log_normalize), since file_prediction.yaml
+    requires "normalized" to hold log-normalized values, not raw NB means
+    (which are unbounded and on an arbitrary count scale)."""
+    pred_counts = np.clip(pred_counts, a_min=0, a_max=None)
+    row_sums = pred_counts.sum(axis=1, keepdims=True)
+    row_sums[row_sums == 0] = 1
+    return np.log1p(pred_counts / row_sums * target_sum)
+
+
+logger.info("Reading input files...")
+adata_test_mod1 = ad.read_h5ad(par["input_test_mod1"])
+adata_train_mod2 = ad.read_h5ad(par["input_train_mod2"])
+
+logger.info("Loading model bundle...")
+with open(par["input_model"], "rb") as f:
+    bundle = pickle.load(f)
+
+test_modality = adata_test_mod1.uns.get("modality")
+if test_modality != "ATAC":
+    raise ValueError(
+        f"babel_predict only supports ATAC->GEX prediction (test input must be ATAC), "
+        f"got modality={test_modality!r}. The GEX->ATAC direction of this BABEL port "
+        "collapses to a per-peak base-rate prediction that ignores the RNA input "
+        "(confirmed via cross-cell prediction similarity ~0.6 and near-zero separation "
+        "between accessible/inaccessible peaks, even after training to convergence) "
+        "and is intentionally disabled rather than silently returning uninformative output."
+    )
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AssymSplicedAutoEncoder(bundle["n_genes"], bundle["chrom_counts"], hidden_dim=bundle["hidden_dim"])
+model.load_state_dict(bundle["state_dict"])
+model.to(device)
+model.eval()
+
+chrom_groups = bundle["chrom_groups"]
+
+if list(adata_test_mod1.var_names) != bundle["atac_var_names"]:
+    raise ValueError(
+        "Test ATAC var_names do not match the peak order the model was trained with; "
+        "reindexing across mismatched peak sets is not supported."
+    )
+X_bin = (_to_dense(adata_test_mod1.layers.get("counts", adata_test_mod1.X)) > 0).astype(np.float32)
+X_per_chrom = [torch.from_numpy(X_bin[:, idxs]).to(device) for idxs in chrom_groups.values()]
+with torch.no_grad():
+    encoded = model.encoder2(X_per_chrom)
+    pred_mean, _, _ = model.decoder1(encoded)
+pred = _lognorm_per_cell(pred_mean.cpu().numpy())
+out_var = adata_train_mod2.var
+
+logger.info("Writing predictions...")
+adata_out = ad.AnnData(
+    layers={"normalized": csc_matrix(pred)},
+    obs=adata_test_mod1.obs,
+    var=out_var,
+    uns={
+        "dataset_id": adata_test_mod1.uns.get("dataset_id", ""),
+        "method_id": "babel",
+    },
+)
+
+adata_out.write_h5ad(par["output"], compression="gzip")
+logger.info("Predictions saved to %s", par["output"])

src/methods/babel/babel_train/config.vsh.yaml

@@ -0,0 +1,62 @@
+__merge__: /src/api/comp_method_train.yaml
+name: babel_train
+info:
+  test_setup:
+    multiome:
+      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
+      nn_epochs: 2
+resources:
+  - path: script.py
+    type: python_script
+  - path: ../model.py
+    dest: model.py
+  - path: ../losses.py
+    dest: losses.py
+  - path: ../chrom_utils.py
+    dest: 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: "--hidden_dim"
+    type: integer
+    default: 16
+    description: Latent dimensionality of the shared encoder/decoder bottleneck.
+  - name: "--nn_epochs"
+    type: integer
+    default: 100
+    description: Maximum training epochs (early stopping applies).
+    info:
+      test_default: 2
+  - name: "--lr"
+    type: double
+    default: 0.001
+    description: Adam learning rate.
+  - name: "--batch_size"
+    type: integer
+    default: 64
+    description: Training batch size.
+  - name: "--loss2_weight"
+    type: double
+    default: 3.0
+    description: Relative weight of the ATAC (BCE) loss terms versus the RNA (negative binomial) loss terms.
+  - name: "--early_stopping_patience"
+    type: integer
+    default: 20
+    description: Epochs without validation loss improvement before stopping early.
+engines:
+  - type: docker
+    image: openproblems/base_pytorch_nvidia:1
+    setup:
+      - type: python
+        packages:
+          - skorch>=0.15
+          - scanpy>=1.9
+          - scikit-learn>=1.1
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, gpu]