get.py

"""Helper functions for accessing data."""

from __future__ import annotations

from typing import TYPE_CHECKING, TypedDict

import numpy as np
import pandas as pd
from anndata import AnnData
from numpy.typing import NDArray

from .._compat import CSBase

if TYPE_CHECKING:
    from collections.abc import Collection, Iterable
    from typing import Any, Literal, Unpack

    from anndata._core.sparse_dataset import BaseCompressedSparseDataset
    from anndata._core.views import ArrayView

    from .._compat import DaskArray


# --------------------------------------------------------------------------------
# Plotting data helpers
# --------------------------------------------------------------------------------


# TODO: implement diffxpy method, make singledispatch
def rank_genes_groups_df(
    adata: AnnData,
    group: str | Iterable[str] | None,
    *,
    key: str = "rank_genes_groups",
    pval_cutoff: float | None = None,
    log2fc_min: float | None = None,
    log2fc_max: float | None = None,
    gene_symbols: str | None = None,
) -> pd.DataFrame:
    """Get :func:`scanpy.tl.rank_genes_groups` results in the form of a :class:`~pandas.DataFrame`.

    Params
    ------
    adata
        Object to get results from.
    group
        Which group (as in :func:`scanpy.tl.rank_genes_groups`'s `groupby`
        argument) to return results from. Can be a list. All groups are
        returned if groups is `None`.
    key
        Key differential expression groups were stored under.
    pval_cutoff
        Return only adjusted p-values below the  cutoff.
    log2fc_min
        Minimum logfc to return.
    log2fc_max
        Maximum logfc to return.
    gene_symbols
        Column name in `.var` DataFrame that stores gene symbols. Specifying
        this will add that column to the returned dataframe.

    Example
    -------
    >>> import scanpy as sc
    >>> pbmc = sc.datasets.pbmc68k_reduced()
    >>> sc.tl.rank_genes_groups(pbmc, groupby="louvain", use_raw=True)
    >>> dedf = sc.get.rank_genes_groups_df(pbmc, group="0")

    """
    if isinstance(group, str):
        group = [group]
    if group is None:
        group = list(adata.uns[key]["names"].dtype.names)
    method = adata.uns[key]["params"]["method"]
    if method == "logreg":
        colnames = ["names", "scores"]
    else:
        colnames = ["names", "scores", "logfoldchanges", "pvals", "pvals_adj"]

    d = [pd.DataFrame(adata.uns[key][c])[group] for c in colnames]
    d = pd.concat(d, axis=1, names=[None, "group"], keys=colnames)
    d = d.stack(level=1, future_stack=True).reset_index()
    d["group"] = pd.Categorical(d["group"], categories=group)
    d = d.sort_values(["group", "level_0"]).drop(columns="level_0")

    if method != "logreg":
        if pval_cutoff is not None:
            d = d[d["pvals_adj"] < pval_cutoff]
        if log2fc_min is not None:
            d = d[d["logfoldchanges"] > log2fc_min]
        if log2fc_max is not None:
            d = d[d["logfoldchanges"] < log2fc_max]
    if gene_symbols is not None:
        d = d.join(adata.var[gene_symbols], on="names")

    for pts, name in {"pts": "pct_nz_group", "pts_rest": "pct_nz_reference"}.items():
        if pts in adata.uns[key]:
            pts_df = (
                adata
                .uns[key][pts][group]
                .rename_axis(index="names")
                .reset_index()
                .melt(id_vars="names", var_name="group", value_name=name)
            )
            d = d.merge(pts_df)

    # remove group column for backward compat if len(group) == 1
    if len(group) == 1:
        d.drop(columns="group", inplace=True)

    return d.reset_index(drop=True)


