ezmsg-sigproc 1.7.0__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

@@ -1,34 +1,35 @@
-import typing
+"""
+Take the absolute value of the data.
 
-import numpy as np
-import ezmsg.core as ez
-from ezmsg.util.generator import consumer
+.. 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
 
-from ..base import GenAxisArray
 
+class AbsSettings:
+    pass
 
-@consumer
-def abs() -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Take the absolute value of the data. See :obj:`np.abs` for more details.
 
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the absolute value of the input :obj:`AxisArray` data.
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = replace(msg_in, data=np.abs(msg_in.data))
+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 AbsSettings(ez.Settings):
-    pass
+class Abs(BaseTransformerUnit[None, AxisArray, AxisArray, AbsTransformer]): ...  # SETTINGS = None
 
 
-class Abs(GenAxisArray):
-    SETTINGS = AbsSettings
+def abs() -> AbsTransformer:
+    """
+    Take the absolute value of the data. See :obj:`np.abs` for more details.
 
-    def construct_generator(self):
-        self.STATE.gen = abs()
+    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

@@ -1,40 +1,48 @@
-import typing
+"""
+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 numpy as np
 import ezmsg.core as ez
-from ezmsg.util.generator import consumer
+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
 
-from ..base import GenAxisArray
 
+class ClipSettings(ez.Settings):
+    min: float | None = None
+    """Lower clip bound. If None, no lower clipping is applied."""
 
-@consumer
-def clip(a_min: float, a_max: float) -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Clips the data to be within the specified range. See :obj:`np.clip` for more details.
+    max: float | None = None
+    """Upper clip bound. If None, no upper clipping is applied."""
 
-    Args:
-        a_min: Lower clip bound
-        a_max: Upper clip bound
 
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the clipped version of the input :obj:`AxisArray` data.
+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),
+        )
 
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = replace(msg_in, data=np.clip(msg_in.data, a_min, a_max))
 
+class Clip(BaseTransformerUnit[ClipSettings, AxisArray, AxisArray, ClipTransformer]):
+    SETTINGS = ClipSettings
 
-class ClipSettings(ez.Settings):
-    a_min: float
-    a_max: float
 
+def clip(min: float | None = None, max: float | None = None) -> ClipTransformer:
+    """
+    Clips the data to be within the specified range.
 
-class Clip(GenAxisArray):
-    SETTINGS = ClipSettings
+    Args:
+        min: Lower clip bound. If None, no lower clipping is applied.
+        max: Upper clip bound. If None, no upper clipping is applied.
 
-    def construct_generator(self):
-        self.STATE.gen = clip(a_min=self.SETTINGS.a_min, a_max=self.SETTINGS.a_max)
+    Returns:
+        :obj:`ClipTransformer`.
+    """
+    return ClipTransformer(ClipSettings(min=min, max=max))