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

ezmsg/sigproc/kaiser.py

@@ -6,10 +6,10 @@ import numpy.typing as npt
 import scipy.signal
 
 from .filter import (
-    FilterBaseSettings,
-    FilterByDesignTransformer,
     BACoeffs,
     BaseFilterByDesignTransformerUnit,
+    FilterBaseSettings,
+    FilterByDesignTransformer,
 )
 
 
@@ -20,30 +20,30 @@ class KaiserFilterSettings(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.
     """
 
     ripple: float | None = None
     """
-    Upper bound for the deviation (in dB) of the magnitude of the filter's frequency response from that of 
+    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 
+    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 
+    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’}
     """
@@ -88,9 +88,7 @@ def kaiser_design_fun(
     return (taps, np.array([1.0]))
 
 
-class KaiserFilterTransformer(
-    FilterByDesignTransformer[KaiserFilterSettings, BACoeffs]
-):
+class KaiserFilterTransformer(FilterByDesignTransformer[KaiserFilterSettings, BACoeffs]):
     def get_design_function(
         self,
     ) -> typing.Callable[[float], BACoeffs | None]:
@@ -104,7 +102,5 @@ class KaiserFilterTransformer(
         )
 
 
-class KaiserFilter(
-    BaseFilterByDesignTransformerUnit[KaiserFilterSettings, KaiserFilterTransformer]
-):
+class KaiserFilter(BaseFilterByDesignTransformerUnit[KaiserFilterSettings, KaiserFilterTransformer]):
     SETTINGS = KaiserFilterSettings

ezmsg/sigproc/math/abs.py

@@ -14,9 +14,7 @@ class AbsTransformer(BaseTransformer[None, AxisArray, AxisArray]):
         return replace(message, data=np.abs(message.data))
 
 
-class Abs(
-    BaseTransformerUnit[None, AxisArray, AxisArray, AbsTransformer]
-): ...  # SETTINGS = None
+class Abs(BaseTransformerUnit[None, AxisArray, AxisArray, AbsTransformer]): ...  # SETTINGS = None
 
 
 def abs() -> AbsTransformer:

ezmsg/sigproc/math/add.py

@@ -0,0 +1,121 @@
+"""Signal addition utilities."""
+
+import asyncio
+import typing
+from dataclasses import dataclass, field
+
+import ezmsg.core as ez
+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 BaseTransformer, BaseTransformerUnit
+
+# --- 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,5 +1,5 @@
-import numpy as np
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 

ezmsg/sigproc/math/difference.py

@@ -1,4 +1,9 @@
+import asyncio
+import typing
+from dataclasses import dataclass, field
+
 import ezmsg.core as ez
+from ezmsg.baseproc.util.asio import run_coroutine_sync
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -10,12 +15,11 @@ class ConstDifferenceSettings(ez.Settings):
     """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."""
+    """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]
-):
+class ConstDifferenceTransformer(BaseTransformer[ConstDifferenceSettings, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
         return replace(
             message,
@@ -25,17 +29,11 @@ class ConstDifferenceTransformer(
         )
 
 
-class ConstDifference(
-    BaseTransformerUnit[
-        ConstDifferenceSettings, AxisArray, AxisArray, ConstDifferenceTransformer
-    ]
-):
+class ConstDifference(BaseTransformerUnit[ConstDifferenceSettings, AxisArray, AxisArray, ConstDifferenceTransformer]):
     SETTINGS = ConstDifferenceSettings
 
 
-def const_difference(
-    value: float = 0.0, subtrahend: bool = True
-) -> ConstDifferenceTransformer:
+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
@@ -47,27 +45,91 @@ def const_difference(
 
     Returns: :obj:`ConstDifferenceTransformer`.
     """
-    return ConstDifferenceTransformer(
-        ConstDifferenceSettings(value=value, subtrahend=subtrahend)
-    )
-
-
-# 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
+    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

@@ -9,9 +9,7 @@ class InvertTransformer(BaseTransformer[None, AxisArray, AxisArray]):
         return replace(message, data=1 / message.data)
 
 
-class Invert(
-    BaseTransformerUnit[None, AxisArray, AxisArray, InvertTransformer]
-): ...  # SETTINGS = None
+class Invert(BaseTransformerUnit[None, AxisArray, AxisArray, InvertTransformer]): ...  # SETTINGS = None
 
 
 def invert() -> InvertTransformer:

ezmsg/sigproc/math/log.py

@@ -1,5 +1,5 @@
-import numpy as np
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -17,11 +17,7 @@ class LogSettings(ez.Settings):
 class LogTransformer(BaseTransformer[LogSettings, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
         data = message.data
-        if (
-            self.settings.clip_zero
-            and np.any(data <= 0)
-            and np.issubdtype(data.dtype, np.floating)
-        ):
+        if self.settings.clip_zero and np.any(data <= 0) and np.issubdtype(data.dtype, np.floating):
             data = np.clip(data, a_min=np.finfo(data.dtype).tiny, a_max=None)
         return replace(message, data=np.log(data) / np.log(self.settings.base))
 

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

@@ -1,5 +1,5 @@
-import numpy as np
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
 from .base import BaseTransformer, BaseTransformerUnit
@@ -65,7 +65,5 @@ class QuantizeTransformer(BaseTransformer[QuantizeSettings, AxisArray, AxisArray
         return replace(message, data=data)
 
 
-class QuantizerUnit(
-    BaseTransformerUnit[QuantizeSettings, AxisArray, AxisArray, QuantizeTransformer]
-):
+class QuantizerUnit(BaseTransformerUnit[QuantizeSettings, AxisArray, AxisArray, QuantizeTransformer]):
     SETTINGS = QuantizeSettings

ezmsg/sigproc/resample.py

@@ -2,15 +2,15 @@ import asyncio
 import math
 import time
 
+import ezmsg.core as ez
 import numpy as np
 import scipy.interpolate
-import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray, LinearAxis
 from ezmsg.util.messages.util import replace
 
 from .base import (
-    BaseStatefulProcessor,
     BaseConsumerUnit,
+    BaseStatefulProcessor,
     processor_state,
 )
 from .util.axisarray_buffer import HybridAxisArrayBuffer, HybridAxisBuffer
@@ -29,7 +29,7 @@ class ResampleSettings(ez.Settings):
     fill_value: str = "extrapolate"
     """
     Value to use for out-of-bounds samples.
-    If 'extrapolate', the transformer will extrapolate. 
+    If 'extrapolate', the transformer will extrapolate.
     If 'last', the transformer will use the last sample.
     See scipy.interpolate.interp1d for more options.
     """
@@ -57,9 +57,9 @@ class ResampleState:
     """
     The buffer for the reference axis (usually a time axis). The interpolation function
     will be evaluated at the reference axis values.
-    When resample_rate is None, this buffer will be filled with the axis from incoming 
+    When resample_rate is None, this buffer will be filled with the axis from incoming
     _reference_ messages.
-    When resample_rate is not None (i.e., prescribed float resample_rate), this buffer 
+    When resample_rate is not None (i.e., prescribed float resample_rate), this buffer
     is filled with a synthetic axis that is generated from the incoming signal messages.
     """
 
@@ -67,7 +67,7 @@ class ResampleState:
     """
     The last value of the reference axis that was returned. This helps us to know
     what the _next_ returned value should be, and to avoid returning the same value.
-    TODO: We can eliminate this variable if we maintain "by convention" that the 
+    TODO: We can eliminate this variable if we maintain "by convention" that the
     reference axis always has 1 value at its start that we exclude from the resampling.
     """
 
@@ -79,9 +79,7 @@ class ResampleState:
     """
 
 
-class ResampleProcessor(
-    BaseStatefulProcessor[ResampleSettings, AxisArray, AxisArray, ResampleState]
-):
+class ResampleProcessor(BaseStatefulProcessor[ResampleSettings, AxisArray, AxisArray, ResampleState]):
     def _hash_message(self, message: AxisArray) -> int:
         ax_idx: int = message.get_axis_idx(self.settings.axis)
         sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
@@ -135,17 +133,11 @@ class ResampleProcessor(
         ax_idx = message.get_axis_idx(self.settings.axis)
         if self.settings.resample_rate is not None and message.data.shape[ax_idx] > 0:
             in_ax = message.axes[self.settings.axis]
-            in_t_end = (
-                in_ax.data[-1]
-                if hasattr(in_ax, "data")
-                else in_ax.value(message.data.shape[ax_idx] - 1)
-            )
+            in_t_end = in_ax.data[-1] if hasattr(in_ax, "data") else in_ax.value(message.data.shape[ax_idx] - 1)
             out_gain = 1 / self.settings.resample_rate
             prev_t_end = self.state.last_ref_ax_val
             n_synth = math.ceil((in_t_end - prev_t_end) * self.settings.resample_rate)
-            synth_ref_axis = LinearAxis(
-                unit="s", gain=out_gain, offset=prev_t_end + out_gain
-            )
+            synth_ref_axis = LinearAxis(unit="s", gain=out_gain, offset=prev_t_end + out_gain)
             self.state.ref_axis_buffer.write(synth_ref_axis, n_samples=n_synth)
 
         self.state.last_write_time = time.time()
@@ -193,11 +185,7 @@ class ResampleProcessor(
         # Get source to train interpolation
         src_axarr = src.peek()
         src_axis = src_axarr.axes[self.settings.axis]
-        x = (
-            src_axis.data
-            if hasattr(src_axis, "data")
-            else src_axis.value(np.arange(src_axarr.data.shape[0]))
-        )
+        x = src_axis.data if hasattr(src_axis, "data") else src_axis.value(np.arange(src_axarr.data.shape[0]))
 
         # Only resample at reference values that have not been interpolated over previously.
         b_ref = ref_xvec > self.state.last_ref_ax_val
@@ -208,11 +196,7 @@ class ResampleProcessor(
 
         if len(ref_idx) == 0:
             # Nothing to interpolate over; return empty data
-            null_ref = (
-                replace(ref_ax, data=ref_ax.data[:0])
-                if hasattr(ref_ax, "data")
-                else ref_ax
-            )
+            null_ref = replace(ref_ax, data=ref_ax.data[:0]) if hasattr(ref_ax, "data") else ref_ax
             return replace(
                 src_axarr,
                 data=src_axarr.data[:0, ...],
@@ -222,17 +206,12 @@ class ResampleProcessor(
         xnew = ref_xvec[ref_idx]
 
         # Identify source data indices around ref tvec with some padding for better interpolation.
-        src_start_ix = max(
-            0, np.where(x > xnew[0])[0][0] - 2 if np.any(x > xnew[0]) else 0
-        )
+        src_start_ix = max(0, np.where(x > xnew[0])[0][0] - 2 if np.any(x > xnew[0]) else 0)
 
         x = x[src_start_ix:]
         y = src_axarr.data[src_start_ix:]
 
-        if (
-            isinstance(self.settings.fill_value, str)
-            and self.settings.fill_value == "last"
-        ):
+        if isinstance(self.settings.fill_value, str) and self.settings.fill_value == "last":
             fill_value = (y[0], y[-1])
         else:
             fill_value = self.settings.fill_value