def _check_indices(
    dim_df: pd.DataFrame,
    alt_index: pd.Index,
    *,
    dim: Literal["obs", "var"],
    keys: Iterable[str],
    alias_index: pd.Index | None = None,
    use_raw: bool = False,
) -> tuple[list[str], list[str], list[str]]:
    """Check indices for `obs_df` and `var_df`."""
    alt_repr = "adata.raw" if use_raw else "adata"

    alt_dim = ("obs", "var")[dim == "obs"]

    alias_name = None
    if alias_index is not None:
        alt_names = pd.Series(alt_index, index=alias_index)
        alias_name = alias_index.name
        alt_search_repr = f"{alt_dim}['{alias_name}']"
    else:
        alt_names = pd.Series(alt_index, index=alt_index)
        alt_search_repr = f"{alt_dim}_names"

    col_keys = []
    index_keys = []
    index_aliases = []
    not_found = []

    # check that adata.obs does not contain duplicated columns
    # if duplicated columns names are present, they will
    # be further duplicated when selecting them.
    if not dim_df.columns.is_unique:
        dup_cols = dim_df.columns[dim_df.columns.duplicated()].tolist()
        msg = (
            f"adata.{dim} contains duplicated columns. Please rename or remove "
            "these columns first.\n`"
            f"Duplicated columns {dup_cols}"
        )
        raise ValueError(msg)

    if not alt_index.is_unique:
        msg = (
            f"{alt_repr}.{alt_dim}_names contains duplicated items\n"
            f"Please rename these {alt_dim} names first for example using "
            f"`adata.{alt_dim}_names_make_unique()`"
        )
        raise ValueError(msg)

    # use only unique keys, otherwise duplicated keys will
    # further duplicate when reordering the keys later in the function
    for key in dict.fromkeys(keys):
        if key in dim_df.columns:
            col_keys.append(key)
            if key in alt_names.index:
                msg = f"The key {key!r} is found in both adata.{dim} and {alt_repr}.{alt_search_repr}."
                raise KeyError(msg)
        elif key in alt_names.index:
            val = alt_names[key]
            if isinstance(val, pd.Series):
                # while var_names must be unique, adata.var[gene_symbols] does not
                # It's still ambiguous to refer to a duplicated entry though.
                assert alias_index is not None
                msg = f"Found duplicate entries for {key!r} in {alt_repr}.{alt_search_repr}."
                raise KeyError(msg)
            index_keys.append(val)
            index_aliases.append(key)
        else:
            not_found.append(key)
    if len(not_found) > 0:
        msg = (
            f"Could not find keys {not_found!r} in columns of `adata.{dim}` or in"
            f" {alt_repr}.{alt_search_repr}."
        )
        raise KeyError(msg)

    return col_keys, index_keys, index_aliases


def _get_array_values(
    x,
    /,
    dim_names: pd.Index,
    keys: Iterable[str],
    *,
    axis: Literal[0, 1],
    backed: bool,
):
    # TODO: This should be made easier on the anndata side
    mutable_idxer = [slice(None), slice(None)]
    idx = dim_names.get_indexer(keys)

    # for backed AnnData is important that the indices are ordered
    if backed:
        idx_order = np.argsort(idx)
        rev_idxer = mutable_idxer.copy()
        mutable_idxer[axis] = idx[idx_order]
        rev_idxer[axis] = np.argsort(idx_order)
        matrix = x[tuple(mutable_idxer)][tuple(rev_idxer)]
    else:
        mutable_idxer[axis] = idx
        matrix = x[tuple(mutable_idxer)]

    if isinstance(matrix, CSBase):
        matrix = matrix.toarray()

    return matrix


