Declare chunk_dim, prime channel fingerprints, adopt the axis-aware state hash

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

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

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:

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:

src/ezmsg/simbiophys/dnss/spike.py

@@ -305,20 +305,18 @@ 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_),
                 data=np.array([], dtype=np.int_),
                 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:

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

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)

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:

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

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]