ezmsg-sigproc 1.2.2__py3-none-any.whl → 2.10.0__py3-none-any.whl

ezmsg/sigproc/gaussiansmoothing.py

@@ -0,0 +1,89 @@
+import warnings
+from typing import Callable
+
+import numpy as np
+
+from .filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+)
+
+
+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,106 @@
+import functools
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
+
+from .filter import (
+    BACoeffs,
+    BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
+)
+
+
+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

ezmsg/sigproc/linear.py

@@ -0,0 +1,120 @@
+"""
+Apply a linear transformation: output = scale * input + offset.
+
+Supports per-element scale and offset along a specified axis.
+For full matrix transformations, use :obj:`AffineTransformTransformer` instead.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseStatefulTransformer, BaseTransformerUnit, processor_state
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class LinearTransformSettings(ez.Settings):
+    scale: float | list[float] | npt.ArrayLike = 1.0
+    """Scale factor(s). Can be a scalar (applied to all elements) or an array
+    matching the size of the specified axis for per-element scaling."""
+
+    offset: float | list[float] | npt.ArrayLike = 0.0
+    """Offset value(s). Can be a scalar (applied to all elements) or an array
+    matching the size of the specified axis for per-element offset."""
+
+    axis: str | None = None
+    """Axis along which to apply per-element scale/offset. If None, scalar
+    scale/offset are broadcast to all elements."""
+
+
+@processor_state
+class LinearTransformState:
+    scale: npt.NDArray = None
+    """Prepared scale array for broadcasting."""
+
+    offset: npt.NDArray = None
+    """Prepared offset array for broadcasting."""
+
+
+class LinearTransformTransformer(
+    BaseStatefulTransformer[LinearTransformSettings, AxisArray, AxisArray, LinearTransformState]
+):
+    """Apply linear transformation: output = scale * input + offset.
+
+    This transformer is optimized for element-wise linear operations with
+    optional per-channel (or per-axis) coefficients. For full matrix
+    transformations, use :obj:`AffineTransformTransformer` instead.
+
+    Examples:
+        # Uniform scaling and offset
+        >>> transformer = LinearTransformTransformer(LinearTransformSettings(scale=2.0, offset=1.0))
+
+        # Per-channel scaling (e.g., for 3-channel data along "ch" axis)
+        >>> transformer = LinearTransformTransformer(LinearTransformSettings(
+        ...     scale=[0.5, 1.0, 2.0],
+        ...     offset=[0.0, 0.1, 0.2],
+        ...     axis="ch"
+        ... ))
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        """Hash based on shape and axis to detect when broadcast shapes need recalculation."""
+        axis = self.settings.axis
+        if axis is not None:
+            axis_idx = message.get_axis_idx(axis)
+            return hash((message.data.ndim, axis_idx, message.data.shape[axis_idx]))
+        return hash(message.data.ndim)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        """Prepare scale/offset arrays with proper broadcast shapes."""
+        xp = get_namespace(message.data)
+        ndim = message.data.ndim
+
+        scale = self.settings.scale
+        offset = self.settings.offset
+
+        # Convert settings to arrays
+        if isinstance(scale, (list, np.ndarray)):
+            scale = xp.asarray(scale, dtype=xp.float64)
+        else:
+            # Scalar: create a 0-d array
+            scale = xp.asarray(float(scale), dtype=xp.float64)
+
+        if isinstance(offset, (list, np.ndarray)):
+            offset = xp.asarray(offset, dtype=xp.float64)
+        else:
+            # Scalar: create a 0-d array
+            offset = xp.asarray(float(offset), dtype=xp.float64)
+
+        # If axis is specified and we have 1-d arrays, reshape for proper broadcasting
+        if self.settings.axis is not None and ndim > 0:
+            axis_idx = message.get_axis_idx(self.settings.axis)
+
+            if scale.ndim == 1:
+                # Create shape for broadcasting: all 1s except at axis_idx
+                broadcast_shape = [1] * ndim
+                broadcast_shape[axis_idx] = scale.shape[0]
+                scale = xp.reshape(scale, broadcast_shape)
+
+            if offset.ndim == 1:
+                broadcast_shape = [1] * ndim
+                broadcast_shape[axis_idx] = offset.shape[0]
+                offset = xp.reshape(offset, broadcast_shape)
+
+        self._state.scale = scale
+        self._state.offset = offset
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        result = message.data * self._state.scale + self._state.offset
+        return replace(message, data=result)
+
+
+class LinearTransform(BaseTransformerUnit[LinearTransformSettings, AxisArray, AxisArray, LinearTransformTransformer]):
+    """Unit wrapper for LinearTransformTransformer."""
+
+    SETTINGS = LinearTransformSettings

ezmsg/sigproc/math/abs.py

