ezmsg-sigproc 2.8.0__tar.gz → 2.10.0__tar.gz

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.8.0
+Version: 2.10.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

ezmsg_sigproc-2.10.0/docs/source/guides/explanations/array_api.rst

@@ -0,0 +1,156 @@
+Array API Support
+=================
+
+ezmsg-sigproc provides support for the `Python Array API standard
+<https://data-apis.org/array-api/>`_, enabling many transformers to work with
+arrays from different backends such as NumPy, CuPy, PyTorch, and JAX.
+
+What is the Array API?
+----------------------
+
+The Array API is a standardized interface for array operations across different
+Python array libraries. By coding to this standard, ezmsg-sigproc transformers
+can process data regardless of which array library created it, enabling:
+
+- **GPU acceleration** via CuPy or PyTorch tensors
+- **Framework interoperability** for integration with ML pipelines
+- **Hardware flexibility** without code changes
+
+How It Works
+------------
+
+Compatible transformers use `array-api-compat <https://github.com/data-apis/array-api-compat>`_
+to detect the input array's namespace and use the appropriate operations:
+
+.. code-block:: python
+
+    from array_api_compat import get_namespace
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)  # numpy, cupy, torch, etc.
+        result = xp.abs(message.data)     # Uses the correct backend
+        return replace(message, data=result)
+
+Usage Example
+-------------
+
+Using Array API compatible transformers with CuPy for GPU acceleration:
+
+.. code-block:: python
+
+    import cupy as cp
+    from ezmsg.util.messages.axisarray import AxisArray
+    from ezmsg.sigproc.math.abs import AbsTransformer
+    from ezmsg.sigproc.math.clip import ClipTransformer, ClipSettings
+
+    # Create data on GPU
+    gpu_data = cp.random.randn(1000, 64).astype(cp.float32)
+    message = AxisArray(gpu_data, dims=["time", "ch"])
+
+    # Process entirely on GPU - no data transfer!
+    abs_transformer = AbsTransformer()
+    clip_transformer = ClipTransformer(ClipSettings(min=0.0, max=1.0))
+
+    result = clip_transformer(abs_transformer(message))
+    # result.data is still a CuPy array on GPU
+
+Compatible Modules
+------------------
+
+The following transformers fully support the Array API standard:
+
+Math Operations
+^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.math.abs`
+     - Absolute value
+   * - :mod:`ezmsg.sigproc.math.clip`
+     - Clip values to a range
+   * - :mod:`ezmsg.sigproc.math.log`
+     - Logarithm with configurable base
+   * - :mod:`ezmsg.sigproc.math.scale`
+     - Multiply by a constant
+   * - :mod:`ezmsg.sigproc.math.invert`
+     - Compute 1/x
+   * - :mod:`ezmsg.sigproc.math.difference`
+     - Subtract a constant (ConstDifferenceTransformer)
+
+Signal Processing
+^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.diff`
+     - Compute differences along an axis
+   * - :mod:`ezmsg.sigproc.transpose`
+     - Transpose/permute array dimensions
+   * - :mod:`ezmsg.sigproc.linear`
+     - Per-channel linear transform (scale + offset)
+   * - :mod:`ezmsg.sigproc.aggregate`
+     - Aggregate operations (AggregateTransformer only)
+
+Coordinate Transforms
+^^^^^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.coordinatespaces`
+     - Cartesian/polar coordinate conversions
+
+Limitations
+-----------
+
+Some operations remain NumPy-only due to lack of Array API equivalents:
+
+- **Random number generation**: Modules using ``np.random`` (e.g., ``denormalize``)
+- **SciPy operations**: Filtering (``scipy.signal.lfilter``), FFT, wavelets
+- **Advanced indexing**: Some slicing operations for metadata handling
+- **Memory layout**: ``np.require`` for contiguous array optimization (NumPy only)
+
+Metadata arrays (axis labels, coordinates) typically remain as NumPy arrays
+since they are not performance-critical.
+
+Adding Array API Support
+------------------------
+
+When contributing new transformers, follow this pattern:
+
+.. code-block:: python
+
+    from array_api_compat import get_namespace
+    from ezmsg.baseproc import BaseTransformer
+    from ezmsg.util.messages.axisarray import AxisArray
+    from ezmsg.util.messages.util import replace
+
+    class MyTransformer(BaseTransformer[MySettings, AxisArray, AxisArray]):
+        def _process(self, message: AxisArray) -> AxisArray:
+            xp = get_namespace(message.data)
+
+            # Use xp instead of np for array operations
+            result = xp.sqrt(xp.abs(message.data))
+
+            return replace(message, data=result)
+
+Key guidelines:
+
+1. Call ``get_namespace(message.data)`` at the start of ``_process``
+2. Use ``xp.function_name`` instead of ``np.function_name``
+3. Note that some functions have different names:
+   - ``np.concatenate`` → ``xp.concat``
+   - ``np.transpose`` → ``xp.permute_dims``
+4. Keep metadata operations (axis labels, etc.) as NumPy
+5. Use in-place operations (``/=``, ``*=``) where possible for efficiency

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/docs/source/guides/sigproc/content-sigproc.rst

@@ -4,7 +4,8 @@ ezmsg-sigproc
 Timeseries signal processing implementations in ezmsg, leveraging numpy and scipy.
 Most of the methods and classes in this extension are intended to be used in building signal processing pipelines.
 They use :class:`ezmsg.util.messages.axisarray.AxisArray` as the primary data structure for passing signals between components.
-The message's data are expected to be a numpy array.
+The message's data are typically NumPy arrays, though many transformers support the
+:doc:`Array API standard <../explanations/array_api>` for use with CuPy, PyTorch, and other backends.
 
 .. note:: Some generators might yield valid :class:`AxisArray` messages with ``.data`` size of 0.
 This may occur when the generator receives inadequate data to produce a valid output, such as when windowing or buffering.
@@ -21,3 +22,4 @@ This may occur when the generator receives inadequate data to produce a valid ou
     base
     units
     processors
+    ../explanations/array_api

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/docs/source/guides/tutorials/signalprocessing.rst

@@ -109,7 +109,7 @@ First, we need to install the `ezmsg-sigproc` package if we haven't already. Thi
 
 .. code-block:: bash
 
-    pip install "ezmsg[sigproc]"
+    pip install "ezmsg-sigproc"
 
 
 |ezmsg_logo_small| Creating the `Downsample` signal processor
@@ -237,7 +237,7 @@ The first two methods deal with the state of the processor (and are only require
 
 In order to implement these methods, we need to understand our preferred message format: `AxisArray`. This is a flexible and powerful class for handling multi-dimensional arrays with named axes, which is particularly useful for signal processing applications. I have already used `AxisArray` in our code as the input message and output message types.
 
-A detailed explanation of the `AxisArray` class is beyond the scope of this tutorial, but you can refer to the :doc:`AxisArray explainer <../explanations/axisarray>` as well as the :doc:`API reference <../reference/API/axisarray>` for more information.
+A detailed explanation of the `AxisArray` class is beyond the scope of this tutorial, but you can refer to the `AxisArray explainer <https://www.ezmsg.org/explanations/axisarray.html>`_ as well as the `API reference <https://www.ezmsg.org/ezmsg/reference/API/axisarray.html>`_ for more information.
 
 Brief Aside on AxisArray
 =================================
