diff --git a/pyproject.toml b/pyproject.toml index d128850..890bb1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,8 +10,11 @@ requires-python = ">=3.10" dynamic = ["version"] dependencies = [ "array-api-compat>=1.11.0", - "ezmsg>=3.9.0", - "ezmsg-baseproc>=1.6.1", + # 3.10.0b2 for AxisArray.chunk_dim, which every producer here declares, and + # CoordinateAxis.fingerprint, which they prime so downstream consumers do + # not each recompute it. + "ezmsg>=3.10.0b2", + "ezmsg-baseproc>=1.12.0", # axis-aware default state hash + hash witness "ezmsg-sigproc>=2.23.0", "ezmsg-event>=0.9.0", "numba>=0.59.0", diff --git a/src/ezmsg/simbiophys/baseline_drift.py b/src/ezmsg/simbiophys/baseline_drift.py index f69209f..80a79cb 100644 --- a/src/ezmsg/simbiophys/baseline_drift.py +++ b/src/ezmsg/simbiophys/baseline_drift.py @@ -166,11 +166,10 @@ class BaselineDriftTransformer( channel of the input. """ - def _hash_message(self, message: AxisArray) -> int: - time_axis = message.axes.get("time") - gain = time_axis.gain if time_axis is not None else 0.0 - n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 - return hash((n_channels, gain)) + # No `_hash_message`: the default already folds in the channel count and the + # chunk axis's gain, which is all this hashed before, and additionally the + # channel *fingerprint* -- one drift process is warmed up per channel, so a + # relabel at a fixed count leaves each channel wearing another's drift. def _reset_state(self, message: AxisArray) -> None: n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 diff --git a/src/ezmsg/simbiophys/cosine_encoder.py b/src/ezmsg/simbiophys/cosine_encoder.py index eab1928..95e7f78 100644 --- a/src/ezmsg/simbiophys/cosine_encoder.py +++ b/src/ezmsg/simbiophys/cosine_encoder.py @@ -152,6 +152,8 @@ def load_from_file( # Create channel axis for output messages ch_labels = np.array([f"ch{i}" for i in range(len(baseline))]) self.ch_axis = AxisArray.CoordinateAxis(data=ch_labels, dims=["ch"]) + # Primed once per model -- see noise.py for why. + self.ch_axis.fingerprint self.validate() @@ -183,6 +185,8 @@ def init_random( # Create channel axis for output messages ch_labels = np.array([f"ch{i}" for i in range(output_ch)]) self.ch_axis = AxisArray.CoordinateAxis(data=ch_labels, dims=["ch"]) + # Primed once per model -- see noise.py for why. + self.ch_axis.fingerprint self.validate() @@ -208,6 +212,15 @@ class CosineEncoderTransformer( - Any other cosine-tuning based encoding """ + def _hash_message(self, message: AxisArray) -> int: + # The tuning parameters come from `settings` alone -- a file or a seeded + # draw -- and never from the message, so the axis-aware default would + # redraw every channel's preferred direction the first time anything + # upstream relabelled, silently changing the simulated population + # mid-stream. The one thing `_reset_state` does read is the array + # backend, to move the parameters onto it. + return hash(get_namespace(message.data).__name__) + def _reset_state(self, message: AxisArray) -> None: """Initialize encoder parameters.""" if self.settings.model_file is not None: diff --git a/src/ezmsg/simbiophys/dnss/lfp.py b/src/ezmsg/simbiophys/dnss/lfp.py index feb7d2c..cd30025 100644 --- a/src/ezmsg/simbiophys/dnss/lfp.py +++ b/src/ezmsg/simbiophys/dnss/lfp.py @@ -240,16 +240,14 @@ def _reset_state(self, time_axis: LinearAxis) -> None: next(self._state.lfp_gen) # Pre-construct template AxisArray with channel axis + ch_axis = AxisArray.CoordinateAxis(data=np.arange(self.settings.n_ch), dims=["ch"]) + # Primed once for the stream -- see noise.py for why. + ch_axis.fingerprint self._state.template = AxisArray( data=np.zeros((0, self.settings.n_ch), dtype=np.float64), dims=["time", "ch"], - axes={ - "time": time_axis, - "ch": AxisArray.CoordinateAxis( - data=np.arange(self.settings.n_ch), - dims=["ch"], - ), - }, + axes={"time": time_axis, "ch": ch_axis}, + chunk_dim="time", ) def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: diff --git a/src/ezmsg/simbiophys/dnss/spike.py b/src/ezmsg/simbiophys/dnss/spike.py index 986c269..9a884e0 100644 --- a/src/ezmsg/simbiophys/dnss/spike.py +++ b/src/ezmsg/simbiophys/dnss/spike.py @@ -305,6 +305,9 @@ def _reset_state(self, time_axis: LinearAxis) -> None: next(self._state.spike_gen) # Pre-construct template AxisArray with channel axis + ch_axis = AxisArray.CoordinateAxis(data=np.arange(self.settings.n_ch), dims=["ch"]) + # Primed once for the stream -- see noise.py for why. + ch_axis.fingerprint self._state.template = AxisArray( data=sparse.COO( coords=np.array([[], []], dtype=np.int_), @@ -312,13 +315,8 @@ def _reset_state(self, time_axis: LinearAxis) -> None: shape=(0, self.settings.n_ch), ), dims=["time", "ch"], - axes={ - "time": time_axis, - "ch": AxisArray.CoordinateAxis( - data=np.arange(self.settings.n_ch), - dims=["ch"], - ), - }, + axes={"time": time_axis, "ch": ch_axis}, + chunk_dim="time", ) def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: diff --git a/src/ezmsg/simbiophys/dynamic_colored_noise.py b/src/ezmsg/simbiophys/dynamic_colored_noise.py index 8632561..0fde279 100644 --- a/src/ezmsg/simbiophys/dynamic_colored_noise.py +++ b/src/ezmsg/simbiophys/dynamic_colored_noise.py @@ -219,14 +219,10 @@ class DynamicColoredNoiseTransformer( >>> noise_output = transformer(beta_input) # Output at 30 kHz """ - def _hash_message(self, message: AxisArray) -> int: - """Hash based on number of channels and sample rate to detect stream changes.""" - time_axis = message.axes.get("time") - # LinearAxis has gain (1/fs) rather than fs directly - gain = time_axis.gain if time_axis is not None else 0.0 - # Number of channels is dim 1 for 2D data, or 1 for 1D data - n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 - return hash((n_channels, gain)) + # No `_hash_message`: the default already folds in the channel count and the + # chunk axis's gain, which is all this hashed before, and additionally the + # channel *fingerprint* -- the delay lines and per-channel coefficients are + # tied to specific channels, so a relabel at a fixed count has to reset. def _reset_state(self, message: AxisArray) -> None: """Initialize filter states and compute timing parameters.""" diff --git a/src/ezmsg/simbiophys/line_noise.py b/src/ezmsg/simbiophys/line_noise.py index c71332c..3b314d5 100644 --- a/src/ezmsg/simbiophys/line_noise.py +++ b/src/ezmsg/simbiophys/line_noise.py @@ -21,7 +21,7 @@ BaseTransformerUnit, processor_state, ) -from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.messages.axisarray import AxisArray, AxisBase from ezmsg.util.messages.util import replace from .oscillator import advance_drifting_sine, freq_drift_step_std @@ -87,15 +87,23 @@ class LineNoiseTransformer(BaseStatefulTransformer[LineNoiseSettings, AxisArray, channels may vary. When ``freq`` is None the input passes through unchanged. """ + def _chunk_axis(self, message: AxisArray) -> AxisBase | None: + """The axis the stream grows along, however the producer named it.""" + dim = message.chunk_dim or next((d for d in self.STREAMING_DIMS if d in message.dims), None) + return message.axes.get(dim) + def _hash_message(self, message: AxisArray) -> int: - time_axis = message.axes.get("time") - gain = time_axis.gain if time_axis is not None else 0.0 - return hash(gain) + # Deliberately narrower than the default. Every state array is (1, 1) and + # broadcasts across however many channels arrive, so the channel count and + # fingerprint the default folds in would only restart the phase + # accumulator for a sinusoid that did not change. The sample period is the + # one thing this depends on. + return hash(getattr(self._chunk_axis(message), "gain", None)) def _reset_state(self, message: AxisArray) -> None: if self.settings.freq is None: return - time_axis = message.axes.get("time") + time_axis = self._chunk_axis(message) self._state.dt = time_axis.gain if time_axis is not None else 1.0 self._state.ang_freq = np.array([[2.0 * np.pi * self.settings.freq]], dtype=np.float64) self._state.amp = np.array([[self.settings.amp]], dtype=np.float64) diff --git a/src/ezmsg/simbiophys/noise.py b/src/ezmsg/simbiophys/noise.py index b1eff63..3b14bd5 100644 --- a/src/ezmsg/simbiophys/noise.py +++ b/src/ezmsg/simbiophys/noise.py @@ -48,16 +48,20 @@ class WhiteNoiseProducer(BaseClockDrivenProducer[WhiteNoiseSettings, WhiteNoiseS def _reset_state(self, time_axis: LinearAxis) -> None: """Initialize template with channel axis.""" n_ch = self.settings.n_ch + ch_axis = AxisArray.CoordinateAxis(data=np.arange(n_ch), dims=["ch"]) + # Compute the channel fingerprint once, now. It is cached on the axis and + # pickled with it, and every message reuses this same axis object, so one + # checksum covers the whole stream. Left cold it would be computed by the + # first stateful consumer in this process -- and, since unpickling builds + # a new axis object per message, by the first consumer in every other + # process, on every message. + ch_axis.fingerprint self._state.template = AxisArray( data=np.zeros((0, n_ch)), dims=["time", "ch"], - axes={ - "time": time_axis, - "ch": AxisArray.CoordinateAxis( - data=np.arange(n_ch), - dims=["ch"], - ), - }, + axes={"time": time_axis, "ch": ch_axis}, + # Messages append along `time`; `ch` describes the stream. + chunk_dim="time", ) def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: diff --git a/src/ezmsg/simbiophys/oscillator.py b/src/ezmsg/simbiophys/oscillator.py index 76d8435..01028bd 100644 --- a/src/ezmsg/simbiophys/oscillator.py +++ b/src/ezmsg/simbiophys/oscillator.py @@ -125,16 +125,14 @@ class SpiralProducer(BaseClockDrivenProducer[SpiralGeneratorSettings, SpiralGene def _reset_state(self, time_axis: LinearAxis) -> None: """Initialize template.""" + ch_axis = AxisArray.CoordinateAxis(data=np.array(["x", "y"]), dims=["ch"]) + # Primed once for the stream -- see noise.py for why. + ch_axis.fingerprint self._state.template = AxisArray( data=np.zeros((0, 2)), dims=["time", "ch"], - axes={ - "time": time_axis, - "ch": AxisArray.CoordinateAxis( - data=np.array(["x", "y"]), - dims=["ch"], - ), - }, + axes={"time": time_axis, "ch": ch_axis}, + chunk_dim="time", ) def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: @@ -237,16 +235,14 @@ def _reset_state(self, time_axis: LinearAxis) -> None: n_ch = self.settings.n_ch # Create template + ch_axis = AxisArray.CoordinateAxis(data=np.arange(n_ch), dims=["ch"]) + # Primed once for the stream -- see noise.py for why. + ch_axis.fingerprint self._state.template = AxisArray( data=np.zeros((0, n_ch)), dims=["time", "ch"], - axes={ - "time": time_axis, - "ch": AxisArray.CoordinateAxis( - data=np.arange(n_ch), - dims=["ch"], - ), - }, + axes={"time": time_axis, "ch": ch_axis}, + chunk_dim="time", ) # Convert settings to arrays and validate diff --git a/tests/unit/test_cosine_encoder.py b/tests/unit/test_cosine_encoder.py index 64871e1..99966f4 100644 --- a/tests/unit/test_cosine_encoder.py +++ b/tests/unit/test_cosine_encoder.py @@ -159,17 +159,16 @@ def test_directional_tuning(self): transformer._state.pd = np.array([[0.0]]) # Preferred direction = 0 (rightward) transformer._state.speed_modulation = np.array([[0.0]]) transformer._state.ch_axis = AxisArray.CoordinateAxis(data=np.array(["ch0"]), dims=["ch"]) - transformer._hash = 0 time_axis = AxisArray.TimeAxis(fs=100.0, offset=0.0) # Polar: magnitude=1, angle=0 (aligned with pd) aligned = AxisArray(np.array([[1.0, 0.0]]), dims=["time", "ch"], axes={"time": time_axis}) + # Adopt the hash these messages produce, so the hand-set parameters above + # are treated as already initialized rather than redrawn from `seed`. + transformer._hash = transformer._hash_message(aligned) output_aligned = transformer(aligned).data[0, 0] - # Reset hash to reuse state - transformer._hash = 0 - # Polar: magnitude=1, angle=pi (opposite to pd) opposite = AxisArray(np.array([[1.0, np.pi]]), dims=["time", "ch"], axes={"time": time_axis}) output_opposite = transformer(opposite).data[0, 0] @@ -191,16 +190,16 @@ def test_speed_modulation(self): transformer._state.pd = np.array([[0.0]]) transformer._state.speed_modulation = np.array([[5.0]]) transformer._state.ch_axis = AxisArray.CoordinateAxis(data=np.array(["ch0"]), dims=["ch"]) - transformer._hash = 0 time_axis = AxisArray.TimeAxis(fs=100.0, offset=0.0) # Different magnitudes slow = AxisArray(np.array([[1.0, 0.0]]), dims=["time", "ch"], axes={"time": time_axis}) + # Adopt the hash these messages produce, so the hand-set parameters above + # are treated as already initialized rather than redrawn from `seed`. + transformer._hash = transformer._hash_message(slow) output_slow = transformer(slow).data[0, 0] - transformer._hash = 0 - fast = AxisArray(np.array([[2.0, 0.0]]), dims=["time", "ch"], axes={"time": time_axis}) output_fast = transformer(fast).data[0, 0] diff --git a/tests/unit/test_state_reset_semantics.py b/tests/unit/test_state_reset_semantics.py new file mode 100644 index 0000000..98b856a --- /dev/null +++ b/tests/unit/test_state_reset_semantics.py @@ -0,0 +1,169 @@ +"""What each stateful stage here treats as "a different stream". + +Two facts these simulators now rely on, neither of which anything else notices +if it stops being true: + +* Every producer declares :attr:`AxisArray.chunk_dim`, so a downstream processor + knows which dimension grows rather than guessing ``"time"``, and hands its + channel axis over with the fingerprint already computed, so nobody downstream + pays the checksum on every message. +* The stages that hold per-channel state let the axis-aware default hash decide + when to reset, and the two that genuinely do not hold per-channel state say so + explicitly. Getting that backwards is silent either way: too eager and a + simulated population is redrawn mid-run, too lazy and one channel's drift is + handed to another. +""" + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.simbiophys import ( + BaselineDriftSettings, + BaselineDriftTransformer, + CosineEncoderSettings, + CosineEncoderTransformer, + DynamicColoredNoiseSettings, + DynamicColoredNoiseTransformer, + LineNoiseSettings, + LineNoiseTransformer, +) +from ezmsg.simbiophys.dnss.lfp import DNSSLFPProducer, DNSSLFPSettings +from ezmsg.simbiophys.dnss.spike import DNSSSpikeProducer, DNSSSpikeSettings +from ezmsg.simbiophys.noise import WhiteNoiseProducer, WhiteNoiseSettings +from ezmsg.simbiophys.oscillator import ( + SinGeneratorSettings, + SinProducer, + SpiralGeneratorSettings, + SpiralProducer, +) + + +def signal(labels: list[str], fs: float = 100.0, n_time: int = 16, dim: str = "time") -> AxisArray: + return AxisArray( + np.zeros((n_time, len(labels))), + dims=[dim, "ch"], + axes={ + dim: AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key="dev", + chunk_dim=dim, + ) + + +PRODUCERS = [ + ("white_noise", WhiteNoiseProducer, WhiteNoiseSettings(n_ch=4, fs=100.0), 100.0), + ("spiral", SpiralProducer, SpiralGeneratorSettings(fs=100.0), 100.0), + ("sin", SinProducer, SinGeneratorSettings(n_ch=4, fs=100.0), 100.0), + ("dnss_lfp", DNSSLFPProducer, DNSSLFPSettings(fs=100.0), 100.0), + # Spike patterns are tabulated at 30 kHz and cannot be resampled. + ("dnss_spike", DNSSSpikeProducer, DNSSSpikeSettings(fs=30000.0), 30000.0), +] + + +@pytest.mark.parametrize("name,cls,settings,fs", PRODUCERS, ids=[p[0] for p in PRODUCERS]) +class TestProducersDescribeTheirStream: + def _produce(self, cls, settings, fs): + producer = cls(settings) + time_axis = AxisArray.TimeAxis(fs=fs, offset=0.0) + producer._reset_state(time_axis) + return producer._produce(10, time_axis) + + def test_declares_the_chunk_dim(self, name, cls, settings, fs): + """Without this a windowing stage downstream has to guess, and ``"time"`` + is present-but-wrong the moment the message becomes ``(win, time, ch)``.""" + assert self._produce(cls, settings, fs).chunk_dim == "time" + + def test_hands_over_a_primed_channel_axis(self, name, cls, settings, fs): + """The axis is built once per stream, so one checksum covers every + message and every consumer -- including across a process boundary, where + unpickling otherwise hands out a cold axis per message, forever.""" + msg = self._produce(cls, settings, fs) + cold = sorted( + dim + for dim, axis in msg.axes.items() + if isinstance(axis, CoordinateAxis) and "_fingerprint" not in axis.__dict__ + ) + assert not cold, f"axes published without a fingerprint: {cold}" + + +class TestPerChannelStateResetsOnRelabel: + """Same channel *count*, different channels. The delay lines and drift + anchors are one-per-channel, so reusing them hands each new channel the + history of whoever used to sit in its column.""" + + @pytest.mark.parametrize( + "cls,settings", + [ + (BaselineDriftTransformer, BaselineDriftSettings(seed=0)), + (DynamicColoredNoiseTransformer, DynamicColoredNoiseSettings(seed=0)), + ], + ids=["baseline_drift", "dynamic_colored_noise"], + ) + def test_a_relabel_at_a_fixed_count_is_a_new_stream(self, cls, settings): + proc = cls(settings) + assert proc._hash_message(signal(["a", "b"])) != proc._hash_message(signal(["c", "d"])) + + @pytest.mark.parametrize( + "cls,settings", + [ + (BaselineDriftTransformer, BaselineDriftSettings(seed=0)), + (DynamicColoredNoiseTransformer, DynamicColoredNoiseSettings(seed=0)), + ], + ids=["baseline_drift", "dynamic_colored_noise"], + ) + def test_a_longer_chunk_is_the_same_stream(self, cls, settings): + """The chunk dimension is excluded, so ordinary chunk-size jitter must + not throw away a warmed-up filter.""" + proc = cls(settings) + assert proc._hash_message(signal(["a", "b"], n_time=16)) == proc._hash_message(signal(["a", "b"], n_time=64)) + + +class TestLineNoiseIsCommonMode: + """Its state is all ``(1, 1)`` and broadcasts, so it deliberately ignores + what the default hash cares most about.""" + + def test_a_relabel_does_not_restart_the_phase(self): + proc = LineNoiseTransformer(LineNoiseSettings(freq=60.0)) + assert proc._hash_message(signal(["a", "b"])) == proc._hash_message(signal(["c", "d", "e"])) + + def test_a_sample_rate_change_does(self): + proc = LineNoiseTransformer(LineNoiseSettings(freq=60.0)) + assert proc._hash_message(signal(["a", "b"], fs=100.0)) != proc._hash_message(signal(["a", "b"], fs=500.0)) + + def test_it_reads_the_declared_chunk_dim_not_the_name_time(self): + """A stream that grows along ``samp`` has no ``time`` axis at all; the + old hard-coded lookup silently fell back to a period of zero.""" + proc = LineNoiseTransformer(LineNoiseSettings(freq=60.0)) + slow = signal(["a", "b"], fs=100.0, dim="samp") + fast = signal(["a", "b"], fs=500.0, dim="samp") + assert proc._hash_message(slow) != proc._hash_message(fast) + + proc(slow) + assert proc.state.dt == pytest.approx(1 / 100.0) + + +class TestCosineEncoderKeepsItsPopulation: + """The tuning parameters come from settings -- a file, or a seeded draw -- + and never from the message.""" + + def test_an_upstream_relabel_does_not_redraw(self): + # seed=None deliberately: with a fixed seed a redraw reproduces the same + # preferred directions and the assertion below could not tell the two + # apart. Here a reset is visible. + proc = CosineEncoderTransformer(CosineEncoderSettings(output_ch=8, seed=None)) + polar = signal(["mag", "ang"]) + proc(polar) + pd_before = np.array(proc.state.pd, copy=True) + + proc(signal(["magnitude", "angle"])) + np.testing.assert_array_equal(proc.state.pd, pd_before) + + def test_the_output_channel_axis_is_primed_and_reused(self): + proc = CosineEncoderTransformer(CosineEncoderSettings(output_ch=8, seed=42)) + first = proc(signal(["mag", "ang"])) + assert "_fingerprint" in first.axes["ch"].__dict__ + # The same object every message, which is what lets a downstream + # processor settle its own hash on one pointer comparison. + assert proc(signal(["mag", "ang"])).axes["ch"] is first.axes["ch"]