Documentation for the model

def configure_optimizers(self):
    """@see pl.LightningModule"""
    # https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam
    # not working because of poor weight decay implem
    if self.optim == "adam":
        optimizer = optim.Adam(
            self.parameters(),
            lr=self.hparams.lr,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=self.weight_decay,
            amsgrad=False,
            fused=self.fused_adam,
        )
    elif self.optim == "adamW":
        optimizer = optim.AdamW(
            self.parameters(),
            lr=self.hparams.lr,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=self.weight_decay,
            amsgrad=False,
            fused=self.fused_adam,
        )
    elif self.optim == "galore":
        raise NotImplementedError("Galore optimizer not implemented")
        # param_groups = [
        #    {
        #        "params": [
        #            v for k, v in self.named_parameters() if "transformer" not in k
        #        ]
        #    },
        #    {
        #        "params": [
        #            v for k, v in self.named_parameters() if "transformer" in k
        #        ],
        #        "rank": 128,
        #        "update_proj_gap": 200,
        #        "scale": 0.25,
        #        "proj_type": "std",
        #    },
        # ]
        # optimizer = GaLoreAdamW(param_groups, lr=self.hparams.lr)
    else:
        raise ValueError(f"Unknown optimizer: {self.optim}")
    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        patience=self.lr_reduce_patience,
        factor=self.lr_reduce_factor,
        verbose=True,
    )
    lr_dict = {
        "scheduler": lr_scheduler,
        # The unit of the scheduler's step size, could also be 'step'.
        # 'epoch' updates the scheduler on epoch end whereas 'step'
        # updates it after a optimizer update.
        "interval": "epoch",
        # How many epochs/steps should pass between calls to
        # `scheduler.step()`. 1 corresponds to updating the learning
        # rate after every epoch/step.
        "frequency": 1,
        # Metric to to monitor for schedulers like `ReduceLROnPlateau`
        "monitor": self.lr_reduce_monitor,
    }
    self.lrfinder_steps = 0
    for val in self.trainer.callbacks:
        if type(val) is _LRCallback:
            self.lrfinder_steps = val.num_training
        if type(val) is LearningRateFinder:
            self.lrfinder_steps = val._num_training_steps
    return [optimizer], [lr_dict]

