labelarray 0.1.0__py3-none-any.whl

labelarray/collection.py

@@ -0,0 +1,267 @@
+"""A collection of labelled arrays sharing axes.
+
+A ``Collection`` holds variables with different axes; each axis has one size
+and at most one label array, shared by every variable on it. Selection,
+assignment, renaming, joins under the exact label rule, and nearest-label
+matching are methods. Conversion to and from xarray needs xarray, imported
+only when asked for.
+
+This module needs NumPy only.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Hashable, Mapping, Sequence
+from typing import Any
+
+import numpy as np
+
+from labelarray import labels as exact
+from labelarray.array import (
+    Array,
+    _checked_coords,
+    _cut,
+    _freeze,
+    _take,
+    label_index,
+    plain,
+    positions_of,
+)
+from labelarray.errors import ShapeError
+
+
+class Collection:
+    """Variables with different axes; each axis has one size and at most one label array."""
+
+    __slots__ = ("attrs", "coord_attrs", "coords", "sizes", "variables")
+
+    def __init__(self, variables: Mapping[Hashable, Array], coords: Mapping[Hashable, Any] | None = None,
+                 attrs: Mapping[Hashable, Any] | None = None, coord_attrs: Mapping[Hashable, Any] | None = None):
+        sizes: dict[Hashable, int] = {}
+        for name, value in variables.items():
+            for dim, size in zip(value.dims, value.data.shape):
+                if sizes.setdefault(dim, size) != size:
+                    raise ShapeError(f"{name}: axis {dim} has {size} positions, other variables have {sizes[dim]}")
+        coords = dict(coords or {})
+        for dim, labels in coords.items():
+            sizes.setdefault(dim, len(labels))
+        self.variables = {name: Array._make(v.data, v.dims, {}, v.attrs, name) for name, v in variables.items()}
+        self.coords = _checked_coords(coords, sizes)
+        self.coord_attrs = {d: dict(a) for d, a in (coord_attrs or {}).items() if d in self.coords}
+        self.attrs = dict(attrs or {})
+        self.sizes = sizes
+
+    @property
+    def dims(self) -> tuple[Hashable, ...]:
+        return tuple(self.sizes)
+
+    def array(self, name: Hashable) -> Array:
+        """One variable, with the labels of its axes."""
+        value = self.variables[name]
+        return Array._make(value.data, value.dims, {d: self.coords[d] for d in value.dims if d in self.coords},
+                           value.attrs, name)
+
+    def evolve(self, variables: Mapping | None = None, coords: Mapping | None = None, **changes: Any) -> Collection:
+        """A copy with parts replaced (``variables``, ``coords``, ``attrs``, ``coord_attrs``)."""
+        return Collection(
+            self.variables if variables is None else variables,
+            self.coords if coords is None else coords,
+            changes.get("attrs", self.attrs), changes.get("coord_attrs", self.coord_attrs),
+        )
+
+    # ── Selection ───────────────────────────────────────────────────────────
+
+    def isel(self, indexers: Mapping[Hashable, Any] | None = None, **more: Any) -> Collection:
+        """Positions, a slice, or a mask per axis, on every variable with that axis; an integer removes it."""
+        variables, coords = dict(self.variables), dict(self.coords)
+        for dim, index in {**(indexers or {}), **more}.items():
+            if getattr(index, "dtype", None) == bool:
+                index = np.flatnonzero(index)  # scan the mask once; every variable reuses the positions
+            for name, value in variables.items():
+                if dim in value.dims:
+                    variables[name] = _take(value, dim, index)
+            if dim in coords:
+                if isinstance(index, (int, np.integer)):
+                    del coords[dim]
+                else:
+                    coords[dim] = _cut(coords[dim], index)
+        return self.evolve(variables, coords)
+
+    def sel(self, indexers: Mapping[Hashable, Any] | None = None, **more: Any) -> Collection:
+        """Labels per axis, looked up exactly; one label removes the axis, a list keeps it."""
+        wanted = {**(indexers or {}), **more}
+        return self.isel({dim: positions_of(self.coords[dim], labels, dim) for dim, labels in wanted.items()})
+
+    def select(self, names: Sequence[Hashable]) -> Collection:
+        """Only the variables *names*, with the labels of the axes they use."""
+        if not names:
+            return self.evolve(variables={})
+        variables = {name: self.variables[name] for name in names}
+        used = {d for v in variables.values() for d in v.dims}
+        return self.evolve(variables, {d: c for d, c in self.coords.items() if d in used})
+
+    # ── Building ────────────────────────────────────────────────────────────
+
+    def assign(self, name: Hashable, value: Any) -> Collection:
+        """*value* as the variable *name*; its labels are added for axes that have none yet."""
+        if not isinstance(value, Array):
+            value = Array(value)
+        variables = dict(self.variables)
+        variables[name] = value
+        coords = dict(self.coords)
+        for dim, labels in value.coords.items():
+            coords.setdefault(dim, labels)
+        return self.evolve(variables, coords)
+
+    def rename(self, names: Mapping[Hashable, Hashable]) -> Collection:
+        """Variables renamed."""
+        return self.evolve({names.get(n, n): v for n, v in self.variables.items()})
+
+    # ── Joins ───────────────────────────────────────────────────────────────
+
+    def join(self, other: Collection, how: str) -> Collection:
+        """Both collections' variables, on shared axes' labels joined *how* (left, inner, or outer).
+
+        Labels are paired under the exact rule; positions a side lacks are
+        missing, with integers and booleans promoted as xarray promotes them.
+        Attributes that conflict are dropped, as xarray's ``drop_conflicts``.
+        """
+        left, right = self, other
+        for dim in [d for d in left.coords if d in right.coords]:
+            ours, theirs = left.coords[dim], right.coords[dim]
+            if exact.equal(ours, theirs):
+                continue
+            target = _joined_labels(ours, theirs, how, dim)
+            left = left._reindexed(dim, label_index(ours, "join", dim).get_indexer(target), target)
+            right = right._reindexed(dim, label_index(theirs, "join", dim).get_indexer(target), target)
+        variables = {**{n: left.array(n) for n in left.variables}, **{n: right.array(n) for n in right.variables}}
+        coords = {**right.coords, **left.coords}
+        coord_attrs = {d: _combine_attrs(left.coord_attrs.get(d, {}), right.coord_attrs.get(d, {})) for d in coords}
+        return Collection(variables, coords, _combine_attrs(left.attrs, right.attrs), coord_attrs)
+
+    def nearest(self, dim: Hashable, labels: np.ndarray, tolerance: Any) -> Collection:
+        """This collection on the labels *labels* for *dim*, each matched to its nearest label within *tolerance*.
+
+        A label with none near enough gives missing values; a tie goes to the lower label.
+        """
+        theirs = self.coords[dim]
+        order = np.argsort(theirs, kind="stable")
+        ordered = theirs[order]
+        upper = np.clip(np.searchsorted(ordered, labels), 0, len(ordered) - 1)
+        lower = np.clip(upper - 1, 0, len(ordered) - 1)
+        pick_lower = np.abs(labels - ordered[lower]) <= np.abs(ordered[upper] - labels)
+        nearest = np.where(pick_lower, lower, upper)
+        positions = np.where(np.abs(ordered[nearest] - labels) <= tolerance, order[nearest], -1)
+        return self._reindexed(dim, positions, labels)
+
+    def _reindexed(self, dim: Hashable, positions: np.ndarray, labels: np.ndarray) -> Collection:
+        variables = {}
+        for name in self.variables:
+            value = self.array(name)
+            variables[name] = _reindex(value, dim, positions, labels) if dim in value.dims else value
+        coords = dict(self.coords)
+        coords[dim] = labels
+        return self.evolve(variables, coords)
+
+    # ── Conversion ──────────────────────────────────────────────────────────
+
+    def to_dict(self) -> dict:
+        """As xarray's ``Dataset.to_dict(data="array")`` gives it, without xarray."""
+        return {
+            "coords": {d: {"dims": (d,), "attrs": dict(self.coord_attrs.get(d, {})), "data": plain(labels)}
+                       for d, labels in self.coords.items()},
+            "attrs": dict(self.attrs),
+            "dims": dict(self.sizes),
+            "data_vars": {n: {"dims": v.dims, "attrs": dict(v.attrs), "data": v.data} for n, v in self.variables.items()},
+        }
+
+    def to_xarray(self) -> Any:
+        """An xarray ``Dataset`` sharing the values; needs xarray."""
+        import xarray as xr
+
+        coords = {d: (d, plain(labels), dict(self.coord_attrs.get(d, {}))) for d, labels in self.coords.items()}
+        return xr.Dataset({n: (v.dims, v.data, dict(v.attrs)) for n, v in self.variables.items()},
+                          coords=coords, attrs=dict(self.attrs))
+
+    @classmethod
+    def from_xarray(cls, dataset: Any) -> Collection:
+        """A collection from an in-memory xarray ``Dataset`` whose coordinates are its axes' labels."""
+        variables = {name: Array._make(np.asarray(value.data), tuple(value.dims), {}, dict(value.attrs), name)
+                     for name, value in dataset.data_vars.items()}
+        coords = {name: dataset[name].values for name in dataset.coords}
+        coord_attrs = {name: dict(dataset[name].attrs) for name in dataset.coords}
+        return cls(variables, coords, dict(dataset.attrs), coord_attrs)
+
+
+def _joined_labels(ours: np.ndarray, theirs: np.ndarray, how: str, dim: Hashable) -> np.ndarray:
+    """The labels a join keeps, in the order xarray's joins give them."""
+    if how == "left":
+        return ours
+    found = label_index(theirs, "join", dim).get_indexer(ours) >= 0
+    if how == "inner":
+        return _freeze(ours[found])
+    extra = theirs[label_index(ours, "join", dim).get_indexer(theirs) < 0]
+    union = np.concatenate([ours, extra])
+    try:
+        return _freeze(np.sort(union, kind="stable"))
+    except TypeError:  # labels that cannot be ordered keep their first-seen order
+        return _freeze(union)
+
+
+def _fill_for(dtype: np.dtype) -> tuple[np.dtype, Any]:
+    """The type and missing value a reindexed variable needs, as xarray promotes them."""
+    if dtype.kind in "iu":
+        # As xarray promotes: integers of up to 16 bits to float32, wider ones to float64.
+        return np.dtype(np.float32 if dtype.itemsize <= 2 else np.float64), np.nan
+    if dtype.kind == "b":
+        return np.dtype(np.float64), np.nan
+    if dtype.kind == "f" or dtype.kind == "c":
+        return dtype, np.nan
+    if dtype.kind in "mM":
+        return dtype, np.datetime64("NaT") if dtype.kind == "M" else np.timedelta64("NaT")
+    return np.dtype(object), np.nan
+
+
+def _reindex(value: Array, dim: Hashable, positions: np.ndarray, labels: np.ndarray) -> Array:
+    """*value* on new labels for *dim*: position -1 becomes missing."""
+    axis = value.axis(dim)
+    missing = positions < 0
+    if not value.data.shape[axis]:  # nothing to take from: every position is missing
+        dtype, fill = _fill_for(value.data.dtype) if len(positions) else (value.data.dtype, None)
+        shape = value.data.shape[:axis] + (len(positions),) + value.data.shape[axis + 1:]
+        taken = np.full(shape, fill, dtype=dtype) if len(positions) else np.empty(shape, dtype=dtype)
+        coords = dict(value.coords)
+        coords[dim] = labels
+        return Array._make(taken, value.dims, coords, value.attrs, value.name)
+    taken = np.take(value.data, np.where(missing, 0, positions), axis=axis)
+    if missing.any():
+        dtype, fill = _fill_for(value.data.dtype)
+        taken = taken.astype(dtype, copy=True) if taken.dtype != dtype else taken.copy()
+        index = (slice(None),) * axis + (missing,)
+        taken[index] = fill
+    coords = dict(value.coords)
+    coords[dim] = labels
+    return Array._make(taken, value.dims, coords, value.attrs, value.name)
+
+
+def _combine_attrs(*mappings: Mapping) -> dict:
+    """xarray's ``drop_conflicts``: keep keys whose values agree wherever they appear."""
+    result: dict = {}
+    dropped: set = set()
+    for mapping in mappings:
+        for key, value in mapping.items():
+            if key in dropped:
+                continue
+            if key in result and not _same(result[key], value):
+                dropped.add(key)
+                del result[key]
+            elif key not in result:
+                result[key] = value
+    return result
+
+
+def _same(a: Any, b: Any) -> bool:
+    try:
+        return bool(np.array_equal(a, b)) if isinstance(a, np.ndarray) or isinstance(b, np.ndarray) else bool(a == b)
+    except Exception:  # noqa: BLE001 - values that cannot be compared count as different
+        return False

