diff --git a/src/ezmsg/simbiophys/_cosine_encoder_mlx.py b/src/ezmsg/simbiophys/_cosine_encoder_mlx.py new file mode 100644 index 0000000..6d869c0 --- /dev/null +++ b/src/ezmsg/simbiophys/_cosine_encoder_mlx.py @@ -0,0 +1,17 @@ +"""MLX implementation detail for :mod:`ezmsg.simbiophys.cosine_encoder`. + +Kept separate so importing ezmsg-simbiophys never requires MLX. The public +transformer imports this module lazily only after receiving an MLX array. +""" + +import mlx.core as mx + + +def _cosine_encode(polar, baseline, modulation, preferred_direction, speed_modulation): + magnitude = polar[:, 0:1] + angle = polar[:, 1:2] + return baseline + modulation * magnitude * mx.cos(angle - preferred_direction) + speed_modulation * magnitude + + +cosine_encode = mx.compile(_cosine_encode) +"""Shape-specialized compiled cosine encoder shared by transformer instances.""" diff --git a/src/ezmsg/simbiophys/cosine_encoder.py b/src/ezmsg/simbiophys/cosine_encoder.py index 8d6411f..eab1928 100644 --- a/src/ezmsg/simbiophys/cosine_encoder.py +++ b/src/ezmsg/simbiophys/cosine_encoder.py @@ -240,17 +240,30 @@ def _process(self, message: AxisArray) -> AxisArray: if polar.ndim != 2 or polar.shape[1] != 2: raise ValueError(f"Expected polar coords with shape (n_samples, 2), got {polar.shape}") - # Extract polar components (from CART2POL: magnitude, angle) - magnitude = polar[:, 0:1] # (n_samples, 1) - angle = polar[:, 1:2] # (n_samples, 1) - # Compute output: baseline + modulation * magnitude * cos(angle - pd) + speed_mod * magnitude # State arrays are pre-shaped to (1, output_ch) for broadcasting - output = ( - self.state.baseline - + self.state.modulation * magnitude * xp.cos(angle - self.state.pd) - + self.state.speed_modulation * magnitude - ) + if xp.__name__ == "mlx.core": + # Import lazily: MLX is optional and unavailable on non-Apple hosts. + # Compilation fuses this elementwise expression and caches the + # shape-specialized graph for the stable chunks used by simulators. + from ._cosine_encoder_mlx import cosine_encode + + output = cosine_encode( + polar, + self.state.baseline, + self.state.modulation, + self.state.pd, + self.state.speed_modulation, + ) + else: + # Extract polar components (from CART2POL: magnitude, angle). + magnitude = polar[:, 0:1] # (n_samples, 1) + angle = polar[:, 1:2] # (n_samples, 1) + output = ( + self.state.baseline + + self.state.modulation * magnitude * xp.cos(angle - self.state.pd) + + self.state.speed_modulation * magnitude + ) return replace( message, diff --git a/src/ezmsg/simbiophys/line_noise.py b/src/ezmsg/simbiophys/line_noise.py index 1604884..c71332c 100644 --- a/src/ezmsg/simbiophys/line_noise.py +++ b/src/ezmsg/simbiophys/line_noise.py @@ -10,6 +10,8 @@ continuous across chunks and as the frequency changes. """ +import math + import ezmsg.core as ez import numpy as np import numpy.typing as npt @@ -24,6 +26,14 @@ from .oscillator import advance_drifting_sine, freq_drift_step_std +MLX_BROADCAST_MIN_ELEMENTS = 200_000 +"""Minimum signal size for keeping line-noise addition on MLX. + +Below this crossover, MLX dispatch costs more than the existing fused NumPy +round-trip. Above it, transferring only the single generated sine column and +broadcasting on-device avoids converting the full sample-by-channel signal. +""" + class LineNoiseSettings(ez.Settings): freq: float | None = None @@ -99,11 +109,8 @@ def _process(self, message: AxisArray) -> AxisArray: return message # pass-through xp = get_namespace(message.data) - data = np.asarray(message.data, dtype=np.float64) - was_1d = data.ndim == 1 - if was_1d: - data = data[:, np.newaxis] - n_samples = data.shape[0] + was_1d = message.data.ndim == 1 + n_samples = message.data.shape[0] sine, self._state.phase, self._state.freq_off = advance_drifting_sine( n_samples, @@ -117,10 +124,22 @@ def _process(self, message: AxisArray) -> AxisArray: self._state.rng, ) - out = data + sine # sine is (n_samples, 1), broadcasts over channels + use_mlx_broadcast = xp.__name__ == "mlx.core" and math.prod(message.data.shape) >= MLX_BROADCAST_MIN_ELEMENTS + if use_mlx_broadcast: + # Keep the large signal on MLX and transfer only the common-mode + # (n_samples, 1) sine. This preserves the NumPy RNG/state semantics + # while avoiding a round-trip of every signal channel. + data = message.data[:, np.newaxis] if was_1d else message.data + out = data + xp.asarray(sine) + else: + data = np.asarray(message.data, dtype=np.float64) + if was_1d: + data = data[:, np.newaxis] + out = data + sine # sine broadcasts over channels + if was_1d: out = out[:, 0] - out_data = xp.asarray(out) if xp is not np else out + out_data = out if use_mlx_broadcast else xp.asarray(out) return replace(message, data=out_data) diff --git a/tests/benchmark/bench_mlx.py b/tests/benchmark/bench_mlx.py new file mode 100644 index 0000000..e81234b --- /dev/null +++ b/tests/benchmark/bench_mlx.py @@ -0,0 +1,129 @@ +"""Manual NumPy/MLX microbenchmarks for the simulator hot paths. + +Run on Apple Silicon from the repository root: + + .venv/bin/python tests/benchmark/bench_mlx.py + +MLX is evaluated after every chunk. This mirrors a streaming pipeline and avoids +timing a growing collection of lazy graphs or including their retained memory. +""" + +from __future__ import annotations + +import platform +import time +from collections.abc import Callable + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray + +from ezmsg.simbiophys import ( + CosineEncoderSettings, + CosineEncoderTransformer, + DynamicColoredNoiseSettings, + DynamicColoredNoiseTransformer, + LineNoiseSettings, + LineNoiseTransformer, +) + +N_CHUNKS = 200 + + +def _message(data, fs: float) -> AxisArray: + dims = ["time"] if data.ndim == 1 else ["time", "ch"] + return AxisArray(data, dims=dims, axes={"time": AxisArray.TimeAxis(fs=fs, offset=0.0)}) + + +def _time_chunks( + transformer: Callable[[AxisArray], AxisArray], + chunks: list[AxisArray], + evaluate: Callable[[object], None] | None = None, +) -> tuple[float, AxisArray]: + # Warm up compilation/JIT and state initialization outside the timed region. + warm = transformer(chunks[0]) + if evaluate is not None: + evaluate(warm.data) + + start = time.perf_counter() + for chunk in chunks[1:]: + output = transformer(chunk) + if evaluate is not None: + evaluate(output.data) + elapsed = time.perf_counter() - start + return elapsed / (len(chunks) - 1), output + + +def _report(name: str, numpy_seconds: float, mlx_seconds: float) -> None: + print( + f"{name:36s} NumPy {numpy_seconds * 1e6:9.1f} us/chunk | " + f"MLX {mlx_seconds * 1e6:9.1f} us/chunk | " + f"speedup {numpy_seconds / mlx_seconds:5.2f}x" + ) + + +def benchmark_cosine(mx) -> None: + rng = np.random.default_rng(1) + polar = np.column_stack( + ( + np.abs(rng.standard_normal(500)), + rng.uniform(-np.pi, np.pi, 500), + ) + ).astype(np.float32) + numpy_message = _message(polar, 100.0) + mlx_message = _message(mx.array(polar), 100.0) + mx.eval(mlx_message.data) + numpy_chunks = [numpy_message] * N_CHUNKS + mlx_chunks = [mlx_message] * N_CHUNKS + + settings = CosineEncoderSettings(output_ch=256, baseline=10.0, modulation=20.0, seed=42) + numpy_seconds, numpy_out = _time_chunks(CosineEncoderTransformer(settings), numpy_chunks) + mlx_seconds, mlx_out = _time_chunks(CosineEncoderTransformer(settings), mlx_chunks, mx.eval) + np.testing.assert_allclose(np.asarray(mlx_out.data), numpy_out.data, rtol=5e-3, atol=1e-4) + _report("Cosine encoder (500 x 2 -> 256)", numpy_seconds, mlx_seconds) + + +def benchmark_dynamic_colored_noise(mx) -> None: + beta = np.linspace(0.5, 1.95, 50, dtype=np.float32)[:, np.newaxis] + beta = np.broadcast_to(beta, (50, 8)).copy() + numpy_message = _message(beta, 100.0) + mlx_message = _message(mx.array(beta), 100.0) + mx.eval(mlx_message.data) + numpy_chunks = [numpy_message] * N_CHUNKS + mlx_chunks = [mlx_message] * N_CHUNKS + + settings = DynamicColoredNoiseSettings(output_fs=30000.0, n_poles=5, seed=42) + numpy_seconds, numpy_out = _time_chunks(DynamicColoredNoiseTransformer(settings), numpy_chunks) + mlx_seconds, mlx_out = _time_chunks(DynamicColoredNoiseTransformer(settings), mlx_chunks, mx.eval) + np.testing.assert_allclose(np.asarray(mlx_out.data), numpy_out.data, rtol=1e-4, atol=1e-5) + _report("Colored noise (50 x 8 -> 15000)", numpy_seconds, mlx_seconds) + + +def benchmark_line_noise(mx) -> None: + rng = np.random.default_rng(2) + signal = rng.standard_normal((15000, 256), dtype=np.float32) + numpy_message = _message(signal, 30000.0) + mlx_message = _message(mx.array(signal), 30000.0) + mx.eval(mlx_message.data) + numpy_chunks = [numpy_message] * N_CHUNKS + mlx_chunks = [mlx_message] * N_CHUNKS + + settings = LineNoiseSettings(freq=60.0, amp=10.0, drift_rate=0.002, seed=42) + numpy_seconds, numpy_out = _time_chunks(LineNoiseTransformer(settings), numpy_chunks) + mlx_seconds, mlx_out = _time_chunks(LineNoiseTransformer(settings), mlx_chunks, mx.eval) + np.testing.assert_allclose(np.asarray(mlx_out.data), numpy_out.data, rtol=1e-5, atol=2e-6) + _report("Line noise (15000 x 256)", numpy_seconds, mlx_seconds) + + +def main() -> None: + if platform.system() != "Darwin" or platform.machine() != "arm64": + raise SystemExit("These benchmarks require MLX on Apple Silicon") + + import mlx.core as mx + + benchmark_cosine(mx) + benchmark_dynamic_colored_noise(mx) + benchmark_line_noise(mx) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_cosine_encoder.py b/tests/unit/test_cosine_encoder.py index 13954c5..64871e1 100644 --- a/tests/unit/test_cosine_encoder.py +++ b/tests/unit/test_cosine_encoder.py @@ -1,7 +1,6 @@ """Unit tests for ezmsg.simbiophys.cosine_encoder module.""" import platform -import time import numpy as np import pytest @@ -241,71 +240,36 @@ def test_multiple_samples(self): @requires_apple_silicon -def test_cosine_encoder_mlx_benchmark(): - """Benchmark CosineEncoderTransformer: numpy vs MLX input.""" +def test_cosine_encoder_mlx_matches_numpy_across_chunk_shapes(): + """MLX preserves the backend and matches NumPy as compiled shapes vary.""" import mlx.core as mx - n_samples = 500 - n_chunks = 200 - output_ch = 256 - fs = 100.0 - settings = CosineEncoderSettings( - output_ch=output_ch, + output_ch=256, baseline=10.0, modulation=20.0, speed_modulation=5.0, seed=42, ) - - # Pre-generate chunks as numpy rng = np.random.default_rng(42) - np_chunks = [] - for i in range(n_chunks + 1): # +1 for warmup + xformer_np = CosineEncoderTransformer(settings) + xformer_mx = CosineEncoderTransformer(settings) + + for n_samples in (20, 37, 100): magnitude = np.abs(rng.standard_normal((n_samples, 1))).astype(np.float32) angle = rng.uniform(-np.pi, np.pi, (n_samples, 1)).astype(np.float32) polar = np.hstack([magnitude, angle]) - np_chunks.append( - AxisArray( - polar, - dims=["time", "ch"], - axes={"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * n_samples / fs)}, - ) - ) - - # MLX versions - mx_chunks = [AxisArray(data=mx.array(chunk.data), dims=chunk.dims, axes=chunk.axes) for chunk in np_chunks] - - # --- Numpy --- - xformer_np = CosineEncoderTransformer(settings) - xformer_np(np_chunks[0]) # Warmup - - t0 = time.perf_counter() - np_outputs = [xformer_np(chunk) for chunk in np_chunks[1:]] - t_numpy = time.perf_counter() - t0 - - # --- MLX --- - xformer_mx = CosineEncoderTransformer(settings) - xformer_mx(mx_chunks[0]) # Warmup - mx.eval(xformer_mx(mx_chunks[0]).data) - - t0 = time.perf_counter() - mx_outputs = [xformer_mx(chunk) for chunk in mx_chunks[1:]] - for out in mx_outputs: - mx.eval(out.data) - t_mlx = time.perf_counter() - t0 - - # Verify output is MLX array - last_mx = mx_outputs[-1] - assert isinstance(last_mx.data, mx.array), f"Expected mx.array, got {type(last_mx.data)}" - - # Correctness: compare outputs - for np_out, mx_out in zip(np_outputs, mx_outputs): - np.testing.assert_allclose(np.asarray(mx_out.data), np_out.data, rtol=5e-3, atol=1e-4) - - print( - f"\n CosineEncoder benchmark ({n_chunks} chunks, {n_samples}×2 → {output_ch} ch):" - f"\n numpy: {t_numpy:.4f}s ({t_numpy / n_chunks * 1000:.2f} ms/chunk)" - f"\n mlx: {t_mlx:.4f}s ({t_mlx / n_chunks * 1000:.2f} ms/chunk)" - f"\n ratio (mlx/numpy): {t_mlx / t_numpy:.2f}x" - ) + axes = {"time": AxisArray.LinearAxis(gain=0.01, offset=0.0)} + out_np = xformer_np(AxisArray(polar, dims=["time", "ch"], axes=axes)) + out_mx = xformer_mx(AxisArray(mx.array(polar), dims=["time", "ch"], axes=axes)) + mx.eval(out_mx.data) + + assert isinstance(out_mx.data, mx.array) + for parameter in ( + xformer_mx.state.baseline, + xformer_mx.state.modulation, + xformer_mx.state.pd, + xformer_mx.state.speed_modulation, + ): + assert isinstance(parameter, mx.array) + np.testing.assert_allclose(np.asarray(out_mx.data), out_np.data, rtol=5e-3, atol=1e-4) diff --git a/tests/unit/test_dynamic_colored_noise.py b/tests/unit/test_dynamic_colored_noise.py index 8d5535b..c728e20 100644 --- a/tests/unit/test_dynamic_colored_noise.py +++ b/tests/unit/test_dynamic_colored_noise.py @@ -1,7 +1,6 @@ """Unit tests for ezmsg.simbiophys.dynamic_colored_noise module.""" import platform -import time import numpy as np import pytest @@ -637,73 +636,27 @@ def test_empty_output_accumulation(self): @requires_apple_silicon -def test_dynamic_colored_noise_mlx_benchmark(): - """Benchmark DynamicColoredNoiseTransformer: numpy vs MLX input.""" +def test_dynamic_colored_noise_mlx_preserves_backend_and_matches_numpy(): + """The host-native recurrence returns MLX data without advancing state twice.""" import mlx.core as mx - n_input = 50 - n_chunks = 200 - n_channels = 16 - fs = 100.0 - output_fs = 1000.0 - settings = DynamicColoredNoiseSettings( - output_fs=output_fs, + output_fs=1000.0, n_poles=5, smoothing_tau=0.01, initial_beta=1.0, seed=42, ) - - # Pre-generate chunks as numpy rng = np.random.default_rng(42) - np_chunks = [] - for i in range(n_chunks + 1): # +1 for warmup - beta = rng.uniform(0.5, 2.0, (n_input, n_channels)).astype(np.float64) - np_chunks.append( - AxisArray( - beta, - dims=["time", "ch"], - axes={"time": AxisArray.LinearAxis(gain=1.0 / fs, offset=i * n_input / fs)}, - ) - ) - - # MLX versions - mx_chunks = [AxisArray(data=mx.array(chunk.data), dims=chunk.dims, axes=chunk.axes) for chunk in np_chunks] - - # --- Numpy --- xformer_np = DynamicColoredNoiseTransformer(settings) - xformer_np(np_chunks[0]) # Warmup + xformer_mx = DynamicColoredNoiseTransformer(settings) - t0 = time.perf_counter() - np_outputs = [xformer_np(chunk) for chunk in np_chunks[1:]] - t_numpy = time.perf_counter() - t0 + for chunk_idx in range(3): + beta = rng.uniform(0.5, 1.95, (50, 16)).astype(np.float64) + axes = {"time": AxisArray.LinearAxis(gain=0.01, offset=chunk_idx * 0.5)} + out_np = xformer_np(AxisArray(beta, dims=["time", "ch"], axes=axes)) + out_mx = xformer_mx(AxisArray(mx.array(beta), dims=["time", "ch"], axes=axes)) + mx.eval(out_mx.data) - # --- MLX --- - xformer_mx = DynamicColoredNoiseTransformer(settings) - xformer_mx(mx_chunks[0]) # Warmup - mx.eval(xformer_mx(mx_chunks[0]).data) - - t0 = time.perf_counter() - mx_outputs = [xformer_mx(chunk) for chunk in mx_chunks[1:]] - for out in mx_outputs: - mx.eval(out.data) - t_mlx = time.perf_counter() - t0 - - # Verify output is MLX array - last_mx = mx_outputs[-1] - assert isinstance(last_mx.data, mx.array), f"Expected mx.array, got {type(last_mx.data)}" - - # Correctness: compare first chunk outputs (seeds differ due to warmup count, - # but shapes should match and values should be finite) - for np_out, mx_out in zip(np_outputs[:5], mx_outputs[:5]): - mx_data = np.asarray(mx_out.data) - assert mx_data.shape == np_out.data.shape - assert np.all(np.isfinite(mx_data)) - - print( - f"\n DynamicColoredNoise benchmark ({n_chunks} chunks, {n_input}×{n_channels} → {output_fs}Hz):" - f"\n numpy: {t_numpy:.4f}s ({t_numpy / n_chunks * 1000:.2f} ms/chunk)" - f"\n mlx: {t_mlx:.4f}s ({t_mlx / n_chunks * 1000:.2f} ms/chunk)" - f"\n ratio (mlx/numpy): {t_mlx / t_numpy:.2f}x" - ) + assert isinstance(out_mx.data, mx.array) + np.testing.assert_allclose(np.asarray(out_mx.data), out_np.data, rtol=1e-4, atol=1e-5) diff --git a/tests/unit/test_line_noise.py b/tests/unit/test_line_noise.py index 021b3d4..b7882cc 100644 --- a/tests/unit/test_line_noise.py +++ b/tests/unit/test_line_noise.py @@ -1,11 +1,19 @@ """Unit tests for ezmsg.simbiophys.line_noise module.""" +import platform + import numpy as np +import pytest from ezmsg.util.messages.axisarray import AxisArray from numpy.fft import rfft, rfftfreq from ezmsg.simbiophys import LineNoiseSettings, LineNoiseTransformer +requires_apple_silicon = pytest.mark.skipif( + platform.machine() != "arm64" or platform.system() != "Darwin", + reason="Requires Apple Silicon for MLX", +) + def _msg(n, fs, offset, n_ch=4, val=0.0): return AxisArray( @@ -71,3 +79,24 @@ def test_drift_rate_magnitude(self): offs = np.array(offs) one_sec_drift = np.std(np.diff(offs)) # std of 1 s increments ~ drift_rate assert 0.5 * rate < one_sec_drift < 1.5 * rate + + @requires_apple_silicon + @pytest.mark.parametrize("shape", [(300, 256), (1000, 256), (15000, 256), (250000,)]) + def test_mlx_matches_numpy_on_both_sides_of_dispatch_crossover(self, shape): + import mlx.core as mx + + rng = np.random.default_rng(20) + data = rng.standard_normal(shape).astype(np.float32) + settings = LineNoiseSettings(freq=60.0, amp=10.0, drift_rate=0.002, seed=42) + numpy_tr = LineNoiseTransformer(settings) + mlx_tr = LineNoiseTransformer(settings) + + for chunk_idx in range(2): + axes = {"time": AxisArray.TimeAxis(fs=30000.0, offset=chunk_idx * shape[0] / 30000.0)} + dims = ["time"] if data.ndim == 1 else ["time", "ch"] + out_np = numpy_tr(AxisArray(data, dims=dims, axes=axes)) + out_mx = mlx_tr(AxisArray(mx.array(data), dims=dims, axes=axes)) + mx.eval(out_mx.data) + + assert isinstance(out_mx.data, mx.array) + np.testing.assert_allclose(np.asarray(out_mx.data), out_np.data, rtol=1e-5, atol=2e-6)