ezmsg-sigproc 2.2.0__py3-none-any.whl → 2.4.0__py3-none-any.whl

ezmsg/sigproc/filterbankdesign.py

@@ -0,0 +1,136 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+
+from ezmsg.util.messages.util import replace
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .base import (
+    BaseStatefulTransformer,
+    processor_state,
+)
+
+from .filterbank import (
+    FilterbankTransformer,
+    FilterbankSettings,
+    FilterbankMode,
+    MinPhaseMode,
+)
+
+from .kaiser import KaiserFilterSettings, kaiser_design_fun
+
+
+class FilterbankDesignSettings(ez.Settings):
+    filters: typing.Iterable[KaiserFilterSettings]
+
+    mode: FilterbankMode = FilterbankMode.CONV
+    """
+    "conv", "fft", or "auto". If "auto", the mode is determined by the size of the input data.
+      fft mode is more efficient for long kernels. However, fft mode uses non-overlapping windows and will
+      incur a delay equal to the window length, which is larger than the largest kernel.
+      conv mode is less efficient but will return data for every incoming chunk regardless of how small it is
+      and thus can provide shorter latency updates.
+    """
+
+    min_phase: MinPhaseMode = MinPhaseMode.NONE
+    """
+    If not None, convert the kernels to minimum-phase equivalents. Valid options are
+      'hilbert', 'homomorphic', and 'homomorphic-full'. Complex filters not supported.
+      See `scipy.signal.minimum_phase` for details.
+    """
+
+    axis: str = "time"
+    """The name of the axis to operate on. This should usually be "time"."""
+
+    new_axis: str = "kernel"
+    """The name of the new axis corresponding to the kernel index."""
+
+
+@processor_state
+class FilterbankDesignState:
+    filterbank: FilterbankTransformer | None = None
+    needs_redesign: bool = False
+
+
+class FilterbankDesignTransformer(
+    BaseStatefulTransformer[
+        FilterbankDesignSettings, AxisArray, AxisArray, FilterbankDesignState
+    ],
+):
+    """
+    Transformer that designs and applies a filterbank based on Kaiser windowed FIR filters.
+    """
+
+    @classmethod
+    def get_message_type(cls, dir: str) -> type[AxisArray]:
+        if dir in ("in", "out"):
+            return AxisArray
+        else:
+            raise ValueError(f"Invalid direction: {dir}. Must be 'in' or 'out'.")
+
+    def update_settings(
+        self, new_settings: typing.Optional[FilterbankDesignSettings] = None, **kwargs
+    ) -> None:
+        """
+        Update settings and mark that filter coefficients need to be recalculated.
+
+        Args:
+            new_settings: Complete new settings object to replace current settings
+            **kwargs: Individual settings to update
+        """
+        # Update settings
+        if new_settings is not None:
+            self.settings = new_settings
+        else:
+            self.settings = replace(self.settings, **kwargs)
+
+        # Set flag to trigger recalculation on next message
+        if self.state.filterbank is not None:
+            self.state.needs_redesign = True
+
+    def _calculate_kernels(self, fs: float) -> list[npt.NDArray]:
+        kernels = []
+        for filter in self.settings.filters:
+            output = kaiser_design_fun(
+                fs,
+                cutoff=filter.cutoff,
+                ripple=filter.ripple,
+                width=filter.width,
+                pass_zero=filter.pass_zero,
+                wn_hz=filter.wn_hz,
+            )
+
+            kernels.append(np.array([1.0]) if output is None else output[0])
+        return kernels
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.state.filterbank is not None and self.state.needs_redesign:
+            self._reset_state(message)
+            self.state.needs_redesign = False
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1
+        axis_idx = message.get_axis_idx(axis)
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((message.key, samp_shape, gain))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis_obj = message.axes[self.settings.axis]
+        assert isinstance(axis_obj, AxisArray.LinearAxis)
+        fs = 1 / axis_obj.gain
+        kernels = self._calculate_kernels(fs)
+        new_settings = FilterbankSettings(
+            kernels=kernels,
+            mode=self.settings.mode,
+            min_phase=self.settings.min_phase,
+            axis=self.settings.axis,
+            new_axis=self.settings.new_axis,
+        )
+        self.state.filterbank = FilterbankTransformer(settings=new_settings)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        return self.state.filterbank(message)

ezmsg/sigproc/firfilter.py

