Merge pull request #94 from ezmsg-org/dev

Add trapezoid to list of aggregation functions and expose function choice in BandPower

Author: cboulay

src/ezmsg/sigproc/aggregate.py

@@ -36,6 +36,7 @@ class AggregationFunction(OptionsEnum):
     NANSUM = "nansum"
     ARGMIN = "argmin"
     ARGMAX = "argmax"
+    TRAPEZOID = "trapezoid"
 
 
 AGGREGATORS = {
@@ -54,6 +55,9 @@ class AggregationFunction(OptionsEnum):
     AggregationFunction.NANSUM: np.nansum,
     AggregationFunction.ARGMIN: np.argmin,
     AggregationFunction.ARGMAX: np.argmax,
+    # Note: Some methods require x-coordinates and
+    #  are handled specially in `_process`.
+    AggregationFunction.TRAPEZOID: np.trapezoid,
 }
 
 
@@ -144,10 +148,23 @@ def _process(self, message: AxisArray) -> AxisArray:
         ax_idx = message.get_axis_idx(axis)
         agg_func = AGGREGATORS[self.settings.operation]
 
-        out_data = [
-            agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx)
-            for sl in self._state.slices
-        ]
+        if self.settings.operation in [
+            AggregationFunction.TRAPEZOID,
+        ]:
+            # Special handling for methods that require x-coordinates.
+            out_data = [
+                agg_func(
+                    slice_along_axis(message.data, sl, axis=ax_idx),
+                    x=self._state.ax_vec[sl],
+                    axis=ax_idx,
+                )
+                for sl in self._state.slices
+            ]
+        else:
+            out_data = [
+                agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx)
+                for sl in self._state.slices
+            ]
 
         msg_out = replace(
             message,

src/ezmsg/sigproc/bandpower.py

@@ -36,6 +36,9 @@ class BandPowerSettings(ez.Settings):
     (min, max) tuples of band limits in Hz.
     """
 
+    aggregation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply to each band."""
+
 
 class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, AxisArray]):
     @staticmethod
@@ -50,7 +53,7 @@ def _initialize_processors(
                 settings=RangedAggregateSettings(
                     axis="freq",
                     bands=settings.bands,
-                    operation=AggregationFunction.MEAN,
+                    operation=settings.aggregation,
                 )
             ),
         }
@@ -68,6 +71,7 @@ def bandpower(
         (17, 30),
         (70, 170),
     ],
+    aggregation: AggregationFunction = AggregationFunction.MEAN,
 ) -> BandPowerTransformer:
     """
     Calculate the average spectral power in each band.
@@ -77,6 +81,8 @@ def bandpower(
     """
     return BandPowerTransformer(
         settings=BandPowerSettings(
-            spectrogram_settings=spectrogram_settings, bands=bands
+            spectrogram_settings=spectrogram_settings,
+            bands=bands,
+            aggregation=aggregation,
         )
     )

tests/unit/test_aggregate.py

@@ -110,6 +110,30 @@ def test_arg_aggregate(agg_func: AggregationFunction):
     assert np.array_equal(out_dat, expected_dat)
 
 
+def test_trapezoid():
+    bands = [(5.0, 20.0), (30.0, 50.0)]
+    in_msgs = [_ for _ in get_msg_gen()]
+    gen = ranged_aggregate(
+        axis="freq", bands=bands, operation=AggregationFunction.TRAPEZOID
+    )
+    out_msgs = [gen.send(_) for _ in in_msgs]
+
+    out_dat = AxisArray.concatenate(*out_msgs, dim="time").data
+
+    # Calculate expected data using trapezoidal integration
+    in_data = AxisArray.concatenate(*in_msgs, dim="time").data
+    targ_ax = in_msgs[0].axes["freq"]
+    targ_ax_vec = targ_ax.value(np.arange(in_data.shape[-1]))
+    expected = []
+    for start, stop in bands:
+        inds = np.logical_and(targ_ax_vec >= start, targ_ax_vec <= stop)
+        expected.append(np.trapezoid(in_data[..., inds], x=targ_ax_vec[inds], axis=-1))
+    expected = np.stack(expected, axis=-1)
+
+    assert out_dat.shape == expected.shape
+    assert np.allclose(out_dat, expected)
+
+
 @pytest.mark.parametrize("change_ax", ["ch", "freq"])
 def test_aggregate_handle_change(change_ax: str):
     """

tests/unit/test_bandpower.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.sigproc.bandpower import bandpower, SpectrogramSettings
+from ezmsg.sigproc.bandpower import bandpower, SpectrogramSettings, AggregationFunction
 
 from tests.helpers.util import (
     create_messages_with_periodic_signal,
@@ -45,6 +45,7 @@ def test_bandpower():
             window_shift=0.1,
         ),
         bands=bands,
+        aggregation=AggregationFunction.MEAN,
     )
     results = [gen.send(_) for _ in messages]