labelarray/contract.py

@@ -0,0 +1,47 @@
+"""Contraction over named axes with ``np.matmul``: the sum of products over some axes, keeping the others.
+
+Operands are arranged as ``(batch, rows, contracted)`` and ``(batch,
+contracted, columns)`` C-ordered arrays, whatever their memory layout, and
+multiplied with ``np.matmul``; batch axes are those both operands keep. The
+result has the kept axes in order of first appearance, left then right (as
+``np.einsum`` and ``xarray.dot`` order them), and is C-ordered. Arranging
+the operands the same way for every layout makes the result independent of
+layout. On 10M points rotated by a 3 x 3 matrix this takes 6.5 ms where
+``np.einsum`` took 13.3 (68.7 on row-major input). Floating-point warnings
+are off for the multiplication: BLAS libraries do not set the flags NumPy
+reads reliably (Accelerate under NumPy 2.0 sets them on ordinary finite
+data), so an overflow shows as infinity in the result, without a warning.
+This module needs NumPy only.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Hashable, Sequence
+
+import numpy as np
+
+
+def contract(left: np.ndarray, left_dims: Sequence[Hashable], right: np.ndarray, right_dims: Sequence[Hashable],
+             over: Sequence[Hashable]) -> tuple[np.ndarray, tuple[Hashable, ...]]:
+    """The sum over *over* of products of *left* and *right*, and its axis names."""
+    left_dims, right_dims, over = tuple(left_dims), tuple(right_dims), tuple(over)
+    sizes = {**dict(zip(right_dims, right.shape)), **dict(zip(left_dims, left.shape))}
+    batch = tuple(d for d in left_dims if d in right_dims and d not in over)
+    rows = tuple(d for d in left_dims if d not in right_dims and d not in over)
+    columns = tuple(d for d in right_dims if d not in left_dims and d not in over)
+    kept = tuple(dict.fromkeys(d for d in (*left_dims, *right_dims) if d not in over))
+
+    def arranged(values, dims, first, second):
+        order = [dims.index(d) for d in (*batch, *first, *second)]
+        shape = (math.prod(sizes[d] for d in batch), math.prod(sizes[d] for d in first),
+                 math.prod(sizes[d] for d in second))
+        return np.ascontiguousarray(np.transpose(values, order)).reshape(shape)
+
+    with np.errstate(all="ignore"):
+        product = np.matmul(arranged(left, left_dims, rows, over), arranged(right, right_dims, over, columns))
+    result = product.reshape(tuple(sizes[d] for d in (*batch, *rows, *columns)))
+    produced = (*batch, *rows, *columns)
+    result = np.transpose(result, [produced.index(d) for d in kept])
+    # np.ascontiguousarray would make a 0-d result 1-d; copy only what is not C-ordered already.
+    return (result if result.flags.c_contiguous else result.copy(order="C")), kept

labelarray/errors.py

@@ -0,0 +1,21 @@
+"""labelarray's errors: one family of ``ValueError``, so callers catching bad input keep working."""
+
+
+class LabelArrayError(ValueError):
+    """Base class for labelarray's errors."""
+
+
+class LabelError(LabelArrayError):
+    """Labels that cannot be ordered, compared, or looked up, or a label that is absent."""
+
+
+class AlignError(LabelArrayError):
+    """Labels on a shared axis do not match."""
+
+
+class GroupError(LabelArrayError):
+    """Key values or a key transform's settings that cannot give groups."""
+
+
+class ShapeError(LabelArrayError):
+    """Axis names, sizes, or ranks that do not fit together."""

