ezmsg-sigproc 1.8.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

ezmsg/sigproc/butterworthfilter.py

@@ -2,20 +2,22 @@ import functools
 import typing
 
 import scipy.signal
-from ezmsg.util.messages.axisarray import AxisArray
 from scipy.signal import normalize
 
 from .filter import (
     FilterBaseSettings,
-    FilterCoefsMultiType,
-    FilterBase,
-    filter_gen_by_design,
+    BACoeffs,
+    SOSCoeffs,
+    FilterByDesignTransformer,
+    BaseFilterByDesignTransformerUnit,
 )
 
 
 class ButterworthFilterSettings(FilterBaseSettings):
     """Settings for :obj:`ButterworthFilter`."""
 
+    # axis and coef_type are inherited from FilterBaseSettings
+
     order: int = 0
     """
     Filter order
@@ -72,7 +74,7 @@ def butter_design_fun(
     cutoff: float | None = None,
     coef_type: str = "ba",
     wn_hz: bool = True,
-) -> FilterCoefsMultiType | None:
+) -> BACoeffs | SOSCoeffs | None:
     """
     See :obj:`ButterworthFilterSettings.filter_specs` for an explanation of specifying different
     filter types (lowpass, highpass, bandpass, bandstop) from the parameters.
@@ -109,52 +111,54 @@ def butter_design_fun(
     return coefs
 
 
-class ButterworthFilter(FilterBase):
-    SETTINGS = ButterworthFilterSettings
-
-    def design_filter(
+class ButterworthFilterTransformer(
+    FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]
+):
+    def get_design_function(
         self,
-    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
+    ) -> 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,
+            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
+
+
 def butter(
     axis: str | None,
     order: int = 0,
     cuton: float | None = None,
     cutoff: float | None = None,
     coef_type: str = "ba",
-) -> typing.Generator[AxisArray, AxisArray, None]:
+    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.
 
-    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.
-        coef_type: "ba" or "sos"
-
     Returns:
-        A primed generator object which accepts an :obj:`AxisArray` via .send(axis_array)
-         and yields an :obj:`AxisArray` with filtered data.
-
+        :obj:`ButterworthFilterTransformer`
     """
-    design_fun = functools.partial(
-        butter_design_fun,
-        order=order,
-        cuton=cuton,
-        cutoff=cutoff,
-        coef_type=coef_type,
+    return ButterworthFilterTransformer(
+        ButterworthFilterSettings(
+            axis=axis,
+            order=order,
+            cuton=cuton,
+            cutoff=cutoff,
+            coef_type=coef_type,
+            wn_hz=wn_hz,
+        )
     )
-    return filter_gen_by_design(axis, coef_type, design_fun)

ezmsg/sigproc/cheby.py

@@ -6,13 +6,17 @@ from scipy.signal import normalize
 
 from .filter import (
     FilterBaseSettings,
-    FilterCoefsMultiType,
-    FilterBase,
+    FilterByDesignTransformer,
+    BACoeffs,
+    SOSCoeffs,
+    BaseFilterByDesignTransformerUnit,
 )
 
 
 class ChebyshevFilterSettings(FilterBaseSettings):
-    """Settings for :obj:`ButterworthFilter`."""
+    """Settings for :obj:`ChebyshevFilter`."""
+
+    # axis and coef_type are inherited from FilterBaseSettings
 
     order: int = 0
     """
@@ -63,7 +67,7 @@ def cheby_design_fun(
     coef_type: str = "ba",
     cheby_type: str = "cheby1",
     wn_hz: bool = True,
-) -> FilterCoefsMultiType:
+) -> BACoeffs | SOSCoeffs | None:
     """
     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.
@@ -100,20 +104,28 @@ def cheby_design_fun(
     return coefs
 
 
-class ChebyshevFilter(FilterBase):
-    SETTINGS = ChebyshevFilterSettings
-
-    def design_filter(
+class ChebyshevFilterTransformer(
+    FilterByDesignTransformer[ChebyshevFilterSettings, BACoeffs | SOSCoeffs]
+):
+    def get_design_function(
         self,
-    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | 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,
+            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,
         )
+
+
+class ChebyshevFilter(
+    BaseFilterByDesignTransformerUnit[
+        ChebyshevFilterSettings, ChebyshevFilterTransformer
+    ]
+):
+    SETTINGS = ChebyshevFilterSettings

ezmsg/sigproc/combfilter.py

@@ -0,0 +1,163 @@
+import functools
+import typing
+import numpy as np
+import scipy.signal
+from scipy.signal import normalize
+
+from .filter import (
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+    BACoeffs,
+    SOSCoeffs,
+    BaseFilterByDesignTransformerUnit,
+)
+
+
+class CombFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`CombFilter`."""
+
+    # axis and coef_type are inherited from FilterBaseSettings
+
+    fundamental_freq: float = 60.0
+    """
+    Fundamental frequency in Hz
+    """
+
+    num_harmonics: int = 3
+    """
+    Number of harmonics to include (including fundamental)
+    """
+
+    q_factor: float = 35.0
+    """
+    Quality factor (Q) for each peak/notch
+    """
+
+    filter_type: str = "notch"
+    """
+    Type of comb filter: 'notch' removes harmonics, 'peak' passes harmonics at the expense of others.
+    """
+
+    quality_scaling: str = "constant"
+    """
+    'constant': same quality for all harmonics results in wider bands at higher frequencies,
+    'proportional': quality proportional to frequency results in constant bandwidths.
+    """
+
+
+def comb_design_fun(
+    fs: float,
+    fundamental_freq: float = 60.0,
+    num_harmonics: int = 3,
+    q_factor: float = 35.0,
+    filter_type: str = "notch",
+    coef_type: str = "sos",
+    quality_scaling: str = "constant",
+) -> BACoeffs | SOSCoeffs | None:
+    """
+    Design a comb filter as cascaded second-order sections targeting a fundamental frequency and its harmonics.
+
+    Returns:
+        The filter coefficients as SOS (recommended) or (b, a) for finite precision stability.
+    """
+    if coef_type != "sos" and coef_type != "ba":
+        raise ValueError("Comb filter only supports 'sos' or 'ba' coefficient types")
+
+    # Generate all SOS sections
+    all_sos = []
+
+    for i in range(1, num_harmonics + 1):
+        freq = fundamental_freq * i
+
+        # Skip if frequency exceeds Nyquist
+        if freq >= fs / 2:
+            continue
+
+        # Adjust Q factor based on scaling method
+        current_q = q_factor
+        if quality_scaling == "proportional":
+            current_q = q_factor * i
+
+        if filter_type == "notch":
+            sos = scipy.signal.iirnotch(w0=freq, Q=current_q, fs=fs)
+        else:  # peak filter
+            sos = scipy.signal.iirpeak(w0=freq, Q=current_q, fs=fs)
+        # Though .iirnotch and .iirpeak return b, a pairs, these are second order so
+        #  we can use them directly as SOS sections.
+        #  Check:
+        # assert np.allclose(scipy.signal.tf2sos(sos[0], sos[1])[0], np.hstack(sos))
+
+        all_sos.append(np.hstack(sos))
+
+    if not all_sos:
+        return None
+
+    # Combine all SOS sections
+    combined_sos = np.vstack(all_sos)
+
+    if coef_type == "ba":
+        # Convert to transfer function form
+        b, a = scipy.signal.sos2tf(combined_sos)
+        return normalize(b, a)
+
+    return combined_sos
+
+
+class CombFilterTransformer(
+    FilterByDesignTransformer[CombFilterSettings, BACoeffs | SOSCoeffs]
+):
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        return functools.partial(
+            comb_design_fun,
+            fundamental_freq=self.settings.fundamental_freq,
+            num_harmonics=self.settings.num_harmonics,
+            q_factor=self.settings.q_factor,
+            filter_type=self.settings.filter_type,
+            coef_type=self.settings.coef_type,
+            quality_scaling=self.settings.quality_scaling,
+        )
+
+
+class CombFilterUnit(
+    BaseFilterByDesignTransformerUnit[CombFilterSettings, CombFilterTransformer]
+):
+    SETTINGS = CombFilterSettings
+
+
+def comb(
+    axis: str | None,
+    fundamental_freq: float = 50.0,
+    num_harmonics: int = 3,
+    q_factor: float = 35.0,
+    filter_type: str = "notch",
+    coef_type: str = "sos",
+    quality_scaling: str = "constant",
+) -> CombFilterTransformer:
+    """
+    Create a comb filter for enhancing or removing a fundamental frequency and its harmonics.
+
+    Args:
+        axis: Axis to filter along
+        fundamental_freq: Base frequency in Hz
+        num_harmonics: Number of harmonic peaks/notches (including fundamental)
+        q_factor: Quality factor for peak/notch width
+        filter_type: 'notch' to remove or 'peak' to enhance harmonics
+        coef_type: Coefficient type ('sos' recommended for stability)
+        quality_scaling: How to handle bandwidths across harmonics
+
+    Returns:
+        CombFilterTransformer
+    """
+    return CombFilterTransformer(
+        CombFilterSettings(
+            axis=axis,
+            fundamental_freq=fundamental_freq,
+            num_harmonics=num_harmonics,
+            q_factor=q_factor,
+            filter_type=filter_type,
+            coef_type=coef_type,
+            quality_scaling=quality_scaling,
+        )
+    )

