ezmsg-sigproc 2.1.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

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/gaussiansmoothing.py

@@ -0,0 +1,93 @@
+from typing import Callable
+import warnings
+
+import numpy as np
+
+from .filter import (
+    FilterBaseSettings,
+    BACoeffs,
+    FilterByDesignTransformer,
+    BaseFilterByDesignTransformerUnit,
+)
+
+
+class GaussianSmoothingSettings(FilterBaseSettings):
+    sigma: float | None = 1.0
+    """
+    sigma : float
+        Standard deviation of the Gaussian kernel.
+    """
+
+    width: int | None = 4
+    """
+    width : int
+        Number of standard deviations covered by the kernel window if kernel_size is not provided.
+    """
+
+    kernel_size: int | None = None
+    """
+    kernel_size : int | None
+        Length of the kernel in samples. If provided, overrides automatic calculation.
+    """
+
+
+def gaussian_smoothing_filter_design(
+    sigma: float = 1.0,
+    width: int = 4,
+    kernel_size: int | None = None,
+) -> BACoeffs | None:
+    # Parameter checks
+    if sigma <= 0:
+        raise ValueError(f"sigma must be positive. Received: {sigma}")
+
+    if width <= 0:
+        raise ValueError(f"width must be positive. Received: {width}")
+
+    if kernel_size is not None:
+        if kernel_size < 1:
+            raise ValueError(f"kernel_size must be >= 1. Received: {kernel_size}")
+    else:
+        kernel_size = int(2 * width * sigma + 1)
+
+    # Warn if kernel_size is smaller than recommended but don't fail
+    expected_kernel_size = int(2 * width * sigma + 1)
+    if kernel_size < expected_kernel_size:
+        ## TODO: Either add a warning or determine appropriate kernel size and raise an error
+        warnings.warn(
+            f"Provided kernel_size {kernel_size} is smaller than recommended "
+            f"size {expected_kernel_size} for sigma={sigma} and width={width}. "
+            "The kernel may be truncated."
+        )
+
+    from scipy.signal.windows import gaussian
+
+    b = gaussian(kernel_size, std=sigma)
+    b /= np.sum(b)  # Ensure normalization
+    a = np.array([1.0])
+
+    return b, a
+
+
+class GaussianSmoothingFilterTransformer(
+    FilterByDesignTransformer[GaussianSmoothingSettings, BACoeffs]
+):
+    def get_design_function(
+        self,
+    ) -> Callable[[float], BACoeffs]:
+        # Create a wrapper function that ignores fs parameter since gaussian smoothing doesn't need it
+        def design_wrapper(fs: float) -> BACoeffs:
+            return gaussian_smoothing_filter_design(
+                sigma=self.settings.sigma,
+                width=self.settings.width,
+                kernel_size=self.settings.kernel_size,
+            )
+
+        return design_wrapper
+
+
+class GaussianSmoothingFilter(
+    BaseFilterByDesignTransformerUnit[
+        GaussianSmoothingSettings, GaussianSmoothingFilterTransformer
+    ]
+):
+    SETTINGS = GaussianSmoothingSettings

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