@@ -0,0 +1,35 @@
+"""
+Take the absolute value of the data.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class AbsSettings:
+    pass
+
+
+class AbsTransformer(BaseTransformer[None, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        return replace(message, data=xp.abs(message.data))
+
+
+class Abs(BaseTransformerUnit[None, AxisArray, AxisArray, AbsTransformer]): ...  # SETTINGS = None
+
+
+def abs() -> AbsTransformer:
+    """
+    Take the absolute value of the data. See :obj:`np.abs` for more details.
+
+    Returns: :obj:`AbsTransformer`.
+
+    """
+    return AbsTransformer()

ezmsg/sigproc/math/add.py

@@ -0,0 +1,120 @@
+"""Add 2 signals or add a constant to a signal."""
+
+import asyncio
+import typing
+from dataclasses import dataclass, field
+
+import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.baseproc.util.asio import run_coroutine_sync
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+# --- Constant Addition (single input) ---
+
+
+class ConstAddSettings(ez.Settings):
+    value: float = 0.0
+    """Number to add to the input data."""
+
+
+class ConstAddTransformer(BaseTransformer[ConstAddSettings, AxisArray, AxisArray]):
+    """Add a constant value to input data."""
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        return replace(message, data=message.data + self.settings.value)
+
+
+class ConstAdd(BaseTransformerUnit[ConstAddSettings, AxisArray, AxisArray, ConstAddTransformer]):
+    """Unit wrapper for ConstAddTransformer."""
+
+    SETTINGS = ConstAddSettings
+
+
+# --- Two-input Addition ---
+
+
+@dataclass
+class AddState:
+    """State for Add processor with two input queues."""
+
+    queue_a: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+    queue_b: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+
+
+class AddProcessor:
+    """Processor that adds two AxisArray signals together.
+
+    This processor maintains separate queues for two input streams and
+    adds corresponding messages element-wise. It assumes both inputs
+    have compatible shapes and aligned time spans.
+    """
+
+    def __init__(self):
+        self._state = AddState()
+
+    @property
+    def state(self) -> AddState:
+        return self._state
+
+    @state.setter
+    def state(self, state: AddState | bytes | None) -> None:
+        if state is not None:
+            # TODO: Support hydrating state from bytes
+            # if isinstance(state, bytes):
+            #     self._state = pickle.loads(state)
+            # else:
+            self._state = state
+
+    def push_a(self, msg: AxisArray) -> None:
+        """Push a message to queue A."""
+        self._state.queue_a.put_nowait(msg)
+
+    def push_b(self, msg: AxisArray) -> None:
+        """Push a message to queue B."""
+        self._state.queue_b.put_nowait(msg)
+
+    async def __acall__(self) -> AxisArray:
+        """Await and add the next messages from both queues."""
+        a = await self._state.queue_a.get()
+        b = await self._state.queue_b.get()
+        return replace(a, data=a.data + b.data)
+
+    def __call__(self) -> AxisArray:
+        """Synchronously get and add the next messages from both queues."""
+        return run_coroutine_sync(self.__acall__())
+
+    # Aliases for legacy interface
+    async def __anext__(self) -> AxisArray:
+        return await self.__acall__()
+
+    def __next__(self) -> AxisArray:
+        return self.__call__()
+
+
+class Add(ez.Unit):
+    """Add two signals together.
+
+    Assumes compatible/similar axes/dimensions and aligned time spans.
+    Messages are paired by arrival order (oldest from each queue).
+    """
+
+    INPUT_SIGNAL_A = ez.InputStream(AxisArray)
+    INPUT_SIGNAL_B = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    async def initialize(self) -> None:
+        self.processor = AddProcessor()
+
+    @ez.subscriber(INPUT_SIGNAL_A)
+    async def on_a(self, msg: AxisArray) -> None:
+        self.processor.push_a(msg)
+
+    @ez.subscriber(INPUT_SIGNAL_B)
+    async def on_b(self, msg: AxisArray) -> None:
+        self.processor.push_b(msg)
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def output(self) -> typing.AsyncGenerator:
+        while True:
+            yield self.OUTPUT_SIGNAL, await self.processor.__acall__()

ezmsg/sigproc/math/clip.py

@@ -0,0 +1,48 @@
+"""
+Clips the data to be within the specified range.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+import ezmsg.core as ez
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class ClipSettings(ez.Settings):
+    min: float | None = None
+    """Lower clip bound. If None, no lower clipping is applied."""
+
+    max: float | None = None
+    """Upper clip bound. If None, no upper clipping is applied."""
+
+
+class ClipTransformer(BaseTransformer[ClipSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        return replace(
+            message,
+            data=xp.clip(message.data, self.settings.min, self.settings.max),
+        )
+
+
+class Clip(BaseTransformerUnit[ClipSettings, AxisArray, AxisArray, ClipTransformer]):
+    SETTINGS = ClipSettings
+
+
+def clip(min: float | None = None, max: float | None = None) -> ClipTransformer:
+    """
+    Clips the data to be within the specified range.
+
+    Args:
+        min: Lower clip bound. If None, no lower clipping is applied.
+        max: Upper clip bound. If None, no upper clipping is applied.
+
+    Returns:
+        :obj:`ClipTransformer`.
+    """
+    return ClipTransformer(ClipSettings(min=min, max=max))

ezmsg/sigproc/math/difference.py

@@ -0,0 +1,143 @@
+"""
+Take the difference between 2 signals or between a signal and a constant value.
+
+.. note::
+    :obj:`ConstDifferenceTransformer` supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    :obj:`DifferenceProcessor` (two-input difference) currently requires NumPy arrays.
+"""
+
+import asyncio
+import typing
+from dataclasses import dataclass, field
+
+import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.baseproc.util.asio import run_coroutine_sync
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class ConstDifferenceSettings(ez.Settings):
+    value: float = 0.0
+    """number to subtract or be subtracted from the input data"""
+
+    subtrahend: bool = True
+    """If True (default) then value is subtracted from the input data. If False, the input data
+    is subtracted from value."""
+
+
+class ConstDifferenceTransformer(BaseTransformer[ConstDifferenceSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        return replace(
+            message,
+            data=(message.data - self.settings.value)
+            if self.settings.subtrahend
+            else (self.settings.value - message.data),
+        )
+
+
+class ConstDifference(BaseTransformerUnit[ConstDifferenceSettings, AxisArray, AxisArray, ConstDifferenceTransformer]):
+    SETTINGS = ConstDifferenceSettings
+
+
+def const_difference(value: float = 0.0, subtrahend: bool = True) -> ConstDifferenceTransformer:
+    """
+    result = (in_data - value) if subtrahend else (value - in_data)
+    https://en.wikipedia.org/wiki/Template:Arithmetic_operations
+
+    Args:
+        value: number to subtract or be subtracted from the input data
+        subtrahend: If True (default) then value is subtracted from the input data.
+         If False, the input data is subtracted from value.
+
+    Returns: :obj:`ConstDifferenceTransformer`.
+    """
+    return ConstDifferenceTransformer(ConstDifferenceSettings(value=value, subtrahend=subtrahend))
+
+
+# --- Two-input Difference ---
+
+
+@dataclass
+class DifferenceState:
+    """State for Difference processor with two input queues."""
+
+    queue_a: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+    queue_b: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue)
+
+
+class DifferenceProcessor:
+    """Processor that subtracts two AxisArray signals (A - B).
+
+    This processor maintains separate queues for two input streams and
+    subtracts corresponding messages element-wise. It assumes both inputs
+    have compatible shapes and aligned time spans.
+    """
+
+    def __init__(self):
+        self._state = DifferenceState()
+
+    @property
+    def state(self) -> DifferenceState:
+        return self._state
+
+    @state.setter
+    def state(self, state: DifferenceState | bytes | None) -> None:
+        if state is not None:
+            self._state = state
+
+    def push_a(self, msg: AxisArray) -> None:
+        """Push a message to queue A (minuend)."""
+        self._state.queue_a.put_nowait(msg)
+
+    def push_b(self, msg: AxisArray) -> None:
+        """Push a message to queue B (subtrahend)."""
+        self._state.queue_b.put_nowait(msg)
+
+    async def __acall__(self) -> AxisArray:
+        """Await and subtract the next messages (A - B)."""
+        a = await self._state.queue_a.get()
+        b = await self._state.queue_b.get()
+        return replace(a, data=a.data - b.data)
+
+    def __call__(self) -> AxisArray:
+        """Synchronously get and subtract the next messages."""
+        return run_coroutine_sync(self.__acall__())
+
+    # Aliases for legacy interface
+    async def __anext__(self) -> AxisArray:
+        return await self.__acall__()
+
+    def __next__(self) -> AxisArray:
+        return self.__call__()
+
+
+class Difference(ez.Unit):
+    """Subtract two signals (A - B).
+
+    Assumes compatible/similar axes/dimensions and aligned time spans.
+    Messages are paired by arrival order (oldest from each queue).
+
+    OUTPUT = INPUT_SIGNAL_A - INPUT_SIGNAL_B
+    """
+
+    INPUT_SIGNAL_A = ez.InputStream(AxisArray)
+    INPUT_SIGNAL_B = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    async def initialize(self) -> None:
+        self.processor = DifferenceProcessor()
+
+    @ez.subscriber(INPUT_SIGNAL_A)
+    async def on_a(self, msg: AxisArray) -> None:
+        self.processor.push_a(msg)
+
+    @ez.subscriber(INPUT_SIGNAL_B)
+    async def on_b(self, msg: AxisArray) -> None:
+        self.processor.push_b(msg)
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def output(self) -> typing.AsyncGenerator:
+        while True:
+            yield self.OUTPUT_SIGNAL, await self.processor.__acall__()

ezmsg/sigproc/math/invert.py

@@ -0,0 +1,28 @@
+"""
+Compute the multiplicative inverse (1/x) of the data.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class InvertTransformer(BaseTransformer[None, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        return replace(message, data=1 / message.data)
+
+
+class Invert(BaseTransformerUnit[None, AxisArray, AxisArray, InvertTransformer]): ...  # SETTINGS = None
+
+
+def invert() -> InvertTransformer:
+    """
+    Take the inverse of the data.
+
+    Returns: :obj:`InvertTransformer`.
+    """
+    return InvertTransformer()