From 68b54e0bc0428518d8466bebb91500a593c8ffb2 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Mon, 3 Aug 2026 22:26:13 -0400 Subject: [PATCH 1/2] BinnedAggregate: allow a tuple of operations, stacked on a metric axis Reducing a high-rate signal for display wants each bin's min *and* max -- stride decimation lands between samples and clips the peaks, so a spike shrinks or vanishes depending on where the stride falls, whereas the two extremes preserve peak amplitude exactly. That is the only thing a dedicated min/max decimator would do that this transformer did not already do: bin on a shared schedule, carry the open partial bin across message boundaries, aggregate, label the output axis. Rather than write a second transformer with its own copy of that arithmetic, `operation` now accepts a tuple. Each function is applied to the same bins and the results stack on a trailing axis, `newaxis` (default "metric"), coordinate-labelled with the AggregationFunction values -- "min", "max" -- so a consumer reads names rather than relying on positional convention. Trailing rather than in place, so the binned axis keeps its position. Computing them together is not just less code: the bins are cut once and sliced once, so the results are identical by construction rather than by two transformers agreeing. There is a test comparing a tuple against separate single-op transformers over the same chunking to hold that. The output shape follows the *type* of `operation`, not the count. A scalar behaves exactly as before -- no trailing axis, no new key in `axes` -- so existing streams are untouched, and a one-element tuple does produce the axis, so a caller assembling its tuple programmatically gets a stable shape either way. The empty-payload path has to build the trailing axis at full width too; a zero-length message of the wrong rank will not concatenate with the messages around it, and at a high input rate most chunks close no bin, so that path is the common one rather than an edge case. --- src/ezmsg/sigproc/binned_aggregate.py | 98 ++++++++++++++++--- tests/unit/test_binned_aggregate.py | 134 ++++++++++++++++++++++++++ 2 files changed, 216 insertions(+), 16 deletions(-) diff --git a/src/ezmsg/sigproc/binned_aggregate.py b/src/ezmsg/sigproc/binned_aggregate.py index 91cd065..f61b9b6 100644 --- a/src/ezmsg/sigproc/binned_aggregate.py +++ b/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 @@ -99,6 +116,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: @@ -114,8 +137,18 @@ def _reset_state(self, message: AxisArray) -> None: self._state.schedule = schedule 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) @@ -123,15 +156,50 @@ def _aggregate(self, xp, segment, axis_idx: int): 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 _metric_axis(self) -> AxisArray.CoordinateAxis: + """Label the trailing axis so consumers read names, not positions.""" + return AxisArray.CoordinateAxis( + data=np.array([op.value for op in self._operations]), + dims=[self.settings.newaxis], + ) + + 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._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 +244,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), ) diff --git a/tests/unit/test_binned_aggregate.py b/tests/unit/test_binned_aggregate.py index f795cde..edfc740 100644 --- a/tests/unit/test_binned_aggregate.py +++ b/tests/unit/test_binned_aggregate.py @@ -255,3 +255,137 @@ 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]) From 13fcc30c59cabfdd7d053ce5159763d3e6627220 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Mon, 3 Aug 2026 23:26:32 -0400 Subject: [PATCH 2/2] Build the metric axis once instead of per message It depends only on settings -- the operations and the axis name -- so rebuilding it in _out_axes on every message meant constructing a numpy string array per message on the hot path, for a value that cannot have changed. Moved to _reset_state and held in state. Every output now carries the same object, so a downstream identity check on the axis is a pointer comparison rather than an array comparison. --- src/ezmsg/sigproc/binned_aggregate.py | 24 ++++++++++++++++-------- tests/unit/test_binned_aggregate.py | 11 +++++++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/src/ezmsg/sigproc/binned_aggregate.py b/src/ezmsg/sigproc/binned_aggregate.py index f61b9b6..dc7687b 100644 --- a/src/ezmsg/sigproc/binned_aggregate.py +++ b/src/ezmsg/sigproc/binned_aggregate.py @@ -106,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] @@ -135,6 +142,14 @@ 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 @property @@ -166,13 +181,6 @@ def _aggregate(self, xp, segment, axis_idx: int): 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 _metric_axis(self) -> AxisArray.CoordinateAxis: - """Label the trailing axis so consumers read names, not positions.""" - return AxisArray.CoordinateAxis( - data=np.array([op.value for op in self._operations]), - dims=[self.settings.newaxis], - ) - def _out_dims(self, message: AxisArray) -> list[str]: dims = list(message.dims) return dims + [self.settings.newaxis] if self._multi else dims @@ -184,7 +192,7 @@ def _out_axes(self, message: AxisArray, step: BinStep) -> dict: self.settings.axis: replace(axis_info, gain=step.output_gain, offset=step.output_offset), } if self._multi: - axes[self.settings.newaxis] = self._metric_axis() + axes[self.settings.newaxis] = self._state.metric_axis return axes def _empty_like(self, message: AxisArray, axis_idx: int, step: BinStep) -> AxisArray: diff --git a/tests/unit/test_binned_aggregate.py b/tests/unit/test_binned_aggregate.py index edfc740..62a496f 100644 --- a/tests/unit/test_binned_aggregate.py +++ b/tests/unit/test_binned_aggregate.py @@ -389,3 +389,14 @@ def test_multi_op_non_time_axis(): 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)