babel: fix the atac loss target ordering, and trim peak memory

[rcannood]

Hi @benjaminfreyuu !

The sparse rewrite in the PR you created looks good to me :).

Three follow-ups on top of your branch, one commit each, so you can drop any of them independently.

Changes

• c8fc2fc fix permuted atac target in babel loss. ChromDecoder concatenates its per-chromosome heads in chrom_groups order, but target2_bin was sliced out of y, which is in original var_names order. Those two only agree if the peaks happen to be grouped by chromosome in lexicographic order -- on the 2021 fixture 1499/1500 columns are permuted. Checked against the real PairedDataset/_collate:

  fix   torch.cat(X['x2_per_chrom']) == X_atac[:, chrom_order] : True
  bug   y-sliced target              == X_atac[:, chrom_order] : False
  

  Taking the target from X["x2_per_chrom"], which get_loss already receives, is order-independent so no reindex is needed. It also drops y2 entirely: one less full ATAC copy in RAM, and ~17% off the per-cell densification.

• 60ec731 free the redundant atac copies before training. At net.fit() five objects each held a full copy of the ATAC non-zeros: both input AnnDatas, X_atac_bin, X_atac_bin_csc, and the per-chromosome slices. The CSC view and the inputs are pure temporaries that never went out of scope. Extract the var_names up front, then del.

• aa19067 size babel_predict chunks by dense bytes, not cell count. The per-chromosome tensors together span the full peak width and exist on the host and the device at once, so chunk_size = 4096 is ~3.7 GB at 229k peaks -- more than it needs to be on a 16 GB T4 (midgpu). Capping one chunk's dense footprint at 512 MiB gives 586 cells/chunk at 229k peaks, and leaves the 1500-peak fixture at 4096, i.e. unchanged.

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 -- the committed output_model.pkl still loads, the bundle keys are unchanged

Not changed

It looks like reindex_to_chrom_groups() is no longer used, and could be removed.

[benjaminfreyuu]

Looks good!