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

ezmsg/sigproc/aggregate.py

@@ -0,0 +1,276 @@
+"""
+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
+from array_api_compat import get_namespace
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    AxisBase,
+    replace,
+    slice_along_axis,
+)
+
+from .spectral import OptionsEnum
+
+
+class AggregationFunction(OptionsEnum):
+    """Enum for aggregation functions available to be used in :obj:`ranged_aggregate` operation."""
+
+    NONE = "None (all)"
+    MAX = "max"
+    MIN = "min"
+    MEAN = "mean"
+    MEDIAN = "median"
+    STD = "std"
+    SUM = "sum"
+    NANMAX = "nanmax"
+    NANMIN = "nanmin"
+    NANMEAN = "nanmean"
+    NANMEDIAN = "nanmedian"
+    NANSTD = "nanstd"
+    NANSUM = "nansum"
+    ARGMIN = "argmin"
+    ARGMAX = "argmax"
+    TRAPEZOID = "trapezoid"
+
+
+AGGREGATORS = {
+    AggregationFunction.NONE: np.all,
+    AggregationFunction.MAX: np.max,
+    AggregationFunction.MIN: np.min,
+    AggregationFunction.MEAN: np.mean,
+    AggregationFunction.MEDIAN: np.median,
+    AggregationFunction.STD: np.std,
+    AggregationFunction.SUM: np.sum,
+    AggregationFunction.NANMAX: np.nanmax,
+    AggregationFunction.NANMIN: np.nanmin,
+    AggregationFunction.NANMEAN: np.nanmean,
+    AggregationFunction.NANMEDIAN: np.nanmedian,
+    AggregationFunction.NANSTD: np.nanstd,
+    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,
+}
+
+
+class RangedAggregateSettings(ez.Settings):
+    """
+    Settings for ``RangedAggregate``.
+    """
+
+    axis: str | None = None
+    """The name of the axis along which to apply the bands."""
+
+    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.
+    """
+
+    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
+
+
+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:
+            self._state.ax_vec = target_axis.value(np.arange(message.data.shape[ax_idx]))
+
+        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"):
+                if self._state.ax_vec.dtype.type is np.str_:
+                    sl_dat = f"{self._state.ax_vec[start]} - {self._state.ax_vec[stop]}"
+                else:
+                    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,
+        )
+
+    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 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,
+            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)
+
+        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:
+    """
+    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."""
+
+    operation: AggregationFunction = AggregationFunction.MEAN
+    """:obj:`AggregationFunction` to apply."""
+
+
+class AggregateTransformer(BaseTransformer[AggregateSettings, AxisArray, AxisArray]):
+    """
+    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.
+    """
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        axis_idx = message.get_axis_idx(self.settings.axis)
+        op = 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

@@ -0,0 +1,80 @@
+from dataclasses import field
+
+import ezmsg.core as ez
+from ezmsg.baseproc import (
+    BaseProcessor,
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .aggregate import (
+    AggregationFunction,
+    RangedAggregateSettings,
+    RangedAggregateTransformer,
+)
+from .spectrogram import SpectrogramSettings, SpectrogramTransformer
+
+
+class BandPowerSettings(ez.Settings):
+    """
+    Settings for ``BandPower``.
+    """
+
+    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 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 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

@@ -0,0 +1,149 @@
+"""
+Backwards-compatible re-exports from ezmsg.baseproc.
+
+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.
+"""
+
+import warnings
+
+warnings.warn(
+    "Importing from 'ezmsg.sigproc.base' is deprecated. Please import from 'ezmsg.baseproc' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+# 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,
+)
+
+__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",
+]

ezmsg/sigproc/butterworthfilter.py

@@ -1,18 +1,59 @@
-import ezmsg.core as ez
+import functools
+import typing
+
 import scipy.signal
-import numpy as np
+from scipy.signal import normalize
+
+from .filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+    SOSCoeffs,
+)
 
-from .filter import Filter, FilterState, FilterSettingsBase
 
-from typing import Optional, Tuple, Union
+class ButterworthFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`ButterworthFilter`."""
 
+    # axis and coef_type are inherited from FilterBaseSettings
 
-class ButterworthFilterSettings(FilterSettingsBase):
     order: int = 0
-    cuton: Optional[float] = None  # Hz
-    cutoff: Optional[float] = None  # Hz
+    """
+    Filter order
+    """
+
+    cuton: float | None = None
+    """
+    Cuton frequency (Hz). If `cutoff` is not specified then this is the highpass corner. Otherwise,
+    if this is lower than `cutoff` then this is the beginning of the bandpass
+    or if this is greater than `cutoff` then this is the end of the bandstop.
+    """
+
+    cutoff: float | None = None
+    """
+    Cutoff frequency (Hz). If `cuton` is not specified then this is the lowpass corner. Otherwise,
+    if this is greater than `cuton` then this is the end of the bandpass,
+    or if this is less than `cuton` then this is the beginning of the bandstop.
+    """
+
+    wn_hz: bool = True
+    """
+    Set False if provided Wn are normalized from 0 to 1, where 1 is the Nyquist frequency
+    """
+
+    def filter_specs(
+        self,
+    ) -> tuple[str, float | tuple[float, float]] | None:
+        """
+        Determine the filter type given the corner frequencies.
 
