diff --git a/docs/api/get.md b/docs/api/get.md index 82f74ed411..039769d0a6 100644 --- a/docs/api/get.md +++ b/docs/api/get.md @@ -19,5 +19,6 @@ useful formats. get.obs_df get.var_df get.rank_genes_groups_df + get.aggregate ``` diff --git a/docs/release-notes/1.10.0.md b/docs/release-notes/1.10.0.md index 542ee15fd9..0266b2bdf9 100644 --- a/docs/release-notes/1.10.0.md +++ b/docs/release-notes/1.10.0.md @@ -15,6 +15,7 @@ * Enhanced dask support for some internal utilities, paving the way for more extensive dask support {pr}`2696` {smaller}`P Angerer` * {func}`scanpy.pp.pca`, {func}`scanpy.pp.scale`, {func}`scanpy.pl.embedding`, and {func}`scanpy.experimental.pp.normalize_pearson_residuals_pca` now support a `mask` parameter {pr}`2272` {smaller}`C Bright, T Marcella, & P Angerer` +* New function {func}`scanpy.get.aggregate` which allows grouped aggregations over your data. Useful for pseudobulking! {pr}`2590` {smaller}`Isaac Virshup` {smaller}`Ilan Gold` {smaller}`Jon Bloom` * {func}`scanpy.tl.rank_genes_groups` no longer warns that it's default was changed from t-test_overestim_var to t-test {pr}`2798` {smaller}`L Heumos` * {func}`scanpy.tl.leiden` now offers `igraph`'s implementation of the leiden algorithm via via `flavor` when set to `igraph`. `leidenalg`'s implementation is still default, but discouraged. {pr}`2815` {smaller}`I Gold` * {func}`scanpy.pp.highly_variable_genes` has new flavor `seurat_v3_paper` that is in its implementation consistent with the paper description in Stuart et al 2018. {pr}`2792` {smaller}`E Roellin` diff --git a/scanpy/_utils/__init__.py b/scanpy/_utils/__init__.py index 625e3b71e3..3ef16300aa 100644 --- a/scanpy/_utils/__init__.py +++ b/scanpy/_utils/__init__.py @@ -874,3 +874,13 @@ def _choose_graph(adata, obsp, neighbors_key): "to compute a neighborhood graph." ) return neighbors["connectivities"] + + +def _resolve_axis( + axis: Literal["obs", 0, "var", 1], +) -> tuple[Literal[0], Literal["obs"]] | tuple[Literal[1], Literal["var"]]: + if axis in {0, "obs"}: + return (0, "obs") + if axis in {1, "var"}: + return (1, "var") + raise ValueError(f"`axis` must be either 0, 1, 'obs', or 'var', was {axis!r}") diff --git a/scanpy/get/__init__.py b/scanpy/get/__init__.py index 08567cfc6e..56c0d3c130 100644 --- a/scanpy/get/__init__.py +++ b/scanpy/get/__init__.py @@ -1,7 +1,6 @@ -# Public -# Private from __future__ import annotations +from ._aggregated import aggregate from .get import ( _check_mask, _get_obs_rep, @@ -15,6 +14,7 @@ "_check_mask", "_get_obs_rep", "_set_obs_rep", + "aggregate", "obs_df", "rank_genes_groups_df", "var_df", diff --git a/scanpy/get/_aggregated.py b/scanpy/get/_aggregated.py new file mode 100644 index 0000000000..159f7572e3 --- /dev/null +++ b/scanpy/get/_aggregated.py @@ -0,0 +1,387 @@ +from __future__ import annotations + +from functools import singledispatch +from typing import TYPE_CHECKING, Literal, Union, get_args + +import numpy as np +import pandas as pd +from anndata import AnnData, utils +from scipy import sparse + +from .._utils import _resolve_axis +from .get import _check_mask + +if TYPE_CHECKING: + from collections.abc import Collection, Iterable + + from numpy.typing import NDArray + +Array = Union[np.ndarray, sparse.csc_matrix, sparse.csr_matrix] +AggType = Literal["count_nonzero", "mean", "sum", "var"] + + +class Aggregate: + """\ + Functionality for generic grouping and aggregating. + + There is currently support for count_nonzero, sum, mean, and variance. + + **Implementation** + + Moments are computed using weighted sum aggregation of data by some feature + via multiplication by a sparse coordinate matrix A. + + Runtime is effectively computation of the product `A @ X`, i.e. the count of (non-zero) + entries in X with multiplicity the number of group memberships for that entry. + This is `O(data)` for partitions (each observation belonging to exactly one group), + independent of the number of groups. + + Params + ------ + groupby + :class:`~pandas.Categorical` containing values for grouping by. + data + Data matrix for aggregation. + mask + Mask to be used for aggregation. + """ + + def __init__( + self, + groupby: pd.Categorical, + data: Array, + *, + mask: NDArray[np.bool_] | None = None, + ) -> None: + self.groupby = groupby + self.indicator_matrix = sparse_indicator(groupby, mask=mask) + self.data = data + + groupby: pd.Categorical + indicator_matrix: sparse.coo_matrix + data: Array + + def count_nonzero(self) -> NDArray[np.integer]: + """\ + Count the number of observations in each group. + + Returns + ------- + Array of counts. + """ + # pattern = self.data._with_data(np.broadcast_to(1, len(self.data.data))) + # return self.indicator_matrix @ pattern + return self.indicator_matrix @ (self.data != 0) + + def sum(self) -> Array: + """\ + Compute the sum per feature per group of observations. + + Returns + ------- + Array of sum. + """ + return utils.asarray(self.indicator_matrix @ self.data) + + def mean(self) -> Array: + """\ + Compute the mean per feature per group of observations. + + Returns + ------- + Array of mean. + """ + return ( + utils.asarray(self.indicator_matrix @ self.data) + / np.bincount(self.groupby.codes)[:, None] + ) + + def mean_var(self, dof: int = 1) -> tuple[np.ndarray, np.ndarray]: + """\ + Compute the count, as well as mean and variance per feature, per group of observations. + + The formula `Var(X) = E(X^2) - E(X)^2` suffers loss of precision when the variance is a + very small fraction of the squared mean. In particular, when X is constant, the formula may + nonetheless be non-zero. By default, our implementation resets the variance to exactly zero + when the computed variance, relative to the squared mean, nears limit of precision of the + floating-point significand. + + Params + ------ + dof + Degrees of freedom for variance. + + Returns + ------- + Object with `count`, `mean`, and `var` attributes. + """ + assert dof >= 0 + + group_counts = np.bincount(self.groupby.codes) + mean_ = self.mean() + # sparse matrices do not support ** for elementwise power. + mean_sq = ( + utils.asarray(self.indicator_matrix @ _power(self.data, 2)) + / group_counts[:, None] + ) + sq_mean = mean_**2 + var_ = mean_sq - sq_mean + # TODO: Why these values exactly? Because they are high relative to the datatype? + # (unchanged from original code: https://github.com/scverse/anndata/pull/564) + precision = 2 << (42 if self.data.dtype == np.float64 else 20) + # detects loss of precision in mean_sq - sq_mean, which suggests variance is 0 + var_[precision * var_ < sq_mean] = 0 + if dof != 0: + var_ *= (group_counts / (group_counts - dof))[:, np.newaxis] + return mean_, var_ + + +def _power(X: Array, power: float | int) -> Array: + """\ + Generate elementwise power of a matrix. + + Needed for non-square sparse matrices because they do not support `**` so the `.power` function is used. + + Params + ------ + X + Matrix whose power is to be raised. + power + Integer power value + + Returns + ------- + Matrix whose power has been raised. + """ + return X**power if isinstance(X, np.ndarray) else X.power(power) + + +@singledispatch +def aggregate( + adata: AnnData, + by: str | Collection[str], + func: AggType | Iterable[AggType], + *, + axis: Literal["obs", 0, "var", 1] | None = None, + mask: NDArray[np.bool_] | str | None = None, + dof: int = 1, + layer: str | None = None, + obsm: str | None = None, + varm: str | None = None, +) -> AnnData: + """\ + Aggregate data matrix based on some categorical grouping. + + This function is useful for pseudobulking as well as plotting. + + Aggregation to perform is specified by `func`, which can be a single metric or a + list of metrics. Each metric is computed over the group and results in a new layer + in the output `AnnData` object. + + If none of `layer`, `obsm`, or `varm` are passed in, `X` will be used for aggregation data. + If `func` only has length 1 or is just an `AggType`, then aggregation data is written to `X`. + Otherwise, it is written to `layers` or `xxxm` as appropriate for the dimensions of the aggregation data. + + Params + ------ + adata + :class:`~anndata.AnnData` to be aggregated. + by + Key of the column to be grouped-by. + func + How to aggregate. + axis + Axis on which to find group by column. + mask + Boolean mask (or key to column containing mask) to apply along the axis. + dof + Degrees of freedom for variance. Defaults to 1. + layer + If not None, key for aggregation data. + obsm + If not None, key for aggregation data. + varm + If not None, key for aggregation data. + + Returns + ------- + Aggregated :class:`~anndata.AnnData`. + + Examples + -------- + + Calculating mean expression and number of nonzero entries per cluster: + + >>> import scanpy as sc, pandas as pd + >>> pbmc = sc.datasets.pbmc3k_processed().raw.to_adata() + >>> pbmc.shape + (2638, 13714) + >>> aggregated = sc.get.aggregate(pbmc, by="louvain", func=["mean", "count_nonzero"]) + >>> aggregated + AnnData object with n_obs × n_vars = 8 × 13714 + obs: 'louvain' + var: 'n_cells' + layers: 'mean', 'count_nonzero' + + We can group over multiple columns: + + >>> pbmc.obs["percent_mito_binned"] = pd.cut(pbmc.obs["percent_mito"], bins=5) + >>> sc.get.aggregate(pbmc, by=["louvain", "percent_mito_binned"], func=["mean", "count_nonzero"]) + AnnData object with n_obs × n_vars = 40 × 13714 + obs: 'louvain', 'percent_mito_binned' + var: 'n_cells' + layers: 'mean', 'count_nonzero' + + Note that this filters out any combination of groups that wasn't present in the original data. + """ + if axis is None: + axis = 1 if varm else 0 + axis, axis_name = _resolve_axis(axis) + if mask is not None: + mask = _check_mask(adata, mask, axis_name) + data = adata.X + if sum(p is not None for p in [varm, obsm, layer]) > 1: + raise TypeError("Please only provide one (or none) of varm, obsm, or layer") + + if varm is not None: + if axis != 1: + raise ValueError("varm can only be used when axis is 1") + data = adata.varm[varm] + elif obsm is not None: + if axis != 0: + raise ValueError("obsm can only be used when axis is 0") + data = adata.obsm[obsm] + elif layer is not None: + data = adata.layers[layer] + if axis == 1: + data = data.T + elif axis == 1: + # i.e., all of `varm`, `obsm`, `layers` are None so we use `X` which must be transposed + data = data.T + + dim_df = getattr(adata, axis_name) + categorical, new_label_df = _combine_categories(dim_df, by) + # Actual computation + layers = aggregate( + data, + by=categorical, + func=func, + mask=mask, + dof=dof, + ) + result = AnnData( + layers=layers, + obs=new_label_df, + var=getattr(adata, "var" if axis == 0 else "obs"), + ) + + if axis == 1: + return result.T + else: + return result + + +@aggregate.register(np.ndarray) +@aggregate.register(sparse.spmatrix) +def aggregate_array( + data, + by: pd.Categorical, + func: AggType | Iterable[AggType], + *, + mask: NDArray[np.bool_] | None = None, + dof: int = 1, +) -> dict[AggType, np.ndarray]: + groupby = Aggregate(groupby=by, data=data, mask=mask) + result = {} + + funcs = set([func] if isinstance(func, str) else func) + if unknown := funcs - set(get_args(AggType)): + raise ValueError(f"func {unknown} is not one of {get_args(AggType)}") + + if "sum" in funcs: # sum is calculated separately from the rest + agg = groupby.sum() + result["sum"] = agg + # here and below for count, if var is present, these can be calculate alongside var + if "mean" in funcs and "var" not in funcs: + agg = groupby.mean() + result["mean"] = agg + if "count_nonzero" in funcs: + result["count_nonzero"] = groupby.count_nonzero() + if "var" in funcs: + mean_, var_ = groupby.mean_var(dof) + result["var"] = var_ + if "mean" in funcs: + result["mean"] = mean_ + + return result + + +def _combine_categories( + label_df: pd.DataFrame, cols: Collection[str] | str +) -> tuple[pd.Categorical, pd.DataFrame]: + """ + Returns both the result categories and a dataframe labelling each row + """ + from itertools import product + + if isinstance(cols, str): + cols = [cols] + + df = pd.DataFrame( + {c: pd.Categorical(label_df[c]).remove_unused_categories() for c in cols}, + ) + n_categories = [len(df[c].cat.categories) for c in cols] + + # It's like np.concatenate([x for x in product(*[range(n) for n in n_categories])]) + code_combinations = np.indices(n_categories).reshape(len(n_categories), -1) + result_categories = pd.Index( + ["_".join(map(str, x)) for x in product(*[df[c].cat.categories for c in cols])] + ) + + # Dataframe with unique combination of categories for each row + new_label_df = pd.DataFrame( + { + c: pd.Categorical.from_codes(code_combinations[i], df[c].cat.categories) + for i, c in enumerate(cols) + }, + index=result_categories, + ) + + # Calculating result codes + factors = np.ones(len(cols) + 1, dtype=np.int32) # First factor needs to be 1 + np.cumsum(n_categories[::-1], out=factors[1:]) + factors = factors[:-1][::-1] + + code_array = np.zeros((len(cols), df.shape[0]), dtype=np.int32) + for i, c in enumerate(cols): + code_array[i] = df[c].cat.codes + code_array *= factors[:, None] + + result_categorical = pd.Categorical.from_codes( + code_array.sum(axis=0), categories=result_categories + ) + + # Filter unused categories + result_categorical = result_categorical.remove_unused_categories() + new_label_df = new_label_df.loc[result_categorical.categories] + + return result_categorical, new_label_df + + +def sparse_indicator( + categorical: pd.Categorical, + *, + mask: NDArray[np.bool_] | None = None, + weight: NDArray[np.floating] | None = None, +) -> sparse.coo_matrix: + if mask is not None and weight is None: + weight = mask.astype(np.float32) + elif mask is not None and weight is not None: + weight = mask * weight + elif mask is None and weight is None: + weight = np.broadcast_to(1.0, len(categorical)) + A = sparse.coo_matrix( + (weight, (categorical.codes, np.arange(len(categorical)))), + shape=(len(categorical.categories), len(categorical)), + ) + return A diff --git a/scanpy/tests/test_aggregated.py b/scanpy/tests/test_aggregated.py new file mode 100644 index 0000000000..4439f3e0eb --- /dev/null +++ b/scanpy/tests/test_aggregated.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import anndata as ad +import numpy as np +import pandas as pd +import pytest +from scipy.sparse import csr_matrix + +import scanpy as sc +from scanpy._utils import _resolve_axis +from scanpy.testing._helpers import assert_equal +from scanpy.testing._helpers.data import pbmc3k_processed +from scanpy.testing._pytest.params import ARRAY_TYPES_MEM + + +@pytest.fixture +def df_base(): + ax_base = ["A", "B"] + return pd.DataFrame(index=ax_base) + + +@pytest.fixture +def df_groupby(): + ax_groupby = [ + *["v0", "v1", "v2"], + *["w0", "w1"], + *["a1", "a2", "a3"], + *["b1", "b2"], + *["c1", "c2"], + "d0", + ] + + df_groupby = pd.DataFrame(index=pd.Index(ax_groupby, name="cell")) + df_groupby["key"] = pd.Categorical([c[0] for c in ax_groupby]) + df_groupby["key_superset"] = pd.Categorical([c[0] for c in ax_groupby]).map( + {"v": "v", "w": "v", "a": "a", "b": "a", "c": "a", "d": "a"} + ) + df_groupby["key_subset"] = pd.Categorical([c[1] for c in ax_groupby]) + df_groupby["weight"] = 2.0 + return df_groupby + + +@pytest.fixture +def X(): + data = [ + *[[0, -2], [1, 13], [2, 1]], # v + *[[3, 12], [4, 2]], # w + *[[5, 11], [6, 3], [7, 10]], # a + *[[8, 4], [9, 9]], # b + *[[10, 5], [11, 8]], # c + [12, 6], # d + ] + return np.array(data, dtype=np.float32) + + +def gen_adata(data_key, dim, df_base, df_groupby, X): + if (data_key == "varm" and dim == "obs") or (data_key == "obsm" and dim == "var"): + pytest.skip("invalid parameter combination") + + obs_df, var_df = (df_groupby, df_base) if dim == "obs" else (df_base, df_groupby) + data = X.T if dim == "var" and data_key != "varm" else X + if data_key != "X": + data_dict_sparse = {data_key: {"test": csr_matrix(data)}} + data_dict_dense = {data_key: {"test": data}} + else: + data_dict_sparse = {data_key: csr_matrix(data)} + data_dict_dense = {data_key: data} + + adata_sparse = ad.AnnData(obs=obs_df, var=var_df, **data_dict_sparse) + adata_dense = ad.AnnData(obs=obs_df, var=var_df, **data_dict_dense) + return adata_sparse, adata_dense + + +@pytest.mark.parametrize("axis", [0, 1]) +def test_mask(axis): + blobs = sc.datasets.blobs() + mask = blobs.obs["blobs"] == 0 + blobs.obs["mask_col"] = mask + if axis == 1: + blobs = blobs.T + by_name = sc.get.aggregate(blobs, "blobs", "sum", axis=axis, mask="mask_col") + by_value = sc.get.aggregate(blobs, "blobs", "sum", axis=axis, mask=mask) + + assert_equal(by_name, by_value) + + assert np.all(by_name["0"].layers["sum"] == 0) + + +@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM) +@pytest.mark.parametrize("metric", ["sum", "mean", "var", "count_nonzero"]) +def test_aggregate_vs_pandas(metric, array_type): + adata = pbmc3k_processed().raw.to_adata() + adata = adata[ + adata.obs["louvain"].isin(adata.obs["louvain"].cat.categories[:5]), :1_000 + ].copy() + adata.X = array_type(adata.X) + adata.obs["percent_mito_binned"] = pd.cut(adata.obs["percent_mito"], bins=5) + result = sc.get.aggregate(adata, ["louvain", "percent_mito_binned"], metric) + + if metric == "count_nonzero": + expected = ( + (adata.to_df() != 0) + .astype(np.float64) + .join(adata.obs[["louvain", "percent_mito_binned"]]) + .groupby(["louvain", "percent_mito_binned"], observed=True) + .agg("sum") + ) + else: + expected = ( + adata.to_df() + .astype(np.float64) + .join(adata.obs[["louvain", "percent_mito_binned"]]) + .groupby(["louvain", "percent_mito_binned"], observed=True) + .agg(metric) + ) + # TODO: figure out the axis names + expected.index = expected.index.to_frame().apply( + lambda x: "_".join(map(str, x)), axis=1 + ) + expected.index.name = None + expected.columns.name = None + + result_df = result.to_df(layer=metric) + result_df.index.name = None + result_df.columns.name = None + + pd.testing.assert_frame_equal(result_df, expected, check_dtype=False, atol=1e-5) + + +@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM) +@pytest.mark.parametrize("metric", ["sum", "mean", "var", "count_nonzero"]) +def test_aggregate_axis(array_type, metric): + adata = pbmc3k_processed().raw.to_adata() + adata = adata[ + adata.obs["louvain"].isin(adata.obs["louvain"].cat.categories[:5]), :1_000 + ].copy() + adata.X = array_type(adata.X) + expected = sc.get.aggregate(adata, ["louvain"], metric) + actual = sc.get.aggregate(adata.T, ["louvain"], metric, axis=1).T + + assert_equal(expected, actual) + + +def test_aggregate_entry(): + args = ("blobs", ["mean", "var", "count_nonzero"]) + + adata = sc.datasets.blobs() + X_result = sc.get.aggregate(adata, *args) + # layer adata + layer_adata = ad.AnnData( + obs=adata.obs, + var=adata.var, + layers={"test": adata.X.copy()}, + ) + layer_result = sc.get.aggregate(layer_adata, *args, layer="test") + obsm_adata = ad.AnnData( + obs=adata.obs, + var=adata.var, + obsm={"test": adata.X.copy()}, + ) + obsm_result = sc.get.aggregate(obsm_adata, *args, obsm="test") + varm_adata = ad.AnnData( + obs=adata.var, + var=adata.obs, + varm={"test": adata.X.copy()}, + ) + varm_result = sc.get.aggregate(varm_adata, *args, varm="test") + + X_result_min = X_result.copy() + del X_result_min.var + X_result_min.var_names = [str(x) for x in np.arange(X_result_min.n_vars)] + + assert_equal(X_result, layer_result) + assert_equal(X_result_min, obsm_result) + assert_equal(X_result.layers, obsm_result.layers) + assert_equal(X_result.layers, varm_result.T.layers) + + +def test_aggregate_incorrect_dim(): + adata = pbmc3k_processed().raw.to_adata() + + with pytest.raises(ValueError, match="was 'foo'"): + sc.get.aggregate(adata, ["louvain"], "sum", axis="foo") + + +@pytest.mark.parametrize("axis_name", ["obs", "var"]) +def test_aggregate_axis_specification(axis_name): + axis, axis_name = _resolve_axis(axis_name) + by = "blobs" if axis == 0 else "labels" + + adata = sc.datasets.blobs() + adata.var["labels"] = np.tile(["a", "b"], adata.shape[1])[: adata.shape[1]] + + agg_index = sc.get.aggregate(adata, by=by, func="mean", axis=axis) + agg_name = sc.get.aggregate(adata, by=by, func="mean", axis=axis_name) + + np.testing.assert_equal(agg_index.layers["mean"], agg_name.layers["mean"]) + + if axis_name == "obs": + agg_unspecified = sc.get.aggregate(adata, by=by, func="mean") + np.testing.assert_equal(agg_name.layers["mean"], agg_unspecified.layers["mean"]) + + +@pytest.mark.parametrize( + ("matrix", "df", "keys", "metrics", "expected"), + [ + pytest.param( + np.block( + [ + [np.ones((2, 2)), np.zeros((2, 2))], + [np.zeros((2, 2)), np.ones((2, 2))], + ] + ), + pd.DataFrame( + { + "a": ["a", "a", "b", "b"], + "b": ["c", "d", "d", "d"], + } + ), + ["a", "b"], + ["count_nonzero"], # , "sum", "mean"], + ad.AnnData( + obs=pd.DataFrame( + {"a": ["a", "a", "b"], "b": ["c", "d", "d"]}, + index=["a_c", "a_d", "b_d"], + ).astype("category"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(4)]), + layers={ + "count_nonzero": np.array( + [[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 2, 2]] + ), + # "sum": np.array([[2, 0], [0, 2]]), + # "mean": np.array([[1, 0], [0, 1]]), + }, + ), + id="count_nonzero", + ), + pytest.param( + np.block( + [ + [np.ones((2, 2)), np.zeros((2, 2))], + [np.zeros((2, 2)), np.ones((2, 2))], + ] + ), + pd.DataFrame( + { + "a": ["a", "a", "b", "b"], + "b": ["c", "d", "d", "d"], + } + ), + ["a", "b"], + ["sum", "mean", "count_nonzero"], + ad.AnnData( + obs=pd.DataFrame( + {"a": ["a", "a", "b"], "b": ["c", "d", "d"]}, + index=["a_c", "a_d", "b_d"], + ).astype("category"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(4)]), + layers={ + "sum": np.array([[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 2, 2]]), + "mean": np.array([[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 1, 1]]), + "count_nonzero": np.array( + [[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 2, 2]] + ), + }, + ), + id="sum-mean-count_nonzero", + ), + pytest.param( + np.block( + [ + [np.ones((2, 2)), np.zeros((2, 2))], + [np.zeros((2, 2)), np.ones((2, 2))], + ] + ), + pd.DataFrame( + { + "a": ["a", "a", "b", "b"], + "b": ["c", "d", "d", "d"], + } + ), + ["a", "b"], + ["mean"], + ad.AnnData( + obs=pd.DataFrame( + {"a": ["a", "a", "b"], "b": ["c", "d", "d"]}, + index=["a_c", "a_d", "b_d"], + ).astype("category"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(4)]), + layers={ + "mean": np.array([[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 1, 1]]), + }, + ), + id="mean", + ), + ], +) +def test_aggregate_examples(matrix, df, keys, metrics, expected): + adata = ad.AnnData( + X=matrix, + obs=df, + var=pd.DataFrame(index=[f"gene_{i}" for i in range(matrix.shape[1])]), + ) + result = sc.get.aggregate(adata, by=keys, func=metrics) + + print(result) + print(expected) + + assert_equal(expected, result) + + +@pytest.mark.parametrize( + ("label_cols", "cols", "expected"), + [ + pytest.param( + dict( + a=pd.Categorical(["a", "b", "c"]), + b=pd.Categorical(["d", "d", "f"]), + ), + ["a", "b"], + pd.Categorical(["a_d", "b_d", "c_f"]), + id="two_of_two", + ), + pytest.param( + dict( + a=pd.Categorical(["a", "b", "c"]), + b=pd.Categorical(["d", "d", "f"]), + c=pd.Categorical(["g", "h", "h"]), + ), + ["a", "b", "c"], + pd.Categorical(["a_d_g", "b_d_h", "c_f_h"]), + id="three_of_three", + ), + pytest.param( + dict( + a=pd.Categorical(["a", "b", "c"]), + b=pd.Categorical(["d", "d", "f"]), + c=pd.Categorical(["g", "h", "h"]), + ), + ["a", "c"], + pd.Categorical(["a_g", "b_h", "c_h"]), + id="two_of_three-1", + ), + pytest.param( + dict( + a=pd.Categorical(["a", "b", "c"]), + b=pd.Categorical(["d", "d", "f"]), + c=pd.Categorical(["g", "h", "h"]), + ), + ["b", "c"], + pd.Categorical(["d_g", "d_h", "f_h"]), + id="two_of_three-2", + ), + ], +) +def test_combine_categories(label_cols, cols, expected): + from scanpy.get._aggregated import _combine_categories + + label_df = pd.DataFrame(label_cols) + result, result_label_df = _combine_categories(label_df, cols) + + assert isinstance(result, pd.Categorical) + + pd.testing.assert_extension_array_equal(result, expected) + + pd.testing.assert_index_equal( + pd.Index(result), result_label_df.index.astype("category") + ) + + reconstructed_df = pd.DataFrame( + [x.split("_") for x in result], columns=cols, index=result.astype(str) + ).astype("category") + pd.testing.assert_frame_equal(reconstructed_df, result_label_df)