import abc import gc from typing import Callable, Dict, List, Optional, Tuple import os import numpy as np import pickle import lightning.pytorch as pl import torch import torch.nn as nn import torch.nn.functional as F from torch.distributions import ( Bernoulli, ContinuousBernoulli, Poisson, ) from scvi.distributions import NegativeBinomial, ZeroInflatedNegativeBinomial from torchmetrics import ExplainedVariance, MeanSquaredError, MetricCollection from torchmetrics.classification import MulticlassF1Score from self_supervision.models.base.base import MLP from self_supervision.models.contrastive.byol import BYOL from self_supervision.models.contrastive.bt import ( BarlowTwins, Transform, LARS, adjust_learning_rate, ) def _mask_gene_programs_numpy( inputs: torch.Tensor, encoded_gene_program: np.ndarray, masking_rate: float ): """ Randomly choose masking_rate percent of the gene programs (rows) to mask. Mask all the genes in the inputs tensor that are 1s in the chosen row of encoded_gene_program. return a binary mask of the tensor and a fraction of how many genes are set to 0. :param inputs: input data :param encoded_gene_program: encoded gene programs :param masking_rate: rate of gene programs to mask :return: mask, frac: mask is a binary mask of the tensor and frac is a fraction of how many genes are set to 0. """ # Randomly choose masking_rate percent of the gene programs to mask num_gene_programs = encoded_gene_program.shape[0] num_gene_programs_to_mask = int(num_gene_programs * masking_rate) gene_programs_to_mask = np.random.permutation(num_gene_programs)[ :num_gene_programs_to_mask ] # Mask the input tensor and cast it to numpy on cpu mask = np.ones_like(inputs.cpu().numpy()) for ix in gene_programs_to_mask: mask[:, encoded_gene_program[ix, :] == 1] = 0 return mask, 1 - mask.sum() / mask.size def _only_activate_gene_program_numpy( inputs: torch.Tensor, # write that encoded_gene_program is a dict of a tuple of torch.Tensor # encoded_gene_program: Dict[Tuple[torch.Tensor, torch.Tensor]], encoded_gene_program: Dict, ): """ For each row of the inputs tensor randomly choose one transcription factor (key of encoded_gene_program) The first element of encoded_gene_program[key] is a vector (torch.Tensor) of indices of genes in the gene program. The input mask should correspond to that mask, i.e., the mask is 1 where the genes are in the gene program. The second element of encoded_gene_program[key] is a vector (torch.Tensor) of index of the corresponding gene program. The output mask should be 1 that vector is 1. :param inputs: input data :param encoded_gene_program: encoded gene programs :return: (input_mask, output_mask): input_mask is a binary mask of the tensor and output_mask is a binary mask of the tensor. """ # Randomly choose one transcription factor (key of encoded_gene_program) per row of the inputs tensor num_gene_programs = len(encoded_gene_program) gene_programs_to_activate = np.random.permutation(num_gene_programs)[ : inputs.shape[0] ] # Mask the input tensor and cast it to numpy on cpu input_mask = np.zeros_like(inputs.cpu().numpy()) output_mask = np.zeros_like(inputs.cpu().numpy()) for ix, key in enumerate(gene_programs_to_activate): input_mask[ix, encoded_gene_program[key][0]] = 1 output_mask[ix, encoded_gene_program[key][1]] = 1 return input_mask, output_mask def _mask_single_gene_programs( inputs: torch.Tensor, encoded_gene_program: np.ndarray, ): """ Randomly choose 1 gene program, mask everything apart from that gene program both in the input and output. """ # Randomly choose 1 gene program num_gene_programs = encoded_gene_program.shape[0] gene_program_to_mask = np.random.permutation(num_gene_programs)[:1] # Mask the input tensor and cast it to numpy on cpu input_mask = np.zeros_like(inputs.cpu().numpy()) output_mask = np.zeros_like(inputs.cpu().numpy()) input_mask[:, encoded_gene_program[gene_program_to_mask, :] == 1] = 1 output_mask[:, encoded_gene_program[gene_program_to_mask, :] == 1] = 1 return input_mask, output_mask class BaseAutoEncoder(pl.LightningModule, abc.ABC): """ Base class for AutoEncoder models in PyTorch Lightning. Args: gene_dim (int): Dimension of the gene. batch_size (int): Batch size. reconst_loss (str, optional): Reconstruction loss function. Defaults to 'mse'. learning_rate (float, optional): Learning rate. Defaults to 0.005. weight_decay (float, optional): Weight decay. Defaults to 0.1. optimizer (Callable[..., torch.optim.Optimizer], optional): Optimizer class. Defaults to torch.optim.AdamW. lr_scheduler (Callable, optional): Learning rate scheduler. Defaults to None. lr_scheduler_kwargs (Dict, optional): Additional arguments for the learning rate scheduler. Defaults to None. gc_frequency (int, optional): Frequency of garbage collection. Defaults to 1. automatic_optimization (bool, optional): Whether to use automatic optimization. Defaults to True. supervised_subset (Optional[int], optional): Subset size for supervised training. Defaults to None. """ autoencoder: nn.Module # autoencoder mapping von gene_dim to gene_dim def __init__( self, gene_dim: int, batch_size: int, reconst_loss: str = "mse", learning_rate: float = 0.005, weight_decay: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, gc_frequency: int = 1, automatic_optimization: bool = True, supervised_subset: Optional[int] = None, ): super(BaseAutoEncoder, self).__init__() self.automatic_optimization = automatic_optimization self.gene_dim = gene_dim self.batch_size = batch_size self.gc_freq = gc_frequency self.lr = learning_rate self.weight_decay = weight_decay self.optim = optimizer self.reconst_loss = reconst_loss self.lr_scheduler = lr_scheduler self.lr_scheduler_kwargs = lr_scheduler_kwargs self.supervised_subset = supervised_subset metrics = MetricCollection( { "explained_var_weighted": ExplainedVariance( multioutput="variance_weighted" ), "explained_var_uniform": ExplainedVariance( multioutput="uniform_average" ), "mse": MeanSquaredError(), } ) self.train_metrics = metrics.clone(prefix="train_") self.val_metrics = metrics.clone(prefix="val_") self.test_metrics = metrics.clone(prefix="test_") def _calc_reconstruction_loss( self, preds: torch.Tensor, targets: torch.Tensor, reduction: str = "mean" ): """ Calculate the reconstruction loss. Args: preds (torch.Tensor): Predicted values. targets (torch.Tensor): Target values. reduction (str, optional): Reduction method. Defaults to 'mean'. Returns: torch.Tensor: Reconstruction loss. """ if self.reconst_loss == "continuous_bernoulli": loss = -ContinuousBernoulli(probs=preds).log_prob(targets) if reduction == "mean": loss = loss.mean() elif reduction == "sum": loss = loss.sum() elif self.reconst_loss == "bce": loss = F.binary_cross_entropy(preds, targets, reduction=reduction) elif self.reconst_loss == "mae": loss = F.l1_loss(preds, targets, reduction=reduction) else: loss = F.mse_loss(preds, targets, reduction=reduction) return loss @abc.abstractmethod def _step(self, batch): """ Calculate predictions (int64 tensor) and loss. Args: batch: Input batch. Returns: tuple: Tuple containing predictions and loss. """ pass def on_after_batch_transfer(self, batch, dataloader_idx): """ Perform operations on the batch after it has been transferred to the device. Args: batch: Input batch. dataloader_idx: Index of the dataloader. Returns: batch: Processed batch. """ if isinstance(batch, dict): # Case for MultiomicsDataloader return batch else: return batch[0] def forward(self, batch): """ Forward pass of the autoencoder. Args: batch: Input batch. Returns: tuple: Tuple containing the latent representation and the reconstructed output. """ x_in = batch["X"] x_latent = self.encoder(x_in) x_reconst = self.decoder(x_latent) return x_latent, x_reconst def on_train_epoch_end(self) -> None: """ Perform operations at the end of each training epoch. """ gc.collect() def on_validation_epoch_end(self) -> None: """ Perform operations at the end of each validation epoch. """ gc.collect() def configure_optimizers(self): """ Configure the optimizer and learning rate scheduler. Returns: dict: Configuration for the optimizer and learning rate scheduler. """ optimizer_config = { "optimizer": self.optim( self.parameters(), lr=self.lr, weight_decay=self.weight_decay ) } if self.lr_scheduler is not None: lr_scheduler_kwargs = ( {} if self.lr_scheduler_kwargs is None else self.lr_scheduler_kwargs ) interval = lr_scheduler_kwargs.pop("interval", "epoch") monitor = lr_scheduler_kwargs.pop("monitor", "val_loss_epoch") frequency = lr_scheduler_kwargs.pop("frequency", 1) scheduler = self.lr_scheduler( optimizer_config["optimizer"], **lr_scheduler_kwargs ) optimizer_config["lr_scheduler"] = { "scheduler": scheduler, "interval": interval, "monitor": monitor, "frequency": frequency, } return optimizer_config class BaseClassifier(pl.LightningModule, abc.ABC): """ Base classifier module for a PyTorch Lightning model. Args: gene_dim (int): Dimension of the gene feature. type_dim (int): Dimension of the cell type. class_weights (np.ndarray): Array of class weights. child_matrix (np.ndarray): Matrix representing the child relationship between cell types. batch_size (int): Batch size. hvg (bool): Flag indicating whether to use highly variable genes. num_hvgs (int): Number of highly variable genes. supervised_subset (Optional[int], optional): Subset of the data to use for supervised training. Defaults to None. learning_rate (float, optional): Learning rate. Defaults to 0.005. weight_decay (float, optional): Weight decay. Defaults to 0.1. optimizer (Callable[..., torch.optim.Optimizer], optional): Optimizer function. Defaults to torch.optim.AdamW. lr_scheduler (Callable, optional): Learning rate scheduler. Defaults to None. lr_scheduler_kwargs (Dict, optional): Additional arguments for the learning rate scheduler. Defaults to None. gc_frequency (int, optional): Frequency of garbage collection. Defaults to 1. """ classifier: Callable # classifier mapping von gene_dim to type_dim - outputs logits def __init__( self, gene_dim: int, type_dim: int, class_weights: np.ndarray, child_matrix: np.ndarray, batch_size: int, hvg: bool = False, num_hvgs: int = 2000, supervised_subset: Optional[int] = None, learning_rate: float = 0.005, weight_decay: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, gc_frequency: int = 1, ): super(BaseClassifier, self).__init__() self.gene_dim = gene_dim self.type_dim = type_dim self.batch_size = batch_size self.gc_freq = gc_frequency self.num_hvgs = num_hvgs self.supervised_subset = supervised_subset self.lr = learning_rate self.weight_decay = weight_decay self.optim = optimizer self.lr_scheduler = lr_scheduler self.lr_scheduler_kwargs = lr_scheduler_kwargs metrics = MetricCollection( { "f1_micro": MulticlassF1Score(num_classes=type_dim, average="micro"), "f1_macro": MulticlassF1Score(num_classes=type_dim, average="macro"), } ) self.train_metrics = metrics.clone(prefix="train_") self.val_metrics = metrics.clone(prefix="val_") self.test_metrics = metrics.clone(prefix="test_") self.register_buffer("class_weights", torch.tensor(class_weights.astype("f4"))) self.register_buffer("child_lookup", torch.tensor(child_matrix.astype("i8"))) if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) self.hvg_indices = pickle.load( open( root + "/self_supervision/data/hvg_" + str(self.num_hvgs) + "_indices.pickle", "rb", ) ) else: self.hvg_indices = None @abc.abstractmethod def _step(self, batch, training=True) -> Tuple[torch.Tensor, torch.Tensor]: """ Abstract method for performing a step in the training, validation, or testing process. Args: batch: Input batch data. training (bool, optional): Flag indicating whether the model is in training mode. Defaults to True. Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple containing the predictions and the loss. """ pass def hierarchy_correct(self, preds, targets) -> Tuple[torch.Tensor, torch.Tensor]: """ Corrects the predictions and targets based on the hierarchy of cell types. Args: preds (torch.Tensor): Predicted cell types. targets (torch.Tensor): Target cell types. Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple containing the corrected predictions and targets. """ pred_is_child_node_or_node = ( torch.sum( self.child_lookup[targets, :] * F.one_hot(preds, self.type_dim), dim=1 ) > 0 ) return ( torch.where(pred_is_child_node_or_node, targets, preds), # corrected preds torch.where( pred_is_child_node_or_node, preds, targets ), # corrected targets ) def on_after_batch_transfer(self, batch, dataloader_idx): """ Callback function called after transferring the batch to the device. Args: batch: Input batch data. dataloader_idx: Index of the dataloader. Returns: batch: Transformed batch data. """ with torch.no_grad(): batch = batch[0] batch["cell_type"] = torch.squeeze(batch["cell_type"]) return batch def forward(self, x: torch.Tensor): """ Forward pass of the model. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor. """ # If x is a tuple, which can happen in some testing due to lightning, then we need to extract the tensor if isinstance(x, dict): x = x["X"] if isinstance(x, tuple): x = x[0]["X"] if self.hvg_indices: x = x[:, self.hvg_indices] return self.classifier(x) def training_step(self, batch, batch_idx): """ Training step of the model. Args: batch: Input batch data. batch_idx: Index of the batch. Returns: loss: Training loss. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None, None # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} preds, loss = self._step(batch, training=True) self.log("train_loss", loss) f1_macro = self.train_metrics["f1_macro"](preds, batch["cell_type"]) f1_micro = self.train_metrics["f1_micro"](preds, batch["cell_type"]) self.log("train_f1_macro_step", f1_macro) self.log("train_f1_micro_step", f1_micro) if batch_idx % self.gc_freq == 0: gc.collect() return loss def validation_step(self, batch, batch_idx): """ Validation step of the model. Args: batch: Input batch data. batch_idx: Index of the batch. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None, None # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} preds, loss = self._step(batch, training=False) self.log("val_loss", loss) self.val_metrics["f1_macro"].update(preds, batch["cell_type"]) self.val_metrics["f1_micro"].update(preds, batch["cell_type"]) if batch_idx % self.gc_freq == 0: gc.collect() def test_step(self, batch, batch_idx): """ Test step of the model. Args: batch: Input batch data. batch_idx: Index of the batch. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None, None # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} preds, loss = self._step(batch, training=False) self.log("test_loss", loss) self.test_metrics["f1_macro"].update(preds, batch["cell_type"]) self.test_metrics["f1_micro"].update(preds, batch["cell_type"]) if batch_idx % self.gc_freq == 0: gc.collect() def on_train_epoch_end(self) -> None: self.log("train_f1_macro_epoch", self.train_metrics["f1_macro"].compute()) self.train_metrics["f1_macro"].reset() self.log("train_f1_micro_epoch", self.train_metrics["f1_micro"].compute()) self.train_metrics["f1_micro"].reset() gc.collect() def on_validation_epoch_end(self) -> None: f1_macro = self.val_metrics["f1_macro"].compute() self.log("val_f1_macro", f1_macro) self.log("hp_metric", f1_macro) self.val_metrics["f1_macro"].reset() self.log("val_f1_micro", self.val_metrics["f1_micro"].compute()) self.val_metrics["f1_micro"].reset() gc.collect() def on_test_epoch_end(self) -> None: self.log("test_f1_macro", self.test_metrics["f1_macro"].compute()) self.test_metrics["f1_macro"].reset() self.log("test_f1_micro", self.test_metrics["f1_micro"].compute()) self.test_metrics["f1_micro"].reset() gc.collect() def configure_optimizers(self): optimizer_config = { "optimizer": self.optim( self.parameters(), lr=self.lr, weight_decay=self.weight_decay ) } if self.lr_scheduler is not None: lr_scheduler_kwargs = ( {} if self.lr_scheduler_kwargs is None else self.lr_scheduler_kwargs ) interval = lr_scheduler_kwargs.pop("interval", "epoch") monitor = lr_scheduler_kwargs.pop("monitor", "val_loss_epoch") frequency = lr_scheduler_kwargs.pop("frequency", 1) scheduler = self.lr_scheduler( optimizer_config["optimizer"], **lr_scheduler_kwargs ) optimizer_config["lr_scheduler"] = { "scheduler": scheduler, "interval": interval, "monitor": monitor, "frequency": frequency, } return optimizer_config class MLPVariationalAutoEncoder(BaseAutoEncoder): def __init__( self, # fixed params gene_dim: int, units_encoder: List[int], units_decoder: List[int], # params from datamodule batch_size: int, hvg: bool = False, num_hvgs: int = 2000, # model specific params supervised_subset: Optional[int] = None, reconstruction_loss: str = "mse", learning_rate: float = 0.005, weight_decay: float = 0.1, dropout: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, output_activation: Callable[[], torch.nn.Module] = nn.Sigmoid, activation: Callable[[], torch.nn.Module] = nn.SELU, # params for masking masking_rate: Optional[float] = None, masking_strategy: Optional[str] = None, # 'random', 'gene_program' encoded_gene_program: Optional[Dict] = None, ): # check input assert 0.0 <= dropout <= 1.0 assert reconstruction_loss in ["mse", "mae", "continuous_bernoulli", "bce"] if reconstruction_loss in ["continuous_bernoulli", "bce"]: assert output_activation == nn.Sigmoid self.batch_size = batch_size self.supervised_subset = supervised_subset super(MLPVariationalAutoEncoder, self).__init__( gene_dim=gene_dim, batch_size=batch_size, learning_rate=learning_rate, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, supervised_subset=supervised_subset, ) self.encoder = MLP( in_channels=gene_dim, hidden_channels=units_encoder, activation_layer=activation, inplace=False, dropout=dropout, ) # Define decoder network self.decoder = nn.Sequential( MLP( in_channels=units_encoder[-1], hidden_channels=units_decoder + [gene_dim], # norm_layer=_get_norm_layer(batch_norm=batch_norm, layer_norm=layer_norm), activation_layer=activation, inplace=False, dropout=dropout, ), output_activation(), ) self.mean_head = torch.nn.Linear(units_encoder[-1], units_encoder[-1], bias=True) self.logvar_head = torch.nn.Linear(units_encoder[-1], units_encoder[-1], bias=True) self.predict_bottleneck = False metrics = MetricCollection( { "explained_var_weighted": ExplainedVariance( multioutput="variance_weighted" ), "explained_var_uniform": ExplainedVariance( multioutput="uniform_average" ), "mse": MeanSquaredError(), } ) self.train_metrics = metrics.clone(prefix="train_") self.val_metrics = metrics.clone(prefix="val_") self.test_metrics = metrics.clone(prefix="test_") # masking self.masking_rate = masking_rate self.masking_strategy = masking_strategy self.encoded_gene_program = encoded_gene_program self.num_hvgs = num_hvgs if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) self.hvg_indices = pickle.load( open( root + "/self_supervision/data/hvg_" + str(self.num_hvgs) + "_indices.pickle", "rb", ) ) else: self.hvg_indices = None def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) return mean + eps * std def _step(self, batch, training=True): targets = batch["X"] inputs = batch["X"] if self.hvg_indices is not None: inputs = inputs[:, self.hvg_indices] targets = targets[:, self.hvg_indices] if self.masking_rate and self.masking_strategy == "random": mask = ( Bernoulli(probs=1.0 - self.masking_rate) .sample(targets.size()) .to(targets.device) ) # upscale inputs to compensate for masking and convert to same device masked_inputs = 1.0 / (1.0 - self.masking_rate) * (inputs * mask) x_latent = self.encoder(masked_inputs) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) # calculate masked loss on masked part only inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) rec_loss = ( inv_mask * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() kl_div = -0.5 * torch.mean(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) loss = rec_loss + kl_div elif self.masking_rate and self.masking_strategy == "gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs mask, frac = _mask_gene_programs_numpy( inputs, self.encoded_gene_program, self.masking_rate ) mask = torch.tensor(mask).to(inputs.device) # log the fraction of genes masked self.log("frac_genes_masked", frac) # mask, frac = self.mask_gene_programs(inputs, self.gene_program_dict, self.masking_rate) # upscale inputs to compensate for masking masked_inputs = 1.0 / (1.0 - frac) * (inputs * mask.to(inputs.device)) x_latent = self.encoder(masked_inputs) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) # calculate masked loss inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) rec_loss = ( inv_mask.to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() kl_div = -0.5 * torch.mean(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) loss = rec_loss + kl_div elif self.masking_rate and self.masking_strategy == "single_gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs input_mask, output_mask = _mask_single_gene_programs( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) x_latent = self.encoder(masked_inputs) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) rec_loss = ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() kl_div = -0.5 * torch.mean(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) loss = rec_loss + kl_div elif self.masking_rate and self.masking_strategy == "gp_to_tf": with torch.no_grad(): # self.encoded gene program is a Dict of encoded gene programs and the corresponding tf indices input_mask, output_mask = _only_activate_gene_program_numpy( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) x_latent = self.encoder(masked_inputs) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) rec_loss = ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() kl_div = -0.5 * torch.mean(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) loss = rec_loss + kl_div # raise error if masking rate is not none but masking strategy is not implemented elif self.masking_rate and self.masking_strategy not in [ "random", "gene_program", ]: raise ValueError( f"Masking strategy {self.masking_strategy} not implemented." ) # raise error if masking strategy is not none but masking rate is 0.0 elif self.masking_strategy and self.masking_rate == 0.0: raise ValueError( f"Masking rate is 0 ({self.masking_rate})," f" but masking strategy {self.masking_strategy} is not None." ) else: x_latent = self.encoder(inputs.to(targets.device)) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) rec_loss = self._calc_reconstruction_loss(x_reconst, targets, reduction="mean") kl_div = -0.5 * torch.mean(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) loss = rec_loss + kl_div return x_reconst, loss def predict_embedding(self, batch): if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return self.encoder(batch["X"]) def forward(self, x_in): x_latent = self.encoder(x_in) z_mean = self.mean_head(x_latent) z_logvar = self.logvar_head(x_latent) z = self.reparameterize(z_mean, z_logvar) x_reconst = self.decoder(z) return x_latent, x_reconst def training_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] self.log_dict( self.train_metrics(x_reconst, batch["X"]), on_epoch=True, on_step=True ) self.log("train_loss", loss, on_epoch=True, on_step=True) if batch_idx % self.gc_freq == 0: gc.collect() return loss def validation_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch, training=False) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] self.log_dict(self.val_metrics(x_reconst, batch["X"])) self.log("val_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def test_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch, training=False) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] metrics = self.test_metrics(x_reconst, batch["X"]) self.log_dict(metrics) self.log("test_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() return metrics def predict_cell_types(self, x: torch.Tensor): return F.softmax(self(x)[0], dim=1) def predict_step( self, batch, batch_idx, dataloader_idx=None, predict_embedding=False ): if batch_idx % self.gc_freq == 0: gc.collect() # Apply dataset_id filtering only if supervised_subset is set if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} if predict_embedding: return self.encoder(batch["X"]).detach() else: preds_corrected, loss = self._step(batch, training=False) return preds_corrected[mask], batch["cell_type"][mask] else: if predict_embedding: return self.encoder(batch["X"]).detach() else: x_reconst, loss = self._step(batch, training=False) return x_reconst, batch["X"] def get_input(self, batch): if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return batch["X"] class MLPAutoEncoder(BaseAutoEncoder): def __init__( self, # fixed params gene_dim: int, units_encoder: List[int], units_decoder: List[int], # params from datamodule batch_size: int, hvg: bool = False, num_hvgs: int = 2000, # model specific params supervised_subset: Optional[int] = None, reconstruction_loss: str = "mse", learning_rate: float = 0.005, weight_decay: float = 0.1, dropout: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, output_activation: Callable[[], torch.nn.Module] = nn.Sigmoid, activation: Callable[[], torch.nn.Module] = nn.SELU, # params for masking masking_rate: Optional[float] = None, masking_strategy: Optional[str] = None, # 'random', 'gene_program' encoded_gene_program: Optional[Dict] = None, ): # check input assert 0.0 <= dropout <= 1.0 assert reconstruction_loss in ["mse", "mae", "continuous_bernoulli", "bce"] if reconstruction_loss in ["continuous_bernoulli", "bce"]: assert output_activation == nn.Sigmoid self.batch_size = batch_size self.supervised_subset = supervised_subset super(MLPAutoEncoder, self).__init__( gene_dim=gene_dim, batch_size=batch_size, learning_rate=learning_rate, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, supervised_subset=supervised_subset, ) self.encoder = MLP( in_channels=gene_dim, hidden_channels=units_encoder, activation_layer=activation, inplace=False, dropout=dropout, ) # Define decoder network self.decoder = nn.Sequential( MLP( in_channels=units_encoder[-1], hidden_channels=units_decoder + [gene_dim], # norm_layer=_get_norm_layer(batch_norm=batch_norm, layer_norm=layer_norm), activation_layer=activation, inplace=False, dropout=dropout, ), output_activation(), ) self.predict_bottleneck = False metrics = MetricCollection( { "explained_var_weighted": ExplainedVariance( multioutput="variance_weighted" ), "explained_var_uniform": ExplainedVariance( multioutput="uniform_average" ), "mse": MeanSquaredError(), } ) self.train_metrics = metrics.clone(prefix="train_") self.val_metrics = metrics.clone(prefix="val_") self.test_metrics = metrics.clone(prefix="test_") # masking self.masking_rate = masking_rate self.masking_strategy = masking_strategy self.encoded_gene_program = encoded_gene_program self.num_hvgs = num_hvgs if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) self.hvg_indices = pickle.load( open( root + "/self_supervision/data/hvg_" + str(self.num_hvgs) + "_indices.pickle", "rb", ) ) else: self.hvg_indices = None def _step(self, batch, training=True): targets = batch["X"] inputs = batch["X"] if self.hvg_indices is not None: inputs = inputs[:, self.hvg_indices] targets = targets[:, self.hvg_indices] if self.masking_rate and self.masking_strategy == "random": mask = ( Bernoulli(probs=1.0 - self.masking_rate) .sample(targets.size()) .to(targets.device) ) # upscale inputs to compensate for masking and convert to same device masked_inputs = 1.0 / (1.0 - self.masking_rate) * (inputs * mask) x_latent, x_reconst = self(masked_inputs) # calculate masked loss on masked part only inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) loss = ( inv_mask * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() elif self.masking_rate and self.masking_strategy == "gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs mask, frac = _mask_gene_programs_numpy( inputs, self.encoded_gene_program, self.masking_rate ) mask = torch.tensor(mask).to(inputs.device) # log the fraction of genes masked self.log("frac_genes_masked", frac) # mask, frac = self.mask_gene_programs(inputs, self.gene_program_dict, self.masking_rate) # upscale inputs to compensate for masking masked_inputs = 1.0 / (1.0 - frac) * (inputs * mask.to(inputs.device)) x_latent, x_reconst = self(masked_inputs) # calculate masked loss inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) loss = ( inv_mask.to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() elif self.masking_rate and self.masking_strategy == "single_gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs input_mask, output_mask = _mask_single_gene_programs( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) x_latent, x_reconst = self(masked_inputs) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) loss = ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() elif self.masking_rate and self.masking_strategy == "gp_to_tf": with torch.no_grad(): # self.encoded gene program is a Dict of encoded gene programs and the corresponding tf indices input_mask, output_mask = _only_activate_gene_program_numpy( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) x_latent, x_reconst = self(masked_inputs) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) loss = ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(x_reconst, targets, reduction="none") ).mean() # raise error if masking rate is not none but masking strategy is not implemented elif self.masking_rate and self.masking_strategy not in [ "random", "gene_program", ]: raise ValueError( f"Masking strategy {self.masking_strategy} not implemented." ) # raise error if masking strategy is not none but masking rate is 0.0 elif self.masking_strategy and self.masking_rate == 0.0: raise ValueError( f"Masking rate is 0 ({self.masking_rate})," f" but masking strategy {self.masking_strategy} is not None." ) else: x_latent, x_reconst = self(inputs.to(targets.device)) loss = self._calc_reconstruction_loss(x_reconst, targets, reduction="mean") return x_reconst, loss def predict_embedding(self, batch): if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return self.encoder(batch["X"]) def forward(self, x_in): x_latent = self.encoder(x_in) x_reconst = self.decoder(x_latent) return x_latent, x_reconst def training_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] self.log_dict( self.train_metrics(x_reconst, batch["X"]), on_epoch=True, on_step=True ) self.log("train_loss", loss, on_epoch=True, on_step=True) if batch_idx % self.gc_freq == 0: gc.collect() return loss def validation_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch, training=False) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] self.log_dict(self.val_metrics(x_reconst, batch["X"])) self.log("val_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def test_step(self, batch, batch_idx): if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} x_reconst, loss = self._step(batch, training=False) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] metrics = self.test_metrics(x_reconst, batch["X"]) self.log_dict(metrics) self.log("test_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() return metrics def predict_cell_types(self, x: torch.Tensor): return F.softmax(self(x)[0], dim=1) def predict_step( self, batch, batch_idx, dataloader_idx=None, predict_embedding=False ): if batch_idx % self.gc_freq == 0: gc.collect() # Apply dataset_id filtering only if supervised_subset is set if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None # Skip the batch if no items match the condition # Filter the batch based on the mask batch = {key: value[mask] for key, value in batch.items()} if predict_embedding: return self.encoder(batch["X"]).detach() else: preds_corrected, loss = self._step(batch, training=False) return preds_corrected[mask], batch["cell_type"][mask] else: if predict_embedding: return self.encoder(batch["X"]).detach() else: x_reconst, loss = self._step(batch, training=False) return x_reconst, batch["X"] def get_input(self, batch): if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return batch["X"] class MLPNegBin(BaseAutoEncoder): """ Multi-Layer Perceptron (MLP) Autoencoder with Negative Binomial (NegBin) loss. Args: gene_dim (int): Dimensionality of the gene input. units_encoder (List[int]): List of integers specifying the number of hidden units in each encoder layer. units_decoder (List[int]): List of integers specifying the number of hidden units in each decoder layer. batch_size (int): Batch size. hvg (bool): Flag indicating whether to use highly variable genes (HVGs) for training. num_hvgs (int): Number of highly variable genes to use. supervised_subset (Optional[int]): Subset of the data to use for supervised training. Default is None. learning_rate (float): Learning rate for the optimizer. Default is 0.005. weight_decay (float): Weight decay for the optimizer. Default is 0.1. dropout (float): Dropout rate. Default is 0.1. optimizer (Callable[..., torch.optim.Optimizer]): Optimizer class. Default is torch.optim.AdamW. lr_scheduler (Callable): Learning rate scheduler. Default is None. lr_scheduler_kwargs (Dict): Additional keyword arguments for the learning rate scheduler. Default is None. activation (Callable[[], torch.nn.Module]): Activation function. Default is nn.SELU. masking_rate (Optional[float]): Rate of masking for input data. Default is None. masking_strategy (Optional[str]): Masking strategy for input data. Default is None. encoded_gene_program (Optional[Dict]): Encoded gene program. Default is None. Raises: NotImplementedError: If masking is not implemented for MLPNegBin. Attributes: encoder (MLP): Encoder module. decoder (nn.Sequential): Decoder module. masking_rate (Optional[float]): Rate of masking for input data. masking_strategy (Optional[str]): Masking strategy for input data. encoded_gene_program (Optional[Dict]): Encoded gene program. num_hvgs (int): Number of highly variable genes. hvg_indices (Optional[List[int]]): Indices of highly variable genes. """ def __init__( self, gene_dim: int, units_encoder: List[int], units_decoder: List[int], batch_size: int, hvg: bool = False, num_hvgs: int = 2000, supervised_subset: Optional[int] = None, gene_likelihood: str = "nb", # 'zinb', 'nb', 'poisson' learning_rate: float = 0.005, weight_decay: float = 0.1, dropout: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, activation: Callable[[], torch.nn.Module] = nn.SELU, # params for masking masking_rate: Optional[float] = None, masking_strategy: Optional[str] = None, # 'random', 'gene_program' encoded_gene_program: Optional[Dict] = None, ): super(MLPNegBin, self).__init__( gene_dim=gene_dim, batch_size=batch_size, learning_rate=learning_rate, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, supervised_subset=supervised_subset, ) self.masking_rate = masking_rate self.masking_strategy = masking_strategy self.encoded_gene_program = encoded_gene_program self.num_hvgs = num_hvgs self.hvg_indices = self.load_hvg_indices(hvg, num_hvgs) self.gene_likelihood = gene_likelihood self.gene_dim = gene_dim self.encoder = MLP( in_channels=gene_dim, hidden_channels=units_encoder, activation_layer=activation, inplace=False, dropout=dropout, ) # Decoder for gene counts output_dim = self._get_decoder_output_dim() self.decoder = nn.Sequential( MLP( in_channels=units_encoder[-1], hidden_channels=units_decoder + [output_dim], activation_layer=activation, final_activation=torch.nn.Softplus if self.gene_likelihood in ["nb", "zinb"] else None, inplace=False, dropout=dropout, ) ) # Dispersion parameters if gene_likelihood in ["nb", "zinb"]: self.px_r = nn.Parameter(torch.randn(self.gene_dim)) metrics = MetricCollection( { "explained_var_weighted": ExplainedVariance( multioutput="variance_weighted" ), "explained_var_uniform": ExplainedVariance( multioutput="uniform_average" ), "mse": MeanSquaredError(), } ) self.train_metrics = metrics.clone(prefix="train_") self.val_metrics = metrics.clone(prefix="val_") self.test_metrics = metrics.clone(prefix="test_") self.masking_rate = masking_rate self.masking_strategy = masking_strategy self.encoded_gene_program = encoded_gene_program def _get_decoder_output_dim(self): if self.gene_likelihood == "zinb": # Mean, dispersion, dropout logits return 3 * self.gene_dim elif self.gene_likelihood == "nb": # Mean, dispersion return 2 * self.gene_dim elif self.gene_likelihood == "poisson": # Mean only return self.gene_dim else: raise ValueError(f"Invalid gene_likelihood: {self.gene_likelihood}") def load_hvg_indices(self, hvg: bool, num_hvgs: int): """ Load the indices of highly variable genes. Args: hvg (bool): Flag indicating whether to use highly variable genes. num_hvgs (int): Number of highly variable genes. Returns: Optional[List[int]]: List of indices of highly variable genes, or None if hvg is False. """ if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) return pickle.load( open( root + "/self_supervision/data/hvg_" + str(num_hvgs) + "_indices.pickle", "rb", ) ) return None def forward(self, x_in): """ Forward pass of the MLPNegBin model. Args: x_in (torch.Tensor): Input tensor. Returns: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Tuple containing the latent representation, mean, and inverse dispersion. """ x_latent = self.encoder(x_in) output = self.decoder(x_latent) if self.gene_likelihood in ["nb", "zinb"]: theta = self.px_r.exp() # Extract parameters for gene expression likelihood distribution if self.gene_likelihood == "zinb": mu = output[:, : self.gene_dim] pi = output[:, 2 * self.gene_dim :] return x_latent, mu, theta, pi elif self.gene_likelihood == "nb": mu = output[:, : self.gene_dim] return x_latent, mu, theta elif self.gene_likelihood == "poisson": mu = output return x_latent, mu def _calc_reconstruction_loss(self, x, output, reduction="mean"): """ Calculates the reconstruction loss for the autoencoder. Args: x (torch.Tensor): The input data. output (tuple): The output of the autoencoder. reduction (str, optional): Reduction method. Defaults to 'mean'. Returns: torch.Tensor: The reconstruction loss. """ if self.gene_likelihood == "zinb": mu, theta, pi = output px = ZeroInflatedNegativeBinomial(mu=mu, theta=theta, zi_logits=pi) elif self.gene_likelihood == "nb": mu, theta = output px = NegativeBinomial(mu=mu, theta=theta) elif self.gene_likelihood == "poisson": mu = output px = Poisson(mu) if reduction == "mean": return -px.log_prob(x).sum(-1).mean() elif reduction == "none": return -px.log_prob(x) def _calc_reconstruction(self, output): """ Calculates the reconstructed transcriptome from the output of the autoencoder. Args: output (tuple): The output of the autoencoder. Returns: torch.Tensor: The reconstructed transcriptome. """ if self.gene_likelihood == "zinb": mu, theta, pi = output px = ZeroInflatedNegativeBinomial(mu=mu, theta=theta, zi_logits=pi) elif self.gene_likelihood == "nb": mu, theta = output px = NegativeBinomial(mu=mu, theta=theta) elif self.gene_likelihood == "poisson": mu = output px = Poisson(mu) return px.mean def _step(self, batch, training=True): """ Perform a single training/validation step. Args: batch (Dict[str, torch.Tensor]): Batch of input data. training (bool): Flag indicating whether the model is in training mode. Default is True. Returns: Tuple[torch.Tensor, torch.Tensor]: Tuple containing the latent representation and loss. """ targets = batch["X"] inputs = batch["X"] if self.hvg_indices is not None: inputs = inputs[:, self.hvg_indices] targets = targets[:, self.hvg_indices] if self.masking_rate and self.masking_strategy == "random": mask = ( Bernoulli(probs=1.0 - self.masking_rate) .sample(targets.size()) .to(targets.device) ) # upscale inputs to compensate for masking and convert to same device masked_inputs = 1.0 / (1.0 - self.masking_rate) * (inputs * mask) # calculate masked loss on masked part only inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) if self.gene_likelihood == "zinb": x_latent, mu, theta, pi = self(masked_inputs) loss = ( ( inv_mask * self._calc_reconstruction_loss( targets, (mu, theta, pi), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta, pi)) elif self.gene_likelihood == "nb": x_latent, mu, theta = self(masked_inputs) loss = ( ( inv_mask * self._calc_reconstruction_loss( targets, (mu, theta), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta)) elif self.gene_likelihood == "poisson": x_latent, mu = self(masked_inputs) loss = ( ( inv_mask * self._calc_reconstruction_loss(targets, mu, reduction="none") ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction(mu) elif self.masking_rate and self.masking_strategy == "gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs mask, frac = _mask_gene_programs_numpy( inputs, self.encoded_gene_program, self.masking_rate ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(mask).to(inputs.device)) ) # calculate masked loss inv_mask = torch.abs(torch.ones(mask.size()).to(targets.device) - mask) if self.gene_likelihood == "zinb": x_latent, mu, theta, pi = self(masked_inputs) loss = ( ( inv_mask.to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta, pi), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta, pi)) elif self.gene_likelihood == "nb": x_latent, mu, theta = self(masked_inputs) loss = ( ( inv_mask.to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta)) elif self.gene_likelihood == "poisson": x_latent, mu = self(masked_inputs) loss = ( ( inv_mask.to(inputs.device) * self._calc_reconstruction_loss(targets, mu, reduction="none") ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction(mu) elif self.masking_rate and self.masking_strategy == "single_gene_program": with torch.no_grad(): # self.encoded gene program is a numpy array of encoded gene programs input_mask, output_mask = _mask_single_gene_programs( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) if self.gene_likelihood == "zinb": x_latent, mu, theta, pi = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta, pi), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta, pi)) elif self.gene_likelihood == "nb": x_latent, mu, theta = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta)) elif self.gene_likelihood == "poisson": x_latent, mu = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(targets, mu, reduction="none") ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction(mu) elif self.masking_rate and self.masking_strategy == "gp_to_tf": with torch.no_grad(): # self.encoded gene program is a Dict of encoded gene programs and the corresponding tf indices input_mask, output_mask = _only_activate_gene_program_numpy( inputs, self.encoded_gene_program ) # log the fraction of genes masked self.log("frac_genes_masked", frac) # upscale inputs to compensate for masking masked_inputs = ( 1.0 / (1.0 - frac) * (inputs * torch.tensor(input_mask).to(inputs.device)) ) # calculate masked loss only on the output_mask part of the reconstruction inv_output_mask = torch.abs( torch.ones(output_mask.size()).to(targets.device) - output_mask ) if self.gene_likelihood == "zinb": x_latent, mu, theta, pi = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta, pi), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta, pi)) elif self.gene_likelihood == "nb": x_latent, mu, theta = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss( targets, (mu, theta), reduction="none" ) ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction((mu, theta)) elif self.gene_likelihood == "poisson": x_latent, mu = self(masked_inputs) loss = ( ( torch.tensor(inv_output_mask).to(inputs.device) * self._calc_reconstruction_loss(targets, mu, reduction="none") ) .sum(-1) .mean() ) x_reconst = self._calc_reconstruction(mu) # raise error if masking rate is not none but masking strategy is not implemented elif self.masking_rate and self.masking_strategy not in [ "random", "gene_program", ]: raise ValueError( f"Masking strategy {self.masking_strategy} not implemented." ) # raise error if masking strategy is not none but masking rate is 0.0 elif self.masking_strategy and self.masking_rate == 0.0: raise ValueError( f"Masking rate is 0 ({self.masking_rate})," f" but masking strategy {self.masking_strategy} is not None." ) else: # x_latent, mu, theta = self(inputs.to(targets.device)) out = self(inputs.to(targets.device)) if self.gene_likelihood == "zinb": x_latent, mu, theta, pi = out loss = self._calc_reconstruction_loss(targets, (mu, theta, pi)) x_reconst = self._calc_reconstruction((mu, theta, pi)) elif self.gene_likelihood == "nb": x_latent, mu, theta = out loss = self._calc_reconstruction_loss(targets, (mu, theta)) x_reconst = self._calc_reconstruction((mu, theta)) elif self.gene_likelihood == "poisson": x_latent, mu = out loss = self._calc_reconstruction_loss(targets, mu) x_reconst = self._calc_reconstruction(mu) else: raise ValueError(f"Invalid gene_likelihood: {self.gene_likelihood}") return x_latent, loss, x_reconst def predict_embedding(self, batch): """ Predict the embedding for a given batch. Args: batch (Dict[str, torch.Tensor]): Batch of input data. Returns: torch.Tensor: Embedding tensor. """ if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return self.encoder(batch["X"]) def training_step(self, batch, batch_idx): """ Training step. Args: batch (Dict[str, torch.Tensor]): Batch of input data. batch_idx (int): Index of the current batch. Returns: torch.Tensor: Loss value. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return batch = {key: value[mask] for key, value in batch.items()} x_latent, loss, _ = self._step(batch) self.log("train_loss", loss, on_epoch=True, on_step=True) if batch_idx % self.gc_freq == 0: gc.collect() return loss def validation_step(self, batch, batch_idx): """ Validation step. Args: batch (Dict[str, torch.Tensor]): Batch of input data. batch_idx (int): Index of the current batch. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return batch = {key: value[mask] for key, value in batch.items()} x_latent, loss, _ = self._step(batch, training=False) self.log("val_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def test_step(self, batch, batch_idx): """ Test step. Args: batch (Dict[str, torch.Tensor]): Batch of input data. batch_idx (int): Index of the current batch. """ if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset batch = {key: value[mask] for key, value in batch.items()} x_latent, loss, x_reconst = self._step(batch, training=False) if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] metrics = self.test_metrics(x_reconst, batch["X"]) self.log_dict(metrics) self.log("test_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() return metrics class MLPClassifier(BaseClassifier): """ Multi-Layer Perceptron (MLP) classifier for cell type classification. Args: gene_dim (int): Dimension of the gene input. type_dim (int): Dimension of the cell type output. class_weights (np.ndarray): Array of class weights for loss calculation. child_matrix (np.ndarray): Matrix representing the hierarchical relationship between cell types. units (List[int]): List of hidden unit sizes for the MLP. batch_size (int): Batch size for training and inference. hvg (bool): Flag indicating whether to use highly variable genes. num_hvgs (int): Number of highly variable genes to use. supervised_subset (Optional[int], optional): Subset of the data to use for supervised training. Defaults to None. dropout (float, optional): Dropout rate. Defaults to 0.0. learning_rate (float, optional): Learning rate for optimizer. Defaults to 0.005. weight_decay (float, optional): Weight decay for optimizer. Defaults to 0.1. optimizer (Callable[..., torch.optim.Optimizer], optional): Optimizer class. Defaults to torch.optim.AdamW. lr_scheduler (Callable, optional): Learning rate scheduler. Defaults to None. lr_scheduler_kwargs (Dict, optional): Additional keyword arguments for the learning rate scheduler. Defaults to None. """ def __init__( self, gene_dim: int, type_dim: int, class_weights: np.ndarray, child_matrix: np.ndarray, units: List[int], batch_size: int, hvg: bool = False, num_hvgs: int = 2000, supervised_subset: Optional[int] = None, dropout: float = 0.0, learning_rate: float = 0.005, weight_decay: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, ): super(MLPClassifier, self).__init__( gene_dim=gene_dim, type_dim=type_dim, class_weights=class_weights, child_matrix=child_matrix, batch_size=batch_size, hvg=hvg, num_hvgs=num_hvgs, supervised_subset=supervised_subset, learning_rate=learning_rate, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, ) self.classifier = MLP( in_channels=gene_dim, hidden_channels=units + [type_dim], dropout=dropout, inplace=False, ) def _step(self, batch, training=True): """ Perform a forward pass and calculate the loss for a batch. Args: batch (dict): Dictionary containing the input data batch. training (bool, optional): Flag indicating whether the model is in training mode. Defaults to True. Returns: torch.Tensor: Corrected predictions for the batch. torch.Tensor: Loss value for the batch. """ logits = self(batch["X"]) with torch.no_grad(): preds = torch.argmax(logits, dim=1) preds_corrected, targets_corrected = self.hierarchy_correct( preds, batch["cell_type"] ) if training: loss = F.cross_entropy( logits, batch["cell_type"], weight=self.class_weights ) else: loss = F.cross_entropy(logits, targets_corrected) return preds_corrected, loss def predict_step( self, batch, batch_idx, dataloader_idx=None, predict_embedding=False ): """ Perform a prediction step for a batch. Args: batch (tuple or dict): Input data batch. batch_idx (int): Index of the batch. dataloader_idx (int, optional): Index of the dataloader. Defaults to None. predict_embedding (bool, optional): Flag indicating whether to predict the embedding. Defaults to False. Returns: torch.Tensor or None: Predicted values for the batch. torch.Tensor or None: Target values for the batch. """ if isinstance(batch, tuple): batch = batch[0] if batch_idx % self.gc_freq == 0: gc.collect() # Apply dataset_id filtering only if supervised_subset is set if self.supervised_subset is not None: mask = batch["dataset_id"] == self.supervised_subset if not any(mask): return None, None # Skip the batch if no items match the condition filtered_batch = {key: value[mask] for key, value in batch.items()} if predict_embedding: return self.encoder(filtered_batch["X"]).detach(), filtered_batch[ "cell_type" ] else: preds_corrected, loss = self._step(filtered_batch, training=False) return preds_corrected, filtered_batch["cell_type"] else: if predict_embedding: return self.encoder(batch["X"]).detach(), filtered_batch["cell_type"] else: preds_corrected, loss = self._step(batch, training=False) return preds_corrected, batch["cell_type"] def forward_embedding(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass to compute the embedding from the MLP. The computed embedding is the output before the last Linear layer of the MLP. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Computed embedding tensor. """ for i, layer in enumerate(self.classifier): x = layer(x) if i == len(self.classifier) - 2: break return x def predict_embedding(self, batch, batch_idx, dataloader_idx=None) -> torch.Tensor: """ Function to predict the embedding from the given batch. Args: batch (dict): Input data batch. batch_idx (int): Index of the batch. dataloader_idx (int, optional): Index of the dataloader. Defaults to None. Returns: torch.Tensor: Predicted embedding tensor. """ x = batch["X"] return self.forward_embedding(x) class MLPBYOL(BaseAutoEncoder): def __init__( self, # fixed params gene_dim: int, units_encoder: List[int], # params from datamodule batch_size: int, # contrastive learning params backbone: str, # MLP, TabNet augment_type: str, # Gaussian, Uniform augment_intensity: float, use_momentum: bool, # model specific params hvg: bool = False, num_hvgs: int = 2000, lr: float = 0.005, weight_decay: float = 0.1, dropout: float = 0.1, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, lr_scheduler_kwargs: Dict = None, activation: Callable[[], torch.nn.Module] = nn.SELU, ): # check input assert 0.0 <= dropout <= 1.0 self.batch_size = batch_size super(MLPBYOL, self).__init__( gene_dim=gene_dim, batch_size=batch_size, learning_rate=lr, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, ) # assign inner model, that will be trained using the BYOL / SimSiam framework self.backbone = backbone self.gene_dim = gene_dim self.units_encoder = units_encoder self.activation = activation self.dropout = dropout self.inner_model = self._get_inner_model() self.num_hvgs = num_hvgs self.byol = BYOL( net=self.inner_model, image_size=self.gene_dim, augment_type=augment_type, augment_intensity=augment_intensity, batch_size=self.batch_size, use_momentum=use_momentum, ) if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) self.hvg_indices = pickle.load( open( root + "/self_supervision/data/hvg_" + str(self.num_hvgs) + "_indices.pickle", "rb", ) ) else: self.hvg_indices = None def _get_inner_model(self): if self.backbone == "MLP": self.inner_model = MLP( in_channels=self.gene_dim, hidden_channels=self.units_encoder, activation_layer=self.activation, inplace=False, dropout=self.dropout, ) elif self.backbone == "TabNet": raise NotImplementedError else: raise NotImplementedError return self.inner_model def _step(self, batch): loss = self.forward(batch) return loss def predict_embedding(self, batch, batch_idx): if batch_idx % self.gc_freq == 0: gc.collect() return self.encoder(batch["X"]) def forward(self, batch): if self.hvg_indices is not None: batch["X"] = batch["X"][:, self.hvg_indices] return self.byol(batch["X"]) def training_step(self, batch, batch_idx): loss = self._step(batch) self.log("train_loss", loss, on_epoch=True) if batch_idx % self.gc_freq == 0: gc.collect() return loss def validation_step(self, batch, batch_idx): loss = self._step(batch) self.log("val_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def test_step(self, batch, batch_idx): loss = self._step(batch) self.log("test_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def predict_step(self, batch, batch_idx, dataloader_idx: Optional[int] = None): x_reconst, loss = self(batch) if batch_idx % self.gc_freq == 0: gc.collect() return x_reconst, batch # helper class DictAsAttributes: def __init__(self, dictionary): self.__dict__.update(dictionary) class MLPBarlowTwins(BaseAutoEncoder): def __init__( self, # fixed params gene_dim: int, units_encoder: List[int], # params from datamodule CHECKPOINT_PATH: str, augment_intensity: float, train_set_size: int, batch_size: int, hvg: bool = False, num_hvgs: int = 2000, # contrastive learning params backbone: str = 'MLP', # MLP, TabNet learning_rate_weights: float = 0.2, learning_rate_biases: float = 0.0048, lambd: float = 0.0051, lr: float = 0.005, # dummy here weight_decay: float = 1e-6, dropout: float = 0.0, optimizer: Callable[..., torch.optim.Optimizer] = torch.optim.AdamW, lr_scheduler: Callable = None, # dummy here lr_scheduler_kwargs: Dict = None, # dummy here activation: Callable[[], torch.nn.Module] = nn.SELU, ): # check input assert 0.0 <= dropout <= 1.0 super(MLPBarlowTwins, self).__init__( gene_dim=gene_dim, batch_size=batch_size, learning_rate=lr, weight_decay=weight_decay, optimizer=optimizer, lr_scheduler=lr_scheduler, lr_scheduler_kwargs=lr_scheduler_kwargs, automatic_optimization=False, ) self.save_hyperparameters(ignore=["gene_dim"]) self.best_val_loss = np.inf self.best_train_loss = np.inf self.num_hvgs = num_hvgs self.train_set_size = train_set_size self.batch_size = batch_size self.weight_decay = weight_decay self.learning_rate_weights = learning_rate_weights self.learning_rate_biases = learning_rate_biases self.lambd = lambd self.loader_length = self.train_set_size // self.batch_size self.step = 0 # for optimizer # assign inner model, that will be trained using the BYOL / SimSiam framework self.backbone = backbone self.gene_dim = gene_dim self.units_encoder = units_encoder self.activation = activation self.dropout = dropout self.batch_size = batch_size self.hvg = hvg self.transform = Transform(p=augment_intensity) self.CHECKPOINT_PATH = CHECKPOINT_PATH if hvg: root = os.path.dirname( os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) ) self.hvg_indices = pickle.load( open( root + "/self_supervision/data/hvg_" + str(self.num_hvgs) + "_indices.pickle", "rb", ) ) else: self.hvg_indices = None inner_model = self._get_inner_model() # Initialize the model-specific parameters self.dropout = dropout self.activation = activation self.units_encoder = units_encoder self.barlow_twins = self._get_barlow_twins_model(inner_model) def _get_barlow_twins_model(self, inner_model): args = ( self._prepare_barlow_twins_args() ) # Prepare the Barlow Twins arguments (use the relevant params) self.barlow_twins = BarlowTwins(backbone=inner_model, args=args) return self.barlow_twins def _get_inner_model(self): if self.backbone == "MLP": self.inner_model = MLP( in_channels=self.gene_dim, hidden_channels=self.units_encoder, activation_layer=self.activation, inplace=False, dropout=self.dropout, ) else: raise NotImplementedError return self.inner_model def _prepare_barlow_twins_args(self): # Prepare Barlow Twins arguments based on the required parameters args = { # 'data': ..., # Path to your dataset directory (from class B) # 'workers': 8, # number of data loader workers "epochs": 1000, # number of total epochs to run "batch_size": self.batch_size, # mini-batch size "learning_rate_weights": self.learning_rate_weights, # base learning rate for weights "learning_rate_biases": self.learning_rate_biases, # base learning rate for biases and batch norm parameters "weight-decay": self.weight_decay, # weight decay "lambd": self.lambd, # weight on off-diagonal terms "projector": "256-256-512-512", # projector MLP # 'print_freq': 100, # print frequency # 'checkpoint_dir': './checkpoint/' # path to checkpoint directory } self.args = DictAsAttributes(args) return self.args def _step(self, batch): # For Multiomics, batch['X'] has additional dim to squeeze if batch["X"].dim() == 3: batch["X"] = batch["X"].squeeze(1) if self.hvg_indices is not None: x_in = batch["X"][:, self.hvg_indices] else: x_in = batch["X"] # Modify the forward pass to use the Barlow Twins model (from class B) loss = self(x_in) return loss def predict_embedding(self, batch, batch_idx): if batch_idx % self.gc_freq == 0: gc.collect() return self.encoder(batch["X"]) def forward(self, batch): # Apply two different augmentations to the same image y1, y2 = self.transform(batch) # Compute the Barlow Twins loss loss = self.barlow_twins(y1, y2) return loss def training_step(self, batch, batch_idx): # Increment the step counter self.step += 1 # Initialize the optimizer param_weights = [] param_biases = [] for param in self.barlow_twins.parameters(): if param.ndim == 1: param_biases.append(param) else: param_weights.append(param) parameters = [{"params": param_weights}, {"params": param_biases}] opt = LARS( parameters, lr=0, weight_decay=self.weight_decay, weight_decay_filter=True, lars_adaptation_filter=True, ) # Adjust the learning rate adjust_learning_rate(self.args, opt, self.loader_length, self.step) # Zero the gradient before loss computation opt.zero_grad() # Calculate the loss loss = self._step(batch) # Manual backward pass and optimizer step self.manual_backward(loss) opt.step() self.log("train_loss", loss, on_epoch=True) if batch_idx % self.gc_freq == 0: gc.collect() if loss < self.best_train_loss: self.best_train_loss = loss self.barlow_twins.save_model( path=self.CHECKPOINT_PATH + "/best_checkpoint_train.ckpt" ) self.barlow_twins.save_model( path=self.CHECKPOINT_PATH + "/last_checkpoint.ckpt" ) return loss def validation_step(self, batch, batch_idx): loss = self._step(batch) self.log("val_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() if loss < self.best_val_loss: self.best_val_loss = loss self.barlow_twins.save_model( path=self.CHECKPOINT_PATH + "/best_checkpoint_val.ckpt" ) def test_step(self, batch, batch_idx): loss = self._step(batch) self.log("test_loss", loss) if batch_idx % self.gc_freq == 0: gc.collect() def predict_step(self, batch, batch_idx, dataloader_idx: Optional[int] = None): x_reconst, loss = self(batch) if batch_idx % self.gc_freq == 0: gc.collect() return x_reconst, batch