labelarray/grouped.py

@@ -0,0 +1,259 @@
+"""Grouped reductions in one pass over positions sorted by group.
+
+Reducing groups one by one costs a Python step per group, which dominates
+when groups are many and small (points binned into voxels). Here every group
+is reduced at once wherever the result is bit-identical to reducing each
+group on its own, as both backends do otherwise:
+
+- min, max, count, any, and all do not depend on the order values are
+  combined in, nor do integer and boolean sums; ``reduceat`` gives them for
+  every group.
+- Float sums and means depend on the order. NumPy adds fewer than 8 values
+  one after another, starting from 0.0, in every memory layout; from 8 on it
+  adds in blocks. Groups with fewer than 8 positions are therefore added
+  position by position for all groups at once; larger groups are left to the
+  caller, which reduces them one by one.
+
+Other reductions (std, var, median), other types, groups with no positions,
+and reductions over further axes are left to the caller.
+``ONE_PASS = False`` turns it off, for comparison in tests.
+
+``Grouping`` is the object the kernels work from: one integer code per
+position and the group labels in code order, with the group sizes and the
+positions sorted by group computed once and kept. ``reduce``, ``counts``,
+and ``broadcast`` apply a grouping to an ``Array``: every group's result
+along a new group axis, and a result per group back on the positions.
+This module needs NumPy only.
+"""
+
+from __future__ import annotations
+
+import builtins
+from collections.abc import Hashable, Sequence
+from typing import Any
+
+import numpy as np
+
+from labelarray import sums
+from labelarray.array import Array, _freeze, _take, labels_from
+
+ONE_PASS = True
+
+# Up to this many groups, scanning the axis once per group is cheaper than sorting it.
+FEW_GROUPS = 8
+
+
+class Grouping:
+    """Positions assigned to groups: one code per position, and each group's label.
+
+    Codes are integers from 0 to ``len(labels) - 1``; ``labels[code]`` is the
+    group's label. ``counts()`` and ``blocks()`` are computed on first use and
+    kept, so a sort by group serves every reduction of the same grouping.
+    """
+
+    def __init__(self, codes: Any, labels: Sequence[Any]):
+        self.codes = np.asarray(codes)
+        self.labels = list(labels)
+        self._counts: np.ndarray | None = None
+        self._blocks: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None
+
+    def __len__(self) -> int:
+        return len(self.labels)
+
+    def counts(self) -> np.ndarray:
+        """Positions in each group, in code order."""
+        if self._counts is None:
+            self._counts = np.bincount(self.codes, minlength=len(self.labels))
+        return self._counts
+
+    def blocks(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Positions sorted by group, keeping their order within a group, and each group's start and end."""
+        if self._blocks is None:
+            # NumPy's stable sort is a radix sort for 8- and 16-bit integers, far faster than for int64.
+            narrow = np.min_scalar_type(max(len(self.labels) - 1, 0))
+            order = np.argsort(self.codes.astype(narrow, copy=False), kind="stable")
+            ends = np.cumsum(self.counts())
+            self._blocks = (order, np.concatenate(([0], ends[:-1])), ends)
+        return self._blocks
+
+# NumPy adds fewer than this many values one after another; more, in blocks (pairwise summation).
+SMALL = 8
+
+_ORDER_FREE = {"min", "max", "count", "any", "all"}
+
+
+def one_pass(data: np.ndarray, axis: int, op: str, na_rm: bool, min_count: int | None,
+             order: np.ndarray, starts: np.ndarray, ends: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
+    """*op* per group along *axis*, with the group axis first, and which groups it gives.
+
+    Groups are ``order[starts[g]:ends[g]]``. Rows of groups not given are
+    left unset for the caller. None when no group can be reduced here.
+    """
+    if not ONE_PASS or not isinstance(data, np.ndarray):
+        return None
+    kind = data.dtype.kind
+    if kind not in "biuf" or (kind == "f" and data.dtype.itemsize < 4) or op not in (*_ORDER_FREE, "sum", "mean"):
+        return None
+    sizes = ends - starts
+    if not len(sizes) or not sizes.all():
+        return None  # groups with no positions (empty bins) are reduced one by one, as before
+    any_size = op in _ORDER_FREE or (op == "sum" and kind in "biu")
+    done = np.ones(len(sizes), dtype=bool) if any_size else sizes < SMALL
+    if not done.any():
+        return None
+    ordered = np.moveaxis(np.take(data, order, axis=axis), axis, 0)
+    missing = np.isnan(ordered) if kind == "f" else None
+    with np.errstate(all="ignore"):
+        values = _reduce(ordered, missing, kind, op, na_rm, min_count, starts, sizes, done)
+    if done.all():
+        return values, done
+    result = np.empty((len(sizes), *values.shape[1:]), dtype=values.dtype)
+    result[done] = values
+    return result, done
+
+
+def _reduce(ordered, missing, kind, op, na_rm, min_count, starts, sizes, done):
+    """The reduction for the groups in *done*, matching ``Array.reduce`` group by group."""
+    if op in ("min", "max"):
+        skip = na_rm and kind == "f"  # nanmin and nanmax ignore NaN; fmin and fmax do too
+        ufunc = {("min", False): np.minimum, ("min", True): np.fmin,
+                 ("max", False): np.maximum, ("max", True): np.fmax}[op, skip]
+        return ufunc.reduceat(ordered, starts, axis=0)
+    present = ~missing if missing is not None else None
+    if op == "count":
+        present = present if present is not None else np.ones(ordered.shape, dtype=bool)
+        return np.add.reduceat(present, starts, axis=0, dtype=_sum_type(present))
+    if op in ("any", "all"):
+        filled = ordered if missing is None else np.where(missing, op == "all", ordered)
+        ufunc = np.logical_or if op == "any" else np.logical_and
+        result = ufunc.reduceat(filled.astype(bool, copy=False), starts, axis=0)
+        if na_rm or kind in "biu" or missing is None:
+            return result
+        return np.where(~np.logical_or.reduceat(missing, starts, axis=0), result, np.nan)
+    skip = na_rm and kind == "f"
+    source = np.where(missing, 0.0, ordered) if skip else ordered  # nansum and nanmean count NaN as 0
+    if op == "sum":
+        if kind in "biu":
+            total = np.add.reduceat(source, starts, axis=0, dtype=_sum_type(source))
+        else:
+            total = _left_to_right(source, starts[done], sizes[done])
+        if min_count is None:
+            return total
+        counts = _counts(present, ordered.shape, starts, sizes)[done]
+        return np.where(counts >= min_count, total, np.nan)
+    # mean: integers and booleans are added as float64, as np.mean adds them
+    total = _left_to_right(source if kind == "f" else source.astype(np.float64), starts[done], sizes[done])
+    counts = _counts(present if skip else None, ordered.shape, starts, sizes)[done]
+    # np.mean and np.nanmean divide in float64 and store in the sum's type.
+    return (total.astype(np.float64) / counts).astype(total.dtype, copy=False)
+
+
+def _left_to_right(values: np.ndarray, first: np.ndarray, sizes: np.ndarray) -> np.ndarray:
+    """Each group's values added one after another, for groups of fewer than ``SMALL`` positions."""
+    total = values[first] + 0.0  # NumPy's sums start from 0.0, so a group of -0.0 sums to 0.0
+    for rank in range(1, int(sizes.max())):
+        active = np.flatnonzero(sizes > rank)
+        total[active] += values[first[active] + rank]
+    return total
+
+
+def _counts(present, shape, starts, sizes) -> np.ndarray:
+    """Present values per group (every position when *present* is None), shaped like one group's result."""
+    if present is None:
+        return sizes.reshape(-1, *([1] * (len(shape) - 1)))
+    return np.add.reduceat(present, starts, axis=0, dtype=np.intp)
+
+
+def _sum_type(values: np.ndarray) -> np.dtype:
+    """The type ``np.sum`` gives: small integers and booleans add as the platform integer."""
+    return np.sum(values[:0], axis=0).dtype
+
+
+# ── Groupings applied to arrays ─────────────────────────────────────────────
+
+
+def reduce(array: Array, op: str, grouping: Grouping, axis: Hashable, name: Hashable,
+           others: Sequence[Hashable] = (), *, na_rm: bool = False, min_count: int | None = None) -> Array:
+    """*op* per group of positions on *axis*, and over *others*: the groups along a new first axis *name*.
+
+    Each group's result is bit-identical to reducing that group on its own
+    (``Array.reduce`` over a mask), whichever path computes it.
+    """
+    if op in ("sum", "mean") and array.data.dtype == np.float32:
+        axes = (array.axis(axis), *(array.axis(o) for o in others))
+        values = sums.reduce(array.data, op, axes, na_rm=na_rm, min_count=min_count, codes=grouping.codes,
+                             groups=len(grouping), sizes=grouping.counts())
+        dims = tuple(d for d in array.dims if d != axis and d not in others)
+        coords = {d: c for d, c in array.coords.items() if d in dims}
+        return Array._make(values, (name, *dims), {name: labels_from(grouping.labels), **coords})
+    fast = _one_pass_array(array, op, grouping, axis, name, others, na_rm, min_count)
+    if fast is not None:
+        return fast
+    stats = [part.reduce(op, [axis, *others], na_rm=na_rm, min_count=min_count)
+             for part in _group_parts(array, grouping, axis)]
+    labels = labels_from(grouping.labels)
+    if stats:
+        data = np.stack([s.data for s in stats])
+        return Array._make(data, (name, *stats[0].dims), {name: labels, **stats[0].coords})
+    empty = _take(array, axis, builtins.slice(0, 0)).reduce(op, [axis, *others], na_rm=na_rm, min_count=min_count)
+    return Array._make(empty.data[np.newaxis][:0], (name, *empty.dims), {name: labels, **empty.coords})
+
+
+def counts(grouping: Grouping, name: Hashable, per_position: int = 1) -> Array:
+    """Values per group: positions times *per_position*, along the axis *name*."""
+    # Plain NumPy labels, not objects: xarray labels counts this way.
+    return Array._make(grouping.counts() * per_position, (name,), {name: _freeze(np.asarray(grouping.labels))})
+
+
+def broadcast(stats: Array, grouping: Grouping, axis: Hashable, name: Hashable,
+              labels: np.ndarray | None = None) -> Array:
+    """A result per group (along *name*) back on each position of *axis*, labelled *labels* if given."""
+    number = stats.axis(name)
+    # Group results come with the group axis first (reduce). Indexing gathers them
+    # several times faster than np.take, into the same values and the same C-ordered layout.
+    data = stats.data[grouping.codes] if number == 0 else np.take(stats.data, grouping.codes, axis=number)
+    dims = stats.dims[:number] + (axis,) + stats.dims[number + 1:]
+    coords = {d: c for d, c in stats.coords.items() if d != name}
+    if labels is not None:
+        coords[axis] = labels
+    return Array._make(data, dims, coords)
+
+
+def _one_pass_array(array: Array, op: str, grouping: Grouping, axis: Hashable, name: Hashable,
+                    others: Sequence[Hashable], na_rm: bool, min_count: int | None) -> Array | None:
+    """Every group at once where that is exact (``one_pass``); the rest one by one."""
+    if others or len(grouping) <= FEW_GROUPS:
+        return None
+    order, starts, ends = grouping.blocks()
+    fast = one_pass(array.data, array.axis(axis), op, na_rm, min_count, order, starts, ends)
+    if fast is None:
+        return None
+    values, done = fast
+    rest = np.flatnonzero(~done)
+    if len(rest):
+        stats = np.stack([_take(array, axis, order[int(starts[g]):int(ends[g])])
+                          .reduce(op, [axis], na_rm=na_rm, min_count=min_count).data for g in rest])
+        if stats.dtype != values.dtype:
+            return None
+        values[rest] = stats
+    dims = tuple(d for d in array.dims if d != axis)
+    coords = {d: c for d, c in array.coords.items() if d in dims}
+    return Array._make(values, (name, *dims), {name: labels_from(grouping.labels), **coords})
+
+
+def _group_parts(array: Array, grouping: Grouping, axis: Hashable) -> list[Array]:
+    """Each group's positions, in their original order and in a mask's memory layout.
+
+    Indexing an axis other than the first with positions or a mask lays
+    the result out with that axis first in memory; reductions over it then
+    add in another order than over a contiguous copy, which changes float
+    results in the last bit. Each group is therefore indexed by its
+    positions, as a mask would be, so results are bit-identical to xarray
+    and to per-group masks. With many groups, the positions come from one
+    sort instead of one scan of the axis per group.
+    """
+    count = len(grouping)
+    if count <= FEW_GROUPS:
+        return [_take(array, axis, np.flatnonzero(grouping.codes == code)) for code in range(count)]
+    order, starts, ends = grouping.blocks()
+    return [_take(array, axis, order[int(start):int(end)]) for start, end in zip(starts, ends)]