def obs_df(
    adata: AnnData,
    keys: Collection[str] = (),
    obsm_keys: Iterable[tuple[str, int]] = (),
    *,
    layer: str | None = None,
    gene_symbols: str | None = None,
    use_raw: bool = False,
) -> pd.DataFrame:
    """Return values for observations in adata.

    Params
    ------
    adata
        AnnData object to get values from.
    keys
        Keys from either `.var_names`, `.var[gene_symbols]`, or `.obs.columns`.
    obsm_keys
        Tuples of `(key from obsm, column index of obsm[key])`.
    layer
        Layer of `adata` to use as expression values.
    gene_symbols
        Column of `adata.var` to search for `keys` in.
    use_raw
        Whether to get expression values from `adata.raw`.

    Returns
    -------
    A dataframe with `adata.obs_names` as index, and values specified by `keys`
    and `obsm_keys`.

    Examples
    --------
    Getting value for plotting:

    >>> import scanpy as sc
    >>> pbmc = sc.datasets.pbmc68k_reduced()
    >>> plotdf = sc.get.obs_df(
    ...     pbmc, keys=["CD8B", "n_genes"], obsm_keys=[("X_umap", 0), ("X_umap", 1)]
    ... )
    >>> plotdf.columns.astype("string")
    Index(['CD8B', 'n_genes', 'X_umap-0', 'X_umap-1'], dtype='string')
    >>> plotdf.plot.scatter("X_umap-0", "X_umap-1", c="CD8B")  # doctest: +SKIP
    <Axes: xlabel='X_umap-0', ylabel='X_umap-1'>

    Calculating mean expression for marker genes by cluster:

    >>> pbmc = sc.datasets.pbmc68k_reduced()
    >>> marker_genes = ["CD79A", "MS4A1", "CD8A", "CD8B", "LYZ"]
    >>> genedf = sc.get.obs_df(pbmc, keys=["louvain", *marker_genes])
    >>> grouped = genedf.groupby("louvain", observed=True)
    >>> mean, var = grouped.mean(), grouped.var()

    """
    if isinstance(keys, str):
        keys = [keys]
    if use_raw:
        assert layer is None, (
            "Cannot specify use_raw=True and a layer at the same time."
        )
        var = adata.raw.var
    else:
        var = adata.var
    alias_index = pd.Index(var[gene_symbols]) if gene_symbols is not None else None

    obs_cols, var_idx_keys, var_symbols = _check_indices(
        adata.obs,
        var.index,
        dim="obs",
        keys=keys,
        alias_index=alias_index,
        use_raw=use_raw,
    )

    # Make df
    df = pd.DataFrame(index=adata.obs_names)

    # add var values
    if len(var_idx_keys) > 0:
        matrix = _get_array_values(
            _get_obs_rep(adata, layer=layer, use_raw=use_raw),
            var.index,
            var_idx_keys,
            axis=1,
            backed=adata.isbacked,
        )
        df = pd.concat(
            [df, pd.DataFrame(matrix, columns=var_symbols, index=adata.obs_names)],
            axis=1,
        )

    # add obs values
    if len(obs_cols) > 0:
        df = pd.concat([df, adata.obs[obs_cols]], axis=1)

    # reorder columns to given order (including duplicates keys if present)
    if keys:
        df = df[keys]

    for k, idx in obsm_keys:
        added_k = f"{k}-{idx}"
        val = adata.obsm[k]
        if isinstance(val, np.ndarray):
            df[added_k] = np.ravel(val[:, idx])
        elif isinstance(val, CSBase):
            df[added_k] = np.ravel(val[:, idx].toarray())
        elif isinstance(val, pd.DataFrame):
            df[added_k] = val.loc[:, idx]

    return df


def var_df(
    adata: AnnData,
    keys: Collection[str] = (),
    varm_keys: Iterable[tuple[str, int]] = (),
    *,
    layer: str | None = None,
) -> pd.DataFrame:
    """Return values for observations in adata.

    Params
    ------
    adata
        AnnData object to get values from.
    keys
        Keys from either `.obs_names`, or `.var.columns`.
    varm_keys
        Tuples of `(key from varm, column index of varm[key])`.
    layer
        Layer of `adata` to use as expression values.

    Returns
    -------
    A dataframe with `adata.var_names` as index, and values specified by `keys`
    and `varm_keys`.

    """
    # Argument handling
    if isinstance(keys, str):
        keys = [keys]
    var_cols, obs_idx_keys, _ = _check_indices(
        adata.var, adata.obs_names, dim="var", keys=keys
    )

    # initialize df
    df = pd.DataFrame(index=adata.var.index)

    if len(obs_idx_keys) > 0:
        matrix = _get_array_values(
            _get_obs_rep(adata, layer=layer),
            adata.obs_names,
            obs_idx_keys,
            axis=0,
            backed=adata.isbacked,
        ).T
        df = pd.concat(
            [df, pd.DataFrame(matrix, columns=obs_idx_keys, index=adata.var_names)],
            axis=1,
        )

    # add obs values
    if len(var_cols) > 0:
        df = pd.concat([df, adata.var[var_cols]], axis=1)

    # reorder columns to given order
    if keys:
        df = df[keys]

    for k, idx in varm_keys:
        added_k = f"{k}-{idx}"
        val = adata.varm[k]
        if isinstance(val, np.ndarray):
            df[added_k] = np.ravel(val[:, idx])
        elif isinstance(val, CSBase):
            df[added_k] = np.ravel(val[:, idx].toarray())
        elif isinstance(val, pd.DataFrame):
            df[added_k] = val.loc[:, idx]
    return df


class _ObsRep(TypedDict, total=False):
    use_raw: bool
    layer: str | None
    obsm: str | None
    obsp: str | None


