From c8fc2fcb6900efda0f389dec5688976dde952784 Mon Sep 17 00:00:00 2001 From: Robrecht Cannoodt Date: Thu, 13 Aug 2026 11:51:12 +0200 Subject: [PATCH 1/3] fix permuted atac target in babel loss ChromDecoder emits peaks in chrom_groups order, but target2_bin was sliced out of y in original var_names order. --- src/methods/babel/babel_train/script.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/src/methods/babel/babel_train/script.py b/src/methods/babel/babel_train/script.py index eca8bac1..8d8483ec 100644 --- a/src/methods/babel/babel_train/script.py +++ b/src/methods/babel/babel_train/script.py @@ -73,13 +73,12 @@ def _row_tensor(mat, idx): class PairedDataset(Dataset): - def __init__(self, x1, x2_per_chrom, y1, y2, size_factors1): + def __init__(self, x1, x2_per_chrom, y1, size_factors1): # Store the RNA/ATAC matrices sparse (CSR); rows are densified on demand in # __getitem__ so the full dense matrix is never materialized. self.x1 = _as_csr(x1) self.x2_per_chrom = [_as_csr(c) for c in x2_per_chrom] self.y1 = _as_csr(y1) - self.y2 = _as_csr(y2) self.size_factors1 = torch.from_numpy(size_factors1) def __len__(self): @@ -91,7 +90,7 @@ def __getitem__(self, idx): "x2_per_chrom": [_row_tensor(c, idx) for c in self.x2_per_chrom], "size_factors1": self.size_factors1[idx], } - y = torch.cat([_row_tensor(self.y1, idx), _row_tensor(self.y2, idx)]) # dummy combined target, unused directly + y = _row_tensor(self.y1, idx) return X, y @@ -118,16 +117,10 @@ def forward(self, x1, x2_per_chrom, size_factors1): class BabelNet(skorch.NeuralNet): - def __init__(self, *args, n_genes, n_peaks, **kwargs): - self._n_genes = n_genes - self._n_peaks = n_peaks - super().__init__(*args, **kwargs) - def get_loss(self, y_pred, y_true, X=None, training=False): - y_true = y_true.to(self.device) preds11, preds12, preds21, preds22, _, _ = y_pred - target1 = y_true[:, : self._n_genes] - target2_bin = y_true[:, self._n_genes : self._n_genes + self._n_peaks] + target1 = y_true.to(self.device) + target2_bin = torch.cat(X["x2_per_chrom"], dim=1).to(self.device) return self.criterion_(preds11, preds12, preds21, preds22, target1, target2_bin) @@ -172,7 +165,7 @@ def get_loss(self, y_pred, y_true, X=None, training=False): n_genes = X_rna.shape[1] n_peaks = X_atac_bin.shape[1] -dataset = PairedDataset(X_rna, X_atac_per_chrom, Y_rna_counts, X_atac_bin, size_factors) +dataset = PairedDataset(X_rna, X_atac_per_chrom, Y_rna_counts, size_factors) logger.info("Building model (n_genes=%d, n_peaks=%d, %d chromosome groups)...", n_genes, n_peaks, len(chrom_counts)) @@ -181,8 +174,6 @@ def get_loss(self, y_pred, y_true, X=None, training=False): module__input_dim1=n_genes, module__input_dim2=chrom_counts, module__hidden_dim=par["hidden_dim"], - n_genes=n_genes, - n_peaks=n_peaks, criterion=QuadLoss, criterion__loss2_weight=par["loss2_weight"], optimizer=torch.optim.Adam, From 60ec731ce18cfd09e5d4890ff84cccd20b767a55 Mon Sep 17 00:00:00 2001 From: Robrecht Cannoodt Date: Thu, 13 Aug 2026 11:52:15 +0200 Subject: [PATCH 2/3] free the redundant atac copies before training --- src/methods/babel/babel_train/script.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/methods/babel/babel_train/script.py b/src/methods/babel/babel_train/script.py index 8d8483ec..6137ad9e 100644 --- a/src/methods/babel/babel_train/script.py +++ b/src/methods/babel/babel_train/script.py @@ -157,14 +157,22 @@ def get_loss(self, y_pred, y_true, X=None, training=False): X_atac_bin = _atac_binarized(adata_atac) chrom_counts, chrom_groups = parse_chrom_groups(adata_atac.var_names) -# Slice peaks per chromosome on a CSC view (fast column indexing); each group stays -# sparse and is densified per-cell later. -X_atac_bin_csc = X_atac_bin.tocsc() -X_atac_per_chrom = [X_atac_bin_csc[:, idxs] for idxs in chrom_groups.values()] +rna_var_names = list(adata_rna.var_names) +atac_var_names = list(adata_atac.var_names) n_genes = X_rna.shape[1] n_peaks = X_atac_bin.shape[1] +# drop the inputs; everything needed downstream has been extracted above +del adata_mod1_train, adata_mod2_train, adata_atac, adata_rna + +# slice peaks per chromosome on a CSC view (fast column indexing); each group stays +# sparse and is densified per-cell later +X_atac_bin_csc = X_atac_bin.tocsc() +del X_atac_bin +X_atac_per_chrom = [_as_csr(X_atac_bin_csc[:, idxs]) for idxs in chrom_groups.values()] +del X_atac_bin_csc + dataset = PairedDataset(X_rna, X_atac_per_chrom, Y_rna_counts, size_factors) logger.info("Building model (n_genes=%d, n_peaks=%d, %d chromosome groups)...", n_genes, n_peaks, len(chrom_counts)) @@ -210,8 +218,8 @@ def get_loss(self, y_pred, y_true, X=None, training=False): "chrom_groups": chrom_groups, "hidden_dim": par["hidden_dim"], "direction": direction, - "rna_var_names": list(adata_rna.var_names), - "atac_var_names": list(adata_atac.var_names), + "rna_var_names": rna_var_names, + "atac_var_names": atac_var_names, "size_factor_median": float(np.median(np.asarray(Y_rna_counts.sum(axis=1)).ravel())), } From aa1906797ca1e3e186c6b0480615d1f086696372 Mon Sep 17 00:00:00 2001 From: Robrecht Cannoodt Date: Thu, 13 Aug 2026 11:53:01 +0200 Subject: [PATCH 3/3] size babel_predict chunks by dense bytes, not cell count --- src/methods/babel/babel_predict/script.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/methods/babel/babel_predict/script.py b/src/methods/babel/babel_predict/script.py index f9294a2f..cddb4378 100644 --- a/src/methods/babel/babel_predict/script.py +++ b/src/methods/babel/babel_predict/script.py @@ -79,7 +79,10 @@ def _lognorm_per_cell(pred_counts, target_sum=1e4): per_chrom = [X_bin[:, idxs].tocsr() for idxs in chrom_groups.values()] n_cells = X_bin.shape[0] -chunk_size = 4096 +# the per-chromosome tensors together span the full peak width, and exist on the host +# and the device at once -- cap one chunk's dense footprint at ~512 MiB rather than at +# a fixed cell count +chunk_size = int(np.clip(512 * 1024**2 // (X_bin.shape[1] * 4), 1, 4096)) pred_chunks = [] with torch.no_grad(): for start in range(0, n_cells, chunk_size):