From fd01802023b315857b8913bc454c00e55192e224 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Wed, 2 Sep 2026 23:14:11 -0400 Subject: [PATCH 1/4] Add CoordinateAxis.fingerprint, a content-derived axis identity An operation that resolves coordinate *values* into cached state -- channel labels into array indices, or into output labels -- cannot key that cache on shape alone. A source that renames, reorders or swaps channels without changing how many it sends keeps getting the previously resolved answer, and the operation silently emits one channel's samples under another channel's label. Comparing the arrays outright is O(bytes) on every message and in every consumer. `fingerprint` is a small hashable digest, computed on first access and cached on the instance, so the cost is paid once per axis object rather than once per consumer per message. Being derived rather than assigned, there is no counter for a producer to forget to bump; building a new axis is what changes it. The cached value lives in __dict__, so it is pickled with the axis and arrives already computed on the far side of a process boundary. Measured on a 256-channel ChannelMap axis, 150 messages through one process hop with five contents-dependent consumers: 750 digests computed without it, 1 with. crc32 rather than hash(tobytes()) because the copy is not the bottleneck -- on that axis the copy runs at ~94 GB/s while CPython's siphash over the result manages ~5.5 GB/s, and crc32 reads the array's buffer directly at ~29 GB/s. The dtype is stored as the object rather than str(dtype), which numpy builds field by field at roughly ten times the cost of the checksum it annotates. None is returned, and cached, when the contents cannot be digested; callers must treat that as "unknown" rather than as a value that compares equal to another None. fast_replace has to drop the cache. It copies __dict__ straight into the constructor and is called on axes, not only on messages, so a cached _fingerprint reaches CoordinateAxis.__init__ as an unexpected keyword and raises TypeError. Dropping is also the correct semantics: a digest of the old field values must not ride onto a copy that changes them. --- src/ezmsg/util/messages/axisarray.py | 76 +++++++++++++++ src/ezmsg/util/messages/util.py | 10 ++ tests/messages/test_axisarray.py | 139 ++++++++++++++++++++++++++- tests/messages/test_replace.py | 59 ++++++++++++ 4 files changed, 282 insertions(+), 2 deletions(-) create mode 100644 tests/messages/test_replace.py diff --git a/src/ezmsg/util/messages/axisarray.py b/src/ezmsg/util/messages/axisarray.py index 7983fa0a..694677d2 100644 --- a/src/ezmsg/util/messages/axisarray.py +++ b/src/ezmsg/util/messages/axisarray.py @@ -6,6 +6,7 @@ import math import typing import warnings +import zlib import ezmsg.core as ez @@ -117,6 +118,10 @@ def create_time_axis(cls, fs: float, offset: float = 0.0) -> "LinearAxis": return cls(unit="s", gain=1.0 / fs, offset=offset) +# Distinguishes "no fingerprint cached yet" from "cached, and it is None". +_UNSET = object() + + @dataclass class ArrayWithNamedDims: """ @@ -186,6 +191,77 @@ def value(self, x): """ return self.data[x] + @property + def fingerprint(self) -> tuple | None: + """ + A small, hashable stand-in for this axis's *contents*. + + Two axes with equal fingerprints hold equal coordinate values; two axes + with different fingerprints do not. It is derived from the data rather + than assigned, so there is no counter for a producer to forget to bump — + building a new axis is what changes it. + + This exists because a downstream operation that resolves coordinate + *values* into something it then caches — channel labels into array + indices, or into output labels — cannot key that cache on shape alone. + A source that renames, reorders or swaps channels without changing how + many it sends would otherwise keep getting the previous answer, emitting + one channel's samples under another channel's label. Comparing the + arrays outright is O(bytes) on every message and in every consumer; + comparing fingerprints is a tuple compare. + + Computed on first access and cached on the instance, so the cost is paid + once per axis object rather than once per consumer per message. The + cached value is part of ``__dict__``, so it survives pickling and + arrives already computed on the far side of a process boundary. + + ``None`` when the contents cannot be digested (a non-numpy backing + array, or an object dtype holding values with no string form). Callers + must treat ``None`` as "unknown" and fall back to comparing the data — + never as a value that compares equal to another ``None``. + + .. warning:: + Assumes the axis is not mutated after construction. Both + ``ax.data = other`` and ``ax.data[0] = "X"`` leave a stale + fingerprint. Messages fan out to several branches of a graph, so + mutating one in place is already unsafe; this makes it a requirement. + Build a new axis (e.g. via :func:`replace`) instead. + + :return: Hashable content digest, or None if one cannot be computed + :rtype: tuple | None + """ + # _UNSET, not None: a successfully computed fingerprint may itself be + # None, and that result has to be cached too -- recomputing it means + # re-raising out of ascontiguousarray on every access. + cached = self.__dict__.get("_fingerprint", _UNSET) + if cached is _UNSET: + cached = self.__dict__["_fingerprint"] = self._compute_fingerprint() + return cached + + def _compute_fingerprint(self) -> tuple | None: + """Digest the coordinate values, or None if they cannot be digested.""" + try: + data = np.ascontiguousarray(self.data) + except (TypeError, ValueError): + # A device-resident or otherwise non-numpy backing array. Refusing + # to guess is the safe answer; the caller falls back to comparing. + return None + if data.dtype.hasobject: + # An object array's buffer is pointers, so checksumming it directly + # would make two equal axes disagree. Widen to a real dtype first. + try: + data = np.ascontiguousarray(data.astype("U")) + except (TypeError, ValueError): + return None + # crc32 rather than hash(data.tobytes()) because the copy is not the + # bottleneck: on a 256-channel struct axis the copy runs at ~94 GB/s and + # CPython's siphash over the result at ~5.5 GB/s, while crc32 reads the + # array's buffer directly at ~29 GB/s. + # + # dtype goes in as the object, not str(dtype): numpy builds a structured + # dtype's repr field by field, which costs ~10x the checksum it annotates. + return (self.unit, tuple(self.dims), data.dtype, data.shape, zlib.crc32(data)) + @dataclass(eq=False) class AxisArray(ArrayWithNamedDims): diff --git a/src/ezmsg/util/messages/util.py b/src/ezmsg/util/messages/util.py index cd4b492d..61aca2da 100644 --- a/src/ezmsg/util/messages/util.py +++ b/src/ezmsg/util/messages/util.py @@ -5,6 +5,14 @@ T = TypeVar("T") +# Instance attributes that are lazily derived caches rather than dataclass +# fields. fast_replace copies __dict__ wholesale into the constructor, so these +# have to be dropped for two reasons: they are not init parameters (passing one +# raises TypeError), and a value derived from the *old* field values must not be +# carried onto a modified copy. Dropping is always safe -- the copy recomputes +# on next access. Costs ~0.01 us per replace. +_DERIVED_CACHE_ATTRS = ("_fingerprint",) + def fast_replace(arr: T, **kwargs: Any) -> T: """ @@ -30,6 +38,8 @@ def fast_replace(arr: T, **kwargs: Any) -> T: :rtype: T """ out_kwargs = arr.__dict__.copy() # Shallow copy + for name in _DERIVED_CACHE_ATTRS: + out_kwargs.pop(name, None) out_kwargs.update(kwargs) return arr.__class__(**out_kwargs) diff --git a/tests/messages/test_axisarray.py b/tests/messages/test_axisarray.py index 0f79c2b9..21325907 100644 --- a/tests/messages/test_axisarray.py +++ b/tests/messages/test_axisarray.py @@ -6,6 +6,8 @@ from ezmsg.util.messages.axisarray import ( AxisArray, + CoordinateAxis, + replace, shape2d, slice_along_axis, sliding_win_oneaxis, @@ -350,8 +352,8 @@ def test_sliding_win_oneaxis_generic(nwin: int, axis: int, step: int): @pytest.mark.parametrize( "shape,axis,nwin,step", [ - ((100, 64), 0, 50, 1), # (time, channels) — typical EEG window - ((1000, 32), 0, 256, 64), # large time axis with step + ((100, 64), 0, 50, 1), # (time, channels) — typical EEG window + ((1000, 32), 0, 256, 64), # large time axis with step ((8, 1000, 16), 1, 100, 10), # middle axis ], ids=["100x64_win50", "1000x32_win256_step64", "8x1000x16_win100_step10"], @@ -414,3 +416,136 @@ def test_to_xr_dataarray(): assert np.allclose( quality_data.y.data, np.array([-12.6, -12.8, -13.0, -12.4, -12.6]) ) + + +class TestCoordinateAxisFingerprint: + """``CoordinateAxis.fingerprint`` is derived from the contents, not assigned. + + It exists so a consumer that caches something resolved from coordinate + *values* can notice those values changing under a fixed shape, without + paying an O(bytes) comparison per consumer per message. + """ + + @staticmethod + def _axis(labels, **kwargs): + return CoordinateAxis(data=np.array(labels), dims=["ch"], **kwargs) + + def test_equal_contents_agree(self): + """Two separately built but equal axes must agree, or every consumer + would reset on a source that rebuilds its axis per message.""" + assert ( + self._axis(["A", "B", "C"]).fingerprint + == self._axis(["A", "B", "C"]).fingerprint + ) + + def test_different_contents_differ(self): + assert ( + self._axis(["A", "B", "C"]).fingerprint + != self._axis(["X", "Y", "Z"]).fingerprint + ) + + def test_reorder_is_detected(self): + assert ( + self._axis(["A", "B", "C"]).fingerprint + != self._axis(["B", "A", "C"]).fingerprint + ) + + def test_unit_and_dims_are_included(self): + assert ( + self._axis(["A", "B"]).fingerprint + != self._axis(["A", "B"], unit="label").fingerprint + ) + square = np.array([["A", "B"], ["C", "D"]]) + one = CoordinateAxis(data=square, dims=["ch", "x"]) + two = CoordinateAxis(data=square, dims=["ch", "y"]) + assert one.fingerprint != two.fingerprint + + def test_dtype_change_alone_is_detected(self): + wide = CoordinateAxis(data=np.array([1, 2], dtype=np.int64), dims=["ch"]) + narrow = CoordinateAxis(data=np.array([1, 2], dtype=np.int32), dims=["ch"]) + assert wide.fingerprint != narrow.fingerprint + + def test_structured_dtype(self): + dt = np.dtype([("label", "U8"), ("bank", "U2")]) + first = np.array([("e1", "A"), ("e2", "B")], dtype=dt) + same = np.array([("e1", "A"), ("e2", "B")], dtype=dt) + other = np.array([("e1", "A"), ("e2", "C")], dtype=dt) + assert ( + CoordinateAxis(data=first, dims=["ch"]).fingerprint + == CoordinateAxis(data=same, dims=["ch"]).fingerprint + ) + assert ( + CoordinateAxis(data=first, dims=["ch"]).fingerprint + != CoordinateAxis(data=other, dims=["ch"]).fingerprint + ) + + def test_object_dtype_is_content_based(self): + """An object array's buffer holds pointers, so digesting it directly + would make two equal axes disagree and reset consumers every message.""" + one = CoordinateAxis(data=np.array(["c1", "c2"], dtype=object), dims=["ch"]) + two = CoordinateAxis( + data=np.array(["".join(("c", str(i))) for i in (1, 2)], dtype=object), + dims=["ch"], + ) + assert one.fingerprint == two.fingerprint + three = CoordinateAxis(data=np.array(["c1", "c9"], dtype=object), dims=["ch"]) + assert one.fingerprint != three.fingerprint + + def test_non_contiguous_data(self): + """A strided view has no C-contiguous buffer; the digest must gather.""" + strided = np.array(["a", "X", "b", "X", "c", "X"])[::2] + assert not strided.flags["C_CONTIGUOUS"] + assert ( + CoordinateAxis(data=strided, dims=["ch"]).fingerprint + == self._axis(["a", "b", "c"]).fingerprint + ) + + def test_undigestable_contents_report_none(self): + """None means 'unknown', so a caller falls back to comparing rather + than treating two unknowns as equal.""" + + class NoStringForm: + def __str__(self): + raise ValueError("cannot stringify") + + __repr__ = __str__ + + axis = self._axis(["a", "b"]) + axis.data = np.array([NoStringForm(), NoStringForm()], dtype=object) + axis.__dict__.pop("_fingerprint", None) + assert axis.fingerprint is None + + def test_cached_on_the_instance(self): + axis = self._axis([f"e{i:03d}" for i in range(64)]) + assert axis.fingerprint is axis.fingerprint + assert "_fingerprint" in axis.__dict__ + + def test_is_hashable(self): + """Consumers fold it into hash((key, shape, fingerprint)).""" + hash(("key", (30, 3), self._axis(["A", "B", "C"]).fingerprint)) + + def test_survives_pickling(self): + """The cached value crosses a process boundary with the data, so the + far side never recomputes it.""" + import pickle + + axis = self._axis([f"e{i:03d}" for i in range(64)]) + expected = axis.fingerprint + msg = AxisArray( + np.zeros((4, 64)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=10.0), "ch": axis}, + key="k", + ) + restored = pickle.loads(pickle.dumps(msg)).axes["ch"] + assert restored is not axis + assert restored.__dict__.get("_fingerprint") is not None # arrived precomputed + assert restored.fingerprint == expected + + def test_replace_yields_a_fresh_fingerprint(self): + """``replace`` builds a new axis, so the cache cannot leak across.""" + axis = self._axis(["A", "B"]) + _ = axis.fingerprint + updated = replace(axis, data=np.array(["X", "Y"])) + assert updated.fingerprint != axis.fingerprint + diff --git a/tests/messages/test_replace.py b/tests/messages/test_replace.py new file mode 100644 index 00000000..475c64c8 --- /dev/null +++ b/tests/messages/test_replace.py @@ -0,0 +1,59 @@ +"""Tests for ezmsg.util.messages.util.replace.""" + +import numpy as np +import pytest + +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis +from ezmsg.util.messages.util import fast_replace, slow_replace + + +def _axis(labels): + return CoordinateAxis(data=np.array(labels), dims=["ch"]) + + +@pytest.mark.parametrize( + "replace_fn", [fast_replace, slow_replace], ids=["fast", "slow"] +) +class TestReplaceWithDerivedCaches: + """A lazily derived cache on the instance must not reach the constructor. + + ``fast_replace`` copies ``__dict__`` straight into ``__class__(**kwargs)``, + so a cached attribute that is not a dataclass field would raise TypeError. + It must also not be *carried over*: a value derived from the old field + values would be wrong on a copy that changes them. + """ + + def test_replace_after_fingerprint_access(self, replace_fn): + axis = _axis(["A", "B"]) + assert axis.fingerprint is not None # populates the cache + updated = replace_fn(axis, data=np.array(["X", "Y"])) + assert list(updated.data) == ["X", "Y"] + + def test_stale_fingerprint_is_not_carried_over(self, replace_fn): + axis = _axis(["A", "B"]) + before = axis.fingerprint + updated = replace_fn(axis, data=np.array(["X", "Y"])) + assert updated.fingerprint != before + + def test_unrelated_field_change_still_refreshes(self, replace_fn): + """``unit`` is part of the fingerprint, so it has to be recomputed even + though the data did not change.""" + axis = _axis(["A", "B"]) + before = axis.fingerprint + updated = replace_fn(axis, unit="label") + assert updated.fingerprint != before + + def test_axisarray_replace_is_unaffected(self, replace_fn): + axis = _axis(["A", "B"]) + _ = axis.fingerprint + msg = AxisArray( + np.zeros((4, 2)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=10.0), "ch": axis}, + key="k", + ) + updated = replace_fn(msg, data=np.ones((4, 2))) + assert updated.key == "k" + # The axis object is passed through by reference, cache intact. + assert updated.axes["ch"] is axis + assert updated.axes["ch"].fingerprint == axis.fingerprint From b3a9f51481d4d062549d1e0f1047fff34605e1aa Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Wed, 2 Sep 2026 23:14:25 -0400 Subject: [PATCH 2/4] Fix CoordinateAxis.__eq__ comparing only `unit` CoordinateAxis inherits from two dataclasses that each supply an __eq__, and the MRO picks the wrong one: CoordinateAxis -> AxisBase -> ABC -> ArrayWithNamedDims AxisBase is a plain @dataclass, so it generates an __eq__ over its only field, `unit`, which shadows the content comparison in ArrayWithNamedDims. Two coordinate axes therefore compared equal whatever their coordinate values were: axis(["A", "B", "C"]) == axis(["X", "Y", "Z"]) # True and because AxisArray.__eq__ tests `self.axes == other.axes`, so did two messages differing only in their channel labels. Defines __eq__ explicitly on CoordinateAxis, comparing unit, dims and values, and following ArrayWithNamedDims' convention of returning True or NotImplemented so that AxisArray.__eq__'s dict comparison still resolves correctly. LinearAxis equality is unaffected and the class remains unhashable, as before. --- src/ezmsg/util/messages/axisarray.py | 24 +++++++++++ tests/messages/test_axisarray.py | 61 ++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/src/ezmsg/util/messages/axisarray.py b/src/ezmsg/util/messages/axisarray.py index 694677d2..f761c900 100644 --- a/src/ezmsg/util/messages/axisarray.py +++ b/src/ezmsg/util/messages/axisarray.py @@ -191,6 +191,30 @@ def value(self, x): """ return self.data[x] + def __eq__(self, other): + """ + Compare unit, dims and coordinate values. + + Defined explicitly because this class inherits from two dataclasses that + both supply an ``__eq__``, and the MRO picks the wrong one: + ``CoordinateAxis -> AxisBase -> ABC -> ArrayWithNamedDims``. ``AxisBase`` + is a plain ``@dataclass``, so it generates an ``__eq__`` over its only + field, ``unit``, and that shadows the content comparison in + ``ArrayWithNamedDims``. Two coordinate axes sharing a unit therefore + compared equal whatever their coordinate values were -- and, since + ``AxisArray.__eq__`` tests ``self.axes == other.axes``, so did two + messages differing only in their channel labels. + + Follows the same convention as ``ArrayWithNamedDims.__eq__``: True, or + NotImplemented for anything else, which Python resolves to False once + the reflected call also declines. + """ + if self is other: + return True + if other.__class__ is self.__class__ and self.unit == other.unit: + return ArrayWithNamedDims.__eq__(self, other) + return NotImplemented + @property def fingerprint(self) -> tuple | None: """ diff --git a/tests/messages/test_axisarray.py b/tests/messages/test_axisarray.py index 21325907..a1724b1b 100644 --- a/tests/messages/test_axisarray.py +++ b/tests/messages/test_axisarray.py @@ -7,6 +7,7 @@ from ezmsg.util.messages.axisarray import ( AxisArray, CoordinateAxis, + LinearAxis, replace, shape2d, slice_along_axis, @@ -549,3 +550,63 @@ def test_replace_yields_a_fresh_fingerprint(self): updated = replace(axis, data=np.array(["X", "Y"])) assert updated.fingerprint != axis.fingerprint + +class TestCoordinateAxisEquality: + """``CoordinateAxis`` compares its coordinate values, not just its unit. + + It inherits an ``__eq__`` from two dataclasses; the MRO puts ``AxisBase`` + (which compares only ``unit``) ahead of ``ArrayWithNamedDims`` (which + compares contents), so without an explicit ``__eq__`` any two axes sharing a + unit compared equal. + """ + + @staticmethod + def _axis(labels, **kwargs): + return CoordinateAxis(data=np.array(labels), dims=["ch"], **kwargs) + + def test_equal_values(self): + assert self._axis(["A", "B"]) == self._axis(["A", "B"]) + + def test_different_values(self): + assert self._axis(["A", "B"]) != self._axis(["X", "Y"]) + + def test_reordered_values(self): + assert self._axis(["A", "B"]) != self._axis(["B", "A"]) + + def test_different_length(self): + assert self._axis(["A", "B"]) != self._axis(["A", "B", "C"]) + + def test_different_unit(self): + assert self._axis(["A", "B"]) != self._axis(["A", "B"], unit="label") + + def test_different_dims(self): + assert self._axis(["A", "B"]) != CoordinateAxis( + data=np.array(["A", "B"]), dims=["x"] + ) + + def test_identity(self): + axis = self._axis(["A", "B"]) + assert axis == axis # noqa: PLR0124 -- exercises the `self is other` fast path + + def test_other_axis_type(self): + assert self._axis(["A", "B"]) != LinearAxis(gain=1.0) + + def test_linear_axis_equality_is_unaffected(self): + assert LinearAxis(gain=2.0, offset=1.0) == LinearAxis(gain=2.0, offset=1.0) + assert LinearAxis(gain=2.0) != LinearAxis(gain=3.0) + + def test_axisarray_sees_a_relabelled_channel_axis(self): + """The consequence that motivated the fix: ``AxisArray.__eq__`` tests + ``self.axes == other.axes``, so a shadowed axis comparison made two + messages differing only in channel labels compare equal.""" + + def msg(labels): + return AxisArray( + np.zeros((4, 2)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=100.0), "ch": self._axis(labels)}, + key="k", + ) + + assert msg(["A", "B"]) == msg(["A", "B"]) + assert msg(["A", "B"]) != msg(["X", "Y"]) From f4c6342ad0dbba2041a40562062bd6b7b0173517 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Thu, 3 Sep 2026 01:25:57 -0400 Subject: [PATCH 3/4] Add AxisArray.chunk_dim, naming the dimension messages append along Successive messages in a stream concatenate along one dimension, and that one is different in kind from the others: its length is however much arrived this time, a LinearAxis there has an offset that advances every message, and a CoordinateAxis there (irregular event times) has per-message values. Everything else describes the stream's configuration and is stable between reconfigurations. Consumers that cache state keyed on the message layout need that distinction and cannot reliably infer it. It is "time" on a raw signal but "win" downstream of a windowing stage, so a consumer that assumes "time" both thrashes on the window count and stops noticing a change in the window length. Taking dims[0] instead trades that for breaking under transpose, where it would exclude the channel dimension -- the very thing worth watching. Only the producer, which named the dims, reliably knows. Declaring it on the message puts the answer in that one place instead of asking each consumer to guess about a message it did not create. The field defaults to None, meaning "not declared", so existing producers are unaffected and consumers can tell silence from a real declaration. It is placed last, so positional construction is unchanged, and fast_replace carries it across a `replace(msg, data=...)` that keeps the same layout. Defaulting it to "time" instead was tried and rejected: 13 tests here and 91 in ezmsg-sigproc construct messages with no time dimension at all -- (win, freq, ch), (freq, ch), (bin, ch), (epoch, target_freq) -- so every one would have to opt out explicitly. Worse, a post-Window (win, time, ch) message *does* contain "time", so the default would pass validation while naming the wrong dimension, reintroducing the exact failure this field exists to prevent in the one case validation cannot catch. Validation lives in ArrayWithNamedDims.__post_init__ rather than an AxisArray override: reusing the existing frame costs one getattr (+0.036 us on message construction) where a second __post_init__ call costs +0.052, and construction is hot. It makes a rename that forgets to update chunk_dim raise at the site that has to handle it. --- src/ezmsg/util/messages/axisarray.py | 32 +++++++++++++++++ tests/messages/test_axisarray.py | 53 ++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/src/ezmsg/util/messages/axisarray.py b/src/ezmsg/util/messages/axisarray.py index f761c900..eeecd29d 100644 --- a/src/ezmsg/util/messages/axisarray.py +++ b/src/ezmsg/util/messages/axisarray.py @@ -146,6 +146,15 @@ def __post_init__(self): raise ValueError("dims must be same length as data.shape") if len(self.dims) != len(set(self.dims)): raise ValueError("dims contains repeated dim names") + # Checked here rather than in an AxisArray override: reusing this frame + # costs one getattr instead of a second __post_init__ call, and message + # construction is hot. + chunk_dim = getattr(self, "chunk_dim", None) + if chunk_dim is not None and chunk_dim not in self.dims: + raise ValueError( + f"chunk_dim {chunk_dim!r} is not one of dims {self.dims}. " + "An operation that renames this dimension must update chunk_dim too." + ) def __eq__(self, other): if self is other: @@ -307,12 +316,35 @@ class AxisArray(ArrayWithNamedDims): :type attrs: dict[str, typing.Any] :param key: Optional key identifier for this array, typically used to specify source device (default is empty string) :type key: str + :param chunk_dim: Name of the dimension this message is a chunk along, or None if not declared + :type chunk_dim: str | None """ axes: dict[str, AxisBase] = field(default_factory=dict) attrs: dict[str, typing.Any] = field(default_factory=dict) key: str = "" + chunk_dim: str | None = None + """The dimension this message is a *chunk* along. + + Successive messages in a stream append to one another along this dimension, + so it is the one whose extent is arbitrary: its length is however much + arrived this time, a ``LinearAxis`` here has an ``offset`` that advances + every message, and a ``CoordinateAxis`` here (irregular event times) has + per-message *values*. Everything else describes the stream's configuration + and is stable between reconfigurations. + + Consumers that cache state keyed on the message layout need that + distinction, and cannot reliably infer it: it is ``"time"`` on a raw signal + but ``"win"`` downstream of a windowing stage, and taking ``dims[0]`` breaks + under :meth:`transpose`. Declaring it here puts the answer where it is + known -- in the producer -- instead of asking every consumer to guess. + + ``None`` means "not declared", leaving consumers to fall back on their own + convention. Any operation that *renames* this dimension is responsible for + updating it, exactly as it already updates ``dims``. + """ + T = typing.TypeVar("T", bound="AxisArray") def __eq__(self, other): diff --git a/tests/messages/test_axisarray.py b/tests/messages/test_axisarray.py index a1724b1b..31ecfd01 100644 --- a/tests/messages/test_axisarray.py +++ b/tests/messages/test_axisarray.py @@ -610,3 +610,56 @@ def msg(labels): assert msg(["A", "B"]) == msg(["A", "B"]) assert msg(["A", "B"]) != msg(["X", "Y"]) + + +class TestChunkDim: + """`chunk_dim` names the dimension successive messages append along. + + Its extent is whatever arrived this time, so consumers that cache state + keyed on the message layout have to treat it differently from the + dimensions that describe the stream's configuration. + """ + + @staticmethod + def _msg(**kwargs): + return AxisArray( + np.zeros((4, 2)), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=100.0)}, + **kwargs, + ) + + def test_defaults_to_none(self): + """Undeclared, so existing producers are unaffected.""" + assert self._msg().chunk_dim is None + + def test_round_trips(self): + assert self._msg(chunk_dim="time").chunk_dim == "time" + + def test_must_name_an_actual_dim(self): + with pytest.raises(ValueError, match="chunk_dim 'nope' is not one of dims"): + self._msg(chunk_dim="nope") + + def test_replace_carries_it(self): + """A transformer that only changes data keeps the same layout.""" + original = self._msg(chunk_dim="time") + assert replace(original, data=np.ones((4, 2))).chunk_dim == "time" + + def test_a_rename_must_update_it(self): + """Renaming the chunk dimension without updating chunk_dim is caught.""" + original = self._msg(chunk_dim="time") + with pytest.raises(ValueError, match="must update chunk_dim"): + replace(original, dims=["win", "ch"]) + assert replace(original, dims=["win", "ch"], chunk_dim="win").chunk_dim == "win" + + def test_survives_pickling(self): + import pickle + + assert ( + pickle.loads(pickle.dumps(self._msg(chunk_dim="time"))).chunk_dim == "time" + ) + + def test_positional_construction_is_unaffected(self): + """chunk_dim is last, so existing positional calls still work.""" + msg = AxisArray(np.zeros((4, 2)), ["time", "ch"], {}, {}, "key") + assert msg.key == "key" and msg.chunk_dim is None From 3e8bad63096e542a5eda59468b932fd78a088744 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Thu, 3 Sep 2026 09:03:07 -0400 Subject: [PATCH 4/4] Carry chunk_dim through ModifyAxisTransformer `chunk_dim` names a dimension, so a transformer that renames or drops dimensions has to move it along with the dimension it points at. Left alone it is either stale-but-fatal -- `replace()` raises because the old name is no longer in `dims` -- or, on a swap like `{"win": "time", "time": "sample"}`, silently correct-looking while naming a different dimension than before. The first form showed up as 26 failures in ezmsg-sigproc's spectrogram and single-band-power composites once `Window` started declaring `chunk_dim="win"` on its output. Dropping a dimension clears `chunk_dim` rather than leaving it dangling. --- src/ezmsg/util/messages/modify.py | 10 ++++++- tests/messages/test_modify.py | 48 +++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/ezmsg/util/messages/modify.py b/src/ezmsg/util/messages/modify.py index 0f767330..b60d2866 100644 --- a/src/ezmsg/util/messages/modify.py +++ b/src/ezmsg/util/messages/modify.py @@ -47,6 +47,13 @@ def __call__(self, message: AxisArray) -> AxisArray: for ix, old_dim in enumerate(message.dims) if new_dims[ix] is None ] + # A renamed or dropped dimension takes chunk_dim with it: the field + # names a dimension, so leaving it pointing at the old name would either + # raise (the name is gone from dims) or silently describe the wrong one. + chunk_dim = message.chunk_dim + if chunk_dim is not None: + chunk_dim = name_map.get(chunk_dim, chunk_dim) + if len(drop_ax_ix) > 0: new_dims = [d for d in new_dims if d is not None] new_axes.pop(None, None) @@ -55,8 +62,9 @@ def __call__(self, message: AxisArray) -> AxisArray: data=np.squeeze(message.data, axis=tuple(drop_ax_ix)), dims=new_dims, axes=new_axes, + chunk_dim=chunk_dim, ) - return replace(message, dims=new_dims, axes=new_axes) + return replace(message, dims=new_dims, axes=new_axes, chunk_dim=chunk_dim) _send_warned: bool = False diff --git a/tests/messages/test_modify.py b/tests/messages/test_modify.py index 91f3c708..ec8fcb70 100644 --- a/tests/messages/test_modify.py +++ b/tests/messages/test_modify.py @@ -66,3 +66,51 @@ def test_drop_axis(targ_dim_len: int): assert "ch" in res.dims assert "ch" in res.axes assert res.data.shape == (5, 4) + + +class TestChunkDimIsRemapped: + """`chunk_dim` names a dimension, so renaming that dimension has to move it. + + Left alone it would either point at a name no longer in `dims` -- which + AxisArray rejects at construction -- or, worse, silently name whichever + other dimension inherited the old name. + """ + + @staticmethod + def _msg(chunk_dim: str | None, win_len: int = 3): + return AxisArray( + data=np.arange(win_len * 4 * 2).reshape(win_len, 4, 2), + dims=["win", "time", "ch"], + axes={ + "win": AxisArray.TimeAxis(fs=10.0), + "time": AxisArray.TimeAxis(fs=100.0), + "ch": AxisArray.CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"]), + }, + key="test_chunk_dim", + chunk_dim=chunk_dim, + ) + + def test_renaming_the_chunk_dim_moves_it(self): + res = modify_axis({"win": "batch"}).send(self._msg("win")) + assert res.dims == ["batch", "time", "ch"] + assert res.chunk_dim == "batch" + + def test_a_swap_follows_the_dimension_not_the_name(self): + """The case that motivated this: a windowing stage emits `win` as the + chunk dimension and a later stage swaps the two time-like names.""" + res = modify_axis({"win": "time", "time": "sample"}).send(self._msg("win")) + assert res.dims == ["time", "sample", "ch"] + assert res.chunk_dim == "time" + + def test_renaming_another_dim_leaves_it_alone(self): + res = modify_axis({"ch": "channel"}).send(self._msg("win")) + assert res.dims == ["win", "time", "channel"] + assert res.chunk_dim == "win" + + def test_undeclared_stays_undeclared(self): + assert modify_axis({"win": "batch"}).send(self._msg(None)).chunk_dim is None + + def test_dropping_the_chunk_dim_clears_it(self): + res = modify_axis({"win": None}).send(self._msg("win", win_len=1)) + assert res.dims == ["time", "ch"] + assert res.chunk_dim is None