BinnedAggregate: allow a tuple of operations, stacked on a metric axis

src/ezmsg/sigproc/binned_aggregate.py

@@ -16,9 +16,11 @@
 (e.g. 30012 Hz) they diverge in both gain and bin count, so two such streams
 never share a grid.
 
-:obj:`BinnedAggregateTransformer` applies an arbitrary :obj:`AggregationFunction`
-per bin instead of only counting, but it does *not* define its own bin
-boundaries: it drives them through the shared :obj:`ezmsg.sigproc.util.binning.BinSchedule`,
+:obj:`BinnedAggregateTransformer` applies one or more arbitrary
+:obj:`AggregationFunction`\\ s per bin instead of only counting -- a tuple of
+functions stacks its results on a trailing axis, e.g. ``(MIN, MAX)`` for a
+display envelope -- but it does *not* define its own bin boundaries: it drives
+them through the shared :obj:`ezmsg.sigproc.util.binning.BinSchedule`,
 the single source of truth for the grid. Any consumer that goes through the same
 schedule at the same ``bin_duration`` lands on the same grid by construction, so
 two such streams align downstream (e.g. with :obj:`ezmsg.sigproc.merge.Merge`).
@@ -64,8 +66,23 @@ class BinnedAggregateSettings(ez.Settings):
     bin_duration: float = 0.02
     """Output bin duration in seconds."""
 
