diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 510982a5..a4b5f972 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -34,4 +34,4 @@ jobs:
uv tool run ruff check --output-format=github src
- name: Run tests
- run: uv run pytest tests
+ run: uv run pytest tests --benchmark-skip
diff --git a/.gitignore b/.gitignore
index b2b010e7..fdebf0dd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -147,3 +147,4 @@ cython_debug/
src/ezmsg/sigproc/__version__.py
uv.lock
*.local.json
+tmp/
diff --git a/docs/source/guides/explanations/array_api.rst b/docs/source/guides/explanations/array_api.rst
index 8d91312f..500540a6 100644
--- a/docs/source/guides/explanations/array_api.rst
+++ b/docs/source/guides/explanations/array_api.rst
@@ -3,7 +3,7 @@ Array API Support
ezmsg-sigproc provides support for the `Python Array API standard
`_, enabling many transformers to work with
-arrays from different backends such as NumPy, CuPy, PyTorch, and JAX.
+arrays from different backends such as NumPy, CuPy, PyTorch, JAX, and MLX.
What is the Array API?
----------------------
@@ -12,7 +12,8 @@ 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
+- **GPU acceleration** via CuPy, PyTorch, or JAX tensors
+- **Apple Silicon acceleration** via MLX
- **Framework interoperability** for integration with ML pipelines
- **Hardware flexibility** without code changes
@@ -27,7 +28,7 @@ to detect the input array's namespace and use the appropriate operations:
from array_api_compat import get_namespace
def _process(self, message: AxisArray) -> AxisArray:
- xp = get_namespace(message.data) # numpy, cupy, torch, etc.
+ xp = get_namespace(message.data) # numpy, cupy, torch, mlx.core, etc.
result = xp.abs(message.data) # Uses the correct backend
return replace(message, data=result)
@@ -90,14 +91,16 @@ Signal Processing
* - Module
- Description
+ * - :mod:`ezmsg.sigproc.spectrum`
+ - FFT-based spectrum (SpectrumTransformer)
+ * - :mod:`ezmsg.sigproc.aggregate`
+ - Aggregate operations (AggregateTransformer, RangedAggregateTransformer)
* - :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
^^^^^^^^^^^^^^^^^^^^^
@@ -111,17 +114,163 @@ Coordinate Transforms
* - :mod:`ezmsg.sigproc.coordinatespaces`
- Cartesian/polar coordinate conversions
+Composite Pipelines
+^^^^^^^^^^^^^^^^^^^
+
+These ``CompositeProcessor`` pipelines chain Array API-aware steps together.
+When fed non-NumPy arrays, each step in the pipeline preserves the backend:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 70
+
+ * - Module
+ - Description
+ * - :mod:`ezmsg.sigproc.bandpower`
+ - BandPowerTransformer (spectrogram + ranged aggregate)
+ * - :mod:`ezmsg.sigproc.singlebandpow`
+ - RMSBandPowerTransformer (with explicit ``backend`` setting; only after initial IIR filter)
+
+MLX on Apple Silicon
+--------------------
+
+`MLX `_ is an array library for Apple Silicon
+that provides GPU-accelerated operations with a NumPy-like API. ezmsg-sigproc's
+Array API support enables MLX acceleration for spectral analysis and other
+pipelines without code changes to the transformers themselves.
+
+Basic usage
+^^^^^^^^^^^
+
+Pass MLX arrays in your ``AxisArray`` messages:
+
+.. code-block:: python
+
+ import mlx.core as mx
+ import numpy as np
+ from ezmsg.util.messages.axisarray import AxisArray
+ from ezmsg.sigproc.spectrum import SpectrumTransformer, SpectrumSettings
+
+ # Create data as MLX array
+ np_data = np.random.randn(1000, 64).astype(np.float32)
+ message = AxisArray(
+ data=mx.array(np_data),
+ dims=["time", "ch"],
+ axes={"time": AxisArray.TimeAxis(fs=1000.0)},
+ )
+
+ proc = SpectrumTransformer(SpectrumSettings(axis="time"))
+ result = proc(message)
+ # result.data is an mlx.core.array
+
+Lazy evaluation and ``mx.eval``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+MLX uses **lazy evaluation** — computations are not executed until their results
+are needed. This allows MLX to fuse operations and optimize the computation
+graph. However, it means that timing code or downstream consumers may see
+artificially fast "processing" that is actually deferred.
+
+To force evaluation, call ``mx.eval()``:
+
+.. code-block:: python
+
+ result = proc(message)
+ mx.eval(result.data) # Forces computation to complete
+
+For ``CompositeProcessor`` pipelines (like ``BandPowerTransformer``), you can
+override ``_post_process`` to call ``mx.eval()`` automatically so that every
+output is fully materialized:
+
+.. code-block:: python
+
+ class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, AxisArray]):
+ @staticmethod
+ def _initialize_processors(settings):
+ return {
+ "spectrogram": SpectrogramTransformer(settings=settings.spectrogram_settings),
+ "aggregate": RangedAggregateTransformer(...),
+ }
+
+ def _post_process(self, result: AxisArray | None) -> AxisArray | None:
+ if result is not None:
+ try:
+ import mlx.core as mx
+
+ if isinstance(result.data, mx.array):
+ mx.eval(result.data)
+ except ImportError:
+ pass
+ return result
+
+This pattern is used by ``BandPowerTransformer`` and ``RMSBandPowerTransformer``.
+It ensures downstream consumers (ezmsg Units, visualization, logging) receive
+fully evaluated arrays without needing to know about MLX internals.
+It also provides a safety valve so the lazy graph does not accumulate if the graph
+is not evaluated at the right time downstream.
+
+.. note::
+
+ The ``_post_process`` hook is defined on ``CompositeProcessor`` in
+ ezmsg-baseproc. It runs after the entire processor chain completes and
+ receives the final output. The ``try``/``except ImportError`` pattern
+ keeps MLX as an optional dependency.
+
+MLX quirks
+^^^^^^^^^^
+
+MLX's Array API coverage is nearly complete but has a few gaps that
+ezmsg-sigproc works around internally:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 35 40
+
+ * - Feature
+ - MLX status
+ - Workaround
+ * - ``fft(norm=...)``
+ - Not supported
+ - Manual normalization (``/ n``, ``/ sqrt(n)``)
+ * - ``fftshift(axes=int)``
+ - Needs tuple
+ - Always pass ``axes=(idx,)``
+ * - ``fftfreq`` / ``rfftfreq``
+ - Not available
+ - Computed with NumPy (metadata only)
+ * - ``dtype.kind``
+ - No ``.kind`` attribute
+ - ``is_complex_dtype()`` helper in ``ezmsg.sigproc.util.array``
+ * - Window functions
+ - Not available
+ - Computed with NumPy, converted via ``xp.asarray()``
+ * - ``nan*`` functions
+ - Not available
+ - Falls back to NumPy automatically
+ * - Boolean indexing
+ - Not supported
+ - Avoided in hot paths; used only in NumPy metadata code
+ * - Slice with ``np.int64``
+ - Rejected
+ - Slice bounds cast to Python ``int``
+
+These workarounds are handled inside the transformers — user code does not need
+to account for them.
+
Limitations
-----------
Some operations remain NumPy-only due to lack of Array API equivalents:
+- **SciPy operations**: Butterworth filtering (``scipy.signal.sosfilt``) and
+ other scipy-dependent steps. Use ``AsArrayTransformer`` to convert between
+ backends at pipeline boundaries (see ``RMSBandPowerTransformer`` for an example).
- **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)
+- **Trapezoidal integration**: ``np.trapezoid`` has no Array API equivalent.
+ ``RangedAggregateTransformer`` falls back to NumPy transparently.
+- **Memory layout**: ``np.require`` for contiguous array optimization
-Metadata arrays (axis labels, coordinates) typically remain as NumPy arrays
+Metadata arrays (axis labels, coordinates) always remain as NumPy arrays
since they are not performance-critical.
Adding Array API Support
@@ -147,10 +296,27 @@ When contributing new transformers, follow this pattern:
Key guidelines:
-1. Call ``get_namespace(message.data)`` at the start of ``_process``
-2. Use ``xp.function_name`` instead of ``np.function_name``
+1. Call ``get_namespace(message.data)`` at the start of ``_process`` (or
+ ``_reset_state`` for stateful transformers).
+2. Use ``xp.function_name`` instead of ``np.function_name`` for all operations
+ on ``message.data``.
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
+4. Keep metadata operations (axis labels, etc.) as NumPy.
+5. When a backend lacks a function (e.g., MLX has no ``nanmean``), fall back
+ gracefully:
+
+ .. code-block:: python
+
+ func_name = "mean"
+ if hasattr(xp, func_name):
+ result = getattr(xp, func_name)(data, axis=axis_idx)
+ else:
+ result = np.mean(np.asarray(data), axis=axis_idx)
+
+6. For ``CompositeProcessor`` subclasses that may produce MLX output, add
+ a ``_post_process`` override to call ``mx.eval()`` (see the MLX section
+ above).
+7. Use portable helpers from ``ezmsg.sigproc.util.array`` when needed:
+ ``is_complex_dtype``, ``is_float_dtype``, ``xp_asarray``.
diff --git a/pyproject.toml b/pyproject.toml
index 8dca8cfd..0e6ab006 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,12 +9,12 @@ authors = [
]
license = "MIT"
readme = "README.md"
-requires-python = ">=3.10.15"
+requires-python = ">=3.10"
dynamic = ["version"]
dependencies = [
"array-api-compat>=1.11.1",
- "ezmsg-baseproc>=1.3.0",
- "ezmsg>=3.6.0",
+ "ezmsg[axisarray]>=3.7.2",
+ "ezmsg-baseproc>=1.5.0",
"mlx>=0.18.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
"numba>=0.61.0",
"numpy>=1.26.0",
@@ -29,6 +29,7 @@ dev = [
"pre-commit>=4.2.0",
"jupyter>=1.1.1",
"scipy-stubs>=1.15.3.0",
+ "matplotlib>=3.10.8",
{include-group = "lint"},
{include-group = "test"},
{include-group = "docs"},
@@ -39,6 +40,7 @@ lint = [
test = [
"frozendict>=2.4.4",
"pytest-asyncio>=0.24.0",
+ "pytest-benchmark>=5.2.3",
"pytest-cov>=5.0.0",
"pytest>=8.3.3",
]
@@ -70,7 +72,10 @@ packages = ["src/ezmsg"]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
norecursedirs = "tests/helpers"
-addopts = "-p no:warnings"
+addopts = "-p no:warnings --benchmark-group-by=group,param:n_channels -m 'not sosfilt'"
+markers = [
+ "sosfilt: marks sosfilt development benchmarks (deselected by default, run with: -m sosfilt)",
+]
[tool.ruff]
line-length = 120
diff --git a/src/ezmsg/sigproc/aggregate.py b/src/ezmsg/sigproc/aggregate.py
index 6854a0aa..c3687a27 100644
--- a/src/ezmsg/sigproc/aggregate.py
+++ b/src/ezmsg/sigproc/aggregate.py
@@ -2,9 +2,11 @@
Aggregation operations over arrays.
.. note::
- :obj:`AggregateTransformer` supports the :doc:`Array API standard `,
- enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
- :obj:`RangedAggregateTransformer` currently requires NumPy arrays.
+ :obj:`AggregateTransformer` and :obj:`RangedAggregateTransformer` support the
+ :doc:`Array API standard `, enabling use with
+ NumPy, CuPy, PyTorch, and other compatible array libraries.
+ Operations not available on a given backend (nan-variants, trapezoid) fall back
+ to NumPy automatically.
"""
import typing
@@ -131,7 +133,7 @@ def _reset_state(self, message: AxisArray) -> None:
slices = []
for start, stop in self.settings.bands:
inds = np.where(np.logical_and(self._state.ax_vec >= start, self._state.ax_vec <= stop))[0]
- slices.append(np.s_[inds[0] : inds[-1] + 1])
+ slices.append(slice(int(inds[0]), int(inds[-1]) + 1))
if hasattr(target_axis, "data"):
if self._state.ax_vec.dtype.type is np.str_:
sl_dat = f"{self._state.ax_vec[start]} - {self._state.ax_vec[stop]}"
@@ -151,40 +153,59 @@ def _reset_state(self, message: AxisArray) -> None:
def _process(self, message: AxisArray) -> AxisArray:
axis = self.settings.axis or message.dims[0]
ax_idx = message.get_axis_idx(axis)
- agg_func = AGGREGATORS[self.settings.operation]
+ xp = get_namespace(message.data)
+ op = self.settings.operation
- if self.settings.operation in [
- AggregationFunction.TRAPEZOID,
- ]:
- # Special handling for methods that require x-coordinates.
+ if op == AggregationFunction.TRAPEZOID:
+ # Trapezoid requires x-coordinates and has no Array API equivalent;
+ # fall back to numpy.
+ np_data = np.asarray(message.data)
out_data = [
- agg_func(
- slice_along_axis(message.data, sl, axis=ax_idx),
+ np.trapezoid(
+ slice_along_axis(np_data, sl, axis=ax_idx),
x=self._state.ax_vec[sl],
axis=ax_idx,
)
for sl in self._state.slices
]
+ stacked = np.stack(out_data, axis=ax_idx)
+ # Convert back to original backend if needed
+ if xp is not np:
+ stacked = xp.asarray(stacked)
else:
- out_data = [
- agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx) for sl in self._state.slices
- ]
+ # Use Array API function when available, fall back to numpy AGGREGATORS
+ func_name = op.value
+ if hasattr(xp, func_name):
+ agg_func = getattr(xp, func_name)
+ out_data = [
+ agg_func(slice_along_axis(message.data, sl, axis=ax_idx), axis=ax_idx) for sl in self._state.slices
+ ]
+ stacked = xp.stack(out_data, axis=ax_idx)
+ else:
+ # nan-variants etc. — fall back to numpy
+ np_agg = AGGREGATORS[op]
+ np_data = np.asarray(message.data)
+ out_data = [
+ np_agg(slice_along_axis(np_data, sl, axis=ax_idx), axis=ax_idx) for sl in self._state.slices
+ ]
+ stacked = np.stack(out_data, axis=ax_idx)
+ if xp is not np:
+ stacked = xp.asarray(stacked)
msg_out = replace(
message,
- data=np.stack(out_data, axis=ax_idx),
+ data=stacked,
axes={**message.axes, axis: self._state.out_axis},
)
- if self.settings.operation in [
- AggregationFunction.ARGMIN,
- AggregationFunction.ARGMAX,
- ]:
+ if op in (AggregationFunction.ARGMIN, AggregationFunction.ARGMAX):
+ # Post-process: convert indices to axis coordinate values.
+ # ax_vec is always numpy; offsets must be numpy for fancy indexing.
out_data = []
for sl_ix, sl in enumerate(self._state.slices):
- offsets = np.take(msg_out.data, [sl_ix], axis=ax_idx)
+ offsets = np.asarray(slice_along_axis(msg_out.data, sl_ix, axis=ax_idx))
out_data.append(self._state.ax_vec[sl][offsets])
- msg_out.data = np.concatenate(out_data, axis=ax_idx)
+ msg_out.data = np.stack(out_data, axis=ax_idx)
return msg_out
diff --git a/src/ezmsg/sigproc/bandpower.py b/src/ezmsg/sigproc/bandpower.py
index 31bb9a1d..5151b66e 100644
--- a/src/ezmsg/sigproc/bandpower.py
+++ b/src/ezmsg/sigproc/bandpower.py
@@ -52,6 +52,17 @@ def _initialize_processors(
),
}
+ def _post_process(self, result: AxisArray | None) -> AxisArray | None:
+ if result is not None:
+ try:
+ import mlx.core as mx
+
+ if isinstance(result.data, mx.array):
+ mx.eval(result.data)
+ except ImportError:
+ pass
+ return result
+
class BandPower(BaseTransformerUnit[BandPowerSettings, AxisArray, AxisArray, BandPowerTransformer]):
SETTINGS = BandPowerSettings
diff --git a/src/ezmsg/sigproc/filter.py b/src/ezmsg/sigproc/filter.py
index d19ef1b4..3454a1a0 100644
--- a/src/ezmsg/sigproc/filter.py
+++ b/src/ezmsg/sigproc/filter.py
@@ -50,6 +50,103 @@ def _normalize_coefs(
return coef_type, coefs
+def _sosfilt_xp(sos, x, axis_idx, zi, xp):
+ """SOS filtering via parallel prefix scan (direct-form II transposed).
+
+ Solves the IIR linear recurrence z[n+1] = A @ z[n] + B * x[n] using a
+ Hillis-Steele inclusive prefix scan in O(log N) sequential steps instead
+ of O(N), minimizing Python-level loop overhead for lazy-evaluation
+ backends like MLX.
+
+ Args:
+ sos: (n_sections, 6) SOS coefficient array. Each row is [b0, b1, b2, a0, a1, a2].
+ a0 is assumed to be 1.0 (standard for scipy.signal.butter output).
+ x: Input data array.
+ axis_idx: The axis along which to filter.
+ zi: Initial conditions, shape (n_sections, *x.shape[:axis_idx], 2, *x.shape[axis_idx+1:]).
+ xp: Array API namespace.
+
+ Returns:
+ (y, zf) tuple — filtered output and final filter state.
+ """
+ n_sections = sos.shape[0]
+ N = x.shape[axis_idx]
+
+ # Move time to axis 0 for uniform batch handling.
+ x = xp.moveaxis(x, axis_idx, 0) # (N, *batch)
+ zi = xp.moveaxis(zi, axis_idx + 1, 1) # (n_sections, 2, *batch)
+
+ # Flatten batch dims into one.
+ batch_shape = x.shape[1:]
+ batch_size = 1
+ for s in batch_shape:
+ batch_size *= s
+ x = xp.reshape(x, (N, batch_size)) # (N, B)
+ zi = xp.reshape(zi, (n_sections, 2, batch_size)) # (S, 2, B)
+
+ # Pre-allocate output zi.
+ zi_out = xp.zeros((n_sections, 2, batch_size), dtype=x.dtype)
+
+ for s in range(n_sections):
+ _b0 = float(sos[s, 0])
+ _b1 = float(sos[s, 1])
+ _b2 = float(sos[s, 2])
+ _a1 = float(sos[s, 4])
+ _a2 = float(sos[s, 5])
+
+ z_init = zi[s] # (2, B)
+
+ # State recurrence: z[n+1] = A @ z[n] + B_vec * x[n]
+ # A = [[-a1, 1], [-a2, 0]]
+ # B_vec = [b1 - a1*b0, b2 - a2*b0]
+ # Output: y[n] = b0 * x[n] + z[n][0]
+ A_mat = xp_asarray(xp, np.array([[-_a1, 1.0], [-_a2, 0.0]])) # (2, 2)
+ B_vec = xp_asarray(xp, np.array([_b1 - _a1 * _b0, _b2 - _a2 * _b0])) # (2,)
+
+ # Initialize scan elements:
+ # A_scan[n] = A for all n
+ # c_scan[n] = B_vec * x[n]
+ A_scan = xp.zeros((N, 2, 2), dtype=A_mat.dtype)
+ A_scan[:] = A_mat # broadcast A_mat into every row
+ c_scan = B_vec[None, :, None] * x[:, None, :] # (N, 2, B)
+
+ # Hillis-Steele inclusive prefix scan.
+ # Operator: (A_r, c_r) ∘ (A_l, c_l) = (A_r @ A_l, A_r @ c_l + c_r)
+ # After the scan, A_scan[n] = A^(n+1) and
+ # c_scan[n] = Σ_{k=0..n} A^(n-k) @ B_vec * x[k].
+ stride = 1
+ while stride < N:
+ right_A = A_scan[stride:] # (N-stride, 2, 2)
+ left_A = A_scan[:-stride] # (N-stride, 2, 2)
+ right_c = c_scan[stride:] # (N-stride, 2, B)
+ left_c = c_scan[:-stride] # (N-stride, 2, B)
+
+ A_scan[stride:] = right_A @ left_A
+ c_scan[stride:] = right_A @ left_c + right_c
+ stride *= 2
+
+ # Recover all states: z[n+1] = A_scan[n] @ z_init + c_scan[n]
+ z_from_scan = A_scan @ z_init[None, :, :] + c_scan # (N, 2, B)
+
+ # z[0..N-1] for output: prepend z_init, drop z[N].
+ z_needed = xp.zeros((N, 2, batch_size), dtype=x.dtype)
+ z_needed[0] = z_init
+ z_needed[1:] = z_from_scan[:-1]
+
+ # y[n] = b0 * x[n] + z[n][0]; output becomes input for the next section.
+ x = _b0 * x + z_needed[:, 0, :] # (N, B)
+
+ # Final state for this section: z[N]
+ zi_out[s] = z_from_scan[-1]
+
+ # Restore shapes.
+ x = xp.reshape(x, (N,) + batch_shape)
+ zi_out = xp.reshape(zi_out, (n_sections, 2) + batch_shape)
+ x = xp.moveaxis(x, 0, axis_idx)
+ zi_out = xp.moveaxis(zi_out, 1, axis_idx + 1)
+ return x, zi_out
+
+
def _fir_filt_fft(b, data, zi, axis_idx, xp):
"""FIR filtering via FFT convolution with streaming state.
@@ -235,7 +332,11 @@ def _reset_state(self, message: AxisArray) -> None:
zi_expand = (slice(None),) + zi_expand
n_tile = (1,) + n_tile
- self.state.zi = np.tile(zi[zi_expand], n_tile)
+ zi_tiled = np.tile(zi[zi_expand], n_tile)
+ if not is_numpy_array(message.data):
+ xp = get_namespace(message.data)
+ zi_tiled = xp_asarray(xp, zi_tiled)
+ self.state.zi = zi_tiled
self.state.fir_method = None
self.state.fir_b = None
self.state.fir_b_1d = None
@@ -297,7 +398,19 @@ def _process(self, message: AxisArray) -> AxisArray:
else:
_, coefs = _normalize_coefs(self.settings.coefs)
filt_func = {"ba": scipy.signal.lfilter, "sos": scipy.signal.sosfilt}[self.settings.coef_type]
+ input_xp = None if is_numpy_array(message.data) else get_namespace(message.data)
+ if input_xp is not None:
+ # Convert coefs and zi to the input namespace so scipy's
+ # array_namespace sees a single backend and converts back.
+ # NOTE: scipy 1.17 bundles an array_api_compat that does
+ # not recognize MLX, so we also convert the output below.
+ # When scipy's bundled copy gains MLX support, the manual
+ # conversion will become a no-op.
+ coefs = tuple(xp_asarray(input_xp, c) for c in coefs)
dat_out, self.state.zi = filt_func(*coefs, message.data, axis=axis_idx, zi=self.state.zi)
+ if input_xp is not None:
+ dat_out = xp_asarray(input_xp, dat_out)
+ self.state.zi = xp_asarray(input_xp, self.state.zi)
else:
dat_out = message.data
diff --git a/src/ezmsg/sigproc/materialize.py b/src/ezmsg/sigproc/materialize.py
new file mode 100644
index 00000000..79fedb43
--- /dev/null
+++ b/src/ezmsg/sigproc/materialize.py
@@ -0,0 +1,25 @@
+"""
+Materialize (evaluate) lazy array data.
+
+MLX arrays are lazily evaluated — computations are queued but not executed
+until the result is needed. This module provides an explicit evaluation point
+so that downstream processors receive fully-evaluated data.
+"""
+
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+
+
+class MaterializeTransformer(BaseTransformer[None, AxisArray, AxisArray]):
+ def _process(self, message: AxisArray) -> AxisArray:
+ try:
+ import mlx.core as mx
+
+ if isinstance(message.data, mx.array):
+ mx.eval(message.data)
+ except ImportError:
+ pass
+ return message
+
+
+class Materialize(BaseTransformerUnit[None, AxisArray, AxisArray, MaterializeTransformer]): ...
diff --git a/src/ezmsg/sigproc/math/pow.py b/src/ezmsg/sigproc/math/pow.py
index 195df7ae..0438155e 100644
--- a/src/ezmsg/sigproc/math/pow.py
+++ b/src/ezmsg/sigproc/math/pow.py
@@ -7,7 +7,6 @@
"""
import ezmsg.core as ez
-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
@@ -20,8 +19,7 @@ class PowSettings(ez.Settings):
class PowTransformer(BaseTransformer[PowSettings, AxisArray, AxisArray]):
def _process(self, message: AxisArray) -> AxisArray:
- xp = get_namespace(message.data)
- return replace(message, data=xp.pow(message.data, self.settings.exponent))
+ return replace(message, data=message.data**self.settings.exponent)
class Pow(BaseTransformerUnit[PowSettings, AxisArray, AxisArray, PowTransformer]):
diff --git a/src/ezmsg/sigproc/singlebandpow.py b/src/ezmsg/sigproc/singlebandpow.py
index e7155dee..8b6744ec 100644
--- a/src/ezmsg/sigproc/singlebandpow.py
+++ b/src/ezmsg/sigproc/singlebandpow.py
@@ -17,7 +17,7 @@
CompositeProcessor,
)
from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.messages.modify import modify_axis
+from ezmsg.util.messages.modify import ModifyAxisSettings, ModifyAxisTransformer
from .aggregate import AggregateSettings, AggregateTransformer, AggregationFunction
from .butterworthfilter import ButterworthFilterSettings, ButterworthFilterTransformer
@@ -63,12 +63,23 @@ def _initialize_processors(
zero_pad_until="none",
),
"aggregate": AggregateTransformer(AggregateSettings(axis="time", operation=AggregationFunction.MEAN)),
- "rename": modify_axis(name_map={"bin": "time"}),
+ "rename": ModifyAxisTransformer(settings=ModifyAxisSettings(name_map={"bin": "time"})),
}
if settings.apply_sqrt:
procs["sqrt"] = PowTransformer(PowSettings(exponent=0.5))
return procs
+ def _post_process(self, result: AxisArray | None) -> AxisArray | None:
+ if result is not None:
+ try:
+ import mlx.core as mx
+
+ if isinstance(result.data, mx.array):
+ mx.eval(result.data)
+ except ImportError:
+ pass
+ return result
+
class RMSBandPower(BaseTransformerUnit[RMSBandPowerSettings, AxisArray, AxisArray, RMSBandPowerTransformer]):
SETTINGS = RMSBandPowerSettings
diff --git a/src/ezmsg/sigproc/spectrogram.py b/src/ezmsg/sigproc/spectrogram.py
index 41f77137..ad0020ed 100644
--- a/src/ezmsg/sigproc/spectrogram.py
+++ b/src/ezmsg/sigproc/spectrogram.py
@@ -7,7 +7,7 @@
CompositeProcessor,
)
from ezmsg.util.messages.axisarray import AxisArray
-from ezmsg.util.messages.modify import modify_axis
+from ezmsg.util.messages.modify import ModifyAxisSettings, ModifyAxisTransformer
from .spectrum import (
SpectralOutput,
@@ -62,7 +62,7 @@ def _initialize_processors(
transform=settings.transform,
output=settings.output,
),
- "modify_axis": modify_axis(name_map={"win": "time"}),
+ "modify_axis": ModifyAxisTransformer(settings=ModifyAxisSettings(name_map={"win": "time"})),
}
diff --git a/src/ezmsg/sigproc/spectrum.py b/src/ezmsg/sigproc/spectrum.py
index 9fe91729..bdc5c597 100644
--- a/src/ezmsg/sigproc/spectrum.py
+++ b/src/ezmsg/sigproc/spectrum.py
@@ -1,10 +1,11 @@
import enum
+import math
import typing
from functools import partial
import ezmsg.core as ez
import numpy as np
-import numpy.typing as npt
+from array_api_compat import get_namespace
from ezmsg.baseproc import (
BaseStatefulTransformer,
BaseTransformerUnit,
@@ -16,6 +17,8 @@
slice_along_axis,
)
+from .util.array import is_complex_dtype
+
class OptionsEnum(enum.Enum):
@classmethod
@@ -121,9 +124,10 @@ class SpectrumState:
# I would prefer `slice(None)` as f_sl default but this fails because it is mutable.
freq_axis: AxisArray.LinearAxis | None = None
fftfun: typing.Callable | None = None
+ fftshift: typing.Callable | None = None
f_transform: typing.Callable | None = None
new_dims: list[str] | None = None
- window: npt.NDArray | None = None
+ window: typing.Any = None
class SpectrumTransformer(BaseStatefulTransformer[SpectrumSettings, AxisArray, AxisArray, SpectrumState]):
@@ -132,7 +136,7 @@ def _hash_message(self, message: AxisArray) -> int:
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, is_complex_dtype(message.data.dtype), ax_idx, ax_info.gain))
def _reset_state(self, message: AxisArray) -> None:
axis = self.settings.axis or message.dims[0]
@@ -140,34 +144,54 @@ def _reset_state(self, message: AxisArray) -> None:
ax_info = message.axes[axis]
targ_len = message.data.shape[ax_idx]
nfft = self.settings.nfft or targ_len
+ xp = get_namespace(message.data)
- # Pre-calculate windowing
- window = WINDOWS[self.settings.window](targ_len)
- window = window.reshape(
- [1] * ax_idx
- + [
- len(window),
- ]
- + [1] * (message.data.ndim - 1 - ax_idx)
- )
+ # Pre-calculate windowing (always compute with numpy, then convert to backend)
+ window_np = WINDOWS[self.settings.window](targ_len)
+ shape = [1] * ax_idx + [len(window_np)] + [1] * (message.data.ndim - 1 - ax_idx)
+ window = xp.asarray(window_np).reshape(shape)
if self.settings.transform != SpectralTransform.RAW_COMPLEX and not (
self.settings.transform == SpectralTransform.REAL or self.settings.transform == SpectralTransform.IMAG
):
- scale = np.sum(window**2.0) * ax_info.gain
+ scale = float(xp.sum(window**2.0)) * ax_info.gain
if self.settings.window != WindowFunction.NONE:
self.state.window = window
+ # Build FFT closure with manual norm fallback for backends that don't support norm=
+ norm = self.settings.norm
+ if norm == "forward":
+ norm_factor = 1.0 / nfft
+ elif norm == "ortho":
+ norm_factor = 1.0 / math.sqrt(nfft)
+ else:
+ norm_factor = None # backward / None — no scaling
+
+ def _make_fft_closure(raw_fft):
+ """Build a closure that calls *raw_fft* and applies norm manually if needed."""
+
+ def fftfun(x):
+ try:
+ return raw_fft(x, n=nfft, axis=ax_idx, norm=norm)
+ except TypeError:
+ result = raw_fft(x, n=nfft, axis=ax_idx)
+ if norm_factor is not None:
+ result = result * norm_factor
+ return result
+
+ return fftfun
+
# Pre-calculate frequencies and select our fft function.
- b_complex = message.data.dtype.kind == "c"
+ b_complex = is_complex_dtype(message.data.dtype)
self.state.f_sl = slice(None)
+ self.state.fftshift = None
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 = _make_fft_closure(xp.fft.rfft)
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 = _make_fft_closure(xp.fft.fft)
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))
@@ -177,6 +201,13 @@ def _reset_state(self, message: AxisArray) -> None:
elif self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL:
freqs = np.fft.fftshift(freqs, axes=-1)
freqs = freqs[self.state.f_sl]
+
+ # Store fftshift closure if shifting is needed (use tuple for axes — MLX requirement)
+ if (
+ self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL
+ ) or self.settings.output == SpectralOutput.NEGATIVE:
+ self.state.fftshift = partial(xp.fft.fftshift, axes=(ax_idx,))
+
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.new_dims = (
@@ -202,20 +233,18 @@ def f_transform(x):
else:
def f1(x):
- return (np.abs(x) ** 2.0) / scale
+ return (xp.abs(x) ** 2.0) / scale
if self.settings.transform == SpectralTransform.REL_DB:
def f_transform(x):
- return 10 * np.log10(f1(x))
+ return 10 * xp.log10(f1(x))
else:
f_transform = f1
self.state.f_transform = f_transform
def _process(self, message: AxisArray) -> AxisArray:
axis = self.settings.axis or message.dims[0]
- 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[self.settings.out_axis or axis] = self.state.freq_axis
@@ -224,19 +253,11 @@ def _process(self, message: AxisArray) -> AxisArray:
win_dat = message.data * self.state.window
else:
win_dat = message.data
- spec = self.state.fftfun(
- win_dat,
- n=self.settings.nfft or targ_len,
- axis=ax_idx,
- norm=self.settings.norm,
- )
- # Note: norm="forward" equivalent to `/ nfft`
- if (
- self.settings.do_fftshift and self.settings.output == SpectralOutput.FULL
- ) or self.settings.output == SpectralOutput.NEGATIVE:
- spec = np.fft.fftshift(spec, axes=ax_idx)
+ spec = self.state.fftfun(win_dat)
+ if self.state.fftshift is not None:
+ spec = self.state.fftshift(spec)
spec = self.state.f_transform(spec)
- spec = slice_along_axis(spec, self.state.f_sl, ax_idx)
+ spec = slice_along_axis(spec, self.state.f_sl, message.get_axis_idx(axis))
msg_out = replace(message, data=spec, dims=self.state.new_dims, axes=new_axes)
return msg_out
diff --git a/src/ezmsg/sigproc/util/array.py b/src/ezmsg/sigproc/util/array.py
index 4c470753..ef3fe51c 100644
--- a/src/ezmsg/sigproc/util/array.py
+++ b/src/ezmsg/sigproc/util/array.py
@@ -46,6 +46,13 @@ def xp_create(fn, *args, dtype=None, device=None, **extra):
return fn(*args, **kwargs)
+def is_complex_dtype(dtype) -> bool:
+ """Check whether *dtype* is a complex type, portably across backends."""
+ if hasattr(dtype, "kind"):
+ return dtype.kind == "c"
+ return "complex" in str(dtype).lower()
+
+
def is_float_dtype(xp, dtype) -> bool:
"""Check whether *dtype* is a real floating-point type, portably."""
try:
diff --git a/tests/helpers/util.py b/tests/helpers/util.py
index 8f01e1c3..ef509c03 100644
--- a/tests/helpers/util.py
+++ b/tests/helpers/util.py
@@ -1,8 +1,11 @@
+import importlib
import os
+import platform
import tempfile
from pathlib import Path
import numpy as np
+import pytest
from ezmsg.util.messages.axisarray import AxisArray
from frozendict import frozendict
from numpy.lib.stride_tricks import sliding_window_view
@@ -176,3 +179,10 @@ def make_chirp(t, t0, a):
frequency = (a * (t + t0)) ** 2
chirp = np.sin(2 * np.pi * frequency * t)
return chirp, frequency
+
+
+_has_mlx = importlib.util.find_spec("mlx") is not None
+requires_mlx = pytest.mark.skipif(
+ not _has_mlx or platform.machine() != "arm64" or platform.system() != "Darwin",
+ reason="Requires MLX on Apple Silicon",
+)
diff --git a/tests/unit/test_aggregate.py b/tests/unit/test_aggregate.py
index 7e036c2a..e1332bf2 100644
--- a/tests/unit/test_aggregate.py
+++ b/tests/unit/test_aggregate.py
@@ -14,7 +14,7 @@
RangedAggregateTransformer,
)
from tests.helpers.empty_time import check_empty_result, make_empty_msg
-from tests.helpers.util import assert_messages_equal
+from tests.helpers.util import assert_messages_equal, requires_mlx
def get_msg_gen(n_chans=20, n_freqs=100, data_dur=30.0, fs=1024.0, key=""):
@@ -406,3 +406,129 @@ def test_ranged_aggregate_empty_passthrough():
empty = make_empty_msg()
result = proc(empty)
check_empty_result(result)
+
+
+# ============== MLX Tests ==============
+@requires_mlx
+@pytest.mark.parametrize(
+ "operation",
+ [
+ AggregationFunction.MEAN,
+ AggregationFunction.SUM,
+ AggregationFunction.MAX,
+ AggregationFunction.MIN,
+ AggregationFunction.STD,
+ AggregationFunction.MEDIAN,
+ ],
+)
+def test_ranged_aggregate_mlx(operation: AggregationFunction):
+ """RangedAggregateTransformer on MLX arrays should match numpy results."""
+ import mlx.core as mx
+
+ bands = [(5.0, 20.0), (30.0, 50.0)]
+ n_times, n_chans, n_freqs = 100, 8, 100
+
+ np_data = np.arange(n_times * n_chans * n_freqs, dtype=np.float32).reshape(n_times, n_chans, n_freqs)
+ axes = frozendict(
+ {
+ "time": AxisArray.TimeAxis(fs=1024.0, offset=0.0),
+ "freq": AxisArray.LinearAxis(gain=1.0, offset=0.0, unit="Hz"),
+ }
+ )
+ dims = ["time", "ch", "freq"]
+
+ msg_np = AxisArray(data=np_data, dims=dims, axes=axes)
+ msg_mx = AxisArray(data=mx.array(np_data), dims=dims, axes=axes)
+
+ xformer_np = RangedAggregateTransformer(RangedAggregateSettings(axis="freq", bands=bands, operation=operation))
+ xformer_mx = RangedAggregateTransformer(RangedAggregateSettings(axis="freq", bands=bands, operation=operation))
+
+ result_np = xformer_np(msg_np)
+ result_mx = xformer_mx(msg_mx)
+
+ assert isinstance(result_mx.data, mx.array), f"Expected mx.array, got {type(result_mx.data)}"
+ assert result_mx.data.shape == result_np.data.shape
+ assert result_mx.dims == result_np.dims
+ np.testing.assert_allclose(np.asarray(result_mx.data), result_np.data.astype(np.float32), rtol=1e-5, atol=1e-5)
+
+
+@requires_mlx
+def test_ranged_aggregate_mlx_trapezoid():
+ """Trapezoid falls back to numpy but should still work with MLX input."""
+ import mlx.core as mx
+
+ bands = [(5.0, 20.0), (30.0, 50.0)]
+ n_times, n_chans, n_freqs = 50, 4, 100
+ np_data = np.arange(n_times * n_chans * n_freqs, dtype=np.float32).reshape(n_times, n_chans, n_freqs)
+ axes = frozendict(
+ {
+ "time": AxisArray.TimeAxis(fs=1024.0, offset=0.0),
+ "freq": AxisArray.LinearAxis(gain=1.0, offset=0.0, unit="Hz"),
+ }
+ )
+
+ msg_np = AxisArray(data=np_data, dims=["time", "ch", "freq"], axes=axes)
+ msg_mx = AxisArray(data=mx.array(np_data), dims=["time", "ch", "freq"], axes=axes)
+
+ xformer_np = RangedAggregateTransformer(
+ RangedAggregateSettings(axis="freq", bands=bands, operation=AggregationFunction.TRAPEZOID)
+ )
+ xformer_mx = RangedAggregateTransformer(
+ RangedAggregateSettings(axis="freq", bands=bands, operation=AggregationFunction.TRAPEZOID)
+ )
+
+ result_np = xformer_np(msg_np)
+ result_mx = xformer_mx(msg_mx)
+
+ # Trapezoid falls back to numpy then converts back to MLX
+ assert isinstance(result_mx.data, mx.array)
+ np.testing.assert_allclose(np.asarray(result_mx.data), result_np.data.astype(np.float32), rtol=1e-5, atol=1e-5)
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="ranged-aggregate")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["numpy", "mlx"])
+def test_ranged_aggregate_benchmark(backend, n_channels, benchmark):
+ """Benchmark RangedAggregateTransformer: numpy vs MLX."""
+ bands = [(5.0, 50.0), (50.0, 100.0), (100.0, 200.0), (200.0, 400.0)]
+ n_times = 125
+ n_freqs = 501 # Typical positive-freq output for 1000-sample FFT
+
+ rng = np.random.default_rng(42)
+ np_data = rng.standard_normal((n_times, n_channels, n_freqs)).astype(np.float32)
+ axes = frozendict(
+ {
+ "time": AxisArray.LinearAxis(gain=0.5, offset=0.0),
+ "freq": AxisArray.LinearAxis(gain=1.0, offset=0.0, unit="Hz"),
+ }
+ )
+ dims = ["time", "ch", "freq"]
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ msg = AxisArray(data=mx.array(np_data), dims=dims, axes=axes)
+ else:
+ msg = AxisArray(data=np_data, dims=dims, axes=axes)
+
+ settings = RangedAggregateSettings(axis="freq", bands=bands, operation=AggregationFunction.MEAN)
+ xformer = RangedAggregateTransformer(settings)
+
+ # Warmup
+ warmup = xformer(msg)
+ if backend == "mlx":
+ mx.eval(warmup.data)
+
+ def run():
+ result = xformer(msg)
+ if backend == "mlx":
+ mx.eval(result.data)
+ return result
+
+ result = benchmark(run)
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ assert isinstance(result.data, mx.array)
diff --git a/tests/unit/test_bandpower.py b/tests/unit/test_bandpower.py
index 227d71ad..0e23d1b8 100644
--- a/tests/unit/test_bandpower.py
+++ b/tests/unit/test_bandpower.py
@@ -1,6 +1,7 @@
import copy
import numpy as np
+import pytest
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.sigproc.bandpower import (
@@ -9,10 +10,7 @@
BandPowerTransformer,
SpectrogramSettings,
)
-from tests.helpers.util import (
- assert_messages_equal,
- create_messages_with_periodic_signal,
-)
+from tests.helpers.util import assert_messages_equal, create_messages_with_periodic_signal, requires_mlx
def _debug_plot(result):
@@ -72,3 +70,63 @@ def test_bandpower():
mags.append(result.data[ix, 2, 0])
# The sorting of the measured magnitudes should match the sorting of the parameter magnitudes.
assert np.array_equal(np.argsort(mags), np.argsort([_["a"] for _ in sin_params]))
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="bandpower")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["numpy", "mlx"])
+def test_bandpower_benchmark(backend, n_channels, benchmark):
+ """Benchmark BandPowerTransformer end-to-end: numpy vs MLX input."""
+ fs = 1000.0
+ chunk_samples = 500
+ n_chunks = 20
+ win_dur = 0.5
+ win_shift = 0.1
+ bands = [(8, 13), (13, 30), (30, 70), (70, 150)]
+
+ settings = BandPowerSettings(
+ spectrogram_settings=SpectrogramSettings(
+ window_dur=win_dur,
+ window_shift=win_shift,
+ ),
+ bands=bands,
+ aggregation=AggregationFunction.MEAN,
+ )
+
+ # Pre-generate chunk messages as numpy
+ rng = np.random.default_rng(42)
+ np_chunks = []
+ for i in range(n_chunks + 1): # +1 for warmup
+ data = rng.standard_normal((chunk_samples, n_channels)).astype(np.float32)
+ np_chunks.append(
+ AxisArray(
+ data,
+ dims=["time", "ch"],
+ axes={"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * chunk_samples / fs)},
+ )
+ )
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ chunks = [AxisArray(data=mx.array(chunk.data), dims=chunk.dims, axes=chunk.axes) for chunk in np_chunks]
+ else:
+ chunks = np_chunks
+
+ xformer = BandPowerTransformer(settings)
+ xformer(chunks[0]) # Warmup
+
+ def process_all_chunks():
+ outputs = [xformer(chunk) for chunk in chunks[1:]]
+ if backend == "mlx":
+ mx.eval(*[o.data for o in outputs])
+ return outputs
+
+ results = benchmark(process_all_chunks)
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ last_mx = next(o for o in reversed(results) if np.asarray(o.data).size > 0)
+ assert isinstance(last_mx.data, mx.array), f"Expected mx.array, got {type(last_mx.data)}"
diff --git a/tests/unit/test_butter.py b/tests/unit/test_butter.py
index 168537f5..235f8da9 100644
--- a/tests/unit/test_butter.py
+++ b/tests/unit/test_butter.py
@@ -6,6 +6,7 @@
from ezmsg.sigproc.butterworthfilter import ButterworthFilterSettings, ButterworthFilterTransformer
from tests.helpers.empty_time import check_empty_result, check_state_not_corrupted, make_empty_msg, make_msg
+from tests.helpers.util import requires_mlx
@pytest.mark.parametrize(
@@ -298,3 +299,118 @@ def test_butterworth_empty_first():
result = proc(empty)
check_empty_result(result)
check_state_not_corrupted(proc, normal)
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="butterworth")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["mlx", "numpy"])
+def test_butterworth_benchmark(backend, n_channels, benchmark):
+ """Benchmark Butterworth filter: numpy vs MLX input."""
+ fs = 1000.0
+ chunk_samples = 256
+ n_chunks = 20
+ order = 4
+
+ settings = ButterworthFilterSettings(
+ axis="time",
+ order=order,
+ cuton=30.0,
+ cutoff=100.0,
+ coef_type="sos",
+ )
+
+ # Build chunks
+ rng = np.random.default_rng(42)
+ chunks = []
+ for i in range(n_chunks):
+ d = rng.standard_normal((chunk_samples, n_channels)).astype(np.float32)
+ _time_axis = AxisArray.TimeAxis(fs=fs, offset=i * chunk_samples / fs)
+ axes = frozendict(
+ {
+ "time": _time_axis,
+ "ch": AxisArray.CoordinateAxis(data=np.arange(n_channels).astype(str), dims=["ch"]),
+ }
+ )
+ if backend == "mlx":
+ import mlx.core as mx
+
+ chunks.append(AxisArray(mx.array(d), dims=["time", "ch"], axes=axes, key="bench"))
+ else:
+ chunks.append(AxisArray(d, dims=["time", "ch"], axes=axes, key="bench"))
+
+ # Warmup
+ xformer = ButterworthFilterTransformer(settings)
+ warmup = xformer(chunks[0])
+ if backend == "mlx":
+ import mlx.core as mx
+
+ mx.eval(warmup.data)
+
+ def process_all_chunks():
+ outputs = [xformer(chunk) for chunk in chunks[1:]]
+ if backend == "mlx":
+ mx.eval(*[o.data for o in outputs])
+ return outputs
+
+ benchmark(process_all_chunks)
+
+
+@requires_mlx
+@pytest.mark.sosfilt
+@pytest.mark.benchmark(group="sosfilt")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["mlx", "numpy"])
+def test_sosfilt_benchmark(backend, n_channels, benchmark):
+ """Benchmark _sosfilt_xp (MLX) vs scipy.signal.sosfilt (numpy) directly."""
+ from ezmsg.sigproc.filter import _sosfilt_xp
+
+ fs = 1000.0
+ chunk_samples = 256
+ n_chunks = 20
+ order = 4
+
+ # Design filter
+ sos_np = scipy.signal.butter(order, [30.0, 100.0], btype="bandpass", output="sos", fs=fs)
+ zi_template = scipy.signal.sosfilt_zi(sos_np) # (n_sections, 2)
+
+ # Build chunks and per-chunk zi (tiled to channel count)
+ rng = np.random.default_rng(42)
+ chunks = []
+ for i in range(n_chunks):
+ chunks.append(rng.standard_normal((chunk_samples, n_channels)).astype(np.float32))
+
+ # Tile zi: (n_sections, 2) → (n_sections, 2, n_channels)
+ zi_np = np.tile(zi_template[:, :, None], (1, 1, n_channels)) # axis_idx=0, so 2 is at position 1
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ xp = mx
+ chunks_xp = [mx.array(c) for c in chunks]
+ sos_xp = mx.array(sos_np.astype(np.float32))
+ zi_state = mx.array(zi_np.astype(np.float32))
+
+ # Warmup
+ _, zi_state = _sosfilt_xp(sos_xp, chunks_xp[0], axis_idx=0, zi=zi_state, xp=xp)
+ mx.eval(zi_state)
+
+ def run():
+ zi = zi_state
+ for chunk in chunks_xp[1:]:
+ out, zi = _sosfilt_xp(sos_xp, chunk, axis_idx=0, zi=zi, xp=xp)
+ mx.eval(out)
+ return out
+ else:
+ zi_state = zi_np.copy()
+
+ # Warmup
+ _, zi_state = scipy.signal.sosfilt(sos_np, chunks[0], axis=0, zi=zi_state)
+
+ def run():
+ zi = zi_state
+ for chunk in chunks[1:]:
+ out, zi = scipy.signal.sosfilt(sos_np, chunk, axis=0, zi=zi)
+ return out
+
+ benchmark(run)
diff --git a/tests/unit/test_firfilter.py b/tests/unit/test_firfilter.py
index f57bd9c7..d235b3d0 100644
--- a/tests/unit/test_firfilter.py
+++ b/tests/unit/test_firfilter.py
@@ -1,6 +1,3 @@
-import platform
-import time
-
import numpy as np
import pytest
import scipy.signal
@@ -13,11 +10,7 @@
firwin_design_fun,
)
from tests.helpers.empty_time import check_empty_result, check_state_not_corrupted, make_empty_msg, make_msg
-
-requires_apple_silicon = pytest.mark.skipif(
- platform.machine() != "arm64" or platform.system() != "Darwin",
- reason="Requires Apple Silicon for MLX",
-)
+from tests.helpers.util import requires_mlx
@pytest.mark.parametrize(
@@ -317,7 +310,7 @@ def test_fir_empty_first():
check_state_not_corrupted(proc, normal)
-@requires_apple_silicon
+@requires_mlx
@pytest.mark.parametrize(
"cutoff, pass_zero",
[
@@ -421,18 +414,18 @@ def test_firfilter_mlx(cutoff, pass_zero, order, n_dims, time_ax):
np.testing.assert_allclose(mx_result_np, np_result, rtol=5e-3, atol=5e-3)
-@requires_apple_silicon
-def test_firfilter_mlx_benchmark():
+@requires_mlx
+@pytest.mark.benchmark(group="firfilter")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["numpy", "mlx"])
+def test_firfilter_benchmark(backend, n_channels, benchmark):
"""Benchmark FIR filter: numpy (scipy lfilter) vs MLX (FFT convolution)."""
- import mlx.core as mx
-
fs = 1000.0
chunk_samples = 256
- n_channels = 32
- n_chunks = 100
+ n_chunks = 20
order = 51
- np_settings = FIRFilterSettings(
+ settings = FIRFilterSettings(
axis="time",
order=order,
cutoff=100.0,
@@ -445,8 +438,7 @@ def test_firfilter_mlx_benchmark():
# Build chunks
rng = np.random.default_rng(42)
- np_chunks = []
- mx_chunks = []
+ chunks = []
for i in range(n_chunks):
d = rng.standard_normal((chunk_samples, n_channels)).astype(np.float32)
_time_axis = AxisArray.TimeAxis(fs=fs, offset=i * chunk_samples / fs)
@@ -456,42 +448,30 @@ def test_firfilter_mlx_benchmark():
"ch": AxisArray.CoordinateAxis(data=np.arange(n_channels).astype(str), dims=["ch"]),
}
)
- np_chunks.append(AxisArray(d, dims=["time", "ch"], axes=axes, key="bench"))
- mx_chunks.append(AxisArray(mx.array(d), dims=["time", "ch"], axes=axes, key="bench"))
+ if backend == "mlx":
+ import mlx.core as mx
- # Warm up both transformers
- np_xformer = FIRFilterTransformer(settings=np_settings)
- mx_xformer = FIRFilterTransformer(settings=np_settings)
- _ = np_xformer(np_chunks[0])
- _ = mx_xformer(mx_chunks[0])
- mx.eval(mx_xformer(mx_chunks[0]).data) # force MLX compile
+ chunks.append(AxisArray(mx.array(d), dims=["time", "ch"], axes=axes, key="bench"))
+ else:
+ chunks.append(AxisArray(d, dims=["time", "ch"], axes=axes, key="bench"))
- # Numpy timing
- np_xformer = FIRFilterTransformer(settings=np_settings)
- _ = np_xformer(np_chunks[0])
- t0 = time.perf_counter()
- np_outputs = [np_xformer(chunk) for chunk in np_chunks[1:]]
- t_numpy = time.perf_counter() - t0
-
- # MLX timing
- mx_xformer = FIRFilterTransformer(settings=np_settings)
- _ = mx_xformer(mx_chunks[0])
- t0 = time.perf_counter()
- mx_outputs = [mx_xformer(chunk) for chunk in mx_chunks[1:]]
- mx.eval(mx_outputs[-1].data) # force evaluation
- t_mlx = time.perf_counter() - t0
-
- # Correctness check
- for np_out, mx_out in zip(np_outputs, mx_outputs):
- if np_out.data.size > 0:
- np.testing.assert_allclose(np.array(mx_out.data), np_out.data, rtol=1e-4, atol=1e-4)
-
- # Verify MLX output type
- assert isinstance(mx_outputs[0].data, mx.array), "MLX output should remain mx.array"
-
- print(
- f"\n FIR filter benchmark ({n_chunks} chunks, {chunk_samples}x{n_channels}, order={order}):"
- f"\n numpy (scipy lfilter): {t_numpy:.4f}s ({t_numpy / n_chunks * 1000:.2f} ms/chunk)"
- f"\n mlx (FFT convolve): {t_mlx:.4f}s ({t_mlx / n_chunks * 1000:.2f} ms/chunk)"
- f"\n ratio (mlx/numpy): {t_mlx / t_numpy:.2f}x"
- )
+ # Warmup
+ xformer = FIRFilterTransformer(settings=settings)
+ warmup = xformer(chunks[0])
+ if backend == "mlx":
+ import mlx.core as mx
+
+ mx.eval(warmup.data)
+
+ def process_all_chunks():
+ outputs = [xformer(chunk) for chunk in chunks[1:]]
+ if backend == "mlx":
+ mx.eval(*[o.data for o in outputs])
+ return outputs
+
+ outputs = benchmark(process_all_chunks)
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ assert isinstance(outputs[0].data, mx.array), "MLX output should remain mx.array"
diff --git a/tests/unit/test_materialize.py b/tests/unit/test_materialize.py
new file mode 100644
index 00000000..8200b2e6
--- /dev/null
+++ b/tests/unit/test_materialize.py
@@ -0,0 +1,34 @@
+import numpy as np
+import pytest
+from ezmsg.util.messages.axisarray import AxisArray
+
+from ezmsg.sigproc.materialize import MaterializeTransformer
+from tests.helpers.empty_time import check_empty_result, make_empty_msg, make_msg
+from tests.helpers.util import requires_mlx
+
+
+def test_numpy_passthrough():
+ msg = make_msg()
+ xformer = MaterializeTransformer()
+ result = xformer(msg)
+ assert result is msg
+
+
+@requires_mlx
+def test_mlx_evaluates():
+ mx = pytest.importorskip("mlx.core")
+ a = mx.ones((10, 3))
+ b = mx.ones((10, 3))
+ lazy_sum = a + b # lazy — not yet evaluated
+ msg = AxisArray(lazy_sum, dims=["time", "ch"])
+ xformer = MaterializeTransformer()
+ result = xformer(msg)
+ assert isinstance(result.data, mx.array)
+ np.testing.assert_array_equal(np.array(result.data), np.full((10, 3), 2.0))
+
+
+def test_empty_time():
+ msg = make_empty_msg()
+ xformer = MaterializeTransformer()
+ result = xformer(msg)
+ check_empty_result(result)
diff --git a/tests/unit/test_singlebandpow.py b/tests/unit/test_singlebandpow.py
index 2c6b16eb..5f3b4c2b 100644
--- a/tests/unit/test_singlebandpow.py
+++ b/tests/unit/test_singlebandpow.py
@@ -1,4 +1,5 @@
import numpy as np
+import pytest
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.sigproc.butterworthfilter import ButterworthFilterSettings
@@ -9,6 +10,7 @@
SquareLawBandPowerSettings,
SquareLawBandPowerTransformer,
)
+from tests.helpers.util import requires_mlx
def _make_sinusoid(
@@ -178,3 +180,51 @@ def test_squarelaw_bandpower():
assert (
abs(mean_power - expected_ms) < 0.25 * expected_ms
), f"Expected power ~{expected_ms:.3f}, got {mean_power:.3f}"
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="rms-bandpower")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["numpy", "mlx"])
+def test_rms_bandpower_benchmark(backend, n_channels, benchmark):
+ """Benchmark RMSBandPowerTransformer with numpy vs MLX backends."""
+ freq, amplitude, fs = 50.0, 2.0, 30_000.0
+ bin_duration = 0.05
+ chunk_samples = int(0.005 * fs) # 50 samples — one bin per chunk
+ n_chunks = 20
+
+ bandpass = ButterworthFilterSettings(order=4, coef_type="sos", cuton=30.0, cutoff=70.0)
+
+ settings = RMSBandPowerSettings(
+ bandpass=bandpass,
+ bin_duration=bin_duration,
+ apply_sqrt=True,
+ )
+
+ # Pre-generate all chunk messages
+ chunks = []
+ for i in range(n_chunks + 1): # +1 for warmup
+ t = (np.arange(chunk_samples) + i * chunk_samples) / fs
+ data = amplitude * np.sin(2 * np.pi * freq * t[:, None] * np.ones((1, n_channels)))
+ axes = {"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * chunk_samples / fs)}
+ if backend == "mlx":
+ import mlx.core as mx
+
+ data = mx.array(data.astype(np.float32))
+ chunks.append(AxisArray(data, dims=["time", "ch"], axes=axes))
+
+ xformer = RMSBandPowerTransformer(settings)
+ warmup = xformer(chunks[0])
+ if backend == "mlx":
+ import mlx.core as mx
+
+ if warmup is not None and hasattr(warmup.data, "__array__"):
+ mx.eval(warmup.data)
+
+ def process_all_chunks():
+ outputs = [xformer(chunk) for chunk in chunks[1:]]
+ if backend == "mlx":
+ mx.eval(*[o.data for o in outputs if o.data.size > 0])
+ return outputs
+
+ benchmark(process_all_chunks)
diff --git a/tests/unit/test_spectrum.py b/tests/unit/test_spectrum.py
index 7e7ffa7d..7493aa09 100644
--- a/tests/unit/test_spectrum.py
+++ b/tests/unit/test_spectrum.py
@@ -17,6 +17,7 @@
from tests.helpers.util import (
assert_messages_equal,
create_messages_with_periodic_signal,
+ requires_mlx,
)
@@ -239,3 +240,115 @@ def test_spectrum_empty_first():
result2 = proc(normal)
assert result2.data.shape[0] == 5
assert np.all(np.isfinite(result2.data))
+
+
+@requires_mlx
+@pytest.mark.parametrize("transform", [SpectralTransform.REL_DB, SpectralTransform.REL_POWER])
+@pytest.mark.parametrize("output", [SpectralOutput.POSITIVE, SpectralOutput.NEGATIVE, SpectralOutput.FULL])
+def test_spectrum_mlx(transform: SpectralTransform, output: SpectralOutput):
+ """SpectrumTransformer on MLX arrays should produce correct spectral peaks."""
+ import mlx.core as mx
+
+ win_dur = 1.0
+ fs = 1000.0
+ sin_params = [
+ {"a": 1.0, "f": 10.0, "p": 0.0, "dur": 5.0},
+ {"a": 0.5, "f": 20.0, "p": np.pi / 7, "dur": 5.0},
+ {"a": 0.2, "f": 200.0, "p": np.pi / 11, "dur": 5.0},
+ ]
+ win_len = int(win_dur * fs)
+ messages_np = create_messages_with_periodic_signal(sin_params=sin_params, fs=fs, msg_dur=win_dur, win_step_dur=None)
+ msg_np_orig = messages_np[0]
+
+ # Build an MLX message
+ msg_mx = AxisArray(
+ data=mx.array(msg_np_orig.data.astype(np.float32)),
+ dims=msg_np_orig.dims,
+ axes=msg_np_orig.axes,
+ )
+
+ settings = SpectrumSettings(axis="time", window=WindowFunction.HAMMING, transform=transform, output=output)
+
+ proc_mx = SpectrumTransformer(settings)
+ result = proc_mx(msg_mx)
+
+ assert isinstance(result.data, mx.array), f"Expected mx.array, got {type(result.data)}"
+ assert "freq" in result.dims
+ assert "freq" in result.axes
+ assert result.axes["freq"].gain == 1 / win_dur
+
+ # Check correct spectral shape
+ fax_ix = result.get_axis_idx("freq")
+ f_len = win_len if output == SpectralOutput.FULL else (win_len // 2 + 1 - (win_len % 2))
+ assert result.data.shape[fax_ix] == f_len
+
+ # Verify peaks are at the expected frequencies
+ result_np = np.asarray(result.data)
+ f_vec = result.axes["freq"].value(np.arange(f_len))
+ if output == SpectralOutput.NEGATIVE:
+ f_vec = np.abs(f_vec)
+ for s_p in sin_params:
+ f_ix = np.argmin(np.abs(f_vec - s_p["f"]))
+ peak_inds = np.argmax(
+ slice_along_axis(result_np, slice(f_ix - 3, f_ix + 3), axis=fax_ix),
+ axis=fax_ix,
+ )
+ assert np.all(peak_inds == 3), f"Peak not at expected freq {s_p['f']} Hz"
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="spectrum")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["numpy", "mlx"])
+def test_spectrum_benchmark(backend, n_channels, benchmark):
+ """Benchmark SpectrumTransformer: numpy vs MLX on multi-window input."""
+ fs = 1000.0
+ win_dur = 1.0
+ n_windows = 50
+ win_samples = int(win_dur * fs)
+
+ rng = np.random.default_rng(42)
+ np_data = rng.standard_normal((n_windows, win_samples, n_channels)).astype(np.float32)
+
+ settings = SpectrumSettings(axis="time", window=WindowFunction.HAMMING, transform=SpectralTransform.REL_DB)
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ msg = AxisArray(
+ data=mx.array(np_data),
+ dims=["win", "time", "ch"],
+ axes={
+ "win": AxisArray.LinearAxis(gain=win_dur, offset=0.0),
+ "time": AxisArray.TimeAxis(fs=fs),
+ },
+ )
+ else:
+ msg = AxisArray(
+ data=np_data,
+ dims=["win", "time", "ch"],
+ axes={
+ "win": AxisArray.LinearAxis(gain=win_dur, offset=0.0),
+ "time": AxisArray.TimeAxis(fs=fs),
+ },
+ )
+
+ proc = SpectrumTransformer(settings)
+
+ # Warmup
+ warmup = proc(msg)
+ if backend == "mlx":
+ mx.eval(warmup.data)
+
+ def run():
+ result = proc(msg)
+ if backend == "mlx":
+ mx.eval(result.data)
+ return result
+
+ result = benchmark(run)
+
+ if backend == "mlx":
+ import mlx.core as mx
+
+ assert isinstance(result.data, mx.array)
diff --git a/tests/unit/test_window.py b/tests/unit/test_window.py
index d4db1fbe..4e52ee18 100644
--- a/tests/unit/test_window.py
+++ b/tests/unit/test_window.py
@@ -9,7 +9,7 @@
from ezmsg.sigproc.window import WindowTransformer
from tests.helpers.empty_time import check_empty_result, check_state_not_corrupted, make_empty_msg, make_msg
-from tests.helpers.util import assert_messages_equal, calculate_expected_windows
+from tests.helpers.util import assert_messages_equal, calculate_expected_windows, requires_mlx
def test_window_gen_nodur():
@@ -272,3 +272,57 @@ def test_window_empty_with_shift():
result = proc(empty)
assert result.data.size >= 0 # Just check no crash
check_state_not_corrupted(proc, normal, time_dim="time")
+
+
+@requires_mlx
+@pytest.mark.benchmark(group="window")
+@pytest.mark.parametrize("n_channels", [32, 256, 1024])
+@pytest.mark.parametrize("backend", ["mlx", "numpy"])
+def test_window_benchmark(backend, n_channels, benchmark):
+ """Benchmark WindowTransformer: numpy vs MLX input."""
+ fs = 1000.0
+ chunk_samples = 256
+ n_chunks = 20
+ window_dur = 0.5
+ window_shift = 0.1
+
+ rng = np.random.default_rng(42)
+ chunks = []
+ for i in range(n_chunks):
+ d = rng.standard_normal((chunk_samples, n_channels)).astype(np.float32)
+ _time_axis = AxisArray.TimeAxis(fs=fs, offset=i * chunk_samples / fs)
+ axes = frozendict(
+ {
+ "time": _time_axis,
+ "ch": AxisArray.CoordinateAxis(data=np.arange(n_channels).astype(str), dims=["ch"]),
+ }
+ )
+ if backend == "mlx":
+ import mlx.core as mx
+
+ chunks.append(AxisArray(mx.array(d), dims=["time", "ch"], axes=axes, key="bench"))
+ else:
+ chunks.append(AxisArray(d, dims=["time", "ch"], axes=axes, key="bench"))
+
+ xformer = WindowTransformer(
+ axis="time",
+ newaxis="win",
+ window_dur=window_dur,
+ window_shift=window_shift,
+ zero_pad_until="none",
+ )
+
+ # Warmup
+ warmup = xformer(chunks[0])
+ if backend == "mlx":
+ import mlx.core as mx
+
+ mx.eval(warmup.data)
+
+ def process_all_chunks():
+ outputs = [xformer(chunk) for chunk in chunks[1:]]
+ if backend == "mlx":
+ mx.eval(*[o.data for o in outputs])
+ return outputs
+
+ benchmark(process_all_chunks)