ezmsg-sigproc 1.8.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

ezmsg/sigproc/transpose.py

@@ -0,0 +1,131 @@
+from types import EllipsisType
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import (
+    AxisArray,
+    replace,
+)
+
+from .base import (
+    BaseStatefulTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+
+
+class TransposeSettings(ez.Settings):
+    """
+    Settings for :obj:`Transpose` node.
+
+    Fields:
+      axes:
+    """
+
+    axes: tuple[int | str | EllipsisType, ...] | None = None
+    order: str | None = None
+
+
+@processor_state
+class TransposeState:
+    axes_ints: tuple[int, ...] | None = None
+
+
+class TransposeTransformer(
+    BaseStatefulTransformer[TransposeSettings, AxisArray, AxisArray, TransposeState]
+):
+    """
+    Downsampled data simply comprise every `factor`th sample.
+    This should only be used following appropriate lowpass filtering.
+    If your pipeline does not already have lowpass filtering then consider
+    using the :obj:`Decimate` collection instead.
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash(tuple(message.dims))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        if self.settings.axes is None:
+            self._state.axes_ints = None
+        else:
+            ell_ix = [ix for ix, ax in enumerate(self.settings.axes) if ax is Ellipsis]
+            if len(ell_ix) > 1:
+                raise ValueError("Only one Ellipsis is allowed in axes.")
+            ell_ix = ell_ix[0] if len(ell_ix) == 1 else len(message.dims)
+            prefix = []
+            for ax in self.settings.axes[:ell_ix]:
+                if isinstance(ax, int):
+                    prefix.append(ax)
+                else:
+                    if ax not in message.dims:
+                        raise ValueError(f"Axis {ax} not found in message dims.")
+                    prefix.append(message.dims.index(ax))
+            suffix = []
+            for ax in self.settings.axes[ell_ix + 1 :]:
+                if isinstance(ax, int):
+                    suffix.append(ax)
+                else:
+                    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
+            ]
+            re_ix = tuple(prefix + ells + suffix)
+            if re_ix == tuple(range(message.data.ndim)):
+                self._state.axes_ints = None
+            else:
+                self._state.axes_ints = re_ix
+        if self.settings.order is not None and self.settings.order.upper()[0] not in [
+            "C",
+            "F",
+        ]:
+            raise ValueError("order must be 'C' or 'F'.")
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.settings.axes is None and self.settings.order is None:
+            # Passthrough
+            return message
+        return super().__call__(message)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self.state.axes_ints is None:
+            # No transpose required
+            if self.settings.order is None:
+                # No memory relayout required
+                # Note: We should not be able to reach here because it should be shortcutted at passthrough.
+                msg_out = message
+            else:
+                # 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]
+                    ),
+                )
+        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]
+                )
+            msg_out = replace(
+                message,
+                data=data_out,
+                dims=dims_out,
+            )
+        return msg_out
+
+
+class Transpose(
+    BaseTransformerUnit[TransposeSettings, AxisArray, AxisArray, TransposeTransformer]
+):
+    SETTINGS = TransposeSettings
+
+
+def transpose(
+    axes: tuple[int | str | EllipsisType, ...] | None = None, order: str | None = None
+) -> TransposeTransformer:
+    return TransposeTransformer(TransposeSettings(axes=axes, order=order))

ezmsg/sigproc/util/asio.py

@@ -0,0 +1,156 @@
+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)

ezmsg/sigproc/util/message.py

@@ -0,0 +1,31 @@
+import time
+import typing
+from dataclasses import dataclass, field
+
+from ezmsg.util.messages.axisarray import AxisArray
+
+
+@dataclass(unsafe_hash=True)
+class SampleTriggerMessage:
+    timestamp: float = field(default_factory=time.time)
+    """Time of the trigger, in seconds. The Clock depends on the input but defaults to time.time"""
+
+    period: tuple[float, float] | None = None
+    """The period around the timestamp, in seconds"""
+
+    value: typing.Any = None
+    """A value or 'label' associated with the trigger."""
+
+
+@dataclass
+class SampleMessage:
+    trigger: SampleTriggerMessage
+    """The time, window, and value (if any) associated with the trigger."""
+
+    sample: AxisArray
+    """The data sampled around the trigger."""
+
+
+def is_sample_message(message: typing.Any) -> typing.TypeGuard[SampleMessage]:
+    """Check if the message is a SampleMessage."""
+    return hasattr(message, "trigger")

ezmsg/sigproc/util/profile.py

@@ -8,6 +8,9 @@ import typing
 import ezmsg.core as ez
 
 
+HEADER = "Time,Source,Topic,SampleTime,PerfCounter,Elapsed"
+
+
 def get_logger_path() -> Path:
     # Retrieve the logfile name from the environment variable
     logfile = os.environ.get("EZMSG_PROFILE", None)
@@ -26,9 +29,23 @@ def _setup_logger(append: bool = False) -> logging.Logger:
     logpath = get_logger_path()
     logpath.parent.mkdir(parents=True, exist_ok=True)
 
-    if not append:
-        # Remove the file if it exists
-        logpath.unlink(missing_ok=True)
+    write_header = True
+    if logpath.exists() and logpath.is_file():
+        if append:
+            with open(logpath) as f:
+                first_line = f.readline().rstrip()
+            if first_line == HEADER:
+                write_header = False
+            else:
+                # Remove the file if appending, but headers do not match
+                ezmsg_logger = logging.getLogger("ezmsg")
+                ezmsg_logger.warning(
+                    "Profiling header mismatch: please make sure to use the same version of ezmsg for all processes."
+                )
+                logpath.unlink()
+        else:
+            # Remove the file if not appending
+            logpath.unlink()
 
     # Create a logger with the name "ezprofile"
     _logger = logging.getLogger("ezprofile")
@@ -43,13 +60,13 @@ def _setup_logger(append: bool = False) -> logging.Logger:
     # Add the file handler to the logger
     _logger.addHandler(fh)
 
-    # Add the first row without formatting.
-    _logger.debug(",".join(["Time", "Source", "Topic", "SampleTime", "PerfCounter", "Elapsed"]))
+    # Add the header if writing to new file or if header matched header in file.
+    if write_header:
+        _logger.debug(HEADER)
 
     # Set the log message format
     formatter = logging.Formatter(
-        "%(asctime)s,%(message)s",
-        datefmt="%Y-%m-%dT%H:%M:%S%z"
+        "%(asctime)s,%(message)s", datefmt="%Y-%m-%dT%H:%M:%S%z"
     )
     fh.setFormatter(formatter)
 
@@ -89,18 +106,31 @@ def profile_method(trace_oldest: bool = True):
     Returns:
         Callable: The decorated function with profiling.
     """
+
     def profiling_decorator(func: typing.Callable):
         @functools.wraps(func)
         def wrapped_func(caller, *args, **kwargs):
             start = time.perf_counter()
             res = func(caller, *args, **kwargs)
             stop = time.perf_counter()
-            source = '.'.join((caller.__class__.__module__, caller.__class__.__name__))
+            source = ".".join((caller.__class__.__module__, caller.__class__.__name__))
             topic = f"{caller.address}"
             samp_time = _process_obj(res, trace_oldest=trace_oldest)
-            logger.debug(",".join([source, topic, f"{samp_time}", f"{stop}", f"{(stop - start) * 1e3:0.4f}"]))
+            logger.debug(
+                ",".join(
+                    [
+                        source,
+                        topic,
+                        f"{samp_time}",
+                        f"{stop}",
+                        f"{(stop - start) * 1e3:0.4f}",
+                    ]
+                )
+            )
             return res
+
         return wrapped_func if logger.level == logging.DEBUG else func
+
     return profiling_decorator
 
 
@@ -115,17 +145,30 @@ def profile_subpub(trace_oldest: bool = True):
     Returns:
         Callable: The decorated async task with profiling.
     """
+
     def profiling_decorator(func: typing.Callable):
         @functools.wraps(func)
-        async def wrapped_task(unit: ez.Unit, msg: typing.Any = None) -> None:
-            source = '.'.join((unit.__class__.__module__, unit.__class__.__name__))
+        async def wrapped_task(unit: ez.Unit, msg: typing.Any = None):
+            source = ".".join((unit.__class__.__module__, unit.__class__.__name__))
             topic = f"{unit.address}"
             start = time.perf_counter()
             async for stream, obj in func(unit, msg):
                 stop = time.perf_counter()
                 samp_time = _process_obj(obj, trace_oldest=trace_oldest)
-                logger.debug(",".join([source, topic, f"{samp_time}", f"{stop}", f"{(stop - start) * 1e3:0.4f}"]))
+                logger.debug(
+                    ",".join(
+                        [
+                            source,
+                            topic,
+                            f"{samp_time}",
+                            f"{stop}",
+                            f"{(stop - start) * 1e3:0.4f}",
+                        ]
+                    )
+                )
                 start = stop
                 yield stream, obj
+
         return wrapped_task if logger.level == logging.DEBUG else func
+
     return profiling_decorator

ezmsg/sigproc/util/typeresolution.py

@@ -0,0 +1,83 @@
+from types import UnionType
+import typing
+from typing_extensions import get_original_bases
+
+
+def resolve_typevar(cls: type, target_typevar: typing.TypeVar) -> type:
+    """
+    Resolve the concrete type bound to a TypeVar in a class hierarchy.
+    This function traverses the method resolution order (MRO) of the class
+    and checks the original bases of each class in the MRO for the TypeVar.
+    If the TypeVar is found, it returns the concrete type bound to it.
+    If the TypeVar is not found, it raises a TypeError.
+    Args:
+        cls (type): The class to inspect.
+        target_typevar (typing.TypeVar): The TypeVar to resolve.
+    Returns:
+        type: The concrete type bound to the TypeVar.
+    """
+    for base in cls.__mro__:
+        orig_bases = get_original_bases(base)
+        for orig_base in orig_bases:
+            origin = typing.get_origin(orig_base)
+            if origin is None:
+                continue
+            params = getattr(origin, "__parameters__", ())
+            if not params:
+                continue
+            if target_typevar in params:
+                index = params.index(target_typevar)
+                args = typing.get_args(orig_base)
+                try:
+                    return args[index]
+                except IndexError:
+                    pass
+    raise TypeError(f"Could not resolve {target_typevar} in {cls}")
+
+
+TypeLike = typing.Union[type[typing.Any], typing.Any, type(None), None]
+
+
+def check_message_type_compatibility(type1: TypeLike, type2: TypeLike) -> bool:
+    """
+    Check if two types are compatible for message passing.
+    Returns True if:
+    - Both are None/NoneType
+    - Either is typing.Any
+    - type1 is a subclass of type2, which includes
+        - type1 and type2 are concrete types and type1 is a subclass of type2
+        - type1 is None/NoneType and type2 is typing.Optional, or
+        - type1 is subtype of the non-None inner type of type2 if type2 is Optional
+    - type1 is a Union/Optional type and all inner types are compatible with type2
+    Args:
+        type1: First type to compare
+        type2: Second type to compare
+    Returns:
+        bool: True if the types are compatible, False otherwise
+    """
+    # If either is Any, they are compatible
+    if type1 is typing.Any or type2 is typing.Any:
+        return True
+
+    # Handle None as NoneType
+    if type1 is None:
+        type1 = type(None)
+    if type2 is None:
+        type2 = type(None)
+
+    # Handle if type1 is Optional/Union type
+    if typing.get_origin(type1) in {typing.Union, UnionType}:
+        return all(
+            check_message_type_compatibility(inner_type, type2)
+            for inner_type in typing.get_args(type1)
+        )
+
+    # Regular issubclass check. Handles cases like:
+    # - type1 is a subclass of concrete type2
+    # - type1 is a subclass of the inner type of type2 if type2 is Optional
+    # - type1 is a subclass of one of the inner types of type2 if type2 is Union
+    # - type1 is NoneType and type2 is Optional or Union[None, ...] or Union[NoneType, ...]
+    try:
+        return issubclass(type1, type2)
+    except TypeError:
+        return False