babel: keep ATAC/RNA sparse instead of densifying up front

[benjaminfreyuu]

Problem

babel_train densified the full ATAC (and RNA) matrix with a single _to_dense() before building the Dataset. On the full 2022 datasets this OOMs — in the 2026-08-01 run, babel on openproblems_neurips2022/pbmc_multiome/swap crashed with:

X_atac_bin = _atac_binarized(adata_atac) → _to_dense(...) → X.toarray()
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 222. GiB for an
array with shape (130095, 228942) and data type float64

It only passed on the tiny 2021 test fixture (1,500 peaks); the real ATAC data has ~229k peaks.

Fix

Restore the original BABEL approach (wukevin/babel's sc_data_loaders): keep the matrices sparse (CSR) and densify one cell at a time in PairedDataset.__getitem__ — mirroring BABEL's per-cell utils.ensure_arr(X[i]).flatten(). Peaks are sliced per chromosome on a CSC view. babel_predict likewise binarizes sparsely and runs inference in cell chunks rather than loading the whole dense matrix onto the GPU.

Peak memory now scales with a minibatch, not the full matrix.

Verification

• viash test src/methods/babel/babel_train/config.vsh.yaml → 1/1 passed
• viash test src/methods/babel/babel_predict/config.vsh.yaml → 1/1 passed
• Scale probe: on a synthetic 20k × 60k ATAC (9.6 GB dense float64), the sparse path peaks at 1.7 GB, while the old toarray() path is OOM-killed under a 6 GB cap.

Note

This removes the OOM crash. Walltime on the full 2022 data is untested here (no local GPU / full data); per-cell densification matches upstream, but if it proves slow, iterator_train__num_workers and/or the resource label can be tuned as a follow-up.

🤖 Generated with Claude Code