@@ -555,25 +555,6 @@ In a separate Python file in the same directory, you can test the `DownsampleTra
 
 Doing the above is very handy for unit testing your processor as well as for offline processing of data.
 
-.. note:: The `downsample` module in `ezmsg-sigproc` has a utility function for creating a `DownsampleTransformer` instance with the desired settings:
-
-    .. code-block:: python
-
-        def downsample(
-            axis: str = "time",
-            target_rate: float | None = None,
-            factor: int | None = None,
-        ) -> DownsampleTransformer:
-            return DownsampleTransformer(
-                DownsampleSettings(axis=axis, target_rate=target_rate, factor=factor)
-            )
-
-    After importing this utility function, lines 8 and 9 in our code above could now read:
-
-    .. code-block:: python
-
-        downsampler = downsample(axis="time", target_rate=50)
-
 Of course, the real power of `ezmsg` comes from integrating your processor into an `ezmsg` Unit and using it in a processing pipeline. We will see how to do this next.
 
 
@@ -600,15 +581,15 @@ A lot of the behind-the-scenes work is done for you by the `BaseTransformerUnit`
         SETTINGS = DownsampleSettings
 
 
-Connecting it to other `Component`\ s and initialising the transformer are accomplished in the same way that we did in the :doc:`pipeline tutorial <pipeline>`.
+Connecting it to other `Component`\ s and initialising the transformer are accomplished in the same way that we did in the `Pipeline Tutorial <https://www.ezmsg.org/tutorials/pipeline.html>`_
 
 
 |ezmsg_logo_small| See Also
 ************************************
 
 - `Further examples <https://github.com/ezmsg-org/ezmsg/tree/master/examples>`_ can be found in the examples directory in `ezmsg`. These are examples of creating and using `ezmsg` Units and pipelines.