def forward(
    self,
    gene_pos: Tensor,
    expression: Optional[Tensor] = None,
    mask: Optional[Tensor] = None,
    full_depth: Optional[Tensor] = None,
    timepoint: Optional[Tensor] = None,  # (new_minibatch_of_nxt_cells,)
    get_gene_emb: bool = False,
    depth_mult: Optional[Tensor] = None,
    do_sample: bool = False,
    do_mvc: bool = False,
    do_class: bool = False,
    get_attention_layer: list = [],
):
    """
    forward also called on self(), a full forward pass on the model

    Args:
        gene_pos (Tensor): A tensor of shape (minibatch, seq_len)
            representing the genes used for each cell in the minibatch.
        expression (Tensor, optional): A tensor of shape (minibatch, seq_len)
            representing the expression levels of genes in the minibatch. Defaults to None.
        mask (Tensor, optional): A tensor of shape (minibatch, seq_len)
            used to mask certain elements in the sequence during the forward pass. Defaults to None.
        full_depth (Tensor, optional): A tensor of shape (minibatch,)
            representing the full depth of each sequence in the minibatch. Defaults to None.
        timepoint (Tensor, optional): A tensor of shape (minibatch,)
            representing the timepoint associated with each sequence in the minibatch. Defaults to None.
        get_gene_emb (bool, optional): A flag indicating whether to return the gene embeddings.
            If True, the gene embeddings are included in the output. Defaults to False.
        do_sample (bool, optional): A flag indicating whether to sample the expression levels.
            If True, the expression levels are sampled during the forward pass. Defaults to False.
        get_attention_layer (list, optional): A list indicating which attention layers to return.
            If not empty, the specified attention layers are included in the output. Defaults to [].

    Returns:
        dict of output Tensors: A dictionary containing the output tensors from the forward pass.
            The keys of the dictionary depend on the input flags (get_gene_emb, do_sample, get_attention_layer).
            at minima, the dictionary codntains the following:
            - "mean": the mean expression levels
            - "zero_logits": the logits for zero-inflated expression levels
            - "disp": the dispersion parameter
            - "cell_embs": the cell embeddings per class
            - "cell_emb": the main cell embedding
            - "cls_output": the output of the classifier
    """
    encoding = self._encoder(gene_pos, expression, mask, full_depth, timepoint)

    if self.attn_bias != "none":
        if not hasattr(self, "nbias"):
            self.nbias = torch.Tensor(
                load_npz(FILEDIR + "/../../data/bias_sparse.npz").todense()
            ).to(device=gene_pos.device, dtype=torch.float16)
        num = len(self.classes) + 2
        bias = torch.zeros(
            (
                gene_pos.shape[0],
                gene_pos.shape[1] + num,
                gene_pos.shape[1] + num,
            ),
            device=gene_pos.device,
            dtype=torch.float16,
        )
        bias[:, num:, :num] = -10_000  # do not pay attention to the cls embeddings
        bias[:, num:, num:] = self.nbias[gene_pos[:, :, None], gene_pos[:, None, :]]
    transformer_output = self.transformer(
        encoding,
        return_qkv=get_attention_layer,
        bias=bias if self.attn_bias != "none" else None,
        bias_layer=list(range(self.nlayers - 1)),
    )

    depth_mult = expression.sum(1) if depth_mult is None else depth_mult
    if len(get_attention_layer) > 0:
        transformer_output, qkvs = transformer_output
        return (
            self._decoder(
                transformer_output,
                depth_mult,
                get_gene_emb,
                do_sample,
                do_mvc,
                do_class,
            ),
            qkvs,
        )
    else:
        return self._decoder(
            transformer_output,
            depth_mult,
            get_gene_emb,
            do_sample,
            do_mvc,
            do_class,
        )

def get_cell_embs(self, layer_output):
    """
    get_cell_embs

    Args:
        layer_output (Tensor): The output tensor from a layer in the model.

    Raises:
        ValueError: Raised when an unknown cell embedding style is encountered.

    Returns:
        Tensor: The cell embeddings tensor.
    """
    if self.cell_emb_style == "cls" and self.classes is not None:
        # (minibatch, embsize)
        cell_emb = layer_output[:, : 2 + len(self.classes)]
    elif self.cell_emb_style == "avg-pool":
        cell_emb = torch.mean(layer_output, dim=1)
    else:
        raise ValueError(f"Unknown cell_emb_style: {self.cell_emb_style}")
    return cell_emb

def log_adata(self, gtclass=None, name=""):
    """
    log_adata will log an adata from predictions.
    It will log to tensorboard and wandb if available

    see @utils.log_adata
    """
    try:
        mdir = self.logger.save_dir if self.logger.save_dir is not None else "/tmp"
    except:
        mdir = "data/"
    if not os.path.exists(mdir):
        os.makedirs(mdir)
    adata, fig = utils.make_adata(
        self.embs,
        self.classes,
        self.pred if not self.keep_all_cls_pred else None,
        self.attn.get(),
        self.global_step,
        self.label_decoders,
        self.labels_hierarchy,
        gtclass,
        self.name + "_" + name + "_" + str(self.global_rank),
        mdir,
        self.doplot,
    )
    if self.doplot:
        try:
            self.logger.experiment.add_figure(fig)
        except:
            print("couldn't log to tensorboard")
        try:
            self.logger.log_image(key="umaps", images=[fig])
        except:
            print("couldn't log to wandb")

    return adata

def on_fit_start(self):
    """@see pl.LightningModule"""
    if type(self.transformer) is FlashTransformerEncoder:
        for encoder_layers in self.transformer.blocks:
            encoder_layers.set_seq_parallel(True)
    for k, v in self.mat_labels_hierarchy.items():
        self.mat_labels_hierarchy[k] = v.to(self.device)

def on_predict_epoch_end(self):
    """@see pl.LightningModule will"""
    if self.pos.shape[0] < 100:
        return
    if self.pred_log_adata:
        print("adding on disk")
        return self.log_adata(name="predict_part_" + str(self.counter))

