Keeping original categories when subsetting an AnnData object

[LucaMarconato]

If one subsets an AnnData object that has a categorical obs column, the new categories of the column will be only the one that appear in the subset column, an not all the original one.

Is this a bug or the expected behavior?
If it's not a bug, is there an helper function/standard way to keep all the original categories?

[LucaMarconato]

Asking also @giovp since when plotting categorical values in Squidpy could have had to deal with this behavior.

[LucaMarconato]

Please notice that when subsetting a DataFrame all the categories are kept, not just the ones appearing in the subset data.

[LucaMarconato]

Created a simple function to fix this:

def _inplace_fix_subset_categorical_obs(subset_adata: AnnData, original_adata: AnnData) -> None:
    """
    Fix categorical obs columns of subset_adata to match the categories of original_adata.

    Parameters
    ----------
    subset_adata
        The subset AnnData object
    original_adata
        The original AnnData object

    Notes
    -----
    See discussion here: https://github.com/scverse/anndata/issues/997
    """
    obs = subset_adata.obs
    for column in obs.columns:
        is_categorical = pd.api.types.is_categorical_dtype(obs[column])
        if is_categorical:
            c = obs[column].cat.set_categories(original_adata.obs[column].cat.categories)
            obs[column] = c

[giovp]

hat's the purpose of this? cause if it is for the colormap saved in adata.uns[categ_colors] then that is also updated.

[LucaMarconato]

I use it in the spatialdata code for aggregate() when aggregating a categorical column: when performing aggregation on subsets of rows I want to get as a result a matrix always with the same number of columns (one per category). But due to the behavior described in this issue, if I don't call the workaround function that I wrote above, the categories are subset and the output of aggregate() are AnnData objects with different columns.

I also had a related problem in the past. When initializing colors with scanpy for a categorical column, the colors where different after subsetting the AnnData object since some categories resulted being dropped.

[flying-sheep]

This behavior is intentional, as anndata explicitly calls a method named _remove_unused_categories. The question is why the behavior is there, as it happened in a commit that doesn’t have an associated PR with a rationale (8cabf9c)

anndata/anndata/_core/anndata.py

Lines 343 to 346 in e067515

	# fix categories
	uns = copy(adata_ref._uns)
	self._remove_unused_categories(adata_ref.obs, obs_sub, uns)
	self._remove_unused_categories(adata_ref.var, var_sub, uns)

[ivirshup]

@flying-sheep I believe we have this behavior since we were early adopters of categorical values, and pandas threw a ton of errors if you passed categorical arrays where there was no instance of a particular category. However, pandas also did not remove categories when subsetting, nor did it make doing that easy.

I think the situation is much improved, and we should figure out how to deprecate this behavior.

[ivirshup]

@LucaMarconato

I also had a related problem in the past. When initializing colors with scanpy for a categorical column, the colors were different after subsetting the AnnData object since some categories resulted being dropped.

The colors saved in the anndata should be updated so they don't change when unused categoricals are removed.

However, we should probably just save the mapping of categories to colors as a dict, and not need to bother keeping these things synced.

[ivirshup]

This has previously been requested and tracked in #890, so I'm going to close this as duplicated