diff --git a/src/ezmsg/util/messages/axisarray.py b/src/ezmsg/util/messages/axisarray.py index 7983fa0a..271b07c6 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: """ @@ -141,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: @@ -186,6 +200,101 @@ 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: + """ + 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): @@ -207,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): @@ -523,6 +655,13 @@ def iter_over_axis(self: T, axis: str | int) -> Generator[T, None, None]: Yields AxisArray objects for each slice along the given axis, with that dimension removed from the resulting arrays. + Iterating over the chunk dimension consumes it: each slice is a single + element along it, not a chunk of a stream that accumulates there. The + yielded arrays therefore declare no :attr:`chunk_dim` -- keeping the old + one would name a dimension that is no longer in ``dims``, which + :meth:`__post_init__` rejects outright. A caller that knows what the + slices are chunks along should say so on the way out. + :param axis: Dimension name or index to iterate over :type axis: str | int :yields: AxisArray objects for each slice along the axis @@ -533,6 +672,7 @@ def iter_over_axis(self: T, axis: str | int) -> Generator[T, None, None]: dim_name = self.dims[axis_idx] new_dims = [d for i, d in enumerate(self.dims) if i != axis_idx] new_axes = {d: a for d, a in self.axes.items() if d != dim_name} + new_chunk_dim = None if self.chunk_dim == dim_name else self.chunk_dim for it_data in xp.moveaxis(self.data, axis_idx, 0): it_aa = replace( @@ -540,6 +680,7 @@ def iter_over_axis(self: T, axis: str | int) -> Generator[T, None, None]: data=it_data, dims=new_dims, axes=new_axes, + chunk_dim=new_chunk_dim, ) yield it_aa 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/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..84126c20 100644 --- a/tests/messages/test_axisarray.py +++ b/tests/messages/test_axisarray.py @@ -6,6 +6,9 @@ from ezmsg.util.messages.axisarray import ( AxisArray, + CoordinateAxis, + LinearAxis, + replace, shape2d, slice_along_axis, sliding_win_oneaxis, @@ -350,8 +353,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 +417,302 @@ 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 + + +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"]) + + +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 + + +class TestIterOverAxisAndChunkDim: + """Iterating a dimension away has to take ``chunk_dim`` with it. + + ``iter_over_axis`` is the one method that *removes* a dimension -- + ``isel``/``sel`` take along it and ``transpose`` only reorders -- so it is + the one place a stale ``chunk_dim`` can be left naming a dimension that no + longer exists. + """ + + @staticmethod + def _msg(chunk_dim): + return AxisArray( + np.arange(24, dtype=float).reshape(3, 4, 2), + dims=["win", "time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=100.0)}, + key="dev", + chunk_dim=chunk_dim, + ) + + def test_iterating_the_chunk_dim_clears_it(self): + """Each slice is one element along ``win``, not a chunk accumulating + there. Keeping the declaration would name a missing dim, which + ``__post_init__`` rejects -- so this used to raise rather than yield.""" + for sub in self._msg(chunk_dim="win").iter_over_axis("win"): + assert sub.dims == ["time", "ch"] + assert sub.chunk_dim is None + + def test_iterating_another_dim_keeps_it(self): + """``win`` still accumulates, and it is still present, so the + declaration is still true.""" + for sub in self._msg(chunk_dim="win").iter_over_axis("time"): + assert sub.dims == ["win", "ch"] + assert sub.chunk_dim == "win" + + def test_an_undeclared_chunk_dim_stays_undeclared(self): + for sub in self._msg(chunk_dim=None).iter_over_axis("win"): + assert sub.chunk_dim is None + + def test_the_data_is_unchanged(self): + """The fix is bookkeeping only.""" + subs = list(self._msg(chunk_dim="win").iter_over_axis("win")) + assert len(subs) == 3 + assert np.array_equal( + subs[1].data, np.arange(24, dtype=float).reshape(3, 4, 2)[1] + ) + + def test_a_caller_can_declare_what_the_slices_are_chunks_along(self): + """Clearing it is the safe default, not the last word: unbundling + windows yields messages that accumulate along the within-window axis.""" + sub = next(self._msg(chunk_dim="win").iter_over_axis("win")) + assert replace(sub, chunk_dim="time").chunk_dim == "time" 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 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