-- `ezmsg-sigproc` has a large number of already implemented signal processors. More information can be found at the :doc:`ezmsg-sigproc reference <../extensions/sigproc/content-sigproc>`.
-- `Downsample` class reference
+- `ezmsg-sigproc` has a large number of already implemented signal processors. More information can be found at the :doc:`ezmsg-sigproc reference <../sigproc/content-sigproc>`.
+- :doc:`Downsample class reference <../../api/generated/ezmsg.sigproc.downsample>`
 
 .. |ezmsg_logo_small| image:: ../_static/_images/ezmsg_logo.png
   :width: 40

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/src/ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.8.0'
-__version_tuple__ = version_tuple = (2, 8, 0)
+__version__ = version = '2.10.0'
+__version_tuple__ = version_tuple = (2, 10, 0)
 
 __commit_id__ = commit_id = None

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/src/ezmsg/sigproc/aggregate.py

@@ -1,3 +1,12 @@
+"""
+Aggregation operations over arrays.
+
+.. note::
+    :obj:`AggregateTransformer` supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    :obj:`RangedAggregateTransformer` currently requires NumPy arrays.
+"""
+
 import typing
 
 import ezmsg.core as ez

ezmsg_sigproc-2.10.0/src/ezmsg/sigproc/butterworthzerophase.py

