diff --git a/src/ezmsg/event/kernel_activation.py b/src/ezmsg/event/kernel_activation.py index 471129a..84b18fe 100644 --- a/src/ezmsg/event/kernel_activation.py +++ b/src/ezmsg/event/kernel_activation.py @@ -6,13 +6,13 @@ uses a state-based approach that is O(n_events + n_bins) instead of O(n_samples). -Input may be either ``sparse.COO`` (the typical output of -:class:`ezmsg.event.peak.ThresholdCrossingTransformer` in default mode) or a -dense array (from the same transformer with ``output_format=DENSE``). When the -input is dense and the configuration is COUNT + SUM (the rate-computation -case), the binning runs on the input's array namespace and stays on device -(e.g., MLX, CuPy). Other configurations with dense input fall back to -event-extraction and use the same code path as sparse input. +Input may be either ``sparse.COO`` (the default output of +:class:`ezmsg.event.peak.ThresholdCrossingTransformer`) or a dense array from the +same transformer with ``output_format=DENSE``. When the input is dense and the +configuration is COUNT + SUM (the rate-computation case), the binning runs on the +input's array namespace and stays on device (e.g., MLX, CuPy). Other configurations +with dense input fall back to event extraction and use the same code path as sparse +input. """ from enum import Enum @@ -97,6 +97,9 @@ class BinnedKernelActivationState: # Current activation level per channel (for exponential/alpha) activation: npt.NDArray[np.float64] | None = None + dense_carry: object | None = None + """Partial-bin COUNT+SUM state in the dense input's array namespace.""" + # For alpha kernel: auxiliary state variable alpha_aux: npt.NDArray[np.float64] | None = None @@ -145,13 +148,15 @@ def _hash_message(self, message: AxisArray) -> int: if "time" not in message.axes or not hasattr(message.axes["time"], "gain"): raise ValueError("Could not determine sample rate from input message") # str(dtype) works for numpy ('bool', 'float32', ...) and mlx (which doesn't expose dtype.kind). - return hash((message.data.ndim, str(message.data.dtype), n_channels, message.axes["time"].gain)) + backend = "sparse" if isinstance(message.data, sparse.SparseArray) else get_namespace(message.data).__name__ + return hash((message.data.ndim, str(message.data.dtype), n_channels, message.axes["time"].gain, backend)) def _reset_state(self, message: AxisArray) -> None: """Initialize state for new input stream.""" n_channels = message.data.shape[message.get_axis_idx("ch")] if "ch" in message.dims else 1 self._state.activation = np.zeros(n_channels, dtype=np.float64) + self._state.dense_carry = None self._state.samples_since_update = np.zeros(n_channels, dtype=np.int64) # For alpha kernel, we need auxiliary state @@ -420,13 +425,17 @@ def _process_dense_count_sum(self, message: AxisArray) -> AxisArray: else: contrib = (data != 0).astype(xp.float32) - # Pull state into the input namespace for on-device math. - overflow_xp = xp.asarray(self._state.activation.reshape(feature_shape)).astype(xp.float32) + # Keep dense partial-bin state in the input namespace. In particular, do + # not round-trip an MLX carry through np.asarray here: that synchronizes + # the device on every source chunk, including chunks that close no bin. + overflow_xp = ( + xp.zeros(feature_shape, dtype=xp.float32) if self._state.dense_carry is None else self._state.dense_carry + ) if n_bins == 0: # No complete bins this chunk — accumulate everything into the carry-over. new_overflow = overflow_xp + (xp.sum(contrib, axis=0) if n_samples > 0 else overflow_xp * 0) - self._state.activation = np.asarray(new_overflow).reshape(self._state.activation.shape) + self._state.dense_carry = new_overflow return replace( message, data=xp.zeros((0,) + feature_shape, dtype=xp.float32), @@ -467,7 +476,7 @@ def _process_dense_count_sum(self, message: AxisArray) -> AxisArray: new_overflow = xp.sum(contrib[last_bin_end:], axis=0) else: new_overflow = xp.zeros(feature_shape, dtype=cumsum.dtype) - self._state.activation = np.asarray(new_overflow).reshape(self._state.activation.shape) + self._state.dense_carry = new_overflow if self.settings.rate_normalize: output = output / step.output_gain diff --git a/src/ezmsg/event/peak.py b/src/ezmsg/event/peak.py index 1c84017..9da2004 100644 --- a/src/ezmsg/event/peak.py +++ b/src/ezmsg/event/peak.py @@ -154,7 +154,7 @@ def _reset_state(self, message: AxisArray) -> None: self._state.elapsed = np.full((n_features,), self._state.refrac_width + 1, dtype=np.int32) def _can_use_mlx_metal(self, xp, data) -> bool: - """The Metal kernel runs automatically for MLX + DENSE + a config it supports. + """The Metal kernel runs automatically for MLX + dense + a supported config. It cannot recover peak values, align on peaks, enforce a min peak width, or auto-scale, so any of those settings disable the path. @@ -253,6 +253,12 @@ def _process(self, message: AxisArray) -> AxisArray: # (the cpu path's prepended buffer alone yields no new crossings, and the metal # kernel has nothing to scan). Short-circuit before backend dispatch. if message.data.shape[0] == 0: + # _hash_message ignores the sample count, so an empty first chunk still triggers + # _reset_state, which seeds _state.data from data[:1] -> an empty time axis. Force + # a re-reset so the first non-empty chunk seeds the prev-sample reference (the + # metal path indexes _state.data[0], and the cpu path prepends it). + if self._state.data is not None and self._state.data.shape[0] == 0: + self._request_reset() return self._empty_output(message, xp) # MLX-on-device fast path: bypass the numpy event-detection logic and run diff --git a/tests/test_peak.py b/tests/test_peak.py index a54fad1..f871694 100644 --- a/tests/test_peak.py +++ b/tests/test_peak.py @@ -183,6 +183,33 @@ def _require_mlx(): return mx +def test_threshold_crossing_defaults_to_sparse(): + proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0, refrac_dur=0.0)) + msg = AxisArray( + data=np.zeros((10, 2), dtype=np.float32), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=1000.0)}, + ) + + result = proc(msg) + + assert isinstance(result.data, sparse.SparseArray) + + +def test_threshold_crossing_mlx_defaults_to_sparse(): + mx = _require_mlx() + proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0, refrac_dur=0.0)) + msg = AxisArray( + data=mx.zeros((10, 2), dtype=mx.float32), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=1000.0)}, + ) + + result = proc(msg) + + assert isinstance(result.data, sparse.SparseArray) + + @pytest.mark.parametrize( ("threshold", "refrac_dur", "stride"), [ @@ -292,3 +319,33 @@ def test_threshold_crossing_empty_time_first(return_peak_val: bool, auto_scale_t out_normal = proc(msg_normal) assert isinstance(out_normal.data, sparse.SparseArray) assert out_normal.data.shape[1] == N_CH + + +def test_threshold_crossing_empty_time_first_mlx_metal(): + """Empty → normal for MLX inputs: the empty first chunk must not leave the Metal path with + an empty prev-sample reference (regression for `[take] ... from an empty axis`).""" + mx = _require_mlx() + fs = 1000.0 + + proc = ThresholdCrossingTransformer( + ThresholdSettings( + threshold=-1.0, + refrac_dur=0.001, + output_format=OutputFormat.DENSE, + ) + ) + + def mlx_msg(n_time: int) -> AxisArray: + return AxisArray( + data=mx.array(np.random.randn(n_time, N_CH).astype(np.float32)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=fs)}, + ) + + out_empty = proc(mlx_msg(0)) + assert out_empty.data.shape[0] == 0 + + out_normal = proc(mlx_msg(64)) + mx.eval(out_normal.data) + assert not isinstance(out_normal.data, sparse.SparseArray) + assert out_normal.data.shape == (64, N_CH) diff --git a/tests/test_rate.py b/tests/test_rate.py index 05c5641..0864bcc 100644 --- a/tests/test_rate.py +++ b/tests/test_rate.py @@ -157,6 +157,36 @@ def test_rate_dense_input_matches_sparse(fs: float, bin_duration: float, refrac_ np.testing.assert_allclose(sp_msg.data, ds_msg.data) +def test_rate_dense_threshold_preserves_mlx_backend(): + mx = pytest.importorskip("mlx.core") + try: + mx.eval(mx.array([1.0], dtype=mx.float32)) + except RuntimeError as exc: + pytest.skip(f"MLX device unavailable: {exc}") + + fs = 30_000.0 + data = np.zeros((600, 4), dtype=np.float32) + data[10] = -2.0 + data[310] = -2.0 + message = _make_msg(mx.array(data), fs, 0.0) + + threshold = ThresholdCrossingTransformer( + threshold=-1.0, + refrac_dur=0.001, + output_format=OutputFormat.DENSE, + ) + rate = Rate(EventRateSettings(bin_duration=0.02, fractional=False)) + + events = threshold(message) + result = rate(events) + mx.eval(events.data, result.data) + + assert isinstance(events.data, mx.array) + assert isinstance(result.data, mx.array) + assert isinstance(rate.state.dense_carry, mx.array) + np.testing.assert_allclose(np.asarray(result.data), np.full((1, 4), 100.0)) + + def test_rate_empty_time_first(): """Empty → normal: empty first message triggers _reset_state on empty data.""" proc = Rate(EventRateSettings(bin_duration=0.02))