@@ -0,0 +1,119 @@
+import functools
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
+
+from .filter import (
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+)
+
+
+class FIRFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`FIRFilter`.  See scipy.signal.firwin for more details"""
+
+    # axis and coef_type are inherited from FilterBaseSettings
+
+    order: int = 0
+    """
+    Filter order/number of taps
+    """
+
+    cutoff: float | npt.ArrayLike | None = None
+    """
+    Cutoff frequency of filter (expressed in the same units as fs) OR an array of cutoff frequencies 
+    (that is, band edges). In the former case, as a float, the cutoff frequency should correspond with 
+    the half-amplitude point, where the attenuation will be -6dB. In the latter case, the frequencies in 
+    cutoff should be positive and monotonically increasing between 0 and fs/2. The values 0 and fs/2 must 
+    not be included in cutoff.
+    """
+
+    width: float | None = None
+    """
+    If width is not None, then assume it is the approximate width of the transition region (expressed in 
+    the same units as fs) for use in Kaiser FIR filter design. In this case, the window argument is ignored.
+    """
+
+    window: str | None = "hamming"
+    """
+    Desired window to use.  See scipy.signal.get_window for a list of windows and required parameters.
+    """
+
+    pass_zero: bool | str = True
+    """
+    If True, the gain at the frequency 0 (i.e., the “DC gain”) is 1. If False, the DC gain is 0. Can also 
+    be a string argument for the desired filter type (equivalent to btype in IIR design functions).
+    {‘lowpass’, ‘highpass’, ‘bandpass’, ‘bandstop’}
+    """
+
+    scale: bool = True
+    """
+    Set to True to scale the coefficients so that the frequency response is exactly unity at a certain 
+    frequency. That frequency is either:
+    * 0 (DC) if the first passband starts at 0 (i.e. pass_zero is True)
+    * fs/2 (the Nyquist frequency) if the first passband ends at fs/2 
+        (i.e the filter is a single band highpass filter); 
+        center of first passband otherwise
+    """
+
+    wn_hz: bool = True
+    """
+    Set False if provided Wn are normalized from 0 to 1, where 1 is the Nyquist frequency
+    """
+
+
+def firwin_design_fun(
+    fs: float,
+    order: int = 0,
+    cutoff: float | npt.ArrayLike | None = None,
+    width: float | None = None,
+    window: str | None = "hamming",
+    pass_zero: bool | str = True,
+    scale: bool = True,
+    wn_hz: bool = True,
+) -> BACoeffs | None:
+    """
+    Design an `order`th-order FIR filter and return the filter coefficients.
+    See :obj:`FIRFilterSettings` for argument description.
+
+    Returns:
+        The filter taps as designed by firwin
+    """
+    if order > 0:
+        taps = scipy.signal.firwin(
+            numtaps=order,
+            cutoff=cutoff,
+            width=width,
+            window=window,
+            pass_zero=pass_zero,
+            scale=scale,
+            fs=fs if wn_hz else None,
+        )
+        return (taps, np.array([1.0]))
+    return None
+
+
+class FIRFilterTransformer(FilterByDesignTransformer[FIRFilterSettings, BACoeffs]):
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | None]:
+        return functools.partial(
+            firwin_design_fun,
+            order=self.settings.order,
+            cutoff=self.settings.cutoff,
+            width=self.settings.width,
+            window=self.settings.window,
+            pass_zero=self.settings.pass_zero,
+            scale=self.settings.scale,
+            wn_hz=self.settings.wn_hz,
+        )
+
+
+class FIRFilter(
+    BaseFilterByDesignTransformerUnit[FIRFilterSettings, FIRFilterTransformer]
+):
+    SETTINGS = FIRFilterSettings

ezmsg/sigproc/kaiser.py

@@ -0,0 +1,110 @@
+import functools
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
+
+from .filter import (
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+)
+
+
+class KaiserFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`KaiserFilter`"""
+
+    # axis and coef_type are inherited from FilterBaseSettings
+
+    cutoff: float | npt.ArrayLike | None = None
+    """
+    Cutoff frequency of filter (expressed in the same units as fs) OR an array of cutoff frequencies 
+    (that is, band edges). In the former case, as a float, the cutoff frequency should correspond with 
+    the half-amplitude point, where the attenuation will be -6dB. In the latter case, the frequencies in 
+    cutoff should be positive and monotonically increasing between 0 and fs/2. The values 0 and fs/2 must 
+    not be included in cutoff.
+    """
+
+    ripple: float | None = None
+    """
+    Upper bound for the deviation (in dB) of the magnitude of the filter's frequency response from that of 
+    the desired filter (not including frequencies in any transition intervals).
+    See scipy.signal.kaiserord for more information.
+    """
+
+    width: float | None = None
+    """
+    If width is not None, then assume it is the approximate width of the transition region (expressed in 
+    the same units as fs) for use in Kaiser FIR filter design.
+    See scipy.signal.kaiserord for more information.
+    """
+
+    pass_zero: bool | str = True
+    """
+    If True, the gain at the frequency 0 (i.e., the “DC gain”) is 1. If False, the DC gain is 0. Can also 
+    be a string argument for the desired filter type (equivalent to btype in IIR design functions).
+    {‘lowpass’, ‘highpass’, ‘bandpass’, ‘bandstop’}
+    """
+
+    wn_hz: bool = True
+    """
+    Set False if cutoff and width are normalized from 0 to 1, where 1 is the Nyquist frequency
+    """
+
+
+def kaiser_design_fun(
+    fs: float,
+    cutoff: float | npt.ArrayLike | None = None,
+    ripple: float | None = None,
+    width: float | None = None,
+    pass_zero: bool | str = True,
+    wn_hz: bool = True,
+) -> BACoeffs | None:
+    """
+    Design an `order`th-order FIR Kaiser filter and return the filter coefficients.
+    See :obj:`FIRFilterSettings` for argument description.
+
+    Returns:
+        The filter taps as designed by firwin
+    """
+    if ripple is None or width is None or cutoff is None:
+        return None
+
+    width = width / (0.5 * fs) if wn_hz else width
+    n_taps, beta = scipy.signal.kaiserord(ripple, width)
+    if n_taps % 2 == 0:
+        n_taps += 1
+    taps = scipy.signal.firwin(
+        numtaps=n_taps,
+        cutoff=cutoff,
+        window=("kaiser", beta),  # type: ignore
+        pass_zero=pass_zero,  # type: ignore
+        scale=False,
+        fs=fs if wn_hz else None,
+    )
+
+    return (taps, np.array([1.0]))
+
+
+class KaiserFilterTransformer(
+    FilterByDesignTransformer[KaiserFilterSettings, BACoeffs]
+):
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | None]:
+        return functools.partial(
+            kaiser_design_fun,
+            cutoff=self.settings.cutoff,
+            ripple=self.settings.ripple,
+            width=self.settings.width,
+            pass_zero=self.settings.pass_zero,
+            wn_hz=self.settings.wn_hz,
+        )
+
+
+class KaiserFilter(
+    BaseFilterByDesignTransformerUnit[KaiserFilterSettings, KaiserFilterTransformer]
+):
+    SETTINGS = KaiserFilterSettings