@@ -0,0 +1,305 @@
+"""
+Streaming zero-phase Butterworth filter implemented as a two-stage composite processor.
+
+Stage 1: Forward causal Butterworth filter (from ezmsg.sigproc.butterworthfilter)
+Stage 2: Backward acausal filter with buffering (ButterworthBackwardFilterTransformer)
+
+The output is delayed by `pad_length` samples to ensure the backward pass has sufficient
+future context. The pad_length is computed analytically using scipy's heuristic.
+"""
+
+import functools
+import typing
+
+import numpy as np
+import scipy.signal
+from ezmsg.baseproc import BaseTransformerUnit
+from ezmsg.baseproc.composite import CompositeProcessor
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .butterworthfilter import (
+    ButterworthFilterSettings,
+    ButterworthFilterTransformer,
+    butter_design_fun,
+)
+from .filter import BACoeffs, FilterByDesignTransformer, SOSCoeffs
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+
+
+class ButterworthZeroPhaseSettings(ButterworthFilterSettings):
+    """
+    Settings for :obj:`ButterworthZeroPhase`.
+
+    This implements a streaming zero-phase Butterworth filter using forward-backward
+    filtering. The output is delayed by `pad_length` samples to ensure the backward
+    pass has sufficient future context.
+
+    The pad_length is computed by finding where the filter's impulse response decays
+    to `settle_cutoff` fraction of its peak value. This accounts for the filter's
+    actual time constant rather than just its order.
+    """
+
+    # Inherits from ButterworthFilterSettings:
+    # axis, coef_type, order, cuton, cutoff, wn_hz
+
+    settle_cutoff: float = 0.01
+    """
+    Fraction of peak impulse response used to determine settling time.
+    The pad_length is set to the number of samples until the impulse response
+    decays to this fraction of its peak. Default is 0.01 (1% of peak).
+    """
+
+    max_pad_duration: float | None = None
+    """
+    Maximum pad duration in seconds. If set, the pad_length will be capped
+    at this value times the sampling rate. Use this to limit latency for
+    filters with very long impulse responses. Default is None (no limit).
+    """
+
+
+class ButterworthBackwardFilterTransformer(FilterByDesignTransformer[ButterworthFilterSettings, BACoeffs | SOSCoeffs]):
+    """
+    Backward (acausal) Butterworth filter with buffering.
+
+    This transformer buffers its input and applies the filter in reverse,
+    outputting only the "settled" portion where transients have decayed.
+    This introduces a lag of ``pad_length`` samples.
+
+    Intended to be used as stage 2 in a zero-phase filter pipeline, receiving
+    forward-filtered data from a ButterworthFilterTransformer.
+    """
+
+    # Instance attributes (initialized in _reset_state)
+    _buffer: HybridAxisArrayBuffer | None
+    _coefs_cache: BACoeffs | SOSCoeffs | None
+    _zi_tiled: np.ndarray | None
+    _pad_length: int
+
+    def get_design_function(
+        self,
+    ) -> typing.Callable[[float], BACoeffs | SOSCoeffs | None]:
+        return functools.partial(
+            butter_design_fun,
+            order=self.settings.order,
+            cuton=self.settings.cuton,
+            cutoff=self.settings.cutoff,
+            coef_type=self.settings.coef_type,
+            wn_hz=self.settings.wn_hz,
+        )
+
+    def _compute_pad_length(self, fs: float) -> int:
+        """
+        Compute pad length based on the filter's impulse response settling time.
+
+        The pad_length is determined by finding where the impulse response decays
+        to `settle_cutoff` fraction of its peak value. This is then optionally
+        capped by `max_pad_duration`.
+
+        Args:
+            fs: Sampling frequency in Hz.
+
+        Returns:
+            Number of samples for the pad length.
+        """
+        # Design the filter to compute impulse response
+        coefs = self.get_design_function()(fs)
+        if coefs is None:
+            # Filter design failed or is disabled
+            return 0
+
+        # Generate impulse response - use a generous length initially
+        # Start with scipy's heuristic as minimum, then extend if needed
+        if self.settings.coef_type == "ba":
+            min_length = 3 * (self.settings.order + 1)
+        else:
+            n_sections = (self.settings.order + 1) // 2
+            min_length = 3 * n_sections * 2
+
+        # Use 10x the minimum as initial impulse length, or at least 10000 samples
+        # (10000 samples allows for ~333ms at 30kHz, covering most practical cases)
+        impulse_length = max(min_length * 10, 10000)
+
+        # Cap impulse length computation if max_pad_duration is set
+        if self.settings.max_pad_duration is not None:
+            max_samples = int(self.settings.max_pad_duration * fs)
+            impulse_length = min(impulse_length, max_samples + 1)
+
+        impulse = np.zeros(impulse_length)
+        impulse[0] = 1.0
+
+        if self.settings.coef_type == "ba":
+            b, a = coefs
+            h = scipy.signal.lfilter(b, a, impulse)
+        else:
+            h = scipy.signal.sosfilt(coefs, impulse)
+
+        # Find where impulse response settles to settle_cutoff of peak
+        abs_h = np.abs(h)
+        peak = abs_h.max()
+        if peak == 0:
+            return min_length
+
+        threshold = self.settings.settle_cutoff * peak
+        above_threshold = np.where(abs_h > threshold)[0]
+
+        if len(above_threshold) == 0:
+            pad_length = min_length
+        else:
+            pad_length = above_threshold[-1] + 1
+
+        # Ensure at least the scipy heuristic minimum
+        pad_length = max(pad_length, min_length)
+
+        # Apply max_pad_duration cap if set
+        if self.settings.max_pad_duration is not None:
+            max_samples = int(self.settings.max_pad_duration * fs)
+            pad_length = min(pad_length, max_samples)
+
+        return pad_length
+
+    def _reset_state(self, message: AxisArray) -> None:
+        """Reset filter state when stream changes."""
+        self._coefs_cache = None
+        self._zi_tiled = None
+        self._buffer = None
+        # Compute pad_length based on the message's sampling rate
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        fs = 1 / message.axes[axis].gain
+        self._pad_length = self._compute_pad_length(fs)
+        self.state.needs_redesign = True
+
+    def _compute_zi_tiled(self, data: np.ndarray, ax_idx: int) -> None:
+        """Compute and cache the tiled zi for the given data shape.
+
+        Called once per stream (or after filter redesign). The result is
+        broadcast-ready for multiplication by the edge sample on each chunk.
+        """
+        if self.settings.coef_type == "ba":
+            b, a = self._coefs_cache
+            zi_base = scipy.signal.lfilter_zi(b, a)
+        else:  # sos
+            zi_base = scipy.signal.sosfilt_zi(self._coefs_cache)
+
+        n_tail = data.ndim - ax_idx - 1
+
+        if self.settings.coef_type == "ba":
+            zi_expand = (None,) * ax_idx + (slice(None),) + (None,) * n_tail
+            n_tile = data.shape[:ax_idx] + (1,) + data.shape[ax_idx + 1 :]
+        else:  # sos
+            zi_expand = (slice(None),) + (None,) * ax_idx + (slice(None),) + (None,) * n_tail
+            n_tile = (1,) + data.shape[:ax_idx] + (1,) + data.shape[ax_idx + 1 :]
+
+        self._zi_tiled = np.tile(zi_base[zi_expand], n_tile)
+
+    def _initialize_zi(self, data: np.ndarray, ax_idx: int) -> np.ndarray:
+        """Initialize filter state (zi) scaled by edge value."""
+        if self._zi_tiled is None:
+            self._compute_zi_tiled(data, ax_idx)
+        first_sample = np.take(data, [0], axis=ax_idx)
+        return self._zi_tiled * first_sample
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        ax_idx = message.get_axis_idx(axis)
+        fs = 1 / message.axes[axis].gain
+
+        # Check if we need to redesign filter
+        if self._coefs_cache is None or self.state.needs_redesign:
+            self._coefs_cache = self.get_design_function()(fs)
+            self._pad_length = self._compute_pad_length(fs)
+            self._zi_tiled = None  # Invalidate; recomputed on next use.
+            self.state.needs_redesign = False
+
+            # Initialize buffer with duration based on pad_length
+            # Add some margin to handle variable chunk sizes
+            buffer_duration = (self._pad_length + 1) / fs
+            self._buffer = HybridAxisArrayBuffer(duration=buffer_duration, axis=axis)
+
+        # Early exit if filter is effectively disabled
+        if self._coefs_cache is None or self.settings.order <= 0 or message.data.size <= 0:
+            return message
+
+        # Write new data to buffer
+        self._buffer.write(message)
+        n_available = self._buffer.available()
+        n_output = n_available - self._pad_length
+
+        # If we don't have enough data yet, return empty
+        if n_output <= 0:
+            new_shape = list(message.data.shape)
+            new_shape[ax_idx] = 0
+            empty_data = np.empty(new_shape, dtype=message.data.dtype)
+            return replace(message, data=empty_data)
+
+        # Peek all available data from buffer
+        # Note: HybridAxisArrayBuffer moves the target axis to position 0
+        buffered = self._buffer.peek(n_available)
+        combined = buffered.data
+        buffer_ax_idx = 0  # Buffer always puts time axis at position 0
+
+        # Backward filter on reversed data
+        combined_rev = np.flip(combined, axis=buffer_ax_idx)
+        backward_zi = self._initialize_zi(combined_rev, buffer_ax_idx)
+
+        if self.settings.coef_type == "ba":
+            b, a = self._coefs_cache
+            y_bwd_rev, _ = scipy.signal.lfilter(b, a, combined_rev, axis=buffer_ax_idx, zi=backward_zi)
+        else:  # sos
+            y_bwd_rev, _ = scipy.signal.sosfilt(self._coefs_cache, combined_rev, axis=buffer_ax_idx, zi=backward_zi)
+
+        # Reverse back to get output in correct time order
+        y_bwd = np.flip(y_bwd_rev, axis=buffer_ax_idx)
+
+        # Output the settled portion (first n_output samples)
+        y = y_bwd[:n_output]
+
+        # Advance buffer read head to discard output samples, keep pad_length
+        self._buffer.seek(n_output)
+
+        # Build output with adjusted time axis
+        # LinearAxis offset is already correct from the buffer
+        out_axis = buffered.axes[axis]
+
+        # Move axis back to original position if needed
+        if ax_idx != 0:
+            y = np.moveaxis(y, 0, ax_idx)
+
+        return replace(
+            message,
+            data=y,
+            axes={**message.axes, axis: out_axis},
+        )
+
+
+class ButterworthZeroPhaseTransformer(CompositeProcessor[ButterworthZeroPhaseSettings, AxisArray, AxisArray]):
+    """
+    Streaming zero-phase Butterworth filter as a composite of two stages.
+
+    Stage 1 (forward): Standard causal Butterworth filter with state
+    Stage 2 (backward): Acausal Butterworth filter with buffering
+
+    The output is delayed by ``pad_length`` samples.
+    """
+
+    @staticmethod
+    def _initialize_processors(
+        settings: ButterworthZeroPhaseSettings,
+    ) -> dict[str, typing.Any]:
+        # Both stages use the same filter design settings
+        return {
+            "forward": ButterworthFilterTransformer(settings),
+            "backward": ButterworthBackwardFilterTransformer(settings),
+        }
+
+    @classmethod
+    def get_message_type(cls, dir: str) -> type[AxisArray]:
+        if dir in ("in", "out"):
+            return AxisArray
+        raise ValueError(f"Invalid direction: {dir}. Must be 'in' or 'out'.")
+
+
+class ButterworthZeroPhase(
+    BaseTransformerUnit[ButterworthZeroPhaseSettings, AxisArray, AxisArray, ButterworthZeroPhaseTransformer]
+):
+    SETTINGS = ButterworthZeroPhaseSettings

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.10.0}/src/ezmsg/sigproc/coordinatespaces.py