def on_predict_epoch_start(self):
    """@see pl.LightningModule"""
    self.embs = None
    self.attn.data = None
    self.attn.attn = None
    self.counter = 0
    if type(self.transformer) is FlashTransformerEncoder:
        for encoder_layers in self.transformer.blocks:
            encoder_layers.set_seq_parallel(False)

def on_validation_epoch_end(self):
    """@see pl.LightningModule"""
    self.embs = self.all_gather(self.embs).view(-1, self.embs.shape[-1])
    self.info = self.all_gather(self.info).view(-1, self.info.shape[-1])
    self.pred = (
        self.all_gather(self.pred).view(-1, self.pred.shape[-1])
        if self.pred is not None
        else None
    )
    self.pos = self.all_gather(self.pos).view(-1, self.pos.shape[-1])
    if not self.trainer.is_global_zero:
        # print("you are not on the main node. cancelling logging step")
        return
    if self.trainer.state.stage != "sanity_check":
        sch = self.lr_schedulers()
        sch.step(self.trainer.callback_metrics["val_loss"])
        # run the test function on specific dataset
        self.log_adata(
            gtclass=self.info, name="validation_part_" + str(self.counter)
        )
        if (self.current_epoch + 1) % 30 == 0:
            self.on_test_epoch_end()

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    """@see pl.LightningModule"""
    # update params
    optimizer.step(closure=optimizer_closure)

    # manually warm up lr without a scheduler
    # making sure that we don't do this during lrfinder
    for i, pg in enumerate(optimizer.param_groups):
        if (
            self.global_step < self.warmup_duration + self.lrfinder_steps
        ) and self.lrfinder_steps < self.global_step:
            lr_scale = min(1.0, float(self.global_step + 1) / self.warmup_duration)
            pg["lr"] = lr_scale * self.hparams.lr
    for i, pg in enumerate(optimizer.param_groups):
        # if pg["lr"] < 2e-5:
        #    pg["lr"] = 2e-5
        self.log("lr_" + str(i), pg["lr"])

def predict_step(self, batch, batch_idx):
    """
    embed given gene expression, encode the gene embedding and cell embedding.

    Args:
        batch @see training_step

    Returns:
        Tensor: _description_
    """
    return self._predict(
        batch["genes"],
        batch["x"],
        batch["depth"],
        self.predict_mode,
        self.pred_embedding,
        self.get_attention_layer,
        self.predict_depth_mult,
    )

def training_step(
    self,
    batch: Dict[str, Tensor],
    batch_idx,
):
    """
    training_step defines the train loop. It is independent of forward

    @see pl.LightningModule

    Returns:
        _type_: _description_
    """
    # TASK 1 & 2 & 3 (first pass, expression reconstruction, label prediction)
    total_loss, losses = self._full_training(
        batch=batch,
        do_denoise=self.do_denoise,
        noise=self.noise,
        do_next_tp=self.do_next_tp,
        do_cce=self.do_cce,
        cce_sim=self.cce_sim,
        do_ecs=self.do_ecs,
        do_mvc=self.do_mvc,
        do_adv_cls=self.do_adv_cls,
        do_adv_batch=self.do_adv_batch,
        do_cls=self.do_cls,
        do_generate=self.do_generate,
        run_full_forward=self.run_full_forward,
        mask_ratio=self.mask_ratio,
    )
    self.log("train_loss", total_loss, prog_bar=True, sync_dist=True)
    self.log_dict(losses, prog_bar=True, sync_dist=True)
    return total_loss

