ezmsg-sigproc 2.4.1__py3-none-any.whl → 2.6.0__py3-none-any.whl

ezmsg/sigproc/fir_hilbert.py

@@ -0,0 +1,336 @@
+import functools
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import scipy.signal as sps
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from ezmsg.sigproc.base import BaseStatefulTransformer, processor_state
+from ezmsg.sigproc.filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    BaseTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+)
+
+
+class FIRHilbertFilterSettings(FilterBaseSettings):
+    """Settings for :obj:`FIRHilbertFilter`."""
+
+    # axis inherited from FilterBaseSettings
+
+    coef_type: str = "ba"
+    """
+    Coefficient type. Must be 'ba' for FIR.
+    """
+
+    order: int = 170
+    """
+    Filter order (taps = order + 1).
+    Hilbert (type-III) filters require even order (odd taps).
+    If odd order (even taps), order will be incremented by 1.
+    """
+
+    f_lo: float = 1.0
+    """
+    Lower corner of Hilbert “pass” band (Hz).
+    Transition starts at f_lo.
+    """
+
+    f_hi: float | None = None
+    """
+    Upper corner of Hilbert “pass” band (Hz).
+    Transition starts at f_hi.
+    If None, highpass from f_lo to Nyquist.
+    """
+
+    trans_lo: float = 1.0
+    """
+    Transition width (Hz) below f_lo.
+    Decrease to sharpen transition.
+    """
+
+    trans_hi: float = 1.0
+    """
+    Transition width (Hz) at high end.
+    Decrease to sharpen transition.
+    """
+
+    weight_pass: float = 1.0
+    """
+    Weight for Hilbert pass region.
+    """
+
+    weight_stop_lo: float = 1.0
+    """
+    Weight for low stop band.
+    """
+
+    weight_stop_hi: float = 1.0
+    """
+    Weight for high stop band.
+    """
+
+    norm_band: tuple[float, float] | None = None
+    """
+    Optional normalization band (f_lo, f_hi) in Hz for gain normalization.
+    If None, no normalization is applied.
+    """
+
+    norm_freq: float | None = None
+    """
+    Optional normalization frequency in Hz for gain normalization.
+    If None, no normalization is applied.
+    """
+
+
+def fir_hilbert_design_fun(
+    fs: float,
+    order: int = 170,
+    f_lo: float = 1.0,
+    f_hi: float | None = None,
+    trans_lo: float = 1.0,
+    trans_hi: float = 1.0,
+    weight_pass: float = 1.0,
+    weight_stop_lo: float = 1.0,
+    weight_stop_hi: float = 1.0,
+    norm_band: tuple[float, float] | None = None,
+    norm_freq: float | None = None,
+) -> BACoeffs | None:
+    """
+    Hilbert FIR filter design using the Remez exchange algorithm.
+    Design an `order`th-order FIR Hilbert filter and return the filter coefficients.
+    See :obj:`FIRHilbertFilterSettings` for argument description.
+
+    Returns:
+        The filter coefficients as a tuple of (b, a).
+    """
+    if order <= 0:
+        return None
+    if order % 2 == 1:
+        order += 1
+    nyq = fs / 2.0
+    taps = order + 1
+    f1 = max(f_lo, 0.0) + trans_lo
+    f2 = (nyq - trans_hi) if (f_hi is None) else min(f_hi, nyq - trans_hi)
+    if not (0.0 < f1 < f2 < nyq):
+        raise ValueError(
+            f"Hilbert passband collapsed or invalid: "
+            f"f_lo={f_lo}, f_hi={f_hi}, trans_lo={trans_lo}, trans_hi={trans_hi}, fs={fs}"
+        )
+    # Bands: [0, f1-trans_lo] stop ; [f1, f2] pass (Hilbert) ; [f2+trans_hi, nyq] stop
+    bands = [0.0, max(f1 - trans_lo, 0.0), f1, f2, min(f2 + trans_hi, nyq), nyq]
+    desired = [0.0, 1.0, 0.0]
+    weight = [max(weight_stop_lo, 0.0), max(weight_pass, 0.0), max(weight_stop_hi, 0.0)]
+    for i in range(1, len(bands) - 1):
+        if bands[i] <= bands[i - 1]:
+            bands[i] = np.nextafter(bands[i - 1], np.inf)
+    if bands[-2] >= nyq:
+        ez.logger.warning("Hilbert upper stopband collapsed; using 2-band (stop/pass) design.")
+        bands = bands[:-3] + [nyq]
+        desired = desired[:-1]
+        weight = weight[:-1]
+    b = sps.remez(taps, bands, desired, weight=weight, type="hilbert", fs=fs)
+    a = np.array([1.0])
+    g = None
+    if norm_freq is not None:
+        if norm_freq < f1 or norm_freq > f2:
+            ez.logger.warning("Invalid normalization frequency specifications. Skipping normalization.")
+        else:
+            f0 = float(norm_freq)
+            w = 2.0 * np.pi * (np.asarray([f0], dtype=np.float64) / fs)
+            _, H = sps.freqz(b, a, worN=w)
+            g = float(np.abs(H[0]))
+    elif norm_band is not None:
+        lo, hi = norm_band
+        if lo < f1 or hi > f2:
+            lo = max(lo, f1)
+            hi = min(hi, f2)
+            ez.logger.warning("Normalization band outside passband. Clipping to passband for normalization.")
+        if lo >= hi:
+            ez.logger.warning("Invalid normalization band specifications. Skipping normalization.")
+        else:
+            freqs = np.linspace(lo, hi, 2048, dtype=np.float64)
+            w = 2.0 * np.pi * (np.asarray(freqs, dtype=np.float64) / fs)
+            _, H = sps.freqz(b, a, worN=w)
+            g = float(np.median(np.abs(H)))
+    if g is not None and g > 0:
+        b = b / g
+    return (b, a)
+
+
+class FIRHilbertFilterTransformer(FilterByDesignTransformer[FIRHilbertFilterSettings, BACoeffs]):
+    def get_design_function(self) -> typing.Callable[[float], BACoeffs | None]:
+        if self.settings.coef_type != "ba":
+            ez.logger.error("FIRHilbert only supports coef_type='ba'.")
+            raise ValueError("FIRHilbert only supports coef_type='ba'.")
+
+        return functools.partial(
+            fir_hilbert_design_fun,
+            order=self.settings.order,
+            f_lo=self.settings.f_lo,
+            f_hi=self.settings.f_hi,
+            trans_lo=self.settings.trans_lo,
+            trans_hi=self.settings.trans_hi,
+            weight_pass=self.settings.weight_pass,
+            weight_stop_lo=self.settings.weight_stop_lo,
+            weight_stop_hi=self.settings.weight_stop_hi,
+            norm_band=self.settings.norm_band,
+            norm_freq=self.settings.norm_freq,
+        )
+
+    def get_taps(self) -> int | None:
+        if self._state.filter is None:
+            return None
+        b, _ = self._state.filter.settings.coefs
+        return b.size if b is not None else None
+
+
+class FIRHilbertFilterUnit(BaseFilterByDesignTransformerUnit[FIRHilbertFilterSettings, FIRHilbertFilterTransformer]):
+    SETTINGS = FIRHilbertFilterSettings
+
+
+@processor_state
+class FIRHilbertEnvelopeState:
+    filter: FIRHilbertFilterTransformer | None = None
+    delay_buf: np.ndarray | None = None
+    dly: int | None = None
+
+
+class FIRHilbertEnvelopeTransformer(
+    BaseStatefulTransformer[FIRHilbertFilterSettings, AxisArray, AxisArray, FIRHilbertEnvelopeState]
+):
+    """
+    Processor for computing the envelope of a signal using the Hilbert transform.
+
+    This processor applies a Hilbert FIR filter to the input signal to obtain the analytic signal, from which the
+    envelope is computed.
+
+    The processor expects and outputs `AxisArray` messages with a `"time"` (time) axis.
+
+    Settings:
+    ---------
+    order : int
+        Filter order (taps = order + 1).
+        Hilbert (type-III) filters require even order (odd taps).
+        If odd order (even taps), order will be incremented by 1.
+    f_lo : float
+        Lower corner of Hilbert “pass” band (Hz).
+        Transition starts at f_lo.
+    f_hi : float, optional
+        Upper corner of Hilbert “pass” band (Hz).
+        Transition starts at f_hi.
+        If None, highpass from f_lo to Nyquist.
+    trans_lo : float
+        Transition width (Hz) below f_lo.
+        Decrease to sharpen transition.
+    trans_hi : float
+        Transition width (Hz) above f_hi.
+        Decrease to sharpen transition.
+    weight_pass : float
+        Weight for Hilbert pass region.
+    weight_stop_lo : float
+        Weight for low stop band.
+    weight_stop_hi : float
+        Weight for high stop band.
+    norm_band : tuple(float, float), optional
+        Optional normalization band (f_lo, f_hi) in Hz for gain normalization.
+        If None, no normalization is applied.
+    norm_freq : float, optional
+        Optional normalization frequency in Hz for gain normalization.
+        If None, no normalization is applied.
+
+    Example:
+    -----------------------------
+    ```python
+    processor = FIRHilbertEnvelopeTransformer(
+        settings=FIRHilbertFilterSettings(
+            order=170,
+            f_lo=1.0,
+            f_hi=50.0,
+        )
+    )
+    ```
+
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        gain = getattr(self._state.filter, "gain", 0.0)
+        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:
+        self._state.filter = FIRHilbertFilterTransformer(settings=self.settings)
+        self._state.delay_buf = None
+        self._state.dly = None
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        y_imag_msg = self._state.filter(message)
+        y_imag = y_imag_msg.data
+
+        axis_name = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis_name)
+        if self._state.dly is None:
+            taps = self._state.filter.get_taps()
+            self._state.dly = (taps - 1) // 2
+
+        x = message.data
+
+        move_axis = False
+        if axis_idx != x.ndim - 1:
+            x = np.moveaxis(x, axis_idx, -1)
+            y_imag = np.moveaxis(y_imag, axis_idx, -1)
+            move_axis = True
+
+        if self._state.delay_buf is None:
+            lead_shape = x.shape[:-1]
+            self._state.delay_buf = np.zeros(lead_shape + (self._state.dly,), dtype=x.dtype)
+
+        x_cat = np.concatenate([self._state.delay_buf, x], axis=-1)
+        x_delayed_full = x_cat[..., : -self._state.dly]
+        y_real = x_delayed_full[..., -x.shape[-1] :]
+
+        self._state.delay_buf = x_cat[..., -self._state.dly :].copy()
+
+        analytic = y_real.astype(np.complex64) + 1j * y_imag.astype(np.complex64)
+        out = np.abs(analytic)
+
+        if move_axis:
+            out = np.moveaxis(out, -1, axis_idx)
+
+        return replace(message, data=out, axes=message.axes)
+
+
+class FIRHilbertEnvelopeUnit(
+    BaseTransformerUnit[
+        FIRHilbertFilterSettings,
+        AxisArray,
+        AxisArray,
+        FIRHilbertEnvelopeTransformer,
+    ]
+):
+    """
+    Unit wrapper for the `FIRHilbertEnvelopeTransformer`.
+
+    This unit provides a plug-and-play interface for calculating the envelope using the FIR Hilbert transform on a
+    signal in an ezmsg graph-based system. It takes in `AxisArray` inputs and outputs processed data in the same format.
+
+    Example:
+    --------
+    ```python
+    unit = FIRHilbertEnvelopeUnit(
+        settings=FIRHilbertFilterSettings(
+            order=170,
+            f_lo=1.0,
+            f_hi=50.0,
+        )
+    )
+    ```
+    """
+
+    SETTINGS = FIRHilbertFilterSettings

ezmsg/sigproc/fir_pmc.py

@@ -0,0 +1,209 @@
+import functools
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import scipy.signal
+
+from ezmsg.sigproc.filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+)
+
+
+class ParksMcClellanFIRSettings(FilterBaseSettings):
+    """Settings for :obj:`ParksMcClellanFIR`."""
+
+    # axis inherited from FilterBaseSettings
+
+    coef_type: str = "ba"
+    """
+    Coefficient type. Must be 'ba' for FIR.
+    """
+
+    order: int = 0
+    """
+    Filter order (taps = order + 1).
+    PMC FIR filters require even order (odd taps).
+    If odd order (even taps), order will be incremented by 1.
+    """
+
+    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.
+    """
+
+    transition: float = 10.0
+    """
+    Transition bandwidth (Hz) applied to each passband edge.
+    For low/high: single transition. For bands: both edges.
+    """
+
+    weight_pass: float = 1.0
+    """
+    Weight for the passband.
+    Used for both high and low passbands in bandstop filters.
+    """
+
+    weight_stop_lo: float = 1.0
+    """
+    Weight for the lower stopband.
+    Not used for bandstop filters.
+    """
+
+    weight_stop_hi: float = 1.0
+    """
+    Weight for the upper stopband.
+    Used as the central-stop weight for bandstop filters.
+    """
+
+    def filter_specs(
+        self,
+    ) -> tuple[str, tuple[float, float] | float] | None:
+        """
+        Determine the filter type given the corner frequencies.
+
+        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:
+            return "lowpass", self.cutoff
+        elif self.cuton is not None and self.cutoff is None:
+            return "highpass", self.cuton
+        elif self.cuton is not None and self.cutoff is not None:
+            if self.cuton <= self.cutoff:
+                return "bandpass", (self.cuton, self.cutoff)
+            else:
+                return "bandstop", (self.cutoff, self.cuton)
+
+
+def parks_mcclellan_design_fun(
+    fs: float,
+    order: int = 0,
+    cuton: float | None = None,
+    cutoff: float | None = None,
+    transition: float = 10.0,
+    weight_pass: float = 1.0,
+    weight_stop_lo: float = 1.0,
+    weight_stop_hi: float = 1.0,
+) -> BACoeffs | None:
+    """
+    See :obj:`ParksMcClellanFIRSettings.filter_specs` for an explanation of specifying different
+    filter types (lowpass, highpass, bandpass, bandstop) from the parameters.
+
+    Designs a Parks-McClellan FIR filter via the Remez exchange algorithm using the given specifications.
+    PMC filters are equiripple and linear phase.
+
+    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.
+        transition: Transition bandwidth (Hz) applied to each passband edge.
+        weight_pass: Weight for the passband.
+        weight_stop_lo: Weight for the lower stopband.
+        weight_stop_hi: Weight for the upper stopband.
+
+    Returns:
+        The filter coefficients as a tuple of (b, a).
+    """
+    if order <= 0:
+        return None
+    if order % 2 == 1:
+        order += 1
+
+    specs = ParksMcClellanFIRSettings(cuton=cuton, cutoff=cutoff).filter_specs()
+    if specs is None:
+        # Under-specified: no filter
+        return None
+
+    btype, corners = specs
+    nyq = fs / 2.0
+    tw = max(transition, 0.0)
+
+    def clip_hz(x: float) -> float:
+        return float(min(max(x, 0.0), nyq))
+
+    if btype == "lowpass":
+        b = [0.0, clip_hz(corners), clip_hz(corners + tw), nyq]
+        d = [1.0, 0.0]
+        w = [max(weight_pass, 0.0), max(weight_stop_hi, 0.0)]
+
+    elif btype == "highpass":
+        b = [0.0, clip_hz(corners - tw), clip_hz(corners), nyq]
+        d = [0.0, 1.0]
+        w = [max(weight_stop_lo, 0.0), max(weight_pass, 0.0)]
+
+    elif btype == "bandpass":
+        b = [
+            0.0,
+            clip_hz(corners[0] - tw),
+            clip_hz(corners[0]),
+            clip_hz(corners[1]),
+            clip_hz(corners[1] + tw),
+            nyq,
+        ]
+        d = [0.0, 1.0, 0.0]
+        w = [max(weight_stop_lo, 0.0), max(weight_pass, 0.0), max(weight_stop_hi, 0.0)]
+
+    else:
+        b = [
+            0.0,
+            clip_hz(corners[0]),
+            clip_hz(corners[0] + tw),
+            clip_hz(corners[1] - tw),
+            clip_hz(corners[1]),
+            nyq,
+        ]
+        d = [1.0, 0.0, 1.0]
+        # For bandstop we can reuse stop_hi as central-stop weight; stop_lo is the DC-side passband stop weight
+        w = [max(weight_pass, 0.0), max(weight_stop_hi, 0.0), max(weight_pass, 0.0)]
+
+    # Ensure bands strictly increase and have nonzero width per segment
+    # Adjust tiny overlaps due to clipping
+    for i in range(1, len(b)):
+        if b[i] <= b[i - 1]:
+            b[i] = min(b[i - 1] + 1e-6, nyq)
+
+    b = scipy.signal.remez(numtaps=order + 1, bands=b, desired=d, weight=w, fs=fs)
+    return (b, np.array([1.0]))
+
+
+class ParksMcClellanFIRTransformer(FilterByDesignTransformer[ParksMcClellanFIRSettings, BACoeffs]):
+    def get_design_function(self) -> typing.Callable[[float], BACoeffs | None]:
+        if self.settings.coef_type != "ba":
+            ez.logger.error("ParksMcClellanFIR only supports coef_type='ba'.")
+            raise ValueError("ParksMcClellanFIR only supports coef_type='ba'.")
+        return functools.partial(
+            parks_mcclellan_design_fun,
+            order=self.settings.order,
+            cuton=self.settings.cuton,
+            cutoff=self.settings.cutoff,
+            transition=self.settings.transition,
+            weight_pass=self.settings.weight_pass,
+            weight_stop_lo=self.settings.weight_stop_lo,
+            weight_stop_hi=self.settings.weight_stop_hi,
+        )
+
+
+class ParksMcClellanFIR(BaseFilterByDesignTransformerUnit[ParksMcClellanFIRSettings, ParksMcClellanFIRTransformer]):
+    SETTINGS = ParksMcClellanFIRSettings

ezmsg/sigproc/firfilter.py

@@ -6,10 +6,10 @@ import numpy.typing as npt
 import scipy.signal
 
 from .filter import (
-    FilterBaseSettings,
-    FilterByDesignTransformer,
     BACoeffs,
     BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
 )
 
 
@@ -25,16 +25,16 @@ class FIRFilterSettings(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 
+    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 
+    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.
     """
 
