diff --git a/scripts/create_datasets/test_resources.sh b/scripts/create_datasets/test_resources.sh index 00b393d7..5e084c34 100755 --- a/scripts/create_datasets/test_resources.sh +++ b/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 diff --git a/src/methods/babel/babel/config.vsh.yaml b/src/methods/babel/babel/config.vsh.yaml new file mode 100644 index 00000000..8cf998bc --- /dev/null +++ b/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 diff --git a/src/methods/babel/babel/main.nf b/src/methods/babel/babel/main.nf new file mode 100644 index 00000000..9ed90555 --- /dev/null +++ b/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 +} diff --git a/src/methods/babel/babel_predict/config.vsh.yaml b/src/methods/babel/babel_predict/config.vsh.yaml new file mode 100644 index 00000000..d989db25 --- /dev/null +++ b/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] diff --git a/src/methods/babel/babel_predict/script.py b/src/methods/babel/babel_predict/script.py new file mode 100644 index 00000000..837d6b72 --- /dev/null +++ b/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"]) diff --git a/src/methods/babel/babel_train/config.vsh.yaml b/src/methods/babel/babel_train/config.vsh.yaml new file mode 100644 index 00000000..e8bead66 --- /dev/null +++ b/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] \ No newline at end of file diff --git a/src/methods/babel/babel_train/script.py b/src/methods/babel/babel_train/script.py new file mode 100644 index 00000000..f35cedfe --- /dev/null +++ b/src/methods/babel/babel_train/script.py @@ -0,0 +1,214 @@ +import logging +import pickle +import sys + +import anndata as ad +import numpy as np +import scanpy as sc +from scipy.sparse import issparse +import skorch +import torch +from torch import nn +from torch.utils.data import Dataset + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +## VIASH START +par = { + "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", + "output": "output_model.pkl", + "hidden_dim": 16, + "nn_epochs": 100, + "lr": 0.001, + "batch_size": 64, + "loss2_weight": 3.0, + "early_stopping_patience": 20, +} +meta = {"name": "babel", "resources_dir": "src/methods/babel"} +## VIASH END + +sys.path.append(meta["resources_dir"]) + +from chrom_utils import parse_chrom_groups +from model import AssymSplicedAutoEncoder +from losses import QuadLoss + + +def _to_dense(X): + return X.toarray() if issparse(X) else np.asarray(X) + + +def _rna_matrix(adata): + counts = _to_dense(adata.layers.get("counts", adata.X)).astype(np.float32) + if "normalized" in adata.layers: + X = _to_dense(adata.layers["normalized"]).astype(np.float32) + else: + tmp = ad.AnnData(counts.copy()) + sc.pp.normalize_total(tmp) + sc.pp.log1p(tmp) + X = tmp.X.astype(np.float32) + size_factors = counts.sum(axis=1, keepdims=True) + med = np.median(size_factors[size_factors > 0]) if np.any(size_factors > 0) else 1.0 + size_factors = size_factors / (med if med > 0 else 1.0) + return X, counts, size_factors.astype(np.float32) + + +def _atac_binarized(adata): + counts = _to_dense(adata.layers.get("counts", adata.X)) + return (counts > 0).astype(np.float32) + + +class PairedDataset(Dataset): + def __init__(self, x1, x2_per_chrom, y1, y2, size_factors1): + self.x1 = torch.from_numpy(x1) + self.x2_per_chrom = [torch.from_numpy(c) for c in x2_per_chrom] + self.y1 = torch.from_numpy(y1) + self.y2 = torch.from_numpy(y2) + self.size_factors1 = torch.from_numpy(size_factors1) + + def __len__(self): + return self.x1.shape[0] + + def __getitem__(self, idx): + X = { + "x1": self.x1[idx], + "x2_per_chrom": [c[idx] for c in self.x2_per_chrom], + "size_factors1": self.size_factors1[idx], + } + y = torch.cat([self.y1[idx], self.y2[idx]]) # dummy combined target, unused directly + return X, y + + +def _collate(batch): + xs, ys = zip(*batch) + x1 = torch.stack([x["x1"] for x in xs]) + size_factors1 = torch.stack([x["size_factors1"] for x in xs]) + n_chroms = len(xs[0]["x2_per_chrom"]) + x2_per_chrom = [torch.stack([x["x2_per_chrom"][i] for x in xs]) for i in range(n_chroms)] + y = torch.stack(ys) + return {"x1": x1, "x2_per_chrom": x2_per_chrom, "size_factors1": size_factors1}, y + + +class BabelModule(nn.Module): + """Wraps AssymSplicedAutoEncoder so skorch's forward(**X) call signature works + with the dict-of-tensors batch produced by PairedDataset/_collate.""" + + def __init__(self, input_dim1, input_dim2, hidden_dim=16): + super().__init__() + self.autoencoder = AssymSplicedAutoEncoder(input_dim1, input_dim2, hidden_dim=hidden_dim) + + def forward(self, x1, x2_per_chrom, size_factors1): + return self.autoencoder(x1, x2_per_chrom, size_factors1=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] + return self.criterion_(preds11, preds12, preds21, preds22, target1, target2_bin) + + +logger.info("Reading input files...") +adata_mod1_train = ad.read_h5ad(par["input_train_mod1"]) +adata_mod2_train = ad.read_h5ad(par["input_train_mod2"]) + +modality1 = adata_mod1_train.uns.get("modality") +modality2 = adata_mod2_train.uns.get("modality") + +if {modality1, modality2} != {"GEX", "ATAC"}: + raise ValueError( + f"babel only supports GEX<->ATAC translation, got modalities " + f"({modality1!r}, {modality2!r}). This is a direct port of BABEL " + "(Wu et al. 2021), whose architecture is not adapted for other modality pairs." + ) + +if modality1 != "ATAC": + raise ValueError( + f"babel only supports ATAC->GEX prediction (input_train_mod1 must be ATAC, " + f"input_train_mod2 must be GEX), got mod1={modality1!r}, mod2={modality2!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 accessible/inaccessible separation, even after " + "training to convergence) and is intentionally disabled rather than silently " + "producing uninformative output." + ) + +adata_atac, adata_rna = adata_mod1_train, adata_mod2_train +direction = "mod1_is_atac" + +logger.info("Preprocessing RNA and ATAC...") +X_rna, Y_rna_counts, size_factors = _rna_matrix(adata_rna) +X_atac_bin = _atac_binarized(adata_atac) + +chrom_counts, chrom_groups = parse_chrom_groups(adata_atac.var_names) +X_atac_per_chrom = [X_atac_bin[:, idxs] for idxs in chrom_groups.values()] + +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) + +logger.info("Building model (n_genes=%d, n_peaks=%d, %d chromosome groups)...", n_genes, n_peaks, len(chrom_counts)) + +net = BabelNet( + module=BabelModule, + 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, + lr=par["lr"], + max_epochs=par["nn_epochs"], + batch_size=par["batch_size"], + iterator_train__collate_fn=_collate, + iterator_train__shuffle=True, + iterator_valid__collate_fn=_collate, + train_split=skorch.dataset.ValidSplit(0.1), + callbacks=[ + skorch.callbacks.EarlyStopping(patience=par["early_stopping_patience"]), + skorch.callbacks.LRScheduler( + policy=torch.optim.lr_scheduler.ReduceLROnPlateau, + monitor="valid_loss", + mode="min", + factor=0.1, + patience=10, + ), + skorch.callbacks.GradientNormClipping(gradient_clip_value=5), + ], + device="cuda" if torch.cuda.is_available() else "cpu", +) + +logger.info("Training BABEL autoencoder...") +net.fit(dataset, y=None) + +logger.info("Saving model bundle...") +bundle = { + "state_dict": net.module_.autoencoder.state_dict(), + "n_genes": n_genes, + "n_peaks": n_peaks, + "chrom_counts": chrom_counts, + "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), + "size_factor_median": float(np.median(Y_rna_counts.sum(axis=1))), +} + +with open(par["output"], "wb") as f: + pickle.dump(bundle, f, protocol=4) + +logger.info("Training complete. Model saved to %s", par["output"]) diff --git a/src/methods/babel/chrom_utils.py b/src/methods/babel/chrom_utils.py new file mode 100644 index 00000000..22f87408 --- /dev/null +++ b/src/methods/babel/chrom_utils.py @@ -0,0 +1,53 @@ +import logging +import re +from collections import OrderedDict + +import numpy as np + +logger = logging.getLogger(__name__) + +_CHROM_RE = re.compile(r"^chr?([0-9XYM]+)[:\-]", re.IGNORECASE) + + +def parse_chrom_groups(var_names): + """Group feature indices by chromosome, parsed from genomic-coordinate-style + var_names (e.g. "chr1:1000-2000"), matching BABEL's per-chromosome grouping. + + Chromosomes are sorted lexicographically (matching BABEL's `sorted(set(...))`) + and must be reproduced in the same order at predict time. + + Falls back to a single group spanning all features if the names don't match + the expected chromosome-coordinate pattern. + """ + chrom_of_feature = [] + for name in var_names: + m = _CHROM_RE.match(str(name)) + chrom_of_feature.append(m.group(1).upper() if m else None) + + n_parsed = sum(c is not None for c in chrom_of_feature) + unique_chroms = sorted({c for c in chrom_of_feature if c is not None}) + + if n_parsed < len(var_names) or len(unique_chroms) < 2: + logger.warning( + "Could not parse chromosome-of-origin for all %d features from var_names " + "(parsed %d, %d unique chromosomes); falling back to a single feature group.", + len(var_names), n_parsed, len(unique_chroms), + ) + chrom_groups = OrderedDict([("ALL", np.arange(len(var_names)))]) + return [len(var_names)], chrom_groups + + chrom_groups = OrderedDict() + chrom_array = np.array(chrom_of_feature) + for chrom in unique_chroms: + chrom_groups[chrom] = np.where(chrom_array == chrom)[0] + + counts = [len(idxs) for idxs in chrom_groups.values()] + return counts, chrom_groups + + +def reindex_to_chrom_groups(X, chrom_groups): + """Concatenate feature columns of X in chrom_groups order (identity if already sorted). + Used to guarantee predict-time feature order matches the order the model was trained with. + """ + idx = np.concatenate(list(chrom_groups.values())) + return X[:, idx] diff --git a/src/methods/babel/losses.py b/src/methods/babel/losses.py new file mode 100644 index 00000000..63beaeb5 --- /dev/null +++ b/src/methods/babel/losses.py @@ -0,0 +1,76 @@ +"""BABEL-style losses: negative binomial reconstruction for RNA, BCE for binarized ATAC, +combined via QuadLoss with a constant (non-warmup) cross-modality weight, matching the +weighting actually used in BABEL's shipped bin/train_model.py (cross_warmup_delay=0, +link_strength=0, i.e. no alignment term, no warmup schedule). +""" + +import torch +from torch import nn + +_EPS = 1e-8 + + +def negative_binom_loss(y_true, y_pred, theta, mean=True): + """DCA-style negative binomial negative log-likelihood.""" + theta = torch.clamp(theta, max=1e6) + t1 = ( + torch.lgamma(theta + _EPS) + + torch.lgamma(y_true + 1.0) + - torch.lgamma(y_true + theta + _EPS) + ) + t2 = (theta + y_true) * torch.log1p(y_pred / (theta + _EPS)) + y_true * ( + torch.log(theta + _EPS) - torch.log(y_pred + _EPS) + ) + loss = t1 + t2 + return loss.mean() if mean else loss + + +class NegativeBinomialLoss(nn.Module): + def forward(self, preds, target): + mean, theta, _ = preds + return negative_binom_loss(target, mean, theta) + + +class BCELoss(nn.Module): + """Uses only the first decoder output head (the probability/accessibility head).""" + + def __init__(self): + super().__init__() + self._bce = nn.BCELoss() + + def forward(self, preds, target): + prob = preds[0] + return self._bce(torch.clamp(prob, min=_EPS, max=1.0 - _EPS), target) + + +class QuadLoss(nn.Module): + """Combines the four AssymSplicedAutoEncoder paths: + loss11 (RNA->RNA, NB), loss22 (ATAC->ATAC, BCE), + loss21 (ATAC->RNA, NB), loss12 (RNA->ATAC, BCE). + + total = loss11 + loss2_weight*loss22 + cross_weight*(loss21 + loss2_weight*loss12) + + No warmup/link terms, matching BABEL's actual shipped training config. + """ + + def __init__(self, loss2_weight=3.0, cross_weight=1.0): + super().__init__() + self.loss1 = NegativeBinomialLoss() + self.loss2 = BCELoss() + self.loss2_weight = loss2_weight + self.cross_weight = cross_weight + + def get_component_losses(self, preds11, preds12, preds21, preds22, target1, target2_bin): + loss11 = self.loss1(preds11, target1) + loss21 = self.loss1(preds21, target1) + loss12 = self.loss2(preds12, target2_bin) + loss22 = self.loss2(preds22, target2_bin) + return loss11, loss12, loss21, loss22 + + def forward(self, preds11, preds12, preds21, preds22, target1, target2_bin): + loss11, loss12, loss21, loss22 = self.get_component_losses( + preds11, preds12, preds21, preds22, target1, target2_bin + ) + same_modality = loss11 + self.loss2_weight * loss22 + cross_modality = loss21 + self.loss2_weight * loss12 + return same_modality + self.cross_weight * cross_modality diff --git a/src/methods/babel/model.py b/src/methods/babel/model.py new file mode 100644 index 00000000..f9944ca8 --- /dev/null +++ b/src/methods/babel/model.py @@ -0,0 +1,205 @@ +"""BABEL-style spliced autoencoder (Wu et al., PNAS 2021, github.com/wukevin/babel). + +Ported architecture: a dense Encoder/Decoder pair for RNA and a chromosome-split +ChromEncoder/ChromDecoder pair for ATAC, combined into an AssymSplicedAutoEncoder +that produces four translation paths (RNA->RNA, RNA->ATAC, ATAC->RNA, ATAC->ATAC) +from two separate (unshared) per-modality latent spaces. +""" + +import torch +from torch import nn + + +class Exp(nn.Module): + def forward(self, x): + return torch.exp(torch.clamp(x, min=-15, max=15)) + + +class ClippedSoftplus(nn.Module): + def forward(self, x): + return torch.clamp(nn.functional.softplus(x), min=1e-4, max=1e4) + + +def _init_linear(layer): + nn.init.xavier_uniform_(layer.weight) + nn.init.zeros_(layer.bias) + return layer + + +def _as_list(activations, n): + if activations is None: + return [None] * n + if isinstance(activations, nn.Module): + return [activations] + [None] * (n - 1) + activations = list(activations) + return activations + [None] * (n - len(activations)) + + +class Encoder(nn.Module): + """Dense encoder: num_inputs -> 64 -> num_units.""" + + def __init__(self, num_inputs, num_units=32, activation=nn.PReLU): + super().__init__() + self.net = nn.Sequential( + _init_linear(nn.Linear(num_inputs, 64)), + nn.BatchNorm1d(64), + activation(), + _init_linear(nn.Linear(64, num_units)), + nn.BatchNorm1d(num_units), + activation(), + ) + + def forward(self, x): + return self.net(x) + + +class Decoder(nn.Module): + """Dense decoder: num_units -> 64 -> three parallel num_outputs heads.""" + + def __init__(self, num_units, num_outputs, final_activations=None, activation=nn.PReLU): + super().__init__() + self.shared = nn.Sequential( + _init_linear(nn.Linear(num_units, 64)), + nn.BatchNorm1d(64), + activation(), + ) + self.head1 = _init_linear(nn.Linear(64, num_outputs)) + self.head2 = _init_linear(nn.Linear(64, num_outputs)) + self.head3 = _init_linear(nn.Linear(64, num_outputs)) + acts = _as_list(final_activations, 3) + self.act1, self.act2, self.act3 = acts + + def forward(self, x, size_factors=None): + h = self.shared(x) + r1 = self.head1(h) + if self.act1 is not None: + r1 = self.act1(r1) + if size_factors is not None: + r1 = r1 * size_factors + r2 = self.head2(h) + if self.act2 is not None: + r2 = self.act2(r2) + r3 = self.head3(h) + if self.act3 is not None: + r3 = self.act3(r3) + return r1, r2, r3 + + +class ChromEncoder(nn.Module): + """Per-chromosome encoder: each chrom's features -> 32 -> 16, concat, -> latent_dim.""" + + def __init__(self, num_inputs, latent_dim=32, activation=nn.PReLU): + super().__init__() + self.branches = nn.ModuleList([ + nn.Sequential( + _init_linear(nn.Linear(n_i, 32)), + nn.BatchNorm1d(32), + activation(), + _init_linear(nn.Linear(32, 16)), + nn.BatchNorm1d(16), + activation(), + ) + for n_i in num_inputs + ]) + self.chrom_sizes = list(num_inputs) + self.combine = nn.Sequential( + _init_linear(nn.Linear(16 * len(num_inputs), latent_dim)), + nn.BatchNorm1d(latent_dim), + activation(), + ) + + def forward(self, x_per_chrom): + outs = [branch(x_i) for branch, x_i in zip(self.branches, x_per_chrom)] + return self.combine(torch.cat(outs, dim=-1)) + + +class ChromDecoder(nn.Module): + """Inverse of ChromEncoder: latent_dim -> per-chrom chunks -> three parallel + per-chrom output heads, concatenated back to genome-wide vectors.""" + + def __init__(self, num_outputs, latent_dim=32, final_activations=None, activation=nn.PReLU): + super().__init__() + n_chroms = len(num_outputs) + self.chrom_sizes = list(num_outputs) + self.expand = nn.Sequential( + _init_linear(nn.Linear(latent_dim, n_chroms * 16)), + nn.BatchNorm1d(n_chroms * 16), + activation(), + ) + self.shared_branches = nn.ModuleList([ + nn.Sequential( + _init_linear(nn.Linear(16, 32)), + nn.BatchNorm1d(32), + activation(), + ) + for _ in num_outputs + ]) + self.head1 = nn.ModuleList([_init_linear(nn.Linear(32, n_i)) for n_i in num_outputs]) + self.head2 = nn.ModuleList([_init_linear(nn.Linear(32, n_i)) for n_i in num_outputs]) + self.head3 = nn.ModuleList([_init_linear(nn.Linear(32, n_i)) for n_i in num_outputs]) + acts = _as_list(final_activations, 3) + self.act1, self.act2, self.act3 = acts + + def forward(self, x): + h = self.expand(x) + chunks = torch.chunk(h, len(self.chrom_sizes), dim=-1) + r1_parts, r2_parts, r3_parts = [], [], [] + for branch, h1, h2, h3, chunk in zip( + self.shared_branches, self.head1, self.head2, self.head3, chunks + ): + hc = branch(chunk) + r1_parts.append(h1(hc)) + r2_parts.append(h2(hc)) + r3_parts.append(h3(hc)) + r1 = torch.cat(r1_parts, dim=-1) + r2 = torch.cat(r2_parts, dim=-1) + r3 = torch.cat(r3_parts, dim=-1) + if self.act1 is not None: + r1 = self.act1(r1) + if self.act2 is not None: + r2 = self.act2(r2) + if self.act3 is not None: + r3 = self.act3(r3) + return r1, r2, r3 + + +class AssymSplicedAutoEncoder(nn.Module): + """Domain 1 (RNA) uses dense Encoder/Decoder; domain 2 (ATAC) uses chromosome-split + ChromEncoder/ChromDecoder. Four translation paths share the two encoders/decoders: + RNA->RNA, RNA->ATAC, ATAC->RNA, ATAC->ATAC. Latent spaces are per-modality, not tied; + alignment across modalities is enforced only via the training loss (QuadLoss), not + architecturally. + """ + + def __init__( + self, + input_dim1, + input_dim2, + hidden_dim=16, + final_activations1=None, + final_activations2=None, + ): + super().__init__() + final_activations1 = final_activations1 or [Exp(), ClippedSoftplus()] + final_activations2 = final_activations2 or [nn.Sigmoid()] + self.encoder1 = Encoder(input_dim1, num_units=hidden_dim) + self.decoder1 = Decoder(hidden_dim, input_dim1, final_activations=final_activations1) + self.encoder2 = ChromEncoder(input_dim2, latent_dim=hidden_dim) + self.decoder2 = ChromDecoder(input_dim2, latent_dim=hidden_dim, final_activations=final_activations2) + + def encode1(self, x1): + return self.encoder1(x1) + + def encode2(self, x2_per_chrom): + return self.encoder2(x2_per_chrom) + + def forward(self, x1, x2_per_chrom, size_factors1=None): + encoded1 = self.encoder1(x1) + encoded2 = self.encoder2(x2_per_chrom) + + preds11 = self.decoder1(encoded1, size_factors=size_factors1) + preds12 = self.decoder2(encoded1) + preds21 = self.decoder1(encoded2, size_factors=size_factors1) + preds22 = self.decoder2(encoded2) + + return preds11, preds12, preds21, preds22, encoded1, encoded2 diff --git a/src/workflows/run_benchmark/config.vsh.yaml b/src/workflows/run_benchmark/config.vsh.yaml index 514ba8b6..089a0f55 100644 --- a/src/workflows/run_benchmark/config.vsh.yaml +++ b/src/workflows/run_benchmark/config.vsh.yaml @@ -74,6 +74,7 @@ dependencies: - name: methods/guanlab_dengkw_pm - name: methods/novel - name: methods/simple_mlp + - name: methods/babel - name: methods/senkin_tmp - name: metrics/correlation - name: metrics/mse diff --git a/src/workflows/run_benchmark/main.nf b/src/workflows/run_benchmark/main.nf index 890d8b1f..86679faf 100644 --- a/src/workflows/run_benchmark/main.nf +++ b/src/workflows/run_benchmark/main.nf @@ -19,6 +19,7 @@ methods = [ guanlab_dengkw_pm, novel, simple_mlp, + babel, senkin_tmp ]