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

ezmsg/sigproc/math/difference.py

@@ -1,18 +1,47 @@
+"""
+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 numpy as np
 import ezmsg.core as ez
-from ezmsg.util.generator import consumer
+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
 
-from ..base import GenAxisArray
+
+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
 
 
-@consumer
-def const_difference(
-    value: float = 0.0, subtrahend: bool = True
-) -> typing.Generator[AxisArray, AxisArray, None]:
+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
@@ -22,48 +51,93 @@ def const_difference(
         subtrahend: If True (default) then value is subtracted from the input data.
          If False, the input data is subtracted from value.
 
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the difference between the input :obj:`AxisArray` data and the 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.
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = replace(
-            msg_in, data=(msg_in.data - value) if subtrahend else (value - msg_in.data)
-        )
 
+    def __init__(self):
+        self._state = DifferenceState()
 
-class ConstDifferenceSettings(ez.Settings):
-    value: float = 0.0
-    subtrahend: bool = True
+    @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
 
-class ConstDifference(GenAxisArray):
-    SETTINGS = ConstDifferenceSettings
+    def push_a(self, msg: AxisArray) -> None:
+        """Push a message to queue A (minuend)."""
+        self._state.queue_a.put_nowait(msg)
 
-    def construct_generator(self):
-        self.STATE.gen = const_difference(
-            value=self.SETTINGS.value, subtrahend=self.SETTINGS.subtrahend
-        )
+    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)
 
-# class DifferenceSettings(ez.Settings):
-#     pass
-#
-#
-# class Difference(ez.Unit):
-#     SETTINGS = DifferenceSettings
-#
-#     INPUT_SIGNAL_1 = ez.InputStream(AxisArray)
-#     INPUT_SIGNAL_2 = ez.InputStream(AxisArray)
-#     OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
-#
-#     @ez.subscriber(INPUT_SIGNAL_2, zero_copy=True)
-#     @ez.publisher(OUTPUT_SIGNAL)
-#     async def on_input_2(self, message: AxisArray) -> typing.AsyncGenerator:
-#         # TODO: buffer_2
-#         # TODO: take buffer_1 - buffer_2 for ranges that align
-#         # TODO: Drop samples from buffer_1 and buffer_2
-#         if ret is not None:
-#             yield self.OUTPUT_SIGNAL, ret
+    @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

@@ -1,35 +1,28 @@
-import typing
+"""
+Compute the multiplicative inverse (1/x) 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 ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import GenAxisArray
-
-
-@consumer
-def invert() -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Take the inverse of the data.
-
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the inversion 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=1 / msg_in.data)
+class InvertTransformer(BaseTransformer[None, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        return replace(message, data=1 / message.data)
 
 
-class InvertSettings(ez.Settings):
-    pass
+class Invert(BaseTransformerUnit[None, AxisArray, AxisArray, InvertTransformer]): ...  # SETTINGS = None
 
 
-class Invert(GenAxisArray):
-    SETTINGS = InvertSettings
+def invert() -> InvertTransformer:
+    """
+    Take the inverse of the data.
 
-    def construct_generator(self):
-        self.STATE.gen = invert()
+    Returns: :obj:`InvertTransformer`.
+    """
+    return InvertTransformer()

ezmsg/sigproc/math/log.py

@@ -1,19 +1,49 @@
-import typing
+"""
+Take the logarithm 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.
+"""
 
-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 LogSettings(ez.Settings):
+    base: float = 10.0
+    """The base of the logarithm. Default is 10."""
+
+    clip_zero: bool = False
+    """If True, clip the data to the minimum positive value of the data type before taking the log."""
+
+
+class LogTransformer(BaseTransformer[LogSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        data = message.data
+        if self.settings.clip_zero:
+            # Check if any values are <= 0 and dtype is floating point
+            has_non_positive = bool(xp.any(data <= 0))
+            is_floating = xp.isdtype(data.dtype, "real floating")
+            if has_non_positive and is_floating:
+                # Use smallest_normal (Array API equivalent of numpy's finfo.tiny)
+                min_val = xp.finfo(data.dtype).smallest_normal
+                data = xp.clip(data, min_val, None)
+        return replace(message, data=xp.log(data) / xp.log(self.settings.base))
+
+
+class Log(BaseTransformerUnit[LogSettings, AxisArray, AxisArray, LogTransformer]):
+    SETTINGS = LogSettings
 
 
-@consumer
 def log(
     base: float = 10.0,
     clip_zero: bool = False,
-) -> typing.Generator[AxisArray, AxisArray, None]:
+) -> LogTransformer:
     """
     Take the logarithm of the data. See :obj:`np.log` for more details.
 
@@ -21,32 +51,7 @@ def log(
         base: The base of the logarithm. Default is 10.
         clip_zero: If True, clip the data to the minimum positive value of the data type before taking the log.
 
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the logarithm of the input :obj:`AxisArray` data.
+    Returns: :obj:`LogTransformer`.
 
     """
-    msg_out = AxisArray(np.array([]), dims=[""])
-    log_base = np.log(base)
-    while True:
-        msg_in: AxisArray = yield msg_out
-        if (
-            clip_zero
-            and np.any(msg_in.data <= 0)
-            and np.issubdtype(msg_in.data.dtype, np.floating)
-        ):
-            msg_in.data = np.clip(
-                msg_in.data, a_min=np.finfo(msg_in.data.dtype).tiny, a_max=None
-            )
-        msg_out = replace(msg_in, data=np.log(msg_in.data) / log_base)
-
-
-class LogSettings(ez.Settings):
-    base: float = 10.0
-    clip_zero: bool = False
-
-
-class Log(GenAxisArray):
-    SETTINGS = LogSettings
-
-    def construct_generator(self):
-        self.STATE.gen = log(base=self.SETTINGS.base, clip_zero=self.SETTINGS.clip_zero)
+    return LogTransformer(LogSettings(base=base, clip_zero=clip_zero))

ezmsg/sigproc/math/scale.py

@@ -1,40 +1,39 @@
-import typing
+"""
+Scale the data by a constant factor.
+
+.. 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 ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import GenAxisArray
 
+class ScaleSettings(ez.Settings):
+    scale: float = 1.0
+    """Factor by which to scale the data magnitude."""
 
-@consumer
-def scale(scale: float = 1.0) -> typing.Generator[AxisArray, AxisArray, None]:
-    """
-    Scale the data by a constant factor.
 
-    Args:
-        scale: Factor by which to scale the data magnitude.
+class ScaleTransformer(BaseTransformer[ScaleSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        return replace(message, data=self.settings.scale * message.data)
 
-    Returns: A primed generator that, when passed an input message via `.send(msg)`, yields an :obj:`AxisArray`
-     with the data payload containing the input :obj:`AxisArray` data scaled by a constant factor.
 
-    """
-    msg_out = AxisArray(np.array([]), dims=[""])
-    while True:
-        msg_in: AxisArray = yield msg_out
-        msg_out = replace(msg_in, data=scale * msg_in.data)
+class Scale(BaseTransformerUnit[ScaleSettings, AxisArray, AxisArray, ScaleTransformer]):
+    SETTINGS = ScaleSettings
 
 
-class ScaleSettings(ez.Settings):
-    scale: float = 1.0
+def scale(scale: float = 1.0) -> ScaleTransformer:
+    """
+    Scale the data by a constant factor.
 
+    Args:
+        scale: Factor by which to scale the data magnitude.
 
-class Scale(GenAxisArray):
-    SETTINGS = ScaleSettings
+    Returns: :obj:`ScaleTransformer`
 
-    def construct_generator(self):
-        self.STATE.gen = scale(
-            scale=self.SETTINGS.scale,
-        )
+    """
+    return ScaleTransformer(ScaleSettings(scale=scale))

ezmsg/sigproc/messages.py

@@ -1,10 +1,9 @@
-import warnings
 import time
+import warnings
 
 import numpy.typing as npt
 from ezmsg.util.messages.axisarray import AxisArray
 
-
 # UPCOMING: TSMessage Deprecation
 # TSMessage is deprecated because it doesn't handle multiple time axes well.
 # AxisArray has an incompatible API but supports a superset of functionality.

ezmsg/sigproc/quantize.py

@@ -0,0 +1,68 @@
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray, replace
+
+
+class QuantizeSettings(ez.Settings):
+    """
+    Settings for the Quantizer.
+    """
+
+    max_val: float
+    """
+    Clip the data to this maximum value before quantization and map the [min_val max_val] range to the quantized range.
+    """
+
+    min_val: float = 0.0
+    """
+    Clip the data to this minimum value before quantization and map the [min_val max_val] range to the quantized range.
+    Default: 0
+    """
+
+    bits: int = 8
+    """
+    Number of bits for quantization.
+    Note: The data type will be integer of the next power of 2 greater than or equal to this value.
+    Default: 8
+    """
+
+
+class QuantizeTransformer(BaseTransformer[QuantizeSettings, AxisArray, AxisArray]):
+    def _process(
+        self,
+        message: AxisArray,
+    ) -> AxisArray:
+        expected_range = self.settings.max_val - self.settings.min_val
+        scale_factor = 2**self.settings.bits - 1
+        clip_max = self.settings.max_val
+
+        # Determine appropriate integer type based on bits
+        if self.settings.bits <= 1:
+            dtype = bool
+        elif self.settings.bits <= 8:
+            dtype = np.uint8
+        elif self.settings.bits <= 16:
+            dtype = np.uint16
+        elif self.settings.bits <= 32:
+            dtype = np.uint32
+        else:
+            dtype = np.uint64
+            if self.settings.bits == 64:
+                # The practical upper bound before converting to int is: 2**64 - 1025
+                #  Anything larger will wrap around to 0.
+                #
+                clip_max *= 1 - 2e-16
+
+        data = message.data.clip(self.settings.min_val, clip_max)
+        data = (data - self.settings.min_val) / expected_range
+
+        # Scale to the quantized range [0, 2^bits - 1]
+        data = np.rint(scale_factor * data).astype(dtype)
+
+        # Create a new AxisArray with the quantized data
+        return replace(message, data=data)
+
+
+class QuantizerUnit(BaseTransformerUnit[QuantizeSettings, AxisArray, AxisArray, QuantizeTransformer]):
+    SETTINGS = QuantizeSettings