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}")
    if self.lr_reduce_monitor is None:
        print("no lr reduce factor")
        return [optimizer]
    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,
    req_depth: Optional[Tensor] = None,
    timepoint: Optional[Tensor] = None,  # (new_minibatch_of_nxt_cells,)
    get_gene_emb: bool = False,
    metacell_token: Optional[Tensor] = None,  # (minibatch, 1)
    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.
        req_depth (Tensor, optional): A tensor of shape (minibatch,)
            representing the full depth of each sequence in the minibatch. Defaults to None.
        depth_mult (Tensor, optional): A tensor of shape (minibatch,)
            representing the depth multiplier for 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,
        req_depth=req_depth if self.depth_atinput else None,
        timepoint=timepoint,
        metacell_token=metacell_token,
    )
    if self.attn_bias != "none":
        if not hasattr(self, "nbias"):
            bias_path = os.path.join(
                Path(FILEDIR).parent.parent, "data", "bias_sparse.npz"
            )
            self.nbias = torch.Tensor(load_npz(bias_path).todense()).to(
                device=gene_pos.device, dtype=torch.float16
            )
        num = self.cell_embs_count if not self.cell_transformer else 0
        bias = torch.zeros(
            (
                gene_pos.shape[0],
                gene_pos.shape[1] + num,
                gene_pos.shape[1] + num,
            ),
            device=gene_pos.device,
            dtype=torch.float16,
        )
        # fade slowly through the iterations
        fade_factor = 400 / (400 + self.trainer.global_step)
        # bias[:, num:, :num] = -10_000  # do not pay attention to the cls embeddings
        bias[:, num:, num:] = (
            self.nbias[gene_pos[:, :, None], gene_pos[:, None, :]] * fade_factor
        )
    if self.cell_transformer:
        cell_encoding = encoding[:, : self.cell_embs_count, :]
        encoding = encoding[:, self.cell_embs_count :, :]
    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)),
    )
    if len(get_attention_layer) > 0:
        transformer_output, qkvs = transformer_output
    if self.cell_transformer:
        cell_output = self.cell_transformer(cell_encoding, x_kv=transformer_output)
        transformer_output = torch.cat([cell_output, transformer_output], dim=1)
    # if not provided we will mult by the current expression sum
    depth_mult = expression.sum(1) if depth_mult is None else depth_mult
    res = self._decoder(
        transformer_output,
        depth_mult,
        get_gene_emb,
        do_sample,
        do_mvc,
        do_class,
        req_depth=req_depth if not self.depth_atinput else None,
    )
    return (res, qkvs) if len(get_attention_layer) > 0 else res

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(
        pos=self.pos,
        expr_pred=self.expr_pred,
        genes=self.genes,
        embs=self.embs,
        classes=self.classes,
        pred=self.pred,
        attention=self.attn.get(),
        label_decoders=self.label_decoders,
        labels_hierarchy=self.labels_hierarchy,
        gtclass=gtclass,
        doplot=self.doplot,
    )
    adata.write(
        mdir
        + "/step_"
        + str(self.global_step)
        + "_"
        + self.name
        + "_"
        + name
        + "_"
        + str(self.global_rank)
        + ".h5ad"
    )
    if self.doplot:
        logged = False
        try:
            self.logger.experiment.add_figure(fig)
            logged = True
        except:
            print("couldn't log to tensorboard")
        try:
            self.logger.log_image(key="umaps", images=[fig])
            logged = True
        except:
            print("couldn't log to wandb")
        if not logged and fig is not None:
            fig.savefig("" + mdir + "/umap_" + self.name + "_" + name + ".png")

    return adata

def on_fit_start(self):
    """@see pl.LightningModule"""
    if type(self.transformer) is FlashTransformer:
        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 FlashTransformer:
        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 self.trainer.state.stage != "sanity_check":
        if self.trainer.is_global_zero:
            print("logging anndata")
            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) % self.test_every == 0:
                self.on_test_epoch_end()
            # Synchronize all processes with a timeout
        if torch.distributed.is_initialized():
            # Set a timeout that's longer than your test typically takes
            # Write rank to file for debugging
            self.trainer.strategy.barrier()

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    """@see pl.LightningModule"""
    # update params
    # manually warm up lr without a scheduler
    # making sure that we don't do this during lrfinder
    lr_scale = None
    prev_lr = None
    if (
        self.trainer.global_step < self.warmup_duration + self.lrfinder_steps
    ) and self.lrfinder_steps <= self.trainer.global_step:
        for i, pg in enumerate(optimizer.param_groups):
            lr_scale = min(
                1.0, float(self.trainer.global_step + 1) / self.warmup_duration
            )
            prev_lr = pg["lr"]
            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"])
    if optimizer.param_groups[0]["lr"] > self.hparams.lr:
        print(optimizer.param_groups[0]["lr"], self.hparams.lr)
        print(lr_scale, self.warmup_duration, self.trainer.global_step, prev_lr)
        if prev_lr is not None:
            pg["lr"] = prev_lr
        else:
            raise ValueError("OPTIMIZER HAS INCREASED LR. WHYY?")

    optimizer.step(closure=optimizer_closure)

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_
    """
    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_temp=self.cce_temp,
        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_temp=self.cce_temp,
        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"]
    metacell_token = batch.get("is_meta", None)
    # 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=120_000,
                metacell_token=metacell_token,
            )
    else:
        self.info = batch["class"]
        self._predict(
            gene_pos,
            expression,
            depth,
            pred_embedding=self.pred_embedding,
            max_size_in_mem=120_000,
            metacell_token=metacell_token,
        )
    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)