def _get_obs_rep(
    adata: AnnData, **choices: Unpack[_ObsRep]
) -> (
    np.ndarray | CSBase | pd.DataFrame | ArrayView | BaseCompressedSparseDataset | None
):
    """Choose array aligned with obs annotation."""
    # https://github.com/scverse/scanpy/issues/1546
    if not isinstance(use_raw := choices.get("use_raw", False), bool):
        msg = f"use_raw expected to be bool, was {type(use_raw)}."
        raise TypeError(msg)
    assert choices.keys() <= {"layer", "use_raw", "obsm", "obsp"}

    # we do this here so the `case _` branch knows which ones are valid for the
    # respective calling function. E.g. `_get_obs_rep(adata, layer="a", obsm="b")`
    # will say that “Only one of `layer` or `obsm` can be specified.”
    match [(k, v) for k, v in choices.items() if v not in {None, False}]:
        case []:
            return adata.X
        # can’t use {"key": v} as match expression, since they allow additional entries
        case [("layer", layer)]:
            return adata.layers[layer]
        case [("use_raw", True)]:
            return adata.raw.X
        case [("obsm", obsm)]:
            return adata.obsm[obsm]
        case [("obsp", obsp)]:
            return adata.obsp[obsp]
        case _:
            valid = [f"`{k}`" for k in choices]
            valid[-1] = f"or {valid[-1]}"
            msg = f"Only one of {', '.join(valid)} can be specified."
            raise ValueError(msg)


def _set_obs_rep(
    adata: AnnData,
    val: Any,
    *,
    use_raw: bool = False,
    layer: str | None = None,
    obsm: str | None = None,
    obsp: str | None = None,
):
    """Set value for observation rep."""
    is_layer = layer is not None
    is_raw = use_raw is not False
    is_obsm = obsm is not None
    is_obsp = obsp is not None
    choices_made = sum((is_layer, is_raw, is_obsm, is_obsp))
    assert choices_made <= 1
    if choices_made == 0:
        adata.X = val
    elif is_layer:
        adata.layers[layer] = val
    elif use_raw:
        adata.raw.X = val
    elif is_obsm:
        adata.obsm[obsm] = val
    elif is_obsp:
        adata.obsp[obsp] = val
    else:
        msg = (
            "That was unexpected. Please report this bug at:\n\n"
            "\thttps://github.com/scverse/scanpy/issues"
        )
        raise AssertionError(msg)


def _check_mask[M: NDArray[np.bool] | NDArray[np.floating] | pd.Series | None](
    data: AnnData | np.ndarray | CSBase | DaskArray,
    mask: str | M,
    dim: Literal["obs", "var"],
    *,
    allow_probabilities: bool = False,
) -> M:  # Could also be a series, but should be one or the other
    """Validate mask argument.

    Params
    ------
    data
        Annotated data matrix or numpy array.
    mask
        Mask (or probabilities if `allow_probabilities=True`).
        Either an appropriatley sized array, or name of a column.
    dim
        The dimension being masked.
    allow_probabilities
        Whether to allow probabilities as `mask`
    """
    if mask is None:
        return mask
    desc = "mask/probabilities" if allow_probabilities else "mask"

    if isinstance(mask, str):
        if not isinstance(data, AnnData):
            msg = f"Cannot refer to {desc} with string without providing anndata object as argument"
            raise ValueError(msg)

        annot: pd.DataFrame = getattr(data, dim)
        if mask not in annot.columns:
            msg = (
                f"Did not find `adata.{dim}[{mask!r}]`. "
                f"Either add the {desc} first to `adata.{dim}`"
                f"or consider using the {desc} argument with an array."
            )
            raise ValueError(msg)
        mask_array = annot[mask].to_numpy()
    else:
        if len(mask) != data.shape[0 if dim == "obs" else 1]:
            msg = f"The shape of the {desc} do not match the data."
            raise ValueError(msg)
        mask_array = mask

    is_bool = pd.api.types.is_bool_dtype(mask_array.dtype)
    if not allow_probabilities and not is_bool:
        msg = "Mask array must be boolean."
        raise ValueError(msg)
    elif allow_probabilities and not (
        is_bool or pd.api.types.is_float_dtype(mask_array.dtype)
    ):
        msg = f"{desc} array must be boolean or floating point."
        raise ValueError(msg)

    return mask_array