Downstream_tasks/Zero_shot_batch_effect/sc_foundation_evals/utils.py

## Copyright (c) Microsoft Corporation.
## Licensed under the MIT license.

from typing import Dict, Union, List
from anndata import AnnData
import numpy as np
import scanpy as sc
import scib
import torch
import torch.nn.functional as F

from .helpers.custom_logging import log

# MODIFIED wrapper for all scib metrics from 
# https://github.com/bowang-lab/scGPT/blob/5a69912232e214cda1998f78e5b4a7b5ef09fe06/scgpt/utils/util.py#L267
def eval_scib_metrics(
    adata: AnnData,
    batch_key: str = "str_batch",
    label_key: str = "cell_type",
    embedding_key: str = "X_scGPT"
) -> Dict:
    
    # if adata.uns["neighbors"] exists, remove it to make sure the optimal 
    # clustering is calculated for the correct embedding
    # print a warning for the user
    if "neighbors" in adata.uns:        
        log.warning(f"neighbors in adata.uns found \n {adata.uns['neighbors']} "
                    "\nto make sure the optimal clustering is calculated for the "
                    "correct embedding, removing neighbors from adata.uns."
                    "\nOverwriting calculation of neighbors with "
                    f"sc.pp.neighbors(adata, use_rep={embedding_key}).")
        adata.uns.pop("neighbors", None)
        sc.pp.neighbors(adata, use_rep=embedding_key)
        log.info("neighbors in adata.uns removed, new neighbors calculated: "
                 f"{adata.uns['neighbors']}")


    # in case just one batch scib.metrics.metrics doesn't work 
    # call them separately
    results_dict = dict()

    res_max, nmi_max, nmi_all = scib.metrics.clustering.opt_louvain(
            adata,
            label_key=label_key,
            cluster_key="cluster",
            use_rep=embedding_key,
            function=scib.metrics.nmi,
            plot=False,
            verbose=False,
            inplace=True,
            force=True,
    )
    
    results_dict["NMI_cluster/label"] = scib.metrics.nmi(
        adata, 
        "cluster",
        label_key,
        "arithmetic",
        nmi_dir=None
    )

    results_dict["ARI_cluster/label"] = scib.metrics.ari(
        adata, 
        "cluster", 
        label_key
    )

    results_dict["ASW_label"] = scib.metrics.silhouette(
        adata, 
        label_key, 
        embedding_key, 
        "euclidean"
    )   

    results_dict["graph_conn"] = scib.metrics.graph_connectivity(
        adata,
        label_key=label_key
    )
    

    # Calculate this only if there are multiple batches
    if len(adata.obs[batch_key].unique()) > 1:
        results_dict["ASW_batch"] = scib.metrics.silhouette(
            adata,
            batch_key,
            embedding_key,
            "euclidean"
        )

        results_dict["ASW_label/batch"] = scib.metrics.silhouette_batch(
            adata, 
            batch_key,
            label_key, 
            embed=embedding_key, 
            metric="euclidean",
            return_all=False,
            verbose=False
        )

        results_dict["PCR_batch"] = scib.metrics.pcr(
            adata,
            covariate=batch_key,
            embed=embedding_key,
            recompute_pca=True,
            n_comps=50,
            verbose=False
        )

    results_dict["avg_bio"] = np.mean(
        [
            results_dict["NMI_cluster/label"],
            results_dict["ARI_cluster/label"],
            results_dict["ASW_label"],
        ]
    )

    log.debug(
        "\n".join([f"{k}: {v:.4f}" for k, v in results_dict.items()])
    )

    # remove nan value in result_dict
    results_dict = {k: v for k, v in results_dict.items() if not np.isnan(v)}

    return results_dict

