ezmsg-sigproc 1.7.0__py3-none-any.whl → 2.10.0__py3-none-any.whl

ezmsg/sigproc/aggregate.py

@@ -1,18 +1,32 @@
+"""
+Aggregation operations over arrays.
+
+.. note::
+    :obj:`AggregateTransformer` supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    :obj:`RangedAggregateTransformer` currently requires NumPy arrays.
+"""
+
 import typing
 
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
-import ezmsg.core as ez
-from ezmsg.util.generator import consumer
+from array_api_compat import get_namespace
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
     AxisBase,
     replace,
+    slice_along_axis,
 )
 
 from .spectral import OptionsEnum
-from .base import GenAxisArray
 
 
 class AggregationFunction(OptionsEnum):
@@ -33,6 +47,7 @@ class AggregationFunction(OptionsEnum):
     NANSUM = "nansum"
     ARGMIN = "argmin"
     ARGMAX = "argmax"
+    TRAPEZOID = "trapezoid"
 
 
 AGGREGATORS = {
@@ -51,133 +66,211 @@ 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
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
+    [(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."""
 
-    # State variables
+
+@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
+                    ax_dat.append(np.mean(self._state.ax_vec[inds]))
+            else:
+                sl_dat = target_axis.value(np.mean(inds))
+            ax_dat.append(sl_dat)
 
-                # 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.
+        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):
+        return msg_out
+
+
+class RangedAggregate(BaseTransformerUnit[RangedAggregateSettings, AxisArray, AxisArray, RangedAggregateTransformer]):
+    SETTINGS = RangedAggregateSettings
+
+
+def ranged_aggregate(
+    axis: str | None = None,
+    bands: list[tuple[float, float]] | None = None,
+    operation: AggregationFunction = AggregationFunction.MEAN,
+) -> RangedAggregateTransformer:
     """
-    Settings for ``RangedAggregate``.
-    See :obj:`ranged_aggregate` for details.
+    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))
+
+
+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."""
 
-    axis: str | None = None
-    bands: list[tuple[float, float]] | None = None
     operation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply."""
 
 
-class RangedAggregate(GenAxisArray):
+class AggregateTransformer(BaseTransformer[AggregateSettings, AxisArray, AxisArray]):
     """
-    Unit for :obj:`ranged_aggregate`
+    Transformer that aggregates an entire axis using a specified operation.
+
+    Unlike :obj:`RangedAggregateTransformer` which aggregates over specific ranges/bands
+    and preserves the axis (with one value per band), this transformer aggregates the
+    entire axis and removes it from the output, reducing dimensionality by one.
     """
 
-    SETTINGS = RangedAggregateSettings
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        axis_idx = message.get_axis_idx(self.settings.axis)
+        op = self.settings.operation
 
-    def construct_generator(self):
-        self.STATE.gen = ranged_aggregate(
-            axis=self.SETTINGS.axis,
-            bands=self.SETTINGS.bands,
-            operation=self.SETTINGS.operation,
+        if op == AggregationFunction.NONE:
+            raise ValueError("AggregationFunction.NONE is not supported for full-axis aggregation")
+
+        if op == AggregationFunction.TRAPEZOID:
+            # Trapezoid integration requires x-coordinates
+            target_axis = message.get_axis(self.settings.axis)
+            if hasattr(target_axis, "data"):
+                x = target_axis.data
+            else:
+                x = target_axis.value(np.arange(message.data.shape[axis_idx]))
+            agg_data = np.trapezoid(np.asarray(message.data), x=x, axis=axis_idx)
+        else:
+            # Try array-API compatible function first, fall back to numpy
+            func_name = op.value
+            if hasattr(xp, func_name):
+                agg_data = getattr(xp, func_name)(message.data, axis=axis_idx)
+            else:
+                agg_data = AGGREGATORS[op](message.data, axis=axis_idx)
+
+        new_dims = list(message.dims)
+        new_dims.pop(axis_idx)
+
+        new_axes = dict(message.axes)
+        new_axes.pop(self.settings.axis, None)
+
+        return replace(
+            message,
+            data=agg_data,
+            dims=new_dims,
+            axes=new_axes,
         )
+
+
+class AggregateUnit(BaseTransformerUnit[AggregateSettings, AxisArray, AxisArray, AggregateTransformer]):
+    """Unit that aggregates an entire axis using a specified operation."""
+
+    SETTINGS = AggregateSettings

ezmsg/sigproc/bandpower.py

@@ -1,75 +1,80 @@
 from dataclasses import field
-import typing
 
-import numpy as np
 import ezmsg.core as ez
+from ezmsg.baseproc import (
+    BaseProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
 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
+from .aggregate import (
+    AggregationFunction,
+    RangedAggregateSettings,
+    RangedAggregateTransformer,
+)
+from .spectrogram import SpectrogramSettings, SpectrogramTransformer
 
 
-@consumer
-def bandpower(
-    spectrogram_settings: SpectrogramSettings,
-    bands: list[tuple[float, float]] | None = [
-        (17, 30),
-        (70, 170),
-    ],
-) -> typing.Generator[AxisArray, AxisArray, None]:
+class BandPowerSettings(ez.Settings):
     """
-    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.
+    Settings for ``BandPower``.
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
-
-    f_spec = spectrogram(
-        window_dur=spectrogram_settings.window_dur,
-        window_shift=spectrogram_settings.window_shift,
-        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)
 
+    spectrogram_settings: SpectrogramSettings = field(default_factory=SpectrogramSettings)
+    """
+    Settings for spectrogram calculation.
+    """
 
-class BandPowerSettings(ez.Settings):
+    bands: list[tuple[float, float]] | None = field(default_factory=lambda: [(17, 30), (70, 170)])
     """
-    Settings for ``BandPower``.
-    See :obj:`bandpower` for details.
+    (min, max) tuples of band limits in Hz.
     """
 
-    spectrogram_settings: SpectrogramSettings = field(
-        default_factory=SpectrogramSettings
-    )
-    bands: list[tuple[float, float]] | None = field(
-        default_factory=lambda: [(17, 30), (70, 170)]
-    )
+    aggregation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply to each band."""
+
 
+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(GenAxisArray):
-    """:obj:`Unit` for :obj:`bandpower`."""
 
+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,
         )
+    )

ezmsg/sigproc/base.py

@@ -1,39 +1,149 @@
-import math
-import traceback
-import typing
+"""
+Backwards-compatible re-exports from ezmsg.baseproc.
 
-import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.generator import GenState
+This module re-exports all symbols from ezmsg.baseproc to maintain backwards
+compatibility for code that imports from ezmsg.sigproc.base.
 
+New code should import directly from ezmsg.baseproc instead.
+"""
 
-class GenAxisArray(ez.Unit):
-    STATE = GenState
+import warnings
 
-    INPUT_SIGNAL = ez.InputStream(AxisArray)
-    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-    INPUT_SETTINGS = ez.InputStream(ez.Settings)
+warnings.warn(
+    "Importing from 'ezmsg.sigproc.base' is deprecated. Please import from 'ezmsg.baseproc' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
 
-    async def initialize(self) -> None:
-        self.construct_generator()
+# Re-export everything from ezmsg.baseproc for backwards compatibility
+from ezmsg.baseproc import (  # noqa: E402
+    # Protocols
+    AdaptiveTransformer,
+    # Type variables
+    AdaptiveTransformerType,
+    # Stateful classes
+    BaseAdaptiveTransformer,
+    # Unit classes
+    BaseAdaptiveTransformerUnit,
+    BaseAsyncTransformer,
+    # Base processor classes
+    BaseConsumer,
+    BaseConsumerUnit,
+    BaseProcessor,
+    BaseProcessorUnit,
+    BaseProducer,
+    BaseProducerUnit,
+    BaseStatefulConsumer,
+    BaseStatefulProcessor,
+    BaseStatefulProducer,
+    BaseStatefulTransformer,
+    BaseTransformer,
+    BaseTransformerUnit,
+    # Composite classes
+    CompositeProcessor,
+    CompositeProducer,
+    CompositeStateful,
+    Consumer,
+    ConsumerType,
+    GenAxisArray,
+    MessageInType,
+    MessageOutType,
+    Processor,
+    Producer,
+    ProducerType,
+    # Message types
+    SampleMessage,
+    SettingsType,
+    Stateful,
+    StatefulConsumer,
+    StatefulProcessor,
+    StatefulProducer,
+    StatefulTransformer,
+    StateType,
+    Transformer,
+    TransformerType,
+    # Type resolution helpers
+    _get_base_processor_message_in_type,
+    _get_base_processor_message_out_type,
+    _get_base_processor_settings_type,
+    _get_base_processor_state_type,
+    _get_processor_message_type,
+    # Type utilities
+    check_message_type_compatibility,
+    get_base_adaptive_transformer_type,
+    get_base_consumer_type,
+    get_base_producer_type,
+    get_base_transformer_type,
+    is_sample_message,
+    # Decorators
+    processor_state,
+    # Profiling
+    profile_subpub,
+    resolve_typevar,
+)
 
-    # Method to be implemented by subclasses to construct the specific generator
-    def construct_generator(self):
-        raise NotImplementedError
-
-    @ez.subscriber(INPUT_SETTINGS)
-    async def on_settings(self, msg: ez.Settings) -> None:
-        self.apply_settings(msg)
-        self.construct_generator()
-
-    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
-    @ez.publisher(OUTPUT_SIGNAL)
-    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
-        try:
-            ret = self.STATE.gen.send(message)
-            if math.prod(ret.data.shape) > 0:
-                yield self.OUTPUT_SIGNAL, ret
-        except (StopIteration, GeneratorExit):
-            ez.logger.debug(f"Generator closed in {self.address}")
-        except Exception:
-            ez.logger.info(traceback.format_exc())
+__all__ = [
+    # Protocols
+    "Processor",
+    "Producer",
+    "Consumer",
+    "Transformer",
+    "StatefulProcessor",
+    "StatefulProducer",
+    "StatefulConsumer",
+    "StatefulTransformer",
+    "AdaptiveTransformer",
+    # Type variables
+    "MessageInType",
+    "MessageOutType",
+    "SettingsType",
+    "StateType",
+    "ProducerType",
+    "ConsumerType",
+    "TransformerType",
+    "AdaptiveTransformerType",
+    # Decorators
+    "processor_state",
+    # Base processor classes
+    "BaseProcessor",
+    "BaseProducer",
+    "BaseConsumer",
+    "BaseTransformer",
+    # Stateful classes
+    "Stateful",
+    "BaseStatefulProcessor",
+    "BaseStatefulProducer",
+    "BaseStatefulConsumer",
+    "BaseStatefulTransformer",
+    "BaseAdaptiveTransformer",
+    "BaseAsyncTransformer",
+    # Composite classes
+    "CompositeStateful",
+    "CompositeProcessor",
+    "CompositeProducer",
+    # Unit classes
+    "BaseProducerUnit",
+    "BaseProcessorUnit",
+    "BaseConsumerUnit",
+    "BaseTransformerUnit",
+    "BaseAdaptiveTransformerUnit",
+    "GenAxisArray",
+    # Type resolution helpers
+    "get_base_producer_type",
+    "get_base_consumer_type",
+    "get_base_transformer_type",
+    "get_base_adaptive_transformer_type",
+    "_get_base_processor_settings_type",
+    "_get_base_processor_message_in_type",
+    "_get_base_processor_message_out_type",
+    "_get_base_processor_state_type",
+    "_get_processor_message_type",
+    # Message types
+    "SampleMessage",
+    "is_sample_message",
+    # Profiling
+    "profile_subpub",
+    # Type utilities
+    "check_message_type_compatibility",
+    "resolve_typevar",
+]