-    def filter_specs(self) -> Optional[Tuple[str, Union[float, Tuple[float, float]]]]:
+        Returns:
+            A tuple with the first element being a string indicating the filter type
+            (one of "lowpass", "highpass", "bandpass", "bandstop")
+            and the second element being the corner frequency or frequencies.
+
+        """
         if self.cuton is None and self.cutoff is None:
             return None
         elif self.cuton is None and self.cutoff is not None:
@@ -26,41 +67,90 @@ class ButterworthFilterSettings(FilterSettingsBase):
                 return "bandstop", (self.cutoff, self.cuton)
 
 
-class ButterworthFilterState(FilterState):
-    design: ButterworthFilterSettings
+def butter_design_fun(
+    fs: float,
+    order: int = 0,
+    cuton: float | None = None,
+    cutoff: float | None = None,
+    coef_type: str = "ba",
+    wn_hz: bool = True,
+) -> BACoeffs | SOSCoeffs | None:
+    """
+    See :obj:`ButterworthFilterSettings.filter_specs` for an explanation of specifying different
+    filter types (lowpass, highpass, bandpass, bandstop) from the parameters.
+    You are likely to want to use this function with :obj:`filter_by_design`, which only passes `fs` to the design
+    function (this), meaning that you should wrap this function with a lambda or prepare with functools.partial.
 
+    Args:
+        fs: The sampling frequency of the data in Hz.
+        order: Filter order.
+        cuton: Corner frequency of the filter in Hz.
+        cutoff: Corner frequency of the filter in Hz.
+        coef_type: "ba", "sos", or "zpk"
+        wn_hz: Set False if provided Wn are normalized from 0 to 1, where 1 is the Nyquist frequency
 
-class ButterworthFilter(Filter):
-    SETTINGS: ButterworthFilterSettings
-    STATE: ButterworthFilterState
+    Returns:
+        The filter coefficients as a tuple of (b, a) for coef_type "ba", or as a single ndarray for "sos",
+        or (z, p, k) for "zpk".
 
-    INPUT_FILTER = ez.InputStream(ButterworthFilterSettings)
+    """
+    coefs = None
+    if order > 0:
+        btype, cutoffs = ButterworthFilterSettings(order=order, cuton=cuton, cutoff=cutoff).filter_specs()
+        coefs = scipy.signal.butter(
+            order,
+            Wn=cutoffs,
+            btype=btype,
+            fs=fs if wn_hz else None,
+            output=coef_type,
+        )
+    if coefs is not None and coef_type == "ba":
+        coefs = normalize(*coefs)
+    return coefs
 
-    def initialize(self) -> None:
-        self.STATE.design = self.SETTINGS
-        self.STATE.filt_designed = True
-        super().initialize()
 
-    def design_filter(self) -> Optional[Tuple[np.ndarray, np.ndarray]]:
-        specs = self.STATE.design.filter_specs()
-        if self.STATE.design.order > 0 and specs is not None:
-            btype, cut = specs
-            return scipy.signal.butter(
-                self.STATE.design.order,
-                Wn=cut,
-                btype=btype,
-                fs=self.STATE.fs,
-                output="ba",
-            )
-        else:
-            return None
+class ButterworthFilterTransformer(FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]):
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        return functools.partial(
+            butter_design_fun,
+            order=self.settings.order,
+            cuton=self.settings.cuton,
+            cutoff=self.settings.cutoff,
+            coef_type=self.settings.coef_type,
+            wn_hz=self.settings.wn_hz,
+        )
+
+
+class ButterworthFilter(BaseFilterByDesignTransformerUnit[ButterworthFilterSettings, ButterworthFilterTransformer]):
+    SETTINGS = ButterworthFilterSettings
+
 
-    @ez.subscriber(INPUT_FILTER)
-    async def redesign(self, message: ButterworthFilterSettings) -> None:
-        if type(message) is not ButterworthFilterSettings:
-            return
+def butter(
+    axis: str | None,
+    order: int = 0,
+    cuton: float | None = None,
+    cutoff: float | None = None,
+    coef_type: str = "ba",
+    wn_hz: bool = True,
+) -> ButterworthFilterTransformer:
+    """
+    Convenience generator wrapping filter_gen_by_design for Butterworth filters.
+    Apply Butterworth filter to streaming data. Uses :obj:`scipy.signal.butter` to design the filter.
+    See :obj:`ButterworthFilterSettings.filter_specs` for an explanation of specifying different
+    filter types (lowpass, highpass, bandpass, bandstop) from the parameters.
 
-        if self.STATE.design.order != message.order:
-            self.STATE.zi = None
-        self.STATE.design = message
-        self.update_filter()
+    Returns:
+        :obj:`ButterworthFilterTransformer`
+    """
+    return ButterworthFilterTransformer(
+        ButterworthFilterSettings(
+            axis=axis,
+            order=order,
+            cuton=cuton,
+            cutoff=cutoff,
+            coef_type=coef_type,
+            wn_hz=wn_hz,
+        )
+    )