def create_attention_mask_default(vecs: torch.Tensor) -> torch.Tensor:
    """
    Create an attention mask from a vector of positions of unknown genes. 
    Implementation of the mask described in the scGPT v2 preprint, fig S1A
    https://www.biorxiv.org/content/10.1101/2023.04.30.538439v2.full#F7

    Args:
        vecs (torch.Tensor): A bool tensor with position of unknown genes. 
                             shape (batch_size, seq_len)

    Returns:
        torch.Tensor: A bool attention mask. 
                      shape (batch_size, seq_len, seq_len)

    Examples:
    >>> create_attention_mask_default(torch.tensor([[0, 0, 1, 1],
                                                    [0, 0, 0, 1]]).bool())  
    tensor([[[False, False,  True,  True],
             [False, False,  True,  True],
             [False, False, False,  True],
             [False, False,  True, False]],

            [[False, False, False,  True],
             [False, False, False,  True],
             [False, False, False,  True],
             [False, False, False, False]]])

    """
    # check if vecs is a boolean tensor
    if not vecs.dtype == torch.bool:
        # check if vecs is 0 and 1 tensor
        if not torch.all(vecs.eq(0) | vecs.eq(1)):
            raise TypeError("vecs must be a boolean tensor")
        else:
            # convert to boolean tensor
            vecs = vecs.bool()

    # Use broadcasting to expand each vector into a square matrix
    attn_mask = vecs.unsqueeze(-1).repeat(1, 1, vecs.size(1))

    # Create a boolean mask for the diagonal
    diagonal_mask = ~torch.eye(vecs.size(1), device=vecs.device).bool()

    # Use the diagonal mask to set the diagonal of each square matrix to False
    attn_mask &= diagonal_mask

    # Transpose the last two dimensions to make each row of the original vector a column in the matrix
    attn_mask = attn_mask.transpose(-1, -2)

    return attn_mask

def create_attention_mask_modified(unknown_genes: torch.Tensor,
                                   cell_embedding_position: int = 0) -> torch.Tensor:
    """
    Create an attention mask from a vector of positions of unknown genes. 
    This is a modification of the mask described in the scGPT v2 preprint.
    With this attention mask, the cell embedding is always attended to,
    and the unknown genes are only attended to by the cell embedding.

    Args:
        unknown_genes (torch.Tensor): A bool or 0/1 tensor with position of unknown genes. 
                             shape (batch_size, seq_len)

    Returns:
        torch.Tensor: A bool attention mask. 
                      shape (batch_size, seq_len, seq_len)
    Examples:
    >>> create_attention_mask_modified(torch.tensor([[0, 0, 1, 1],
                                                     [0, 0, 0, 1]]).bool())  
    tensor([[[False, False,  True,  True],
             [False, False,  True,  True],
             [False, True, False,  True],
             [False, True,  True, False]],

            [[False, False, False,  True],
             [False, False, False,  True],
             [False, False, False,  True],
             [False, True, True, False]]])

    """
    # if not using GPU show warning
    if unknown_genes.device.type == 'cpu':
        log.warning('create_attention_mask_modified() is not optimized for CPU, '
                    'please use GPU for better performance.' 
                    f'Device: {unknown_genes.device.type} ')
    
    # # change vecs to 0 and 1 tensor
    # if unknown_genes.dtype == torch.bool:
    #     unknown_genes = unknown_genes.float()
    
    if torch.any(unknown_genes[:, cell_embedding_position].eq(True)): # eq(1)
        raise ValueError('Cell embedding position is unknown gene position.')
    
    attn_mask = create_attention_mask_default(unknown_genes)
    
    
    known_genes = ~unknown_genes
    # the cell embedding should always be attended to
    known_genes[:, cell_embedding_position] = False

    for i in range(len(unknown_genes)):
        attn_mask[i][unknown_genes[i].unsqueeze(-1) & known_genes[i].unsqueeze(0)] = True
            
    return attn_mask

def random_mask_value(
    values: Union[torch.Tensor, np.ndarray],
    mask_ratio: float = 0.15,
    mask_value: int = -1,
    pad_value: int = -2,
    mask_cell_embedding: bool = False,
    cell_emb_value: int = 0,
) -> torch.Tensor:
    """
    Randomly mask a batch of data.

    Args:
        values (array-like):
            A batch of tokenized data, with shape (batch_size, n_features).
        mask_ratio (float): The ratio of genes to mask, default to 0.15.
        mask_value (int): The value to mask with, default to -1.
        pad_value (int): The value of padding in the values, will be kept unchanged.
        mask_cell_embedding (bool): Whether to mask the cell embedding, default to False.

    Returns:
        torch.Tensor: A tensor of masked data.
    """
    if isinstance(values, torch.Tensor):
        # it is crucial to clone the tensor, otherwise it changes the original tensor
        values = values.clone().detach().numpy()
    else:
        values = values.copy()

    if not mask_cell_embedding:
        # sanity check that the first element is the cell embedding
        assert np.all(values[:, 0] == cell_emb_value)

    for i in range(len(values)):
        row = values[i]
        non_padding_idx = np.nonzero(row - pad_value)[0]
        if not mask_cell_embedding:
            # remove the first element, which is the cell embedding
            non_padding_idx = non_padding_idx[1:]
        n_mask = int(len(non_padding_idx) * mask_ratio)
        mask_idx = np.random.choice(non_padding_idx, n_mask, replace=False)
        row[mask_idx] = mask_value
    return torch.from_numpy(values).float()

