Add BABEL method (ATAC->GEX prediction)

[benjaminfreyuu]

Ports BABEL (Wu et al. 2021, PNAS) as a train/predict method pair for predicting RNA expression from scATAC-seq accessibility. Uses a dense RNA encoder/decoder and a chromosome-split ATAC encoder/decoder, trained jointly with a negative binomial loss for RNA and BCE for binarized ATAC accessibility.

Verified end-to-end via viash run against real bmmc_multiome/swap data (Docker build, training, prediction): 0.34 mean per-cell correlation with ground truth, in line with the existing knnr_py baseline on the same data.As BABEl is not installable in its original implmentation, code was rewritten here.

Describe your changes

Checklist before requesting a review

• I have performed a self-review of my code

• Check the correct box. Does this PR contain:

  • Breaking changes
  • New functionality
  • Major changes
  • Minor changes
  • Bug fixes

• Proposed changes are described in the CHANGELOG.md

• CI Tests succeed and look good!

[rcannood]

Ran BABEL on a de.NBI Tesla T4. Two blockers, one of which only shows up on a GPU and so can't be caught by CI or by local CPU testing.

1. Not registered at all -- it would silently never run

babel is missing from both places a method has to be listed:

• dependencies: in run_benchmark/config.vsh.yaml
• the methods list in run_benchmark/main.nf

The build stays green, viash test stays green, and the method simply never appears in the results. No error anywhere -- which is the worst version of this, since nothing tells you it didn't run. (#18 has the opposite half of the problem and takes the whole run down instead.)

2. Training crashes on a GPU -- device mismatch

File ".../losses.py", line 17, in negative_binom_loss
    torch.lgamma(theta + _EPS)
RuntimeError: Expected all tensors to be on the same device,
              but found at least two devices, cuda:0 and cpu!

net.fit(dataset, y=None) with

device="cuda" if torch.cuda.is_available() else "cpu",

puts the module and its predictions on cuda:0, but the target coming out of the custom Dataset/collate_fn stays on CPU, so the NB loss gets a mix.

This cannot reproduce on CPU -- there every tensor is on cpu and the check passes -- which is why it survived testing. As submitted, BABEL cannot train on a GPU at all. Verified both directions on the same node:

babel_train (--gpus all)   EXIT=1  11s   CUDA proc 3/10, peak 169 MiB  -> RuntimeError above
babel_train (no --gpus)    EXIT=0        completes fine, model written

Worth noting it is a real GPU path -- the model does move to the device and allocates 169 MiB -- it's a device-handling bug, not a silent CPU fallback.

3. The published method id would be babel_predict

"method_id": meta["name"],

meta["name"] in babel_predict is babel_predict, and both metrics copy the prediction's id verbatim into the score. Confirmed on the output:

prediction: shape=(427, 1500) dtype=float32 uns.method_id='babel_predict'

novel/simple_mlp hardcode the wrapper name for this reason; #20 strips the _predict suffix.

What does work

Using a CPU-trained model, babel_predict runs on the GPU node and produces a valid prediction:

babel_predict  EXIT=0  16s   peak 0% util  159 MiB  CUDA proc 1/15
prediction: shape=(427, 1500) dtype=float32
CONTRACT: OK

It creates a CUDA context but I didn't see sustained GPU compute -- on a test-sized input the predict pass may simply be too short to sample at 1 Hz, so I wouldn't read much into that number either way. Worth re-checking once training works on GPU.

I merged current main into a local copy of this branch -- clean merge, viash ns build green -- and have the registration fix ready; say the word and I'll push it, or leave it to you.