@@ -3,6 +3,10 @@ Coordinate space transformations for streaming data.
 
 This module provides utilities and ezmsg nodes for transforming between
 Cartesian (x, y) and polar (r, theta) coordinate systems.
+
+.. 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 enum import Enum
@@ -11,6 +15,7 @@ from typing import Tuple
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+from array_api_compat import get_namespace, is_array_api_obj
 from ezmsg.baseproc import (
     BaseTransformer,
     BaseTransformerUnit,
@@ -20,14 +25,24 @@ from ezmsg.util.messages.axisarray import AxisArray, replace
 # -- Utility functions for coordinate transformations --
 
 
+def _get_namespace_or_numpy(*args: npt.ArrayLike):
+    """Get array namespace if any arg is an array, otherwise return numpy."""
+    for arg in args:
+        if is_array_api_obj(arg):
+            return get_namespace(arg)
+    return np
+
+
 def polar2z(r: npt.ArrayLike, theta: npt.ArrayLike) -> npt.ArrayLike:
     """Convert polar coordinates to complex number representation."""
-    return r * np.exp(1j * theta)
+    xp = _get_namespace_or_numpy(r, theta)
+    return r * xp.exp(1j * theta)
 
 
 def z2polar(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
     """Convert complex number to polar coordinates (r, theta)."""
-    return np.abs(z), np.angle(z)
+    xp = _get_namespace_or_numpy(z)
+    return xp.abs(z), xp.atan2(xp.imag(z), xp.real(z))
 
 
 def cart2z(x: npt.ArrayLike, y: npt.ArrayLike) -> npt.ArrayLike:
@@ -37,7 +52,8 @@ def cart2z(x: npt.ArrayLike, y: npt.ArrayLike) -> npt.ArrayLike:
 
 def z2cart(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
     """Convert complex number to Cartesian coordinates (x, y)."""
-    return np.real(z), np.imag(z)
+    xp = _get_namespace_or_numpy(z)
+    return xp.real(z), xp.imag(z)
 
 
 def cart2pol(x: npt.ArrayLike, y: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
@@ -90,6 +106,7 @@ class CoordinateSpacesTransformer(BaseTransformer[CoordinateSpacesSettings, Axis
     """
 
     def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)
 
@@ -116,9 +133,9 @@ class CoordinateSpacesTransformer(BaseTransformer[CoordinateSpacesSettings, Axis
             out_a, out_b = pol2cart(component_a, component_b)
 
         # Stack results back along the same axis
-        result = np.stack([out_a, out_b], axis=axis_idx)
+        result = xp.stack([out_a, out_b], axis=axis_idx)
 
-        # Update axis labels if present
+        # Update axis labels if present (use numpy for string labels)
         axes = message.axes
         if axis in axes and hasattr(axes[axis], "data"):
             if self.settings.mode == CoordinateMode.CART2POL: