Source code for spateo.svg.utils

from typing import List, Optional, Union

import dynamo as dyn
import numpy as np
import ot
import pandas as pd
from anndata import AnnData
from scipy.sparse import csr_matrix, issparse
from scipy.sparse.csgraph import floyd_warshall

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

from ..logging import logger_manager as lm


def bin_adata(
    adata: AnnData,
    bin_size: int = 1,
    layer: str = "spatial",
) -> AnnData:
    """Aggregate cell-based adata by bin size. Cells within a bin would be
    aggregated together as one cell.

    Args:
        adata: the input adata.
        bin_size: the size of square to bin adata.

    Returns:
        Aggreated adata.
    """
    adata = adata.copy()
    adata.obsm[layer] = (adata.obsm[layer] // bin_size).astype(np.int32)
    df = (
        pd.DataFrame(adata.X.toarray(), columns=adata.var_names)
        if issparse(adata.X)
        else pd.DataFrame(adata.X, columns=adata.var_names)
    )
    df[["x", "y"]] = adata.obsm[layer]
    df2 = df.groupby(by=["x", "y"]).sum()
    a = AnnData(df2)
    a.uns["__type"] = "UMI"
    a.obs_names = [str(i[0]) + "_" + str(i[1]) for i in df2.index.to_list()]
    a.obsm[layer] = np.array([list(i) for i in df2.index.to_list()], dtype=np.float64)
    return a


def shuffle_adata(
    adata: AnnData,
    seed: int = 0,
    replace: bool = False,
):
    """Shuffle X in anndata object randomly.

    Args:
        adata: AnnData object
        seed: seed for randomly shuffling

    Returns:
        adata: AnnData object
    """
    adata = adata.copy()
    if seed == 0:
        return adata
    np.random.seed(seed)
    if issparse(adata.X):
        tmp = adata.X.toarray()
        if replace:
            tmp = tmp[np.random.randint(len(tmp), size=len(tmp))]
        else:
            np.random.shuffle(tmp)
        adata.X = csr_matrix(tmp)
    else:
        tmp = adata.X
        if replace:
            tmp = tmp[np.random.randint(len(tmp), size=len(tmp))]
        else:
            np.random.shuffle(tmp)
        adata.X = tmp
    return adata


def filter_adata_by_pos_ratio(
    adata,
    pos_ratio,
):
    """Filter out cells with positive ratio lower than a setting value.

    Args:
        adata: AnnData object.
        pos_ratio: Cells with positive ratio lower than this value would be discarded.

    Return:
        AnnData object.

    """
    genes, adata = get_genes_by_pos_ratio(adata, pos_ratio)
    return adata[:, genes].copy()


def get_genes_by_pos_ratio(
    adata: AnnData,
    pos_ratio: float = 0.1,
) -> list:
    """Get genes that have postive ratio higher than a setting value.

    Args:
        adata: AnnData object.
        pos_ratio: The threshold of positive ratio.

    Returns:
        Gene list.
        AnnData object.
    """
    adata = adata.copy()
    adata.var["nCells"] = np.sum(adata.X.toarray() > 0, axis=0) if issparse(adata.X) else np.sum(adata.X > 0, axis=0)
    adata.var["raw_pos_rate"] = adata.var["nCells"] / adata.n_obs
    return adata.var_names[adata.var["nCells"] / adata.n_obs > pos_ratio].to_list(), adata


def add_pos_ratio_to_adata(adata: AnnData, layer: str = None, var_name: str = "raw_pos_rate"):
    """Calculate positive ratios for all genes, and return to AnnData.
    We defind positive ratio of a gene as the percent of cells express this gene.

    Args:
        adata: AnnData object.
        layer: The layer of AnnData, in which the data are used. If not given, we use data in X.
        var_name: The var name for storing positive ratios.

    Return:
        None
    """
    if layer:
        adata.var[var_name] = (
            np.sum(adata.layers[layer].toarray() > 0, axis=0)
            if issparse(adata.layers[layer])
            else np.sum(adata.layers[layer] > 0, axis=0)
        )
    else:
        adata.var[var_name] = (
            np.sum(adata.X.toarray() > 0, axis=0) if issparse(adata.X) else np.sum(adata.X > 0, axis=0)
        )
    adata.var[var_name] = adata.var[var_name] / adata.n_obs


def cal_geodesic_distance(
    adata: AnnData,
    layer: str = "spatial",
    n_neighbors: int = 30,
    min_dis_cutoff: float = 2.0,
    max_dis_cutoff: float = 4.0,
) -> AnnData:
    """Calculate geodesic distance between any pair of genes.

    Args:
        adata: AnnData object.
        layer: The layer of AnnData, in which the data are used.
        n_neighbors: The number of neighbor to connect a cell to its nearest neighbors.
        min_dis_cutoff: Remove cells with minimal distance with its neighbors larger than this value.
                        These cells are like islated cells.
        max_dis_cutoff: Remove cells with maximal distance with its neighbors larger than this value.
                        These cells are like sparse cells.

    Return:
        AnnData object.

    """

    dyn.tl.neighbors(
        adata,
        X_data=adata.obsm[layer],
        n_neighbors=n_neighbors,
        result_prefix="spatial",
    )
    # remove islated one cell
    b = adata[
        np.min(
            adata.obsp["spatial_distances"].toarray(),
            axis=1,
            initial=1e10,
            where=np.array(adata.obsp["spatial_distances"].toarray() > 0),
        )
        <= min_dis_cutoff
    ]
    lm.main_info(f"The cell/buckets number after filtering by min_dis_cutoff is {len(b)}")

    # remove sparse cells
    b = b[np.max(b.obsp["spatial_distances"].toarray(), axis=1) <= max_dis_cutoff]
    lm.main_info(f"The cell/buckets number after filtering by max_dis_cutoff is {len(b)}")

    dyn.tl.neighbors(
        b,
        X_data=b.obsm[layer],
        n_neighbors=n_neighbors,
        result_prefix="spatial",
    )
    # from scipy.sparse import csr_matrix
    # from scipy.sparse.csgraph import floyd_warshall

    conn = b.obsp["spatial_distances"].toarray()
    conn[conn == np.inf] = 0
    conn = csr_matrix(conn)
    dist_matrix, predecessors = floyd_warshall(csgraph=conn, directed=False, return_predecessors=True)
    b.obsp["distance"] = dist_matrix
    return b


def cal_euclidean_distance(
    adata: AnnData,
    layer: str = "spatial",
    min_dis_cutoff: float = np.inf,
    max_dis_cutoff: float = np.inf,
) -> AnnData:
    dyn.tl.neighbors(
        adata,
        X_data=adata.obsm[layer],
        n_neighbors=adata.n_obs,
        result_prefix="spatial",
    )
    # remove islated one cell
    b = adata[
        np.min(
            adata.obsp["spatial_distances"].toarray(),
            axis=1,
            initial=1e10,
            where=np.array(adata.obsp["spatial_distances"].toarray() > 0),
        )
        <= min_dis_cutoff
    ]

    # remove sparse cells
    b = b[np.max(b.obsp["spatial_distances"].toarray(), axis=1) <= max_dis_cutoff]

    # from scipy.sparse import csr_matrix
    # from scipy.sparse.csgraph import floyd_warshall

    conn = b.obsp["spatial_distances"].toarray()
    conn[conn == np.inf] = 0
    conn = csr_matrix(conn)
    dist_matrix, predecessors = floyd_warshall(csgraph=conn, directed=False, return_predecessors=True)
    b.obsp["distance"] = dist_matrix
    return b


def scale_to(
    adata: AnnData,
    to_median: bool = True,
    N: int = 10000,
) -> AnnData:
    """Scale the X array in AnnData.

    Args:
        adata: AnnData object.
        to_median: Whether scale to the median of cell total expressions.
        N: if `to_median` is False, scale data to this value.

    Return:
        AnnData object.

    """
    adata = adata.copy()
    adata.X = adata.X.astype(np.float64)

    if to_median:
        if issparse(adata.X):
            N = np.median(np.sum(adata.X.toarray(), axis=1))
        else:
            N = np.median(np.sum(adata.X, axis=1))

    if issparse(adata.X):
        adata.X = csr_matrix((adata.X.toarray().T / (np.sum(adata.X.toarray(), axis=1) / N)).T)
    else:
        adata.X = (adata.X.T / (np.sum(adata.X, axis=1) / N)).T
    return adata


def cal_wass_dis(M, a, b=[], numItermax=1000000):
    """Computing Wasserstein distance.

    Args:
        M:  (ns,nt) array-like, float – Loss matrix (c-order array in numpy with type float64)
        a: (ns,) array-like, float – Source histogram (uniform weight if empty list)
        b: (nt,) array-like, float – Target histogram (uniform weight if empty list)

    Returns:
        W: (float, array-like) – Optimal transportation loss for the given parameters

    """

    W = ot.emd2(a, b, M, numItermax=numItermax)

    return W


def cal_rank_p(genes, ws, w_df, bin_num=100):
    ws_dict = {}
    for g, w in zip(genes, ws):
        if g not in ws_dict:
            ws_dict[g] = []
        ws_dict[g].append(w)

    sorted_genes = w_df["mean"].sort_values().index.to_list()
    each_bin_gene_num = int(len(sorted_genes) / bin_num) + 1
    each_bin_ws = {}
    bin_of_gene = {}
    for i in range(bin_num):
        each_bin_ws[i] = []
        for g in sorted_genes[i * each_bin_gene_num : (i + 1) * each_bin_gene_num]:
            if np.sum(np.array(ws_dict[g])) > 0:
                each_bin_ws[i].append(ws_dict[g])
            bin_of_gene[g] = i
        each_bin_ws[i] = np.array(each_bin_ws[i])
    rank_p = []
    for g in w_df.index:
        t = each_bin_ws[bin_of_gene[g]].flatten()
        rank_p.append((np.sum(t >= w_df.loc[g, "Wasserstein_distance"]) + 1) / len(t))
    return rank_p, each_bin_ws


def loess_reg(
    adata: AnnData,
    layers: str = "X",
) -> AnnData:
    adata = adata.copy()
    adata.X = adata.X.astype(np.float64)

    if issparse(adata.X):
        adata.X = csr_matrix((adata.X.toarray().T / (np.sum(adata.X.toarray(), axis=1) / N)).T)
    else:
        adata.X = (adata.X.T / (np.sum(adata.X, axis=1) / N)).T
    return adata