Optimize MLX simulator hot paths

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."""

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,

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)
 
 

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()

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)