Source code for campmc.mt.sizes

"""Per-metacell size tables from partition labels."""

from __future__ import annotations

import pandas as pd
from anndata import AnnData




def metacell_sizes(adata: AnnData, *, partition_key: str = "camp") -> pd.DataFrame:
    """Return per-metacell cell counts from ``adata.obs``.

    Parameters
    ----------
    adata
        AnnData with partition labels in ``obs``.
    partition_key
        Column in ``adata.obs`` with metacell labels.

    Returns
    -------
    DataFrame indexed by metacell id with a ``size`` column.

    Examples
    --------

    .. exec-jupyter::

     import campmc as cp
     import scanpy as sc
     adata = sc.datasets.pbmc68k_reduced()
     cp.partition(adata, method="camp3", gamma=50, random_state=0)
     cp.mt.metacell_sizes(adata).head()

    """
    if partition_key not in adata.obs.columns:
        raise ValueError(
            f"adata.obs is missing partition column '{partition_key}'. Run campmc.partition before computing sizes."
        )

    counts = adata.obs[partition_key].value_counts().sort_index()
    return pd.DataFrame({"size": counts}, index=counts.index.rename(partition_key))