"""Aggregate single cells into well- or perturbation-level profiles."""
from __future__ import annotations
from collections.abc import Sequence
from copy import deepcopy
import anndata as ad
import numpy as np
import pandas as pd
from anndata import AnnData
from mantispy._core._numba import MEAN, MEDIAN
from mantispy._core._reduce import group_codes, matrix_dtype, reduce_grouped, reduced_var
from mantispy._core.frames import as_frame, categorize_metadata
from mantispy._core.logging import get_logger
from mantispy._core.provenance import record_params
from mantispy._core.schema import REQUIRED_UNS, get_resolution, resolution_for, resolve_control_column, stamp
FUNCTIONS = {"median": MEDIAN, "mean": MEAN}
#: Keys describing the features or the experiment; result tables are keyed on the input rows, so they are dropped.
_INHERITED = frozenset({"feature_select", "blocklist"})
#: Columns that add up over a group, so they are recomputed for it rather than carried as constant metadata.
_TALLIES = ("Metadata_CellCount", "Metadata_SiteCount", "Metadata_ReplicateCount")
[docs]
def aggregate(
adata: AnnData,
by: Sequence[str] = ("Metadata_Plate", "Metadata_Well"),
func: str = "median",
min_cells: int = 10,
layer: str | None = None,
use_rep: str | None = None,
count_key: str = "Metadata_CellCount",
site_key: str = "Metadata_SiteCount",
) -> AnnData:
"""Aggregate ``adata`` to one profile per group.
Args:
adata: Single cells, or profiles to aggregate further, as its recorded resolution says.
by: Columns defining a profile.
The default is one profile per well.
func: ``"median"`` (the pycytominer default) or ``"mean"``.
min_cells: Groups with fewer cells than this are dropped.
layer: Aggregate this layer instead of ``X``.
use_rep: Aggregate this ``obsm`` representation (e.g. an embedding from ``pp.tvn``/``pp.harmony``) instead of ``X``; the result's ``X`` holds the reduced representation and ``var`` is a plain range index, since the axes are not named features.
Mutually exclusive with ``layer``.
count_key: ``obs`` column the cell count is written to, and read from when ``adata`` holds profiles.
site_key: ``obs`` column the number of fields of view is written to, and read from when ``adata`` holds profiles.
Returns:
A new :class:`~anndata.AnnData` with one row per group.
Rows are named by their group keys joined with ``_``, e.g. ``"plate1_A01"``.
``var`` is carried over unchanged; ``obs`` holds the grouping columns, `count_key`, `site_key` when the fields of view are known, and every other ``obs`` column (not only ``Metadata_`` ones) that is constant within every group.
`count_key` is the number of cells behind a row, so its scope follows ``by``: grouping by site counts the cells of one field of view, grouping by well those of every field.
Profiles contribute the cells they carry rather than one each, and profiles that carry no count give an unknown one.
`site_key` is the number of fields that contributed cells, summed where the rows carry it and counted from ``Metadata_Site`` otherwise.
The resolution recorded is ``"well"`` when ``by`` holds both ``Metadata_Plate`` and ``Metadata_Well``, since a finer grouping such as one row per site is still per-well or finer, and ``"aggregate"`` otherwise.
Raises:
ValueError: ``func`` is not one of ``FUNCTIONS``, or ``use_rep`` and ``layer`` are both given, or ``use_rep`` is not a 2-D representation in ``obsm``.
Notes:
This uses mantispy's own NaN-skipping kernel rather than :func:`scanpy.get.aggregate`, which propagates NaN and is slower on both mean and median.
"""
if func not in FUNCTIONS:
raise ValueError(f"func must be one of {tuple(FUNCTIONS)}, got {func!r}")
# use_rep/layer mutual exclusion and obsm validation are enforced once, in reduce_grouped's _obsm_source.
columns = [by] if isinstance(by, str) else list(by)
codes, keys = group_codes(adata, columns)
values, _, counts = reduce_grouped(
adata, columns, FUNCTIONS[func], layer=layer, use_rep=use_rep, grouping=(codes, keys)
)
frame = as_frame(adata.obs)
tallies = {count_key: counts}
# The recorded resolution, not the column, decides count handling: an object may carry its well's
# count as a covariate. get_resolution raises on an unstamped object rather than assume object-level.
if get_resolution(adata) != "object":
tallies[count_key] = np.full(len(keys), np.nan)
for column in (count_key, site_key):
if column in frame:
tallies[column] = np.bincount(codes, weights=frame[column].to_numpy(dtype=float), minlength=len(keys))
elif "Metadata_Site" in frame:
tallies[site_key] = _site_counts(frame, codes, len(keys))
obs = _group_obs(adata, columns, keys, codes, tallies)
# A group whose count is unknown is kept rather than dropped as too small.
keep = ~(tallies[count_key] < min_cells)
if not keep.any():
get_logger().warning("aggregate: dropped every group, since min_cells=%d exceeds every group size", min_cells)
_warn_dropped_controls(obs, keep, resolve_control_column(adata), min_cells)
obs = obs.loc[keep].reset_index(drop=True)
obs.index = _group_names(obs, columns)
var = reduced_var(adata, use_rep, values.shape[1])
result = ad.AnnData(X=values[keep].astype(matrix_dtype(adata, layer, use_rep=use_rep)), obs=obs, var=var)
store = adata.uns.get("mantispy", {})
resolution = resolution_for(columns)
stamp(result, resolution=resolution, grouped_by=columns)
result.uns["mantispy"].update({key: deepcopy(value) for key, value in store.items() if key in _INHERITED})
dropped = sorted(set(store) - _INHERITED - set(REQUIRED_UNS) - {"params"})
if dropped:
get_logger().debug("aggregate: dropped %s, which describe the input rows", dropped)
call_params = {
"by": columns,
"func": func,
"min_cells": min_cells,
"layer": layer,
"use_rep": use_rep,
"count_key": count_key,
"site_key": site_key,
}
record_params(result, "aggregate", call_params)
return result
def _group_names(obs: pd.DataFrame, columns: list[str]) -> pd.Index:
"""Row names built from the group keys, so aggregates from several screens concatenate without clashes."""
# Columns joined with vectorized string concatenation rather than a row-wise apply over every group.
joined = obs[columns[0]].astype(str)
for column in columns[1:]:
joined = joined + "_" + obs[column].astype(str)
# Keys that join to the same text get a numbered suffix.
return ad.utils.make_index_unique(pd.Index(joined.to_numpy(dtype=str)))
def _warn_dropped_controls(obs: pd.DataFrame, keep: np.ndarray, control_col: str, min_cells: int) -> None:
"""Warn when ``min_cells`` removes every negcon/poscon group of a plate.
A later ``normalize(reference=)`` or ``map(mode="activity")`` then has no reference there, so a silent
drop reads downstream as a confusing "no controls" error far from its cause.
"""
if control_col not in obs:
return
types = obs[control_col].astype(str).to_numpy()
plate_col = "Metadata_Plate" if "Metadata_Plate" in obs else None
plates = obs[plate_col].astype(str).to_numpy() if plate_col else np.zeros(len(obs), dtype=object)
lost = []
for plate in pd.unique(plates):
on_plate = plates == plate
for control in ("negcon", "poscon"):
present = on_plate & (types == control)
if present.any() and not (present & keep).any():
lost.append(f"{control} on plate {plate}" if plate_col else control)
if lost:
get_logger().warning(
"aggregate: min_cells=%d dropped %d control group(s) entirely (%s)",
min_cells,
len(lost),
", ".join(lost[:3]),
deep="normalize(reference=) then has no reference there, so lower min_cells or check the control wells' cell counts",
)
def _site_counts(frame: pd.DataFrame, codes: np.ndarray, n_groups: int) -> np.ndarray:
"""Distinct fields of view among each group's cells, where site 1 of one well and of the next are different fields."""
field = np.zeros(len(frame), dtype=np.int64)
for column in ("Metadata_Plate", "Metadata_Well", "Metadata_Site"):
if column in frame:
level, uniques = pd.factorize(frame[column])
field = field * (len(uniques) + 1) + level + 1
width = int(field.max()) + 1
return np.bincount(np.unique(codes.astype(np.int64) * width + field) // width, minlength=n_groups)
def _group_obs(
adata: AnnData, columns: list[str], keys: pd.Index, codes: np.ndarray, tallies: dict[str, np.ndarray]
) -> pd.DataFrame:
"""Build the aggregated ``obs``: grouping keys, the per-group tallies, every column constant within a group.
The ``Metadata_`` prefix is not a survival filter (spec §13.1): any column constant within every group
is carried, whatever its name; a column that varies within a group is dropped rather than copied.
"""
if len(columns) == 1:
obs = pd.DataFrame({columns[0]: np.asarray(keys)})
else:
obs = pd.DataFrame({name: keys.get_level_values(position) for position, name in enumerate(columns)})
obs = obs.reset_index(drop=True)
for name, values in tallies.items():
obs[name] = values
frame = as_frame(adata.obs)
carried = [column for column in frame.columns if column not in {*columns, *_TALLIES, *tallies}]
if carried:
grouped = frame[carried].groupby(codes, observed=True)
constant = grouped.nunique(dropna=False).le(1).all()
for column in carried:
if constant[column]:
obs[column] = grouped[column].first().to_numpy()
else:
get_logger().debug("aggregate: dropped non-constant column %s", column)
return categorize_metadata(obs)