-    operation: AggregationFunction = AggregationFunction.MEAN
-    """:obj:`AggregationFunction` applied within each bin."""
+    operation: AggregationFunction | tuple[AggregationFunction, ...] = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` applied within each bin.
+
+    A tuple applies several aggregations to the same bins and stacks the results
+    on a new trailing axis named by ``newaxis``, coordinate-labelled with each
+    function's value (``"min"``, ``"max"``, ...). ``(MIN, MAX)`` is the envelope
+    used to put a high-rate signal on a screen without stride-decimation
+    clipping the peaks: it keeps each bin's extremes exactly, so a spike shows
+    at full amplitude no matter where in the bin it fell.
+
+    A one-element tuple still produces the trailing axis -- the output shape
+    follows the *type* of this field, not the number of functions, so a caller
+    that builds its tuple programmatically gets a stable shape.
+    """
+
+    newaxis: str = "metric"
+    """Name of the trailing axis added when ``operation`` is a tuple."""
 
     fractional: bool = True
     """If True (default), bins span a *fractional* ``bin_duration * fs`` samples
@@ -89,6 +106,13 @@ class BinnedAggregateState:
     aggregation works for any operation (not just sums). Its length is kept in
     sync with ``schedule.carry_count``."""
 
+    metric_axis: AxisArray.CoordinateAxis | None = None
+    """The trailing axis attached to every multi-operation output.
+
+    It depends only on settings, so it is built once here rather than rebuilt
+    per message -- and since it is the same object every time, downstream
+    identity checks on the axis stay cheap. ``None`` for a scalar operation."""
+
 
 class BinnedAggregateTransformer(
     BaseStatefulTransformer[BinnedAggregateSettings, AxisArray, AxisArray, BinnedAggregateState]
@@ -99,6 +123,12 @@ class BinnedAggregateTransformer(
     :obj:`RangedAggregateTransformer` (which aggregates static coordinate
     bands), this reduces a high-rate axis to a regularly-binned lower-rate
     axis, carrying the open partial bin across message boundaries.
+
+    ``settings.operation`` may be a tuple, in which case every function is
+    applied to the same bins and the results stack on a trailing ``newaxis``.
+    The bins are computed once and sliced once, so N aggregations cost far less
+    than N copies of this transformer, and -- more importantly -- they are
+    guaranteed to describe the same bins.
     """
 
     def _hash_message(self, message: AxisArray) -> int:
@@ -112,26 +142,72 @@ def _reset_state(self, message: AxisArray) -> None:
         )
         schedule.reset(1.0 / axis_info.gain)
         self._state.schedule = schedule
+        self._state.metric_axis = (
+            AxisArray.CoordinateAxis(
+                data=np.array([op.value for op in self._operations]),
+                dims=[self.settings.newaxis],
+            )
+            if self._multi
+            else None
+        )
         self._state.carry = None
 
-    def _aggregate(self, xp, segment, axis_idx: int):
+    @property
+    def _operations(self) -> tuple[AggregationFunction, ...]:
+        """``operation`` as a tuple, regardless of how it was given."""
         op = self.settings.operation
+        return op if isinstance(op, tuple) else (op,)
+
+    @property
+    def _multi(self) -> bool:
+        """Whether to emit a trailing metric axis."""
+        return isinstance(self.settings.operation, tuple)
+
+    def _apply_one(self, xp, op: AggregationFunction, segment, axis_idx: int):
         func_name = op.value
         if hasattr(xp, func_name):
             return getattr(xp, func_name)(segment, axis=axis_idx)
         # nan-variants etc. are not in the Array API; fall back to numpy.
         result = AGGREGATORS[op](np.asarray(segment), axis=axis_idx)
         return xp.asarray(result) if xp is not np else result
 
-    def _empty_like(self, message: AxisArray, axis_idx: int, step: BinStep) -> AxisArray:
+    def _aggregate(self, xp, segment, axis_idx: int):
+        """Reduce one bin. Multi-op results stack on a new *trailing* axis.
+
+        Trailing rather than in place, so the binned axis keeps its position and
+        every existing consumer of a single-op stream sees an unchanged shape.
+        """
+        if not self._multi:
+            return self._apply_one(xp, self.settings.operation, segment, axis_idx)
+        return xp.stack([self._apply_one(xp, op, segment, axis_idx) for op in self._operations], axis=-1)
+
+    def _out_dims(self, message: AxisArray) -> list[str]:
+        dims = list(message.dims)
+        return dims + [self.settings.newaxis] if self._multi else dims
+
+    def _out_axes(self, message: AxisArray, step: BinStep) -> dict:
         axis_info = message.get_axis(self.settings.axis)
+        axes = {
+            **message.axes,
+            self.settings.axis: replace(axis_info, gain=step.output_gain, offset=step.output_offset),
+        }
+        if self._multi:
+            axes[self.settings.newaxis] = self._state.metric_axis
+        return axes
+
+    def _empty_like(self, message: AxisArray, axis_idx: int, step: BinStep) -> AxisArray:
+        xp = get_namespace(message.data)
+        data = slice_along_axis(message.data, slice(0, 0), axis=axis_idx)
+        if self._multi:
+            # Zero-length along the binned axis, but the metric axis must still
+            # be its full width or the empty message would not match the shape
+            # of the ones around it.
+            data = xp.stack([data] * len(self._operations), axis=-1)
         return replace(
             message,
-            data=slice_along_axis(message.data, slice(0, 0), axis=axis_idx),
-            axes={
-                **message.axes,
-                self.settings.axis: replace(axis_info, gain=step.output_gain, offset=step.output_offset),
-            },
+            data=data,
+            dims=self._out_dims(message),
+            axes=self._out_axes(message, step),
         )
 
     def _process(self, message: AxisArray) -> AxisArray:
@@ -176,10 +252,8 @@ def _process(self, message: AxisArray) -> AxisArray:
         return replace(
             message,
             data=stacked,
-            axes={
-                **message.axes,
-                axis: replace(axis_info, gain=step.output_gain, offset=step.output_offset),
-            },
+            dims=self._out_dims(message),
+            axes=self._out_axes(message, step),
         )
 
 

tests/unit/test_binned_aggregate.py

@@ -255,3 +255,148 @@ async def drive(msg):
     assert len(published) == 1
     _, msg_out = published[0]
     assert msg_out.data.shape[0] > 0
+
+
+# ---- multi-operation (trailing metric axis) --------------------------------
+
+MINMAX = (AggregationFunction.MIN, AggregationFunction.MAX)
+
+
+def test_tuple_operation_stacks_on_a_trailing_axis():
+    fs = 1000.0
+    rng = np.random.default_rng(0)
+    sig = rng.standard_normal((100, 2))
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX, fractional=True)
+    out = proc(_sig_msgs(sig, fs, 100)[0])
+
+    # spb = 20 -> 5 bins, 2 channels, 2 metrics.
+    assert out.dims == ["time", "ch", "metric"]
+    assert out.data.shape == (5, 2, 2)
+    np.testing.assert_allclose(out.data[..., 0], _ref_binned(sig, 20.0, np.min))
+    np.testing.assert_allclose(out.data[..., 1], _ref_binned(sig, 20.0, np.max))
+
+
+def test_metric_axis_is_labelled_from_the_enum():
+    """Consumers read names, not positions -- and the names are the enum values."""
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)
+    out = proc(_sig_msgs(np.ones((100, 2)), 1000.0, 100)[0])
+    metric = out.axes["metric"]
+    assert list(metric.data) == ["min", "max"]
+    assert metric.dims == ["metric"]
+
+
+def test_newaxis_is_configurable():
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX, newaxis="bound")
+    out = proc(_sig_msgs(np.ones((100, 2)), 1000.0, 100)[0])
+    assert out.dims == ["time", "ch", "bound"]
+    assert list(out.axes["bound"].data) == ["min", "max"]
+
+
+def test_scalar_operation_shape_is_unchanged():
+    """The trailing axis appears only for a tuple, so existing streams are
+    untouched."""
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=AggregationFunction.MEAN)
+    out = proc(_sig_msgs(np.ones((100, 2)), 1000.0, 100)[0])
+    assert out.dims == ["time", "ch"]
+    assert out.data.shape == (5, 2)
+    assert "metric" not in out.axes
+
+
+def test_single_element_tuple_still_produces_the_axis():
+    """Shape follows the type of `operation`, not the count, so a caller that
+    builds its tuple programmatically gets a stable shape."""
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=(AggregationFunction.MAX,))
+    out = proc(_sig_msgs(np.ones((100, 2)), 1000.0, 100)[0])
+    assert out.dims == ["time", "ch", "metric"]
+    assert out.data.shape == (5, 2, 1)
+    assert list(out.axes["metric"].data) == ["max"]
+
+
+def test_multi_op_matches_running_each_op_separately():
+    """The whole point of a tuple is that the bins are identical; prove they are
+    by comparing against separate single-op transformers on the same chunking."""
+    fs = 1000.0
+    rng = np.random.default_rng(1)
+    sig = rng.standard_normal((503, 3))
+    msgs = _sig_msgs(sig, fs, 37)  # chunk size coprime with the bin size
+
+    multi = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)
+    only_min = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=AggregationFunction.MIN)
+    only_max = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=AggregationFunction.MAX)
+
+    got = np.concatenate([m.data for m in _run(multi, copy.deepcopy(msgs))], axis=0)
+    exp_min = np.concatenate([m.data for m in _run(only_min, copy.deepcopy(msgs))], axis=0)
+    exp_max = np.concatenate([m.data for m in _run(only_max, copy.deepcopy(msgs))], axis=0)
+
+    np.testing.assert_allclose(got[..., 0], exp_min)
+    np.testing.assert_allclose(got[..., 1], exp_max)
+
+
+def test_multi_op_carries_partial_bins_across_chunks():
+    """A peak in a bin split across two messages must survive.
+
+    This is the property that makes the envelope usable for spike-bearing data:
+    the extremes are exact regardless of where the message boundary falls.
+    """
+    fs = 1000.0
+    sig = np.zeros((60, 1))
+    sig[25, 0] = 9.0  # bin 1 (samples 20..39), and chunk boundary is at 30
+    sig[35, 0] = -4.0
+
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)
+    out = np.concatenate([m.data for m in _run(proc, _sig_msgs(sig, fs, 30))], axis=0)
+
+    assert out.shape == (3, 1, 2)
+    np.testing.assert_allclose(out[1, 0, 0], -4.0)  # min of the split bin
+    np.testing.assert_allclose(out[1, 0, 1], 9.0)  # max of the split bin
+
+
+def test_multi_op_empty_message_keeps_the_metric_width():
+    """An empty payload must still be the right shape, or it will not
+    concatenate with the messages around it."""
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)
+    out = proc(_sig_msgs(np.ones((5, 2)), 1000.0, 5)[0])  # < one bin -> empty
+    assert out.data.shape == (0, 2, 2)
+    assert out.dims == ["time", "ch", "metric"]
+
+
+def test_multi_op_axis_gain_matches_single_op():
+    """The metric axis must not disturb the bin rate."""
+    msgs = _sig_msgs(np.ones((100, 2)), 1000.0, 100)
+    multi = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)(copy.deepcopy(msgs[0]))
+    single = BinnedAggregateTransformer(axis="time", bin_duration=0.02)(copy.deepcopy(msgs[0]))
+    assert multi.axes["time"].gain == single.axes["time"].gain
+    assert multi.axes["time"].offset == single.axes["time"].offset
+
+
+def test_multi_op_non_time_axis():
+    """Binning a non-leading axis must still append the metric axis at the end."""
+    proc = BinnedAggregateTransformer(axis="ch", bin_duration=2.0, operation=MINMAX, fractional=False)
+    msg = AxisArray(
+        data=np.arange(40, dtype=float).reshape(4, 10),
+        dims=["time", "ch"],
+        axes=frozendict(
+            {
+                "time": AxisArray.TimeAxis(fs=10.0, offset=0.0),
+                "ch": AxisArray.TimeAxis(fs=1.0, offset=0.0),
+            }
+        ),
+        key="test_binned_aggregate_ch",
+    )
+    out = proc(msg)
+    assert out.dims == ["time", "ch", "metric"]
+    # 10 channels binned 2 at a time -> 5 bins.
+    assert out.data.shape == (4, 5, 2)
+    np.testing.assert_allclose(out.data[0, :, 0], [0, 2, 4, 6, 8])
+    np.testing.assert_allclose(out.data[0, :, 1], [1, 3, 5, 7, 9])
+
+
+def test_metric_axis_is_built_once_not_per_message():
+    """It depends only on settings, so every output should carry the same
+    object -- both to avoid rebuilding it per message and so a downstream
+    identity check on the axis stays cheap."""
+    proc = BinnedAggregateTransformer(axis="time", bin_duration=0.02, operation=MINMAX)
+    outs = _run(proc, _sig_msgs(np.ones((200, 2)), 1000.0, 40))
+    assert len(outs) > 1
+    first = outs[0].axes["metric"]
+    assert all(o.axes["metric"] is first for o in outs)