diff --git a/pyproject.toml b/pyproject.toml index 68d7962..d693b1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,8 +8,12 @@ readme = "README.md" requires-python = ">=3.10" dynamic = ["version"] dependencies = [ - "ezmsg>=3.9.0", - "ezmsg-baseproc>=1.10.2", + # 3.10.0b2 for CoordinateAxis.fingerprint and AxisArray.chunk_dim, which + # ezmsg-baseproc's default state hash reads. Pinned directly rather than + # left to the transitive requirement: uv only enables pre-releases for a + # package named with a pre-release marker in *this* file. + "ezmsg>=3.10.0b2", + "ezmsg-baseproc>=1.12.0", "ezmsg-sigproc>=2.26.1", "sparse>=0.17.0", "numpy>=2.2.6", diff --git a/src/ezmsg/event/kernel_activation.py b/src/ezmsg/event/kernel_activation.py index 1946344..d99ef5b 100644 --- a/src/ezmsg/event/kernel_activation.py +++ b/src/ezmsg/event/kernel_activation.py @@ -144,12 +144,19 @@ class BinnedKernelActivation( """ def _hash_message(self, message: AxisArray) -> int: - n_channels = message.data.shape[message.get_axis_idx("ch")] if "ch" in message.dims else 1 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). + # The base class covers the key, the dims, the channel count and the + # sample rate. Two things it cannot see are added here: + # dtype -- `str()` rather than `.kind`, which mlx arrays do not expose. + # backend -- the per-channel state is numpy either way, but the code + # paths that fill it differ between sparse, numpy and mlx. + # Folding in the key is new, and is the point: `activation`, + # `dense_carry` and `samples_since_update` are per-channel running state, + # so carrying them into a different stream is the same failure as a + # filter keeping the previous device's history. 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)) + return self._message_hash(message, extra=(str(message.data.dtype), backend)) def _reset_state(self, message: AxisArray) -> None: """Initialize state for new input stream.""" diff --git a/src/ezmsg/event/peak.py b/src/ezmsg/event/peak.py index 9da2004..79b90a8 100644 --- a/src/ezmsg/event/peak.py +++ b/src/ezmsg/event/peak.py @@ -106,11 +106,6 @@ class ThresholdCrossingTransformer( ): """Transformer that detects threshold crossing events.""" - def _hash_message(self, message: AxisArray) -> int: - ax_idx = message.get_axis_idx("time") - sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - return hash((message.key, sample_shape, message.axes["time"].gain)) - def _reset_state(self, message: AxisArray) -> None: """Reset the state variables.""" xp = get_namespace(message.data) diff --git a/src/ezmsg/event/refractory.py b/src/ezmsg/event/refractory.py index 51c4d66..e9306b9 100644 --- a/src/ezmsg/event/refractory.py +++ b/src/ezmsg/event/refractory.py @@ -20,9 +20,6 @@ class Refractory: class RefractoryTransformer(BaseStatefulTransformer[RefractorySettings, AxisArray, AxisArray, Refractory]): - def _hash_message(self, message: AxisArray) -> int: - return super()._hash_message(message) - def _reset_state(self, message: AxisArray) -> None: fs = 1 / message.axes["time"].gain self._state.width = int(self.settings.dur * fs) diff --git a/tests/test_state_reset_semantics.py b/tests/test_state_reset_semantics.py new file mode 100644 index 0000000..a6ce106 --- /dev/null +++ b/tests/test_state_reset_semantics.py @@ -0,0 +1,219 @@ +"""When these transformers rebuild their state, now that the base class decides. + +``ThresholdCrossingTransformer`` and ``RefractoryTransformer`` no longer +implement ``_hash_message``: the base class folds in the message key, the dims, +the length of every dimension except the one the stream is chunked along, the +coordinate values on those dimensions, and the gain and offset of any linear +axis among them. + +The change that matters is the channel *fingerprint*. A source that renames or +reorders its channels without changing how many it sends -- a device +reconfigured mid-session, a montage swapped -- previously looked identical to +these transformers, so per-channel state carried over onto channels it did not +belong to. Nothing about the output announces that; these tests are what would +catch it coming back. + +``BinnedKernelActivation`` still overrides, because dtype and array backend are +things the default cannot see. Its tests here pin both halves: what it adds, and +what it now inherits. +""" + +import numpy as np +import pytest +import sparse +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.event.kernel_activation import ( + ActivationKernelType, + BinnedKernelActivation, + BinnedKernelActivationSettings, +) +from ezmsg.event.peak import ThresholdCrossingTransformer, ThresholdSettings +from ezmsg.event.refractory import RefractorySettings, RefractoryTransformer + +FS = 30_000.0 + + +def dense_msg(labels, n_time=600, fs=FS, key="dev", seed=0): + """A dense message carrying a *labelled* channel axis. + + The conftest helpers omit the ``ch`` axis, which is fine for what they test + but makes a relabel invisible by construction -- there is nothing to relabel. + A real source (blackrock, lsl, a replayed file) always attaches one. + """ + rng = np.random.default_rng(seed) + return AxisArray( + data=rng.standard_normal((n_time, len(labels))), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +def sparse_msg(labels, n_time=600, fs=FS, key="dev", dtype=bool): + coords = np.array( + [[i * 7 % n_time for i in range(len(labels) * 3)], [i % len(labels) for i in range(len(labels) * 3)]] + ) + data = np.ones(coords.shape[1], dtype=dtype) + return AxisArray( + data=sparse.COO(coords, data, shape=(n_time, len(labels))), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +def at_sample(labels, sample_ix, n_time=600, fs=FS, key="dev"): + """One event on every channel, at a chosen sample index.""" + n_ch = len(labels) + coords = np.array([[sample_ix] * n_ch, list(range(n_ch))]) + return AxisArray( + data=sparse.COO(coords, np.ones(n_ch, dtype=bool), shape=(n_time, n_ch)), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +ARM_A = ["armA-1", "armA-2", "armA-3", "armA-4"] +ARM_B = ["armB-1", "armB-2", "armB-3", "armB-4"] + + +class TestARelabelIsNoticed: + """Same channel count, different channels. The case that used to be silent.""" + + def test_threshold_crossing_rebuilds(self): + proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0, auto_scale_tau=0.5)) + proc(dense_msg(ARM_A)) + before = proc._hash + proc(dense_msg(ARM_B, seed=1)) + assert proc._hash != before + + def test_the_adaptive_scaler_does_not_carry_over(self): + """`auto_scale_tau` gives the transformer a per-channel z-scoring state. + Carrying armA's running mean onto armB is the same failure as a filter + keeping the previous device's history.""" + settings = ThresholdSettings(threshold=-1.0, auto_scale_tau=0.5) + + carried = ThresholdCrossingTransformer(settings) + carried(dense_msg(ARM_A, seed=0)) + got = carried(dense_msg(ARM_B, seed=1)) + + clean = ThresholdCrossingTransformer(settings) + want = clean(dense_msg(ARM_B, seed=1)) + + assert (got.data != want.data).nnz == 0, "state survived a channel relabel" + + def test_refractory_rebuilds(self): + proc = RefractoryTransformer(RefractorySettings(dur=0.002)) + proc(sparse_msg(ARM_A)) + before = proc._hash + proc(sparse_msg(ARM_B)) + assert proc._hash != before + + def test_refractory_elapsed_counters_do_not_carry_over(self): + """The refractory window has to be straddled for this to mean anything. + + armA fires on its *last* sample and armB on its *sixth*; at 30 kHz a + 2 ms window is 60 samples, so armA's `elapsed` counters -- if they + survived -- would still be inside the window and would swallow armB's + events entirely. Give both streams the same event positions and the + carried state changes nothing, and the test passes either way. + """ + settings = RefractorySettings(dur=0.002) + n_time = 600 + late = at_sample(ARM_A, n_time - 1, n_time=n_time) + early = at_sample(ARM_B, 5, n_time=n_time) + + carried = RefractoryTransformer(settings) + carried(late) + got = carried(early) + + clean = RefractoryTransformer(settings) + want = clean(at_sample(ARM_B, 5, n_time=n_time)) + + assert want.data.nnz == len(ARM_B), "the reference run should pass every event" + np.testing.assert_array_equal(got.data.todense(), want.data.todense()) + + +class TestWhatMustNotResetStillDoesNot: + def test_chunk_size_jitter_is_ignored(self): + """The chunk dimension's length is whatever arrived; rebuilding on it + would throw away the trailing buffer on every irregular chunk.""" + proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0)) + proc(dense_msg(ARM_A, n_time=600)) + before = proc._hash + proc(dense_msg(ARM_A, n_time=317)) + assert proc._hash == before + + def test_the_same_channels_again_is_not_a_change(self): + proc = RefractoryTransformer(RefractorySettings(dur=0.002)) + proc(sparse_msg(ARM_A)) + before = proc._hash + proc(sparse_msg(ARM_A)) + assert proc._hash == before + + +class TestKernelActivationKeepsItsOverride: + """It adds dtype and backend, which the base class cannot see, and now + inherits the key and fingerprint, which it previously ignored.""" + + @staticmethod + def _proc(): + return BinnedKernelActivation( + BinnedKernelActivationSettings(kernel_type=ActivationKernelType.EXPONENTIAL, tau=0.05, bin_duration=0.02) + ) + + def test_a_dtype_change_still_resets(self): + proc = self._proc() + proc(sparse_msg(ARM_A, dtype=bool)) + before = proc._hash + proc(sparse_msg(ARM_A, dtype=np.float32)) + assert proc._hash != before + + def test_a_relabel_now_resets(self): + """New: the previous hash saw only ndim, dtype, channel count and rate.""" + proc = self._proc() + proc(sparse_msg(ARM_A)) + before = proc._hash + proc(sparse_msg(ARM_B)) + assert proc._hash != before + + def test_a_stream_switch_now_resets(self): + """Also new. `activation`, `dense_carry` and `samples_since_update` are + per-channel running state and must not follow a key change.""" + proc = self._proc() + proc(sparse_msg(ARM_A, key="deviceA")) + before = proc._hash + proc(sparse_msg(ARM_A, key="deviceB")) + assert proc._hash != before + + def test_chunk_size_jitter_is_still_ignored(self): + proc = self._proc() + proc(sparse_msg(ARM_A, n_time=600)) + before = proc._hash + proc(sparse_msg(ARM_A, n_time=317)) + assert proc._hash == before + + def test_a_missing_time_axis_is_still_rejected(self): + proc = self._proc() + msg = AxisArray( + data=sparse.COO(np.zeros((2, 0), dtype=np.int64), np.zeros(0, dtype=bool), shape=(10, 4)), + dims=["time", "ch"], + axes={"ch": CoordinateAxis(data=np.array(ARM_A), dims=["ch"])}, + key="dev", + ) + with pytest.raises(ValueError, match="sample rate"): + proc(msg)