@@ -45,18 +45,18 @@ class FIRFilterSettings(FilterBaseSettings):
 
     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 
+    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 
+    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); 
+    * 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
     """
 
@@ -113,7 +113,5 @@ class FIRFilterTransformer(FilterByDesignTransformer[FIRFilterSettings, BACoeffs
         )
 
 
-class FIRFilter(
-    BaseFilterByDesignTransformerUnit[FIRFilterSettings, FIRFilterTransformer]
-):
+class FIRFilter(BaseFilterByDesignTransformerUnit[FIRFilterSettings, FIRFilterTransformer]):
     SETTINGS = FIRFilterSettings

ezmsg/sigproc/gaussiansmoothing.py

@@ -1,13 +1,13 @@
-from typing import Callable
 import warnings
+from typing import Callable
 
 import numpy as np
 
 from .filter import (
-    FilterBaseSettings,
     BACoeffs,
-    FilterByDesignTransformer,
     BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
 )
 
 
@@ -68,9 +68,7 @@ def gaussian_smoothing_filter_design(
     return b, a
 
 
-class GaussianSmoothingFilterTransformer(
-    FilterByDesignTransformer[GaussianSmoothingSettings, BACoeffs]
-):
+class GaussianSmoothingFilterTransformer(FilterByDesignTransformer[GaussianSmoothingSettings, BACoeffs]):
     def get_design_function(
         self,
     ) -> Callable[[float], BACoeffs]:
@@ -86,8 +84,6 @@ class GaussianSmoothingFilterTransformer(
 
 
 class GaussianSmoothingFilter(
-    BaseFilterByDesignTransformerUnit[
-        GaussianSmoothingSettings, GaussianSmoothingFilterTransformer
-    ]
+    BaseFilterByDesignTransformerUnit[GaussianSmoothingSettings, GaussianSmoothingFilterTransformer]
 ):
     SETTINGS = GaussianSmoothingSettings