def validation_step(
    self,
    batch,
    batch_idx,
):
    """
    validation_step defines the validation loop. It is independent of forward
    @see pl.LightningModule

    Args:
        batch (list[Tensor]): @see training_step
    """
    val_loss, losses = self._full_training(
        batch=batch,
        do_denoise=self.do_denoise,
        noise=self.noise,
        do_next_tp=self.do_next_tp,
        do_cce=self.do_cce,
        cce_sim=self.cce_sim,
        do_ecs=self.do_ecs,
        do_mvc=self.do_mvc,
        do_adv_cls=self.do_adv_cls,
        do_adv_batch=self.do_adv_batch,
        do_cls=self.do_cls,
        do_generate=self.do_generate,
        run_full_forward=self.run_full_forward,
        mask_ratio=self.mask_ratio,
    )
    expression = batch["x"]
    gene_pos = batch["genes"]
    depth = batch["depth"]
    # TODO: make this faster by only calling val loss
    if self.embs is not None:
        if self.embs.shape[0] < 100_000:
            self.info = torch.cat([self.info, batch["class"]])
            self._predict(
                gene_pos,
                expression,
                depth,
                pred_embedding=self.pred_embedding,
                max_size_in_mem=100_000,
            )
    else:
        self.info = batch["class"]
        self._predict(
            gene_pos,
            expression,
            depth,
            pred_embedding=self.pred_embedding,
            max_size_in_mem=100_000,
        )
    self.log("val_loss", val_loss, sync_dist=True)
    self.log_dict(losses, sync_dist=True)
    return val_loss

def forward(self, x: Tensor, batch_labels: Tensor) -> Tensor:
    """
    Args:
        x: Tensor, shape [batch_size, embsize]
    """
    if self.reverse_grad:
        x = grad_reverse(x, lambd=1.0)
    for layer in self.decoder:
        x = layer(x)
    x = self.out_layer(x)
    return F.cross_entropy(x, batch_labels)

def add_attn(
    self, x: List[Tensor], pos: Tensor, expr: Optional[Tensor] = None
) -> None:
    """
    Aggregate the attention or data based on the comp_attn flag.

    Args:
        x (List[Tensor]): List of tensors to aggregate.
        pos (Tensor): Position tensor.
    """
    if self.data is None:
        self.data = torch.zeros(
            [self.gene_dim, self.gene_dim, len(x) * x[0].shape[3]],
            device="cuda",
            dtype=torch.float32,
        )
        self.div = torch.zeros(1, device="cuda", dtype=torch.float32)

    for i, elem in enumerate(x):
        batch, seq_len, _, heads, _ = elem.shape
        if self.apply_softmax:
            attn = torch.nn.functional.softmax(
                elem[:, :, 0, :, :].permute(0, 2, 1, 3)
                @ elem[:, :, 1, :, :].permute(0, 2, 3, 1),
                dim=-1,
            )
            if expr is not None:
                attn = attn * (expr > 0).float()
            self.data[:, :, heads * i : heads * (i + 1)] += (
                attn.sum(0).permute(1, 2, 0) / batch
            )
        else:
            self.data[:, :, heads * i : heads * (i + 1)] += (
                elem[:, :, 0, :, :].permute(0, 2, 1, 3)
                @ elem[:, :, 1, :, :].permute(0, 2, 3, 1)
            ).sum(0).permute(1, 2, 0) / batch
    self.div += 1

def add_qk(
    self, x: List[Tensor], pos: Tensor, expr: Optional[Tensor] = None
) -> None:
    """
    Add data to the internal storage.

    Args:
        x (List[Tensor]): List of tensors to add.
        pos (Tensor): Position tensor.
    """
    if self.data is None:
        self.data = torch.zeros(
            [len(x), self.gene_dim] + list(x[0].shape[2:]), device="cuda"
        )
        self.div = torch.zeros(self.gene_dim, device="cuda")
    for i in range(x[0].shape[0]):
        loc = torch.cat([torch.arange(8, device="cuda"), pos[i] + 8]).int()
        for j in range(len(x)):
            self.data[j, loc, :, :, :] += x[j][i]
        self.div[loc] += 1

def get(self) -> Optional[np.ndarray]:
    """
    Get the aggregated attention or data.

    Returns:
        Optional[np.ndarray]: The aggregated attention or data.
    """
    if self.shared_qk:
        if self.data is None:
            return None
        # shape is (layers, genes, qkv, heads, emb)
        return self.data / self.div.view(1, self.div.shape[0], 1, 1, 1)
    else:
        if self.data is None:
            return None
        self.data.div_(self.div)
        return self.data