ezmsg-sigproc 2.5.0__py3-none-any.whl → 2.7.0__py3-none-any.whl

ezmsg/sigproc/spectrogram.py

@@ -1,20 +1,21 @@
 from typing import Generator
+
 import ezmsg.core as ez
+from ezmsg.baseproc import (
+    BaseStatefulProcessor,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.modify import modify_axis
 
-from .window import Anchor, WindowTransformer
 from .spectrum import (
-    WindowFunction,
-    SpectralTransform,
     SpectralOutput,
+    SpectralTransform,
     SpectrumTransformer,
+    WindowFunction,
 )
-from .base import (
-    CompositeProcessor,
-    BaseStatefulProcessor,
-    BaseTransformerUnit,
-)
+from .window import Anchor, WindowTransformer
 
 
 class SpectrogramSettings(ez.Settings):
@@ -41,9 +42,7 @@ class SpectrogramSettings(ez.Settings):
     """The :obj:`SpectralOutput` format."""
 
 
-class SpectrogramTransformer(
-    CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]
-):
+class SpectrogramTransformer(CompositeProcessor[SpectrogramSettings, AxisArray, AxisArray]):
     @staticmethod
     def _initialize_processors(
         settings: SpectrogramSettings,
@@ -54,9 +53,7 @@ class SpectrogramTransformer(
                 newaxis="win",
                 window_dur=settings.window_dur,
                 window_shift=settings.window_shift,
-                zero_pad_until="shift"
-                if settings.window_shift is not None
-                else "input",
+                zero_pad_until="shift" if settings.window_shift is not None else "input",
                 anchor=settings.window_anchor,
             ),
             "spectrum": SpectrumTransformer(
@@ -69,11 +66,7 @@ class SpectrogramTransformer(
         }
 
 
-class Spectrogram(
-    BaseTransformerUnit[
-        SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer
-    ]
-):
+class Spectrogram(BaseTransformerUnit[SpectrogramSettings, AxisArray, AxisArray, SpectrogramTransformer]):
     SETTINGS = SpectrogramSettings
 
 

ezmsg/sigproc/spectrum.py

@@ -1,21 +1,20 @@
 import enum
-from functools import partial
 import typing
+from functools import partial
 
+import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
-import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import (
-    AxisArray,
-    slice_along_axis,
-    replace,
-)
-
-from .base import (
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+    slice_along_axis,
+)
 
 
 class OptionsEnum(enum.Enum):
@@ -127,17 +126,13 @@ class SpectrumState:
     window: npt.NDArray | None = None
 
 
-class SpectrumTransformer(
-    BaseStatefulTransformer[SpectrumSettings, AxisArray, AxisArray, SpectrumState]
-):
+class SpectrumTransformer(BaseStatefulTransformer[SpectrumSettings, AxisArray, AxisArray, SpectrumState]):
     def _hash_message(self, message: AxisArray) -> int:
         axis = self.settings.axis or message.dims[0]
         ax_idx = message.get_axis_idx(axis)
         ax_info = message.axes[axis]
         targ_len = message.data.shape[ax_idx]
-        return hash(
-            (targ_len, message.data.ndim, message.data.dtype.kind, ax_idx, ax_info.gain)
-        )
+        return hash((targ_len, message.data.ndim, message.data.dtype.kind, ax_idx, ax_info.gain))
 
     def _reset_state(self, message: AxisArray) -> None:
         axis = self.settings.axis or message.dims[0]
@@ -156,8 +151,7 @@ class SpectrumTransformer(
             + [1] * (message.data.ndim - 1 - ax_idx)
         )
         if self.settings.transform != SpectralTransform.RAW_COMPLEX and not (
-            self.settings.transform == SpectralTransform.REAL
-            or self.settings.transform == SpectralTransform.IMAG
+            self.settings.transform == SpectralTransform.REAL or self.settings.transform == SpectralTransform.IMAG
         ):
             scale = np.sum(window**2.0) * ax_info.gain
 
@@ -170,30 +164,21 @@ class SpectrumTransformer(
         if (not b_complex) and self.settings.output == SpectralOutput.POSITIVE:
             # If input is not complex and desired output is SpectralOutput.POSITIVE, we can save some computation
             #  by using rfft and rfftfreq.
-            self.state.fftfun = partial(
-                np.fft.rfft, n=nfft, axis=ax_idx, norm=self.settings.norm
-            )
+            self.state.fftfun = partial(np.fft.rfft, n=nfft, axis=ax_idx, norm=self.settings.norm)
             freqs = np.fft.rfftfreq(nfft, d=ax_info.gain * targ_len / nfft)
         else:
-            self.state.fftfun = partial(
-                np.fft.fft, n=nfft, axis=ax_idx, norm=self.settings.norm
-            )
+            self.state.fftfun = partial(np.fft.fft, n=nfft, axis=ax_idx, norm=self.settings.norm)
             freqs = np.fft.fftfreq(nfft, d=ax_info.gain * targ_len / nfft)
             if self.settings.output == SpectralOutput.POSITIVE:
                 self.state.f_sl = slice(None, nfft // 2 + 1 - (nfft % 2))
             elif self.settings.output == SpectralOutput.NEGATIVE:
                 freqs = np.fft.fftshift(freqs, axes=-1)
                 self.state.f_sl = slice(None, nfft // 2 + 1)
-            elif (
-                self.settings.do_fftshift
-                and self.settings.output == SpectralOutput.FULL
-            ):
+            elif self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL:
                 freqs = np.fft.fftshift(freqs, axes=-1)
             freqs = freqs[self.state.f_sl]
         freqs = freqs.tolist()  # To please type checking
-        self.state.freq_axis = AxisArray.LinearAxis(
-            unit="Hz", gain=freqs[1] - freqs[0], offset=freqs[0]
-        )
+        self.state.freq_axis = AxisArray.LinearAxis(unit="Hz", gain=freqs[1] - freqs[0], offset=freqs[0])
         self.state.new_dims = (
             message.dims[:ax_idx]
             + [
@@ -232,11 +217,7 @@ class SpectrumTransformer(
         ax_idx = message.get_axis_idx(axis)
         targ_len = message.data.shape[ax_idx]
 
-        new_axes = {
-            k: v
-            for k, v in message.axes.items()
-            if k not in [self.settings.out_axis, axis]
-        }
+        new_axes = {k: v for k, v in message.axes.items() if k not in [self.settings.out_axis, axis]}
         new_axes[self.settings.out_axis or axis] = self.state.freq_axis
 
         if self.state.window is not None:
@@ -261,9 +242,7 @@ class SpectrumTransformer(
         return msg_out
 
 
-class Spectrum(
-    BaseTransformerUnit[SpectrumSettings, AxisArray, AxisArray, SpectrumTransformer]
-):
+class Spectrum(BaseTransformerUnit[SpectrumSettings, AxisArray, AxisArray, SpectrumTransformer]):
     SETTINGS = SpectrumSettings
 
 

ezmsg/sigproc/transpose.py

@@ -1,16 +1,16 @@
 from types import EllipsisType
-import numpy as np
-import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import (
-    AxisArray,
-    replace,
-)
 
-from .base import (
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+)
 
 
 class TransposeSettings(ez.Settings):
@@ -30,9 +30,7 @@ class TransposeState:
     axes_ints: tuple[int, ...] | None = None
 
 
-class TransposeTransformer(
-    BaseStatefulTransformer[TransposeSettings, AxisArray, AxisArray, TransposeState]
-):
+class TransposeTransformer(BaseStatefulTransformer[TransposeSettings, AxisArray, AxisArray, TransposeState]):
     """
     Downsampled data simply comprise every `factor`th sample.
     This should only be used following appropriate lowpass filtering.
@@ -67,11 +65,7 @@ class TransposeTransformer(
                     if ax not in message.dims:
                         raise ValueError(f"Axis {ax} not found in message dims.")
                     suffix.append(message.dims.index(ax))
-            ells = [
-                _
-                for _ in range(message.data.ndim)
-                if _ not in prefix and _ not in suffix
-            ]
+            ells = [_ for _ in range(message.data.ndim) if _ not in prefix and _ not in suffix]
             re_ix = tuple(prefix + ells + suffix)
             if re_ix == tuple(range(message.data.ndim)):
                 self._state.axes_ints = None
@@ -100,17 +94,13 @@ class TransposeTransformer(
                 # If the memory is already contiguous in the correct order, np.require won't do anything.
                 msg_out = replace(
                     message,
-                    data=np.require(
-                        message.data, requirements=self.settings.order.upper()[0]
-                    ),
+                    data=np.require(message.data, requirements=self.settings.order.upper()[0]),
                 )
         else:
             dims_out = [message.dims[ix] for ix in self.state.axes_ints]
             data_out = np.transpose(message.data, axes=self.state.axes_ints)
             if self.settings.order is not None:
-                data_out = np.require(
-                    data_out, requirements=self.settings.order.upper()[0]
-                )
+                data_out = np.require(data_out, requirements=self.settings.order.upper()[0])
             msg_out = replace(
                 message,
                 data=data_out,
@@ -119,9 +109,7 @@ class TransposeTransformer(
         return msg_out
 
 
-class Transpose(
-    BaseTransformerUnit[TransposeSettings, AxisArray, AxisArray, TransposeTransformer]
-):
+class Transpose(BaseTransformerUnit[TransposeSettings, AxisArray, AxisArray, TransposeTransformer]):
     SETTINGS = TransposeSettings
 
 

ezmsg/sigproc/util/asio.py

@@ -1,156 +1,25 @@
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
-import contextlib
-import inspect
-import threading
-from typing import Any, Coroutine, TypeVar
-
-T = TypeVar("T")
-
-
-class CoroutineExecutionError(Exception):
-    """Custom exception for coroutine execution failures"""
-
-    pass
-
-
-def run_coroutine_sync(coroutine: Coroutine[Any, Any, T], timeout: float = 30) -> T:
-    """
-    Executes an asyncio coroutine synchronously, with enhanced error handling.
-
-    Args:
-        coroutine: The asyncio coroutine to execute
-        timeout: Maximum time in seconds to wait for coroutine completion (default: 30)
-
-    Returns:
-        The result of the coroutine execution
-
-    Raises:
-        CoroutineExecutionError: If execution fails due to threading or event loop issues
-        TimeoutError: If execution exceeds the timeout period
-        Exception: Any exception raised by the coroutine
-    """
-
-    def run_in_new_loop() -> T:
-        """
-        Creates and runs a new event loop in the current thread.
-        Ensures proper cleanup of the loop.
-        """
-        new_loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(new_loop)
-        try:
-            return new_loop.run_until_complete(
-                asyncio.wait_for(coroutine, timeout=timeout)
-            )
-        finally:
-            with contextlib.suppress(Exception):
-                # Clean up any pending tasks
-                pending = asyncio.all_tasks(new_loop)
-                for task in pending:
-                    task.cancel()
-                new_loop.run_until_complete(
-                    asyncio.gather(*pending, return_exceptions=True)
-                )
-            new_loop.close()
-
-    try:
-        loop = asyncio.get_running_loop()
-    except RuntimeError:
-        try:
-            return asyncio.run(asyncio.wait_for(coroutine, timeout=timeout))
-        except Exception as e:
-            raise CoroutineExecutionError(
-                f"Failed to execute coroutine: {str(e)}"
-            ) from e
-
-    if threading.current_thread() is threading.main_thread():
-        if not loop.is_running():
-            try:
-                return loop.run_until_complete(
-                    asyncio.wait_for(coroutine, timeout=timeout)
-                )
-            except Exception as e:
-                raise CoroutineExecutionError(
-                    f"Failed to execute coroutine in main loop: {str(e)}"
-                ) from e
-        else:
-            with ThreadPoolExecutor() as pool:
-                try:
-                    future = pool.submit(run_in_new_loop)
-                    return future.result(timeout=timeout)
-                except Exception as e:
-                    raise CoroutineExecutionError(
-                        f"Failed to execute coroutine in thread: {str(e)}"
-                    ) from e
-    else:
-        try:
-            future = asyncio.run_coroutine_threadsafe(coroutine, loop)
-            return future.result(timeout=timeout)
-        except Exception as e:
-            raise CoroutineExecutionError(
-                f"Failed to execute coroutine threadsafe: {str(e)}"
-            ) from e
-
-
-class SyncToAsyncGeneratorWrapper:
-    """
-    A wrapper for synchronous generators to be used in an async context.
-    """
-
-    def __init__(self, gen):
-        self._gen = gen
-        self._closed = False
-        # Prime the generator to ready for first send/next call
-        try:
-            is_not_primed = inspect.getgeneratorstate(self._gen) is inspect.GEN_CREATED
-        except AttributeError as e:
-            raise TypeError(
-                "The provided generator is not a valid generator object"
-            ) from e
-        if is_not_primed:
-            try:
-                next(self._gen)
-            except StopIteration:
-                self._closed = True
-            except Exception as e:
-                raise RuntimeError(f"Failed to prime generator: {e}") from e
-
-    async def asend(self, value):
-        if self._closed:
-            raise StopAsyncIteration("Generator is closed")
-        try:
-            return await asyncio.to_thread(self._gen.send, value)
-        except StopIteration as e:
-            self._closed = True
-            raise StopAsyncIteration("Generator is closed") from e
-        except Exception as e:
-            raise RuntimeError(f"Error while sending value to generator: {e}") from e
-
-    async def __anext__(self):
-        if self._closed:
-            raise StopAsyncIteration("Generator is closed")
-        try:
-            return await asyncio.to_thread(self._gen.__next__)
-        except StopIteration as e:
-            self._closed = True
-            raise StopAsyncIteration("Generator is closed") from e
-        except Exception as e:
-            raise RuntimeError(
-                f"Error while getting next value from generator: {e}"
-            ) from e
-
-    async def aclose(self):
-        if self._closed:
-            return
-        try:
-            await asyncio.to_thread(self._gen.close)
-        except Exception as e:
-            raise RuntimeError(f"Error while closing generator: {e}") from e
-        finally:
-            self._closed = True
-
-    def __aiter__(self):
-        return self
-
-    def __getattr__(self, name):
-        return getattr(self._gen, name)
+"""
+Backwards-compatible re-exports from ezmsg.baseproc.util.asio.
+
+New code should import directly from ezmsg.baseproc instead.
+"""
+
+import warnings
+
+warnings.warn(
+    "Importing from 'ezmsg.sigproc.util.asio' is deprecated. Please import from 'ezmsg.baseproc.util.asio' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+from ezmsg.baseproc.util.asio import (  # noqa: E402
+    CoroutineExecutionError,
+    SyncToAsyncGeneratorWrapper,
+    run_coroutine_sync,
+)
+
+__all__ = [
+    "CoroutineExecutionError",
+    "SyncToAsyncGeneratorWrapper",
+    "run_coroutine_sync",
+]

ezmsg/sigproc/util/axisarray_buffer.py

@@ -1,14 +1,15 @@
+"""AxisArray support for .buffer.HybridBuffer."""
+
 import math
 import typing
 
-from array_api_compat import get_namespace
 import numpy as np
-from ezmsg.util.messages.axisarray import AxisArray, LinearAxis, CoordinateAxis
+from array_api_compat import get_namespace
+from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis, LinearAxis
 from ezmsg.util.messages.util import replace
 
 from .buffer import HybridBuffer
 
-
 Array = typing.TypeVar("Array")
 
 
@@ -68,9 +69,7 @@ class HybridAxisBuffer:
         if hasattr(first_axis, "data"):
             # Initialize a CoordinateAxis buffer
             if len(first_axis.data) > 1:
-                _axis_gain = (first_axis.data[-1] - first_axis.data[0]) / (
-                    len(first_axis.data) - 1
-                )
+                _axis_gain = (first_axis.data[-1] - first_axis.data[0]) / (len(first_axis.data) - 1)
             else:
                 _axis_gain = 1.0
             self._coords_gain_estimate = _axis_gain
@@ -107,8 +106,7 @@ class HybridAxisBuffer:
                 )
             if axis.gain != self._linear_axis.gain:
                 raise ValueError(
-                    f"Buffer initialized with gain={self._linear_axis.gain}, "
-                    f"but received gain={axis.gain}."
+                    f"Buffer initialized with gain={self._linear_axis.gain}, but received gain={axis.gain}."
                 )
             if self._linear_n_available + n_samples > self.capacity:
                 # Simulate overflow by advancing the offset and decreasing
@@ -117,16 +115,12 @@ class HybridAxisBuffer:
                 self.seek(n_to_discard)
             # Update the offset corresponding to the oldest sample in the buffer
             #  by anchoring on the new offset and accounting for the samples already available.
-            self._linear_axis.offset = (
-                axis.offset - self._linear_n_available * axis.gain
-            )
+            self._linear_axis.offset = axis.offset - self._linear_n_available * axis.gain
             self._linear_n_available += n_samples
 
     def peek(self, n_samples: int | None = None) -> LinearAxis | CoordinateAxis:
         if self._coords_buffer is not None:
-            return replace(
-                self._coords_template, data=self._coords_buffer.peek(n_samples)
-            )
+            return replace(self._coords_template, data=self._coords_buffer.peek(n_samples))
         else:
             # Return a shallow copy.
             return replace(self._linear_axis, offset=self._linear_axis.offset)
@@ -184,13 +178,9 @@ class HybridAxisBuffer:
         else:
             return None
 
-    def searchsorted(
-        self, values: typing.Union[float, Array], side: str = "left"
-    ) -> typing.Union[int, Array]:
+    def searchsorted(self, values: typing.Union[float, Array], side: str = "left") -> typing.Union[int, Array]:
         if self._coords_buffer is not None:
-            return self._coords_buffer.xp.searchsorted(
-                self._coords_buffer.peek(self.available()), values, side=side
-            )
+            return self._coords_buffer.xp.searchsorted(self._coords_buffer.peek(self.available()), values, side=side)
         else:
             if self.available() == 0:
                 if isinstance(values, float):
@@ -312,9 +302,7 @@ class HybridAxisArrayBuffer:
             axes={**self._template_msg.axes, self._axis: out_axis},
         )
 
-    def peek_axis(
-        self, n_samples: int | None = None
-    ) -> LinearAxis | CoordinateAxis | None:
+    def peek_axis(self, n_samples: int | None = None) -> LinearAxis | CoordinateAxis | None:
         """Retrieves the axis data without advancing the read head."""
         if self._data_buffer is None:
             return None
@@ -369,9 +357,7 @@ class HybridAxisArrayBuffer:
         """
         return self._axis_buffer.gain
 
-    def axis_searchsorted(
-        self, values: typing.Union[float, Array], side: str = "left"
-    ) -> typing.Union[int, Array]:
+    def axis_searchsorted(self, values: typing.Union[float, Array], side: str = "left") -> typing.Union[int, Array]:
         """
         Find the indices into which the given values would be inserted
         into the target axis data to maintain order.