def mask_data(tokenized_data: Dict[str, torch.Tensor],
              mask_ratio: float = 0.15,
              pad_value: int = -2,
              mask_value: int = -1, 
              mask_cell_embedding: bool = False,
              cell_emb_value: int = 0) -> Dict[str, torch.Tensor]:
        """
        Mask the data.
        """
        # this will randomly mask all the values,
        # including or excluding the cell embedding based on the mask_cell_embedding flag
        masked_values = random_mask_value(
            tokenized_data["values"],
            mask_ratio = mask_ratio,
            mask_value = mask_value,
            pad_value = pad_value,
            mask_cell_embedding = mask_cell_embedding,
            cell_emb_value = cell_emb_value
            )
        
        # this will mark the genes which values were masked for loss calculation
        gene_mask = torch.logical_and(tokenized_data["values"].ne(pad_value),
                                      masked_values.eq(mask_value))
    
        out_dict = {
            "gene_ids": tokenized_data["genes"],
            "values": masked_values,
            "target_values": tokenized_data["values"],
            "gene_mask": gene_mask
        }
        
        if "batch_labels" in tokenized_data.keys():
            out_dict["batch_labels"] = tokenized_data["batch_labels"]
        
        return out_dict

def calculate_losses(input: torch.Tensor,
                     output: torch.Tensor, 
                     masked_position: torch.Tensor,
                     non_padded_position: torch.Tensor, #TODO: add this!
                     skip_cell: bool = True,
                     methods: Union[List[str], str] = "all",
                     reduction_method = "mean",
                     cp: int = 0
                     ) -> Dict[str, float]:
    
    if reduction_method not in ["mean", "sum"]:
        msg = f"calculate_loss: reduction_method {reduction_method} is not supported"
        log.error(msg)
        raise ValueError(msg)
    
    implemented_methods = ["mse", "mre", "mae"]
    # check if methods equal to string all 
    if isinstance(methods, str) and methods.lower() == "all":
        methods = implemented_methods

    # if methods is not a list, make it a list    
    methods = [methods] if isinstance(methods, str) else methods

    # make sure methods are lower case
    methods = [method.lower() for method in methods]

    # check if element of methods is implemented
    methods_ = [method for method in methods if method not in implemented_methods]
    if len(methods_) == len(methods):
        msg = f"calculate_loss: methods {methods_} are not implemented"
        log.error(msg)
        raise ValueError(msg)
    
    if len(methods_) > 0:
        msg = f"calculate_loss: methods {methods_} are not implemented"
        log.warning(msg)

    # make sure masked_position is boolean
    if masked_position.dtype != torch.bool:
        # log.warning("evaluate_and_log: masked_position is not boolean")
        masked_position = masked_position.bool()

    if skip_cell:
        # check if cp is valid
        if cp >= input.shape[1]:
            msg = f"evaluate_and_log: cp {cp} is greater than input.shape[1] {input.shape[1]}"
            log.error(msg)
            raise ValueError(msg)
        # remove cp from input and output
        input = torch.concat((input[:,:cp], input[:,cp+1:]), dim=1)
        output = torch.concat((output[:,:cp], output[:,cp+1:]), dim=1)
        masked_position = torch.concat((masked_position[:,:cp], 
                                        masked_position[:,cp+1:]), dim=1)
        non_padded_position = torch.concat((non_padded_position[:,:cp],
                                            non_padded_position[:,cp+1:]), dim=1)
        
    results = dict()
    if "mse" in methods:
        # get the loss for masked values
        loss = F.mse_loss(torch.masked_select(output, masked_position).float(), 
                          torch.masked_select(input, masked_position).float(), 
                          reduction = reduction_method)
        
        # TODO: add masked_select on non padded!
        loss_all = F.mse_loss(torch.masked_select(output, non_padded_position).float(),
                              torch.masked_select(input, non_padded_position).float(), 
                              reduction = reduction_method)

        results["MSE"] = loss
        results["MSE_all"] = loss_all

    if "mre" in methods:
        loss = (torch.abs(torch.masked_select(output, masked_position) - 
                          torch.masked_select(input, masked_position)) /
                          (torch.masked_select(input, masked_position) + 1e-6))
        
        if reduction_method == "mean":
            loss = loss.mean()
        elif reduction_method == "sum":
            loss = loss.sum()

        loss_all = torch.abs((torch.masked_select(output, non_padded_position) -  
                              torch.masked_select(input, non_padded_position)) / 
                              (torch.masked_select(input, non_padded_position) + 1e-6))

        if reduction_method == "mean":
            loss_all = loss_all.mean()
        elif reduction_method == "sum":
            loss_all = loss_all.sum()

        results["MRE"] = loss
        results["MRE_all"] = loss_all

    if "mae" in methods:
        # get the loss for masked values
        loss = F.l1_loss(torch.masked_select(output, masked_position).float(), 
                         torch.masked_select(input, masked_position).float(), 
                         reduction = reduction_method)
        
        loss_all = F.l1_loss(torch.masked_select(output, non_padded_position).float(), 
                             torch.masked_select(input, non_padded_position).float(), 
                             reduction = reduction_method)
        
        results["MAE"] = loss
        results["MAE_all"] = loss_all
    
    return results

