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)