ezmsg/sigproc/decimate.py

@@ -3,31 +3,41 @@ import typing
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray
 
-from .cheby import ChebyshevFilter, ChebyshevFilterSettings
+from .base import BaseTransformerUnit
+from .cheby import ChebyshevFilterTransformer, ChebyshevFilterSettings
 from .downsample import Downsample, DownsampleSettings
-from .filter import FilterCoefsMultiType
+from .filter import BACoeffs, SOSCoeffs
 
 
-class ChebyForDecimate(ChebyshevFilter):
+class ChebyForDecimateTransformer(ChebyshevFilterTransformer[BACoeffs | SOSCoeffs]):
     """
-    A :obj:`ChebyshevFilter` node with a design filter method that additionally accepts a target sampling rate,
+    A :obj:`ChebyshevFilterTransformer` with a design filter method that additionally accepts a target sampling rate,
      and if the target rate cannot be achieved it returns None, else it returns the filter coefficients.
     """
 
-    def design_filter(
+    def get_design_function(
         self,
-    ) -> typing.Callable[[float], FilterCoefsMultiType | None]:
-        def cheby_opt_design_fun(fs: float) -> FilterCoefsMultiType | None:
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        def cheby_opt_design_fun(fs: float) -> BACoeffs | SOSCoeffs | None:
             if fs is None:
                 return None
-            ds_factor = int(fs / (2.5 * self.SETTINGS.Wn))
+            ds_factor = int(fs / (2.5 * self.settings.Wn))
             if ds_factor < 2:
                 return None
-            partial_fun = super(ChebyForDecimate, self).design_filter()
+            partial_fun = super(ChebyForDecimateTransformer, self).get_design_function()
             return partial_fun(fs)
+
         return cheby_opt_design_fun
 
 
+class ChebyForDecimate(
+    BaseTransformerUnit[
+        ChebyshevFilterSettings, AxisArray, AxisArray, ChebyForDecimateTransformer
+    ]
+):
+    SETTINGS = ChebyshevFilterSettings
+
+
 class Decimate(ez.Collection):
     """
     A :obj:`Collection` chaining a :obj:`Filter` node configured as a lowpass Chebyshev filter
@@ -43,7 +53,6 @@ class Decimate(ez.Collection):
     DOWNSAMPLE = Downsample()
 
     def configure(self) -> None:
-
         cheby_settings = ChebyshevFilterSettings(
             order=8,
             ripple_tol=0.05,

ezmsg/sigproc/detrend.py

@@ -0,0 +1,29 @@
+import scipy.signal as sps
+from ezmsg.util.messages.axisarray import AxisArray, replace
+from ezmsg.sigproc.ewma import EWMATransformer, EWMASettings
+from ezmsg.sigproc.base import BaseTransformerUnit
+
+
+class DetrendTransformer(EWMATransformer):
+    """
+    Detrend the data using an exponentially weighted moving average (EWMA)
+     estimate of the mean.
+    """
+
+    def _process(self, message):
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        means, self._state.zi = sps.lfilter(
+            [self._state.alpha],
+            [1.0, self._state.alpha - 1.0],
+            message.data,
+            axis=axis_idx,
+            zi=self._state.zi,
+        )
+        return replace(message, data=message.data - means)
+
+
+class DetrendUnit(
+    BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, DetrendTransformer]
+):
+    SETTINGS = EWMASettings

ezmsg/sigproc/diff.py

@@ -0,0 +1,81 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.sigproc.base import (
+    BaseTransformerUnit,
+    processor_state,
+    BaseStatefulTransformer,
+)
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
+
+
+class DiffSettings(ez.Settings):
+    axis: str | None = None
+    scale_by_fs: bool = False
+
+
+@processor_state
+class DiffState:
+    last_dat: npt.NDArray | None = None
+    last_time: float | None = None
+
+
+class DiffTransformer(
+    BaseStatefulTransformer[DiffSettings, AxisArray, AxisArray, DiffState]
+):
+    def _hash_message(self, message: AxisArray) -> int:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+        return hash((sample_shape, message.key))
+
+    def _reset_state(self, message) -> None:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        self.state.last_dat = slice_along_axis(message.data, slice(0, 1), axis=ax_idx)
+        if self.settings.scale_by_fs:
+            ax_info = message.get_axis(self.settings.axis)
+            if hasattr(ax_info, "data"):
+                if len(ax_info.data) > 1:
+                    self.state.last_time = 2 * ax_info.data[0] - ax_info.data[1]
+                else:
+                    self.state.last_time = ax_info.data[0] - 0.001
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[0]
+        ax_idx = message.get_axis_idx(axis)
+
+        diffs = np.diff(
+            np.concatenate((self.state.last_dat, message.data), axis=ax_idx),
+            axis=ax_idx,
+        )
+        # Prepare last_dat for next iteration
+        self.state.last_dat = slice_along_axis(
+            message.data, slice(-1, None), axis=ax_idx
+        )
+        # Scale by fs if requested. This convers the diff to a derivative. e.g., diff of position becomes velocity.
+        if self.settings.scale_by_fs:
+            ax_info = message.get_axis(axis)
+            if hasattr(ax_info, "data"):
+                dt = np.diff(np.concatenate(([self.state.last_time], ax_info.data)))
+                # Expand dt dims to match diffs
+                exp_sl = (
+                    (None,) * ax_idx
+                    + (Ellipsis,)
+                    + (None,) * (message.data.ndim - ax_idx - 1)
+                )
+                diffs /= dt[exp_sl]
+                self.state.last_time = ax_info.data[-1]  # For next iteration
+            else:
+                diffs /= ax_info.gain
+
+        return replace(message, data=diffs)
+
+
+class DiffUnit(
+    BaseTransformerUnit[DiffSettings, AxisArray, AxisArray, DiffTransformer]
+):
+    SETTINGS = DiffSettings
+
+
+def diff(axis: str = "time", scale_by_fs: bool = False) -> DiffTransformer:
+    return DiffTransformer(DiffSettings(axis=axis, scale_by_fs=scale_by_fs))