import torch

def permute_values(mat: torch.Tensor, 
                   pad_value: int = -2,
                   cell_embedding: bool = True) -> tuple:
    """
    Permute the data in a batch. The data is a 2D matrix with shape 
        (batch_size, seq_len). The non-padded values are permuted and
        pad values are kept at the end of the sequence
    Args:
        - mat: 2D matrix with shape (batch_size, seq_len)
        - pad_value: the value of the pad
        - cell_embedding: if True, the first position is a cell embedding
    Returns:
        - mat_perm: the permuted matrix
        - indx_perm: the indices of the permutation
    """
    if len(mat.shape) > 2:
        raise ValueError("mat should be 2D or 1D")
        
    # reshape the mat if needed so that the mat id 2D
    reshape = False
    if len(mat.shape) == 1:
        reshape = True
        mat = mat.unsqueeze(0)
    
    if cell_embedding:
        cemb = mat[:,0]
        mat = mat[:,1:]
    mat_perm = torch.empty_like(mat)
    indx_perm = torch.empty_like(mat)
    pad_masks = mat == pad_value

    for i in range(mat.shape[0]):
        non_pad = torch.where(~pad_masks[i])[0]
        pad = torch.where(pad_masks[i])[0]
        perm_non_pad = non_pad[torch.randperm(len(non_pad))] 
        perm = torch.cat([perm_non_pad, pad])
        indx_perm[i] = perm
        mat_perm[i] = mat[i][perm]
    
    if cell_embedding: # TODO: test not permuting cell_embedding!
        mat_perm = torch.cat([cemb.unsqueeze(1), mat_perm], dim=1)
        indx_perm = torch.cat((torch.zeros(indx_perm.shape[0], 1).to(indx_perm.device), 
                               indx_perm+1), dim=1)

    if reshape:
        mat_perm = mat_perm.view(-1)
        indx_perm = indx_perm.view(-1)

    # change indx_perm to int
    indx_perm = indx_perm.long()
    return mat_perm, indx_perm

def rearrange(mat: torch.Tensor, 
              indx: torch.Tensor) -> torch.Tensor:
    """
    Rearrange the rows of a 2D tensor A according to the indices indx
    Args:
        - mat: 2D tensor with shape (batch_size, seq_len)
        - indx: 2D tensor with shape (batch_size, seq_len)
    Returns:
        - mat_reordered: 2D tensor with shape (batch_size, seq_len)
    """
    batch_size = mat.shape[0]
    batch_indices = torch.arange(batch_size).view(-1, 1).to(indx.device)
    mat_reordered = mat[batch_indices, indx]
    return mat_reordered

def reverse_permute(mat: torch.Tensor,
                    indx: torch.Tensor) -> torch.Tensor:
    """
    Reverse the permutation of a 2D tensor A according to the indices indx
    Args:
        - mat: 2D tensor with shape (batch_size, seq_len)
        - indx: 2D tensor with shape (batch_size, seq_len)
    Returns:
        - Sorted 2D tensor with shape (batch_size, seq_len)
    """
    dims = torch.arange(indx.size(0)).reshape(-1, 1)
    return mat[dims, indx.argsort()]