"""
Aggregation operations over arrays.
.. note::
:obj:`AggregateTransformer` and :obj:`RangedAggregateTransformer` support the
:doc:`Array API standard </guides/explanations/array_api>`, enabling use with
NumPy, CuPy, PyTorch, and other compatible array libraries.
Operations not available on a given backend (nan-variants, trapezoid) fall back
to NumPy automatically.
"""
import typing
import ezmsg.core as ez
import numpy as np
import numpy.typing as npt
from array_api_compat import get_namespace
from ezmsg.baseproc import (
BaseStatefulTransformer,
BaseTransformer,
BaseTransformerUnit,
processor_state,
)
from ezmsg.util.messages.axisarray import (
AxisArray,
AxisBase,
replace,
slice_along_axis,
)
from .spectral import OptionsEnum
[docs]
class AggregationFunction(OptionsEnum):
"""Enum for aggregation functions available to be used in :obj:`ranged_aggregate` operation."""
NONE = "None (all)"
MAX = "max"
MIN = "min"
MEAN = "mean"
MEDIAN = "median"
STD = "std"
SUM = "sum"
NANMAX = "nanmax"
NANMIN = "nanmin"
NANMEAN = "nanmean"
NANMEDIAN = "nanmedian"
NANSTD = "nanstd"
NANSUM = "nansum"
ARGMIN = "argmin"
ARGMAX = "argmax"
TRAPEZOID = "trapezoid"
AGGREGATORS = {
AggregationFunction.NONE: np.all,
AggregationFunction.MAX: np.max,
AggregationFunction.MIN: np.min,
AggregationFunction.MEAN: np.mean,
AggregationFunction.MEDIAN: np.median,
AggregationFunction.STD: np.std,
AggregationFunction.SUM: np.sum,
AggregationFunction.NANMAX: np.nanmax,
AggregationFunction.NANMIN: np.nanmin,
AggregationFunction.NANMEAN: np.nanmean,
AggregationFunction.NANMEDIAN: np.nanmedian,
AggregationFunction.NANSTD: np.nanstd,
AggregationFunction.NANSUM: np.nansum,
AggregationFunction.ARGMIN: np.argmin,
AggregationFunction.ARGMAX: np.argmax,
# Note: Some methods require x-coordinates and
# are handled specially in `aggregate_slices`.
AggregationFunction.TRAPEZOID: np.trapezoid,
}
# Operations that cannot be evaluated from the values alone.
_NEEDS_COORDINATES = frozenset(
{
AggregationFunction.TRAPEZOID,
AggregationFunction.ARGMIN,
AggregationFunction.ARGMAX,
}
)
[docs]
def needs_coordinates(operation: typing.Union["AggregationFunction", typing.Iterable["AggregationFunction"]]) -> bool:
"""Whether ``operation`` requires the axis's x-coordinates.
Lets a caller skip building a coordinate vector it will not use, which is
worth doing where that vector would have to be constructed per chunk.
Accepts a single function or an iterable of them.
"""
if isinstance(operation, AggregationFunction):
return operation in _NEEDS_COORDINATES
return any(op in _NEEDS_COORDINATES for op in operation)
[docs]
def axis_coordinates(message: AxisArray, axis_name: str) -> npt.NDArray:
"""The coordinate value of every element along ``axis_name``.
Reads a coordinate axis's values directly; evaluates a linear axis over its
own length. Callers whose data does not line up with the message -- anything
carrying samples across message boundaries -- must build their own vector
instead.
"""
target_axis = message.get_axis(axis_name)
if hasattr(target_axis, "data"):
return np.asarray(target_axis.data)
axis_idx = message.get_axis_idx(axis_name)
return target_axis.value(np.arange(message.data.shape[axis_idx]))
def _apply_one(xp, op: AggregationFunction, segment, axis_idx: int):
"""One aggregation over one already-sliced segment."""
func_name = op.value
if hasattr(xp, func_name):
return getattr(xp, func_name)(segment, axis=axis_idx)
# nan-variants and friends are not in the Array API standard.
result = AGGREGATORS[op](np.asarray(segment), axis=axis_idx)
return xp.asarray(result) if xp is not np else result
[docs]
def aggregate_slices(
data,
slices: typing.Sequence[slice],
axis_idx: int,
operation: AggregationFunction,
*,
coordinates: typing.Optional[npt.NDArray] = None,
index_to_coordinate: bool = True,
):
"""Apply ``operation`` to each slice of ``data`` along ``axis_idx``, stacked.
Where the groups come from is the caller's business -- coordinate bands
resolved once, or bin boundaries recomputed per chunk -- but *running* the
aggregation is the same either way, and three operations need more than
``f(segment, axis)`` to do it correctly:
* the nan-variants, which the Array API does not define, so they need a
numpy fallback and a conversion back to the caller's namespace;
* ``TRAPEZOID``, which integrates and so needs the axis's x-coordinates, or
it silently returns an integral in units of *samples* rather than of the
axis;
* ``ARGMIN``/``ARGMAX``, which return a position within the slice. An index
is rarely what anyone wants -- "the peak is at 10.5 Hz" is useful, "the
peak is at offset 7 of this band" is not -- so it is converted back to the
axis coordinate here.
The result has the same rank as ``data``, with ``axis_idx`` reduced to one
entry per slice, so a caller can drop it into the message it came from
having replaced only that axis.
:param data: The array to slice. Any Array API namespace.
:param slices: One slice per output group, in output order. Slices index
``data`` along ``axis_idx``; they need not be contiguous, and an empty
sequence yields a zero-length result.
:param axis_idx: Index of the axis being grouped.
:param operation: The :obj:`AggregationFunction` to apply within each slice.
:param coordinates: The axis's coordinate values, one per element of ``data``
along ``axis_idx``. Required when :func:`needs_coordinates` is True for
``operation``, ignored otherwise.
:param index_to_coordinate: Whether ``ARGMIN``/``ARGMAX`` results are
converted from a within-slice index to an axis coordinate. False leaves
the raw index and needs no ``coordinates``. Every transformer in this
package leaves this True; it exists for a caller that means to index
back into the array it passed in.
:raises ValueError: if ``operation`` needs coordinates and none were given,
or if ``coordinates`` does not span ``data`` along ``axis_idx``. Either
would otherwise produce a plausible-looking but wrong number.
"""
xp = get_namespace(data)
is_arg = operation in (AggregationFunction.ARGMIN, AggregationFunction.ARGMAX)
uses_coordinates = operation == AggregationFunction.TRAPEZOID or (is_arg and index_to_coordinate)
if uses_coordinates:
if coordinates is None:
raise ValueError(f"AggregationFunction.{operation.name} needs the axis coordinates; pass coordinates=...")
coordinates = np.asarray(coordinates)
n = data.shape[axis_idx]
if coordinates.shape[0] != n:
raise ValueError(f"coordinates has {coordinates.shape[0]} entries but data spans {n} along axis {axis_idx}")
if not len(slices):
return slice_along_axis(data, slice(0, 0), axis=axis_idx)
if operation == AggregationFunction.TRAPEZOID:
# No Array API equivalent, and it needs x, so this path is numpy-only.
np_data = np.asarray(data)
stacked = np.stack(
[
np.trapezoid(slice_along_axis(np_data, sl, axis=axis_idx), x=coordinates[sl], axis=axis_idx)
for sl in slices
],
axis=axis_idx,
)
return xp.asarray(stacked) if xp is not np else stacked
stacked = xp.stack(
[_apply_one(xp, operation, slice_along_axis(data, sl, axis=axis_idx), axis_idx) for sl in slices],
axis=axis_idx,
)
if is_arg and index_to_coordinate:
# Indices are relative to each slice, so each is looked up in that
# slice's own coordinates rather than in the whole vector.
out = [
coordinates[sl][np.asarray(slice_along_axis(stacked, sl_ix, axis=axis_idx))]
for sl_ix, sl in enumerate(slices)
]
stacked = np.stack(out, axis=axis_idx)
return xp.asarray(stacked) if xp is not np else stacked
return stacked
[docs]
class RangedAggregateSettings(ez.Settings):
"""
Settings for ``RangedAggregate``.
"""
axis: str | None = None
"""The name of the axis along which to apply the bands."""
bands: list[tuple[float, float]] | None = None
"""
[(band1_min, band1_max), (band2_min, band2_max), ...]
If not set then this acts as a passthrough node.
"""
operation: AggregationFunction = AggregationFunction.MEAN
""":obj:`AggregationFunction` to apply to each band."""
[docs]
@processor_state
class RangedAggregateState:
slices: list[tuple[typing.Any, ...]] | None = None
out_axis: AxisBase | None = None
ax_vec: npt.NDArray | None = None
[docs]
class RangedAggregate(BaseTransformerUnit[RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateTransformer]):
SETTINGS = RangedAggregateSettings
[docs]
def ranged_aggregate(
axis: str | None = None,
bands: list[tuple[float, float]] | None = None,
operation: AggregationFunction = AggregationFunction.MEAN,
) -> RangedAggregateTransformer:
"""
Apply an aggregation operation over one or more bands.
Args:
axis: The name of the axis along which to apply the bands.
bands: [(band1_min, band1_max), (band2_min, band2_max), ...]
If not set then this acts as a passthrough node.
operation: :obj:`AggregationFunction` to apply to each band.
Returns:
:obj:`RangedAggregateTransformer`
"""
return RangedAggregateTransformer(RangedAggregateSettings(axis=axis, bands=bands, operation=operation))
[docs]
class AggregateSettings(ez.Settings):
"""Settings for :obj:`Aggregate`."""
axis: str
"""The name of the axis to aggregate over. This axis will be removed from the output."""
operation: AggregationFunction = AggregationFunction.MEAN
""":obj:`AggregationFunction` to apply.
``ARGMIN``/``ARGMAX`` return the *coordinate* of the extremum along ``axis``
-- "the peak is at 14 Hz" -- not its index, matching
:obj:`RangedAggregateTransformer` and :obj:`BinnedAggregateTransformer`. An
index would be unusable here anyway, since ``axis`` is removed from the
output. Where ``axis`` carries no axis metadata the coordinates are 0, 1,
2, ..., so the result is the index after all."""
[docs]
class AggregateUnit(BaseTransformerUnit[AggregateSettings, AxisArray, AxisArray, AggregateTransformer]):
"""Unit that aggregates an entire axis using a specified operation."""
SETTINGS = AggregateSettings