ezmsg-sigproc 1.4.2__py3-none-any.whl → 1.6.0__py3-none-any.whl

ezmsg/sigproc/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.4.2'
-__version_tuple__ = version_tuple = (1, 4, 2)
+__version__ = version = '1.6.0'
+__version_tuple__ = version_tuple = (1, 6, 0)

ezmsg/sigproc/activation.py

@@ -1,10 +1,10 @@
-from dataclasses import replace
 import typing
 
 import numpy as np
 import scipy.special
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 from .spectral import OptionsEnum
@@ -41,7 +41,7 @@ ACTIVATIONS = {
 
 @consumer
 def activation(
-    function: typing.Union[str, ActivationFunction],
+    function: str | ActivationFunction,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Transform the data with a simple activation function.

ezmsg/sigproc/affinetransform.py

@@ -1,4 +1,3 @@
-from dataclasses import replace
 import os
 from pathlib import Path
 import typing
@@ -6,7 +5,8 @@ import typing
 import numpy as np
 import numpy.typing as npt
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.axisarray import AxisArray, AxisBase
+from ezmsg.util.messages.util import replace
 from ezmsg.util.generator import consumer
 
 from .base import GenAxisArray
@@ -14,8 +14,8 @@ from .base import GenAxisArray
 
 @consumer
 def affine_transform(
-    weights: typing.Union[np.ndarray, str, Path],
-    axis: typing.Optional[str] = None,
+    weights: np.ndarray | str | Path,
+    axis: str | None = None,
     right_multiply: bool = True,
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
@@ -47,7 +47,7 @@ def affine_transform(
 
     # State variables
     # New axis with transformed labels, if required
-    new_axis: typing.Optional[AxisArray.Axis] = None
+    new_axis: AxisBase | None = None
 
     # Reset if any of these change.
     check_input = {"key": None}
@@ -71,10 +71,10 @@ def affine_transform(
             # Determine if we need to modify the transformed axis.
             if (
                 axis in msg_in.axes
-                and hasattr(msg_in.axes[axis], "labels")
+                and hasattr(msg_in.axes[axis], "data")
                 and weights.shape[0] != weights.shape[1]
             ):
-                in_labels = msg_in.axes[axis].labels
+                in_labels = msg_in.axes[axis].data
                 new_labels = []
                 n_in, n_out = weights.shape
                 if len(in_labels) != n_in:
@@ -101,8 +101,8 @@ def affine_transform(
                                 new_labels.append("")
                     elif np.all(b_filled_outputs):
                         # Transform is dropping some of the inputs.
-                        new_labels = np.array(in_labels)[b_used_inputs].tolist()
-                new_axis = replace(msg_in.axes[axis], labels=new_labels)
+                        new_labels = np.array(in_labels)[b_used_inputs]
+                new_axis = replace(msg_in.axes[axis], data=np.array(new_labels))
 
         data = msg_in.data
 
@@ -133,8 +133,8 @@ class AffineTransformSettings(ez.Settings):
     See :obj:`affine_transform` for argument details.
     """
 
-    weights: typing.Union[np.ndarray, str, Path]
-    axis: typing.Optional[str] = None
+    weights: np.ndarray | str | Path
+    axis: str | None = None
     right_multiply: bool = True
 
 
@@ -157,7 +157,7 @@ def zeros_for_noop(data: npt.NDArray, **ignore_kwargs) -> npt.NDArray:
 
 @consumer
 def common_rereference(
-    mode: str = "mean", axis: typing.Optional[str] = None, include_current: bool = True
+    mode: str = "mean", axis: str | None = None, include_current: bool = True
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
     Perform common average referencing (CAR) on streaming data.
@@ -214,7 +214,7 @@ class CommonRereferenceSettings(ez.Settings):
     """
 
     mode: str = "mean"
-    axis: typing.Optional[str] = None
+    axis: str | None = None
     include_current: bool = True
 
 

ezmsg/sigproc/aggregate.py

@@ -1,11 +1,15 @@
-from dataclasses import replace
 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
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    slice_along_axis,
+    AxisBase,
+    replace,
+)
 
 from .spectral import OptionsEnum
 from .base import GenAxisArray
@@ -52,8 +56,8 @@ AGGREGATORS = {
 
 @consumer
 def ranged_aggregate(
-    axis: typing.Optional[str] = None,
-    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None,
+    axis: str | None = None,
+    bands: list[tuple[float, float]] | None = None,
     operation: AggregationFunction = AggregationFunction.MEAN,
 ):
     """
@@ -71,12 +75,12 @@ def ranged_aggregate(
     msg_out = AxisArray(np.array([]), dims=[""])
 
     # State variables
-    slices: typing.Optional[typing.List[typing.Tuple[typing.Any, ...]]] = None
-    out_axis: typing.Optional[AxisArray.Axis] = None
-    ax_vec: typing.Optional[npt.NDArray] = None
+    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}
+    check_inputs = {"gain": None, "offset": None, "len": None, "key": None}
 
     while True:
         msg_in: AxisArray = yield msg_out
@@ -86,35 +90,52 @@ def ranged_aggregate(
             axis = axis or msg_in.dims[0]
             target_axis = msg_in.get_axis(axis)
 
-            b_reset = target_axis.gain != check_inputs["gain"]
-            b_reset = b_reset or target_axis.offset != check_inputs["offset"]
+            # Check if we need to reset state
+            b_reset = msg_in.key != check_inputs["key"]
+            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:
-                check_inputs["gain"] = target_axis.gain
-                check_inputs["offset"] = target_axis.offset
+                # Update check variables
+                check_inputs["key"] = msg_in.key
+                if hasattr(target_axis, "data"):
+                    check_inputs["len"] = len(target_axis.data)
+                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_idx = msg_in.get_axis_idx(axis)
 
-                ax_vec = (
-                    target_axis.offset
-                    + np.arange(msg_in.data.shape[ax_idx]) * target_axis.gain
-                )
+                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 = []
-                mids = []
+                ax_dat = []
                 for start, stop in bands:
                     inds = np.where(np.logical_and(ax_vec >= start, ax_vec <= stop))[0]
-                    mids.append(np.mean(inds) * target_axis.gain + target_axis.offset)
                     slices.append(np.s_[inds[0] : inds[-1] + 1])
-                out_ax_kwargs = {
-                    "unit": target_axis.unit,
-                    "offset": mids[0],
-                    "gain": (mids[1] - mids[0]) if len(mids) > 1 else 1.0,
-                }
-                if hasattr(target_axis, "labels"):
-                    out_ax_kwargs["labels"] = [f"{_[0]} - {_[1]}" for _ in bands]
-                out_axis = replace(target_axis, **out_ax_kwargs)
+                    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,
+                )
 
             agg_func = AGGREGATORS[operation]
             out_data = [
@@ -142,8 +163,8 @@ class RangedAggregateSettings(ez.Settings):
     See :obj:`ranged_aggregate` for details.
     """
 
-    axis: typing.Optional[str] = None
-    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None
+    axis: str | None = None
+    bands: list[tuple[float, float]] | None = None
     operation: AggregationFunction = AggregationFunction.MEAN
 
 

ezmsg/sigproc/bandpower.py

@@ -14,7 +14,7 @@ from .base import GenAxisArray
 @consumer
 def bandpower(
     spectrogram_settings: SpectrogramSettings,
-    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = [
+    bands: list[tuple[float, float]] | None = [
         (17, 30),
         (70, 170),
     ],
@@ -58,7 +58,7 @@ class BandPowerSettings(ez.Settings):
     spectrogram_settings: SpectrogramSettings = field(
         default_factory=SpectrogramSettings
     )
-    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = field(
+    bands: list[tuple[float, float]] | None = field(
         default_factory=lambda: [(17, 30), (70, 170)]
     )
 

ezmsg/sigproc/butterworthfilter.py

@@ -1,15 +1,19 @@
+import functools
 import typing
 
-import ezmsg.core as ez
 import scipy.signal
-import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.generator import consumer
+from scipy.signal import normalize
 
-from .filter import filtergen, Filter, FilterState, FilterSettingsBase
+from .filter import (
+    FilterBaseSettings,
+    FilterCoefsMultiType,
+    FilterBase,
+    filter_gen_by_design,
+)
 
 
-class ButterworthFilterSettings(FilterSettingsBase):
+class ButterworthFilterSettings(FilterBaseSettings):
     """Settings for :obj:`ButterworthFilter`."""
 
     order: int = 0
@@ -17,25 +21,28 @@ class ButterworthFilterSettings(FilterSettingsBase):
     Filter order
     """
 
-    cuton: typing.Optional[float] = None
+    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: typing.Optional[float] = None
+    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,
-    ) -> typing.Optional[
-        typing.Tuple[str, typing.Union[float, typing.Tuple[float, float]]]
-    ]:
+    ) -> tuple[str, float | tuple[float, float]] | None:
         """
         Determine the filter type given the corner frequencies.
 
@@ -58,21 +65,81 @@ class ButterworthFilterSettings(FilterSettingsBase):
                 return "bandstop", (self.cutoff, self.cuton)
 
 
-@consumer
+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,
+) -> FilterCoefsMultiType | 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
+
+    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".
+
+    """
+    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
+
+
+class ButterworthFilter(FilterBase):
+    SETTINGS = ButterworthFilterSettings
+
+    def design_filter(
+        self,
+    ) -> typing.Callable[[float], FilterCoefsMultiType | 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,
+        )
+
+
 def butter(
-    axis: typing.Optional[str],
+    axis: str | None,
     order: int = 0,
-    cuton: typing.Optional[float] = None,
-    cutoff: typing.Optional[float] = None,
+    cuton: float | None = None,
+    cutoff: float | None = None,
     coef_type: str = "ba",
 ) -> typing.Generator[AxisArray, AxisArray, None]:
     """
+    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.
 
     Args:
         axis: The name of the axis to filter.
+            Note: The axis must be represented in the message .axes and be of type AxisArray.LinearAxis.
         order: Filter order.
         cuton: Corner frequency of the filter in Hz.
         cutoff: Corner frequency of the filter in Hz.
@@ -83,79 +150,11 @@ def butter(
          and yields an :obj:`AxisArray` with filtered data.
 
     """
-    # IO
-    msg_out = AxisArray(np.array([]), dims=[""])
-
-    # Check parameters
-    btype, cutoffs = ButterworthFilterSettings(
-        order=order, cuton=cuton, cutoff=cutoff
-    ).filter_specs()
-
-    # State variables
-    # Initialize filtergen as passthrough until we can calculate coefs.
-    filter_gen = filtergen(axis, None, coef_type)
-
-    # Reset if these change.
-    check_input = {"gain": None}
-    # Key not checked because filter_gen will handle resetting if .key changes.
-
-    while True:
-        msg_in: AxisArray = yield msg_out
-        axis = axis or msg_in.dims[0]
-
-        b_reset = msg_in.axes[axis].gain != check_input["gain"]
-        b_reset = b_reset and order > 0  # Not passthrough
-        if b_reset:
-            check_input["gain"] = msg_in.axes[axis].gain
-            coefs = scipy.signal.butter(
-                order,
-                Wn=cutoffs,
-                btype=btype,
-                fs=1 / msg_in.axes[axis].gain,
-                output=coef_type,
-            )
-            filter_gen = filtergen(axis, coefs, coef_type)
-
-        msg_out = filter_gen.send(msg_in)
-
-
-class ButterworthFilterState(FilterState):
-    design: ButterworthFilterSettings
-
-
-class ButterworthFilter(Filter):
-    """:obj:`Unit` for :obj:`butterworth`"""
-
-    SETTINGS = ButterworthFilterSettings
-    STATE = ButterworthFilterState
-
-    INPUT_FILTER = ez.InputStream(ButterworthFilterSettings)
-
-    async def initialize(self) -> None:
-        self.STATE.design = self.SETTINGS
-        self.STATE.filt_designed = True
-        await super().initialize()
-
-    def design_filter(self) -> typing.Optional[typing.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
-
-    @ez.subscriber(INPUT_FILTER)
-    async def redesign(self, message: ButterworthFilterSettings) -> None:
-        if type(message) is not ButterworthFilterSettings:
-            return
-
-        if self.STATE.design.order != message.order:
-            self.STATE.zi = None
-        self.STATE.design = message
-        self.update_filter()
+    design_fun = functools.partial(
+        butter_design_fun,
+        order=order,
+        cuton=cuton,
+        cutoff=cutoff,
+        coef_type=coef_type,
+    )
+    return filter_gen_by_design(axis, coef_type, design_fun)

ezmsg/sigproc/cheby.py

@@ -0,0 +1,119 @@
+import functools
+import typing
+
+import scipy.signal
+from scipy.signal import normalize
+
+from .filter import (
+    FilterBaseSettings,
+    FilterCoefsMultiType,
+    FilterBase,
+)
+
+
+class ChebyshevFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`ButterworthFilter`."""
+
+    order: int = 0
+    """
+    Filter order
+    """
+
+    ripple_tol: float | None = None
+    """
+    The maximum ripple allowed below unity gain in the passband. Specified in decibels, as a positive number.
+    """
+
+    Wn: float | tuple[float, float] | None = None
+    """
+    A scalar or length-2 sequence giving the critical frequencies.
+    For Type I filters, this is the point in the transition band at which the gain first drops below -rp.
+    For digital filters, Wn are in the same units as fs unless wn_hz is False.
+    For analog filters, Wn is an angular frequency (e.g., rad/s).
+    """
+
+    btype: str = "lowpass"
+    """
+    {‘lowpass’, ‘highpass’, ‘bandpass’, ‘bandstop’}
+    """
+
+    analog: bool = False
+    """
+    When True, return an analog filter, otherwise a digital filter is returned.
+    """
+
+    cheby_type: str = "cheby1"
+    """
+    Which type of Chebyshev filter to design. Either "cheby1" or "cheby2".
+    """
+
+    wn_hz: bool = True
+    """
+    Set False if provided Wn are normalized from 0 to 1, where 1 is the Nyquist frequency
+    """
+
+
+def cheby_design_fun(
+    fs: float,
+    order: int = 0,
+    ripple_tol: float | None = None,
+    Wn: float | tuple[float, float] | None = None,
+    btype: str = "lowpass",
+    analog: bool = False,
+    coef_type: str = "ba",
+    cheby_type: str = "cheby1",
+    wn_hz: bool = True,
+) -> FilterCoefsMultiType:
+    """
+    Chebyshev type I and type II digital and analog filter design.
+    Design an `order`th-order digital or analog Chebyshev type I or type II filter and return the filter coefficients.
+    See :obj:`ChebyFilterSettings` for argument description.
+
+    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".
+    """
+    coefs = None
+    if order > 0:
+        if cheby_type == "cheby1":
+            coefs = scipy.signal.cheby1(
+                order,
+                ripple_tol,
+                Wn,
+                btype=btype,
+                analog=analog,
+                output=coef_type,
+                fs=fs if wn_hz else None,
+            )
+        elif cheby_type == "cheby2":
+            coefs = scipy.signal.cheby2(
+                order,
+                ripple_tol,
+                Wn,
+                btype=btype,
+                analog=analog,
+                output=coef_type,
+                fs=fs,
+            )
+    if coefs is not None and coef_type == "ba":
+        coefs = normalize(*coefs)
+    return coefs
+
+
+class ChebyshevFilter(FilterBase):
+    SETTINGS = ChebyshevFilterSettings
+
+    def design_filter(
+        self,
+    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
+        return functools.partial(
+            cheby_design_fun,
+            order=self.SETTINGS.order,
+            ripple_tol=self.SETTINGS.ripple_tol,
+            Wn=self.SETTINGS.Wn,
+            btype=self.SETTINGS.btype,
+            analog=self.SETTINGS.analog,
+            coef_type=self.SETTINGS.coef_type,
+            cheby_type=self.SETTINGS.cheby_type,
+            wn_hz=self.SETTINGS.wn_hz,
+        )

ezmsg/sigproc/decimate.py

@@ -1,9 +1,8 @@
-import scipy.signal
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 
+from .cheby import ChebyshevFilter, ChebyshevFilterSettings
 from .downsample import Downsample, DownsampleSettings
-from .filter import Filter, FilterCoefficients, FilterSettings
 
 
 class Decimate(ez.Collection):
@@ -17,24 +16,21 @@ class Decimate(ez.Collection):
     INPUT_SIGNAL = ez.InputStream(AxisArray)
     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
 
-    FILTER = Filter()
+    FILTER = ChebyshevFilter()
     DOWNSAMPLE = Downsample()
 
     def configure(self) -> None:
+        cheby_settings = ChebyshevFilterSettings(
+            order=8 if self.SETTINGS.factor > 1 else 0,
+            ripple_tol=0.05,
+            Wn=0.8 / self.SETTINGS.factor if self.SETTINGS.factor > 1 else None,
+            btype="lowpass",
+            axis=self.SETTINGS.axis,
+            wn_hz=False,
+        )
+        self.FILTER.apply_settings(cheby_settings)
         self.DOWNSAMPLE.apply_settings(self.SETTINGS)
 
-        if self.SETTINGS.factor < 1:
-            raise ValueError("Decimation factor must be >= 1 (no decimation")
-        elif self.SETTINGS.factor == 1:
-            filt = FilterCoefficients()
-        else:
-            # See scipy.signal.decimate for IIR Filter Condition
-            b, a = scipy.signal.cheby1(8, 0.05, 0.8 / self.SETTINGS.factor)
-            system = scipy.signal.dlti(b, a)
-            filt = FilterCoefficients(b=system.num, a=system.den)  # type: ignore
-
-        self.FILTER.apply_settings(FilterSettings(filt=filt))
-
     def network(self) -> ez.NetworkDefinition:
         return (
             (self.INPUT_SIGNAL, self.FILTER.INPUT_SIGNAL),