ezmsg-sigproc 1.8.2__py3-none-any.whl → 2.1.0__py3-none-any.whl

ezmsg/sigproc/aggregate.py

@@ -3,7 +3,6 @@ import typing
 import numpy as np
 import numpy.typing as npt
 import ezmsg.core as ez
-from ezmsg.util.generator import consumer
 from ezmsg.util.messages.axisarray import (
     AxisArray,
     slice_along_axis,
@@ -12,7 +11,11 @@ from ezmsg.util.messages.axisarray import (
 )
 
 from .spectral import OptionsEnum
-from .base import GenAxisArray
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 
 
 class AggregationFunction(OptionsEnum):
@@ -33,6 +36,7 @@ class AggregationFunction(OptionsEnum):
     NANSUM = "nansum"
     ARGMIN = "argmin"
     ARGMAX = "argmax"
+    TRAPEZOID = "trapezoid"
 
 
 AGGREGATORS = {
@@ -51,133 +55,161 @@ AGGREGATORS = {
     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,
 }
 
 
-@consumer
-def ranged_aggregate(
-    axis: str | None = None,
-    bands: list[tuple[float, float]] | None = None,
-    operation: AggregationFunction = AggregationFunction.MEAN,
-):
+class RangedAggregateSettings(ez.Settings):
+    """
+    Settings for ``RangedAggregate``.
     """
-    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.
+    axis: str | None = None
+    """The name of the axis along which to apply the bands."""
 
-    Returns:
-        A primed generator object ready to yield an :obj:`AxisArray` for each .send(axis_array)
+    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.
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
 
-    # State variables
+    operation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply to each band."""
+
+
+@processor_state
+class RangedAggregateState:
     slices: list[tuple[typing.Any, ...]] | None = None
     out_axis: AxisBase | None = None
     ax_vec: npt.NDArray | None = None
 
-    # Reset if any of these changes. Key not checked because continuity between chunks not required.
-    check_inputs = {"gain": None, "offset": None, "len": None, "key": None}
 
-    while True:
-        msg_in: AxisArray = yield msg_out
-        if bands is None:
-            msg_out = msg_in
+class RangedAggregateTransformer(
+    BaseStatefulTransformer[
+        RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateState
+    ]
+):
+    def __call__(self, message: AxisArray) -> AxisArray:
+        # Override for shortcut passthrough mode.
+        if self.settings.bands is None:
+            return message
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        target_axis = message.get_axis(axis)
+
+        hash_components = (message.key,)
+        if hasattr(target_axis, "data"):
+            hash_components += (len(target_axis.data),)
+        elif isinstance(target_axis, AxisArray.LinearAxis):
+            hash_components += (target_axis.gain, target_axis.offset)
+        return hash(hash_components)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[0]
+        target_axis = message.get_axis(axis)
+        ax_idx = message.get_axis_idx(axis)
+
+        if hasattr(target_axis, "data"):
+            self._state.ax_vec = target_axis.data
         else:
-            axis = axis or msg_in.dims[0]
-            target_axis = msg_in.get_axis(axis)
+            self._state.ax_vec = target_axis.value(
+                np.arange(message.data.shape[ax_idx])
+            )
 
-            # Check if we need to reset state
-            b_reset = msg_in.key != check_inputs["key"]
+        ax_dat = []
+        slices = []
+        for start, stop in self.settings.bands:
+            inds = np.where(
+                np.logical_and(self._state.ax_vec >= start, self._state.ax_vec <= stop)
+            )[0]
+            slices.append(np.s_[inds[0] : inds[-1] + 1])
             if hasattr(target_axis, "data"):
-                b_reset = b_reset or len(target_axis.data) != check_inputs["len"]
-            elif isinstance(target_axis, AxisArray.LinearAxis):
-                b_reset = b_reset or target_axis.gain != check_inputs["gain"]
-                b_reset = b_reset or target_axis.offset != check_inputs["offset"]
-
-            if b_reset:
-                # Update check variables
-                check_inputs["key"] = msg_in.key
-                if hasattr(target_axis, "data"):
-                    check_inputs["len"] = len(target_axis.data)
+                if self._state.ax_vec.dtype.type is np.str_:
+                    sl_dat = f"{self._state.ax_vec[start]} - {self._state.ax_vec[stop]}"
                 else:
-                    check_inputs["gain"] = target_axis.gain
-                    check_inputs["offset"] = target_axis.offset
-
-                # If the axis we are operating on has changed (e.g., "time" or "win" axis always changes),
-                #  or the key has changed, then recalculate slices.
+                    ax_dat.append(np.mean(self._state.ax_vec[inds]))
+            else:
+                sl_dat = target_axis.value(np.mean(inds))
+            ax_dat.append(sl_dat)
+
+        self._state.slices = slices
+        self._state.out_axis = AxisArray.CoordinateAxis(
+            data=np.array(ax_dat),
+            dims=[axis],
+            unit=target_axis.unit,
+        )
 
-                ax_idx = msg_in.get_axis_idx(axis)
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+        agg_func = AGGREGATORS[self.settings.operation]
 
-                if hasattr(target_axis, "data"):
-                    ax_vec = target_axis.data
-                else:
-                    ax_vec = target_axis.value(np.arange(msg_in.data.shape[ax_idx]))
-
-                slices = []
-                ax_dat = []
-                for start, stop in bands:
-                    inds = np.where(np.logical_and(ax_vec >= start, ax_vec <= stop))[0]
-                    slices.append(np.s_[inds[0] : inds[-1] + 1])
-                    if hasattr(target_axis, "data"):
-                        if ax_vec.dtype.type is np.str_:
-                            sl_dat = f"{ax_vec[start]} - {ax_vec[stop]}"
-                        else:
-                            sl_dat = ax_dat.append(np.mean(ax_vec[inds]))
-                    else:
-                        sl_dat = target_axis.value(np.mean(inds))
-                    ax_dat.append(sl_dat)
-
-                out_axis = AxisArray.CoordinateAxis(
-                    data=np.array(ax_dat),
-                    dims=[axis],
-                    unit=target_axis.unit,
+        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,
                 )
-
-            agg_func = AGGREGATORS[operation]
+                for sl in self._state.slices
+            ]
+        else:
             out_data = [
-                agg_func(slice_along_axis(msg_in.data, sl, axis=ax_idx), axis=ax_idx)
-                for sl in slices
+                agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx)
+                for sl in self._state.slices
             ]
 
-            msg_out = replace(
-                msg_in,
-                data=np.stack(out_data, axis=ax_idx),
-                axes={**msg_in.axes, axis: out_axis},
-            )
-            if operation in [AggregationFunction.ARGMIN, AggregationFunction.ARGMAX]:
-                # Convert indices returned by argmin/argmax into the value along the axis.
-                out_data = []
-                for sl_ix, sl in enumerate(slices):
-                    offsets = np.take(msg_out.data, [sl_ix], axis=ax_idx)
-                    out_data.append(ax_vec[sl][offsets])
-                msg_out.data = np.concatenate(out_data, axis=ax_idx)
+        msg_out = replace(
+            message,
+            data=np.stack(out_data, axis=ax_idx),
+            axes={**message.axes, axis: self._state.out_axis},
+        )
 
+        if self.settings.operation in [
+            AggregationFunction.ARGMIN,
+            AggregationFunction.ARGMAX,
+        ]:
+            out_data = []
+            for sl_ix, sl in enumerate(self._state.slices):
+                offsets = np.take(msg_out.data, [sl_ix], axis=ax_idx)
+                out_data.append(self._state.ax_vec[sl][offsets])
+            msg_out.data = np.concatenate(out_data, axis=ax_idx)
 
-class RangedAggregateSettings(ez.Settings):
-    """
-    Settings for ``RangedAggregate``.
-    See :obj:`ranged_aggregate` for details.
-    """
+        return msg_out
 
-    axis: str | None = None
-    bands: list[tuple[float, float]] | None = None
-    operation: AggregationFunction = AggregationFunction.MEAN
 
+class RangedAggregate(
+    BaseTransformerUnit[
+        RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateTransformer
+    ]
+):
+    SETTINGS = RangedAggregateSettings
 
-class RangedAggregate(GenAxisArray):
-    """
-    Unit for :obj:`ranged_aggregate`
+
+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.
 
-    SETTINGS = RangedAggregateSettings
+    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.
 
-    def construct_generator(self):
-        self.STATE.gen = ranged_aggregate(
-            axis=self.SETTINGS.axis,
-            bands=self.SETTINGS.bands,
-            operation=self.SETTINGS.operation,
-        )
+    Returns:
+        :obj:`RangedAggregateTransformer`
+    """
+    return RangedAggregateTransformer(
+        RangedAggregateSettings(axis=axis, bands=bands, operation=operation)
+    )

ezmsg/sigproc/bandpower.py

@@ -1,76 +1,88 @@
 from dataclasses import field
-import typing
 
-import numpy as np
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.generator import consumer, compose
 
-from .spectrogram import spectrogram, SpectrogramSettings
-from .aggregate import ranged_aggregate, AggregationFunction
-from .base import GenAxisArray
-
-
-@consumer
-def bandpower(
-    spectrogram_settings: SpectrogramSettings,
-    bands: list[tuple[float, float]] | None = [
-        (17, 30),
-        (70, 170),
-    ],
-) -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Calculate the average spectral power in each band.
-
-    Args:
-        spectrogram_settings: Settings for spectrogram calculation.
-        bands: (min, max) tuples of band limits in Hz.
-
-    Returns:
-        A primed generator object ready to yield an :obj:`AxisArray` for each .send(axis_array)
-        with the data payload being the average spectral power in each band of the input data.
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
-
-    f_spec = spectrogram(
-        window_dur=spectrogram_settings.window_dur,
-        window_shift=spectrogram_settings.window_shift,
-        window_anchor=spectrogram_settings.window_anchor,
-        window=spectrogram_settings.window,
-        transform=spectrogram_settings.transform,
-        output=spectrogram_settings.output,
-    )
-    f_agg = ranged_aggregate(
-        axis="freq", bands=bands, operation=AggregationFunction.MEAN
-    )
-    pipeline = compose(f_spec, f_agg)
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = pipeline(msg_in)
+from .spectrogram import SpectrogramSettings, SpectrogramTransformer
+from .aggregate import (
+    AggregationFunction,
+    RangedAggregateTransformer,
+    RangedAggregateSettings,
+)
+from .base import (
+    BaseProcessor,
+    CompositeProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+)
 
 
 class BandPowerSettings(ez.Settings):
     """
     Settings for ``BandPower``.
-    See :obj:`bandpower` for details.
     """
 
     spectrogram_settings: SpectrogramSettings = field(
         default_factory=SpectrogramSettings
     )
+    """
+    Settings for spectrogram calculation.
+    """
+
     bands: list[tuple[float, float]] | None = field(
         default_factory=lambda: [(17, 30), (70, 170)]
     )
+    """
+    (min, max) tuples of band limits in Hz.
+    """
+
+    aggregation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply to each band."""
 
 
-class BandPower(GenAxisArray):
-    """:obj:`Unit` for :obj:`bandpower`."""
+class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, AxisArray]):
+    @staticmethod
+    def _initialize_processors(
+        settings: BandPowerSettings,
+    ) -> dict[str, BaseProcessor | BaseStatefulProcessor]:
+        return {
+            "spectrogram": SpectrogramTransformer(
+                settings=settings.spectrogram_settings
+            ),
+            "aggregate": RangedAggregateTransformer(
+                settings=RangedAggregateSettings(
+                    axis="freq",
+                    bands=settings.bands,
+                    operation=settings.aggregation,
+                )
+            ),
+        }
 
+
+class BandPower(
+    BaseTransformerUnit[BandPowerSettings, AxisArray, AxisArray, BandPowerTransformer]
+):
     SETTINGS = BandPowerSettings
 
-    def construct_generator(self):
-        self.STATE.gen = bandpower(
-            spectrogram_settings=self.SETTINGS.spectrogram_settings,
-            bands=self.SETTINGS.bands,
+
+def bandpower(
+    spectrogram_settings: SpectrogramSettings,
+    bands: list[tuple[float, float]] | None = [
+        (17, 30),
+        (70, 170),
+    ],
+    aggregation: AggregationFunction = AggregationFunction.MEAN,
+) -> BandPowerTransformer:
+    """
+    Calculate the average spectral power in each band.
+
+    Returns:
+        :obj:`BandPowerTransformer`
+    """
+    return BandPowerTransformer(
+        settings=BandPowerSettings(
+            spectrogram_settings=spectrogram_settings,
+            bands=bands,
+            aggregation=aggregation,
         )
+    )