From 3a56c4d0ef4378d4deb46bc625055efa651dc290 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Fri, 4 Sep 2026 00:25:00 -0400 Subject: [PATCH] Key state resets on the message layout, via baseproc's axis-aware default ezmsg-baseproc 1.12.0's default `_hash_message` folds in the message key, the dims, the length of every dimension except the chunk dimension, the coordinate *values* on those dimensions, and the gain and offset of any linear axis among them. 23 of the 46 stateful processors here were hashing a strict subset of that and now inherit it; 3 keep an override for something the default cannot see; 1 keeps one for a reason the default gets wrong. What the deletions buy is the channel fingerprint. A source that renames or reorders channels without changing how many it sends -- a device reconfigured mid-session, a montage swapped -- was invisible to these hashes, so per-channel state carried onto channels it did not belong to. For a filter that state is numeric, and the result is not subtle: first 4 samples of the new channel 'armB-1': filter state carried from armA: [-5.919 -4.772 -0.5114 1.064] with a correct reset: [ 0. -0.0013 -0.0023 -0.0026] max |difference| = 11.12 vs new-data amplitude 0.024 Nothing in the output announces it. `tests/unit/test_state_reset_semantics.py` pins the behaviour that fell out, including that same filter case reproduced end to end. Each deletion was checked rather than assumed: a harness ran every override and the default over one stream and compared their reset points, and an override was only removed where the default's set was a superset. The four kept: * `align` -- two alternating input streams; the default's key sensitivity would reset it on every message. * `spectrum` -- the FFT is sized by the chunk dimension, the one length the default deliberately ignores, so it folds it back in via `extra=`. * `filterbank`, `wavelets` -- dtype, which the default cannot see. Three processors gain a *narrower* hash than the default, because their state genuinely depends on less: `AdaptiveStandardScaler` owns no arrays (its two child EWMATransformers hash themselves, verified bit-identical after a relabel), and `Downsample` and `FilterbankDesign` derive everything from one axis' gain. `AxisArray.chunk_dim` is declared where this package creates or consumes one: `Window` names its new axis, `Spectrum` and `Aggregate` clear it when they consume the dimension, `Flatten` follows a renamed preserve axis, `Concat` carries it. Without that the default has to guess, and a wrong guess is not a small error -- it either thrashes on chunk-size jitter or stops noticing real changes. Coordinate axes this package builds are handed over primed, via a new `util.message.with_fingerprint`. The digest is cached on the axis and pickles with it, so computing it at construction spares the first consumer in every receiving process from recomputing it on every message -- unpickling builds a new axis object per message, so a cold axis is re-checksummed forever. Deliberately not primed: coordinate axes along the chunk dimension, whose values are per-message and whose fingerprint no consumer reads. Two tests now assert the opposite of what they did, both deliberately: `test_common_rereference_field_values_change_is_not_detected` and flatten's `test_labels_outside_flatten_axes_do_not_reset`. Both documented a concession made to avoid an O(bytes) per-message cost; `CoordinateAxis.fingerprint` removes that cost by computing the digest once per axis object rather than once per consumer, so the concession is no longer worth making. `tests/helpers/recycled_shm.py` had a latent gap of its own: `_detach` did not copy axes, so it could not have detected a retained axis. Fixed, and verified that it now catches one. Requires ezmsg 3.10.0b2 and ezmsg-baseproc 1.12.0. 4247 passed, up from 4184. --- benchmarks/benchmark_axis_fingerprint.py | 279 ++++++++++ benchmarks/benchmark_concat_fingerprint.py | 224 ++++++++ benchmarks/benchmark_hash_overhead.py | 512 ++++++++++++++++++ benchmarks/benchmark_hash_witness.py | 218 ++++++++ benchmarks/benchmark_state_resets.py | 267 +++++++++ docs/source/conf.py | 1 + .../guides/sigproc/axis_fingerprint.rst | 142 +++++ .../source/guides/sigproc/content-sigproc.rst | 1 + pyproject.toml | 5 +- src/ezmsg/sigproc/adaptive_lattice_notch.py | 5 - src/ezmsg/sigproc/adaptive_lnc.py | 5 - src/ezmsg/sigproc/affinetransform.py | 21 +- src/ezmsg/sigproc/aggregate.py | 25 +- src/ezmsg/sigproc/align.py | 18 +- src/ezmsg/sigproc/binned_aggregate.py | 13 +- src/ezmsg/sigproc/concat.py | 110 +++- src/ezmsg/sigproc/coordinatespaces.py | 4 +- src/ezmsg/sigproc/denormalize.py | 14 +- src/ezmsg/sigproc/diff.py | 5 - src/ezmsg/sigproc/ewma.py | 6 - src/ezmsg/sigproc/fbcca.py | 5 +- src/ezmsg/sigproc/filter.py | 13 - src/ezmsg/sigproc/filterbank.py | 18 +- src/ezmsg/sigproc/filterbankdesign.py | 10 +- src/ezmsg/sigproc/fir_hilbert.py | 7 - src/ezmsg/sigproc/flatten.py | 30 +- src/ezmsg/sigproc/linear.py | 8 - src/ezmsg/sigproc/resample.py | 7 - src/ezmsg/sigproc/rollingscaler.py | 7 - src/ezmsg/sigproc/sampler.py | 7 - src/ezmsg/sigproc/scaler.py | 8 + src/ezmsg/sigproc/signalinjector.py | 5 - src/ezmsg/sigproc/slicer.py | 9 +- src/ezmsg/sigproc/spectrum.py | 28 +- src/ezmsg/sigproc/transpose.py | 3 - src/ezmsg/sigproc/util/channels.py | 245 +++++++++ src/ezmsg/sigproc/util/message.py | 25 + src/ezmsg/sigproc/wavelets.py | 15 +- src/ezmsg/sigproc/window.py | 33 +- tests/helpers/recycled_shm.py | 14 +- tests/unit/test_affine_transform.py | 72 ++- tests/unit/test_axis_fingerprint_priming.py | 182 +++++++ tests/unit/test_buffer_recycling.py | 84 +++ tests/unit/test_concat.py | 134 +++++ tests/unit/test_flatten.py | 70 +++ tests/unit/test_state_reset_semantics.py | 167 ++++++ tests/unit/test_util_channels.py | 137 +++++ 47 files changed, 2997 insertions(+), 221 deletions(-) create mode 100644 benchmarks/benchmark_axis_fingerprint.py create mode 100644 benchmarks/benchmark_concat_fingerprint.py create mode 100644 benchmarks/benchmark_hash_overhead.py create mode 100644 benchmarks/benchmark_hash_witness.py create mode 100644 benchmarks/benchmark_state_resets.py create mode 100644 docs/source/guides/sigproc/axis_fingerprint.rst create mode 100644 tests/unit/test_axis_fingerprint_priming.py create mode 100644 tests/unit/test_state_reset_semantics.py diff --git a/benchmarks/benchmark_axis_fingerprint.py b/benchmarks/benchmark_axis_fingerprint.py new file mode 100644 index 00000000..eb593ad4 --- /dev/null +++ b/benchmarks/benchmark_axis_fingerprint.py @@ -0,0 +1,279 @@ +"""What it costs to notice that a coordinate axis changed, measured. + +A stateful transformer that resolves channel *labels* to array *indices* has to +decide how much of the message to fold into ``_hash_message``. Fold too little +and a source that renames or reorders channels under a fixed key and channel +count keeps getting the previous message's indices -- one channel's samples +emitted under another channel's label (ezmsg-org/ezmsg-sigproc#232). Fold too +much and every message pays for it. + +:mod:`ezmsg.sigproc.util.channels` offers both answers -- +``group_spec_fingerprint`` (O(1), field presence only) and +``coord_value_fingerprint`` (O(bytes), actual values) -- and this script is +where the numbers in their docstrings come from. + +Three results drive how ``coord_value_fingerprint`` is written: + +* **The checksum dominates, not the copy.** ``tobytes()`` runs at ~94 GB/s; + CPython's siphash over the result manages ~5.5 GB/s. ``zlib.crc32`` reads the + array buffer directly at ~29 GB/s, so it is ~5x cheaper end to end. +* **Restricting to one field is not reliably cheaper.** Extracting a field from + a struct array is a strided gather; a wide field (U16 ``label``, 59% of the + itemsize) costs more than checksumming the whole contiguous axis. The + restriction is for invalidation *correctness* -- not resetting when an unread + ``x``/``y`` field churns -- and is only sometimes also a speedup. +* **Never ask numpy for several fields at once.** ``arr[['array', 'bank']]`` + returns a view that keeps the *original* itemsize, so its ``tobytes()`` is the + whole array. Two of eight fields would cost more than all eight. + +Run from the repository root:: + + uv run python benchmarks/benchmark_axis_fingerprint.py + uv run python benchmarks/benchmark_axis_fingerprint.py --n-ch 1024 + uv run python benchmarks/benchmark_axis_fingerprint.py --sections mechanics,scaling +""" + +from __future__ import annotations + +import argparse +import timeit +import zlib + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray + +from ezmsg.sigproc.util.channels import coord_value_fingerprint, group_spec_fingerprint + +# ezmsg-blackrock's ChannelMap ``ch`` axis: the "full metadata" case. 108 B per +# channel, of which the U16 label is 64 B. +CHANNELMAP_DTYPE = np.dtype( + [ + ("label", "U16"), + ("x", " AxisArray.CoordinateAxis: + data = np.zeros(n_ch, dtype=CHANNELMAP_DTYPE) + data["label"] = [f"elec{i:04d}" for i in range(n_ch)] + data["x"] = np.arange(n_ch, dtype=np.float64) + data["y"] = np.arange(n_ch, dtype=np.float64) + data["size"] = 1.0 + data["array"] = np.arange(n_ch) // 128 + data["bank"] = np.array([("A", "B", "C", "D")[i // 64 % 4] for i in range(n_ch)]) + data["elec"] = np.arange(n_ch, dtype=np.int32) + 1 + data["headstage"] = np.arange(n_ch) // 32 + return AxisArray.CoordinateAxis(data=data, dims=["ch"]) + + +def label_axis(n_ch: int) -> AxisArray.CoordinateAxis: + """The plain (unstructured) label axis most sources emit.""" + return AxisArray.CoordinateAxis(data=np.array([f"ch{i:04d}" for i in range(n_ch)]), dims=["ch"]) + + +def make_msg(ch_axis, n_ch: int, n_times: int) -> AxisArray: + return AxisArray( + np.zeros((n_times, n_ch), dtype=np.float32), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=30000.0), "ch": ch_axis}, + key="dev", + attrs={"source": "bench", "session": 3}, + ) + + +def bench(label: str, fn, number: int = 50_000, nbytes: int | None = None) -> float: + fn() # warm + us = timeit.timeit(fn, number=number) / number * 1e6 + rate = f"{nbytes / us / 1000:8.2f} GB/s" if nbytes else "" + print(f" {label:<54} {us:8.3f} us {rate}") + return us + + +def concat_fingerprint(msg: AxisArray, concat_dim: str) -> tuple: + """Verbatim from ``ConcatTransformer._fingerprint`` (concat.py), for comparison. + + Note that concat calls this on *both* of its inputs, so its per-message cost + is twice what is reported here. + """ + ax = msg.axes.get(concat_dim) + ax_hash = hash(ax.data.tobytes()) if ax is not None and hasattr(ax, "data") else None + attrs_fp = frozenset((k, type(v).__name__, repr(v)) for k, v in (msg.attrs or {}).items()) + return (tuple(msg.dims), msg.data.shape, ax_hash, attrs_fp) + + +def section_strategies(struct_msg: AxisArray, plain_msg: AxisArray) -> None: + """Every candidate answer, against the baseline hash it would replace.""" + struct = struct_msg.axes["ch"].data + plain = plain_msg.axes["ch"].data + n_ch = struct_msg.data.shape[1] + + print("\n== strategies ==") + print("\n-- O(1) baselines (what transformers pay today) --") + base = bench("hash((key, n_ch)) [SlicerTransformer today]", lambda: hash((struct_msg.key, n_ch))) + bench("group_spec_fingerprint(msg, 'ch', None) [default]", lambda: group_spec_fingerprint(struct_msg, "ch", None)) + bench("group_spec_fingerprint(msg, 'ch', 'bank')", lambda: group_spec_fingerprint(struct_msg, "ch", "bank")) + + print("\n-- fold the values in: hash(tobytes()) --") + naive = bench("hash(struct.tobytes()) whole ChannelMap", lambda: hash(struct.tobytes())) + bench("hash(labels.tobytes()) plain label axis", lambda: hash(plain.tobytes())) + bench("hash(struct['bank'].tobytes()) one U2 field", lambda: hash(struct["bank"].tobytes())) + bench("hash(struct['label'].tobytes()) one U16 field", lambda: hash(struct["label"].tobytes())) + bench( + "hash(struct[['array','bank']].tobytes()) TRAP: 2 fields", + lambda: hash(struct[["array", "bank"]].tobytes()), + ) + + print("\n-- fold the values in: zlib.crc32 (what util.channels uses) --") + crc = bench("zlib.crc32(struct) whole ChannelMap", lambda: zlib.crc32(struct), nbytes=struct.nbytes) + bench("zlib.crc32(labels) plain label axis", lambda: zlib.crc32(plain), nbytes=plain.nbytes) + + print("\n-- concat.py's _fingerprint (per input; concat calls it twice) --") + concat = bench("_fingerprint(msg) struct axis + 2 attrs", lambda: concat_fingerprint(struct_msg, "ch")) + bench("_fingerprint(msg) plain labels + 2 attrs", lambda: concat_fingerprint(plain_msg, "ch")) + bench( + " ...its attrs frozenset alone", + lambda: frozenset((k, type(v).__name__, repr(v)) for k, v in (struct_msg.attrs or {}).items()), + ) + + print("\n-- reference points --") + guarded = bench("the work being guarded: data[:, :n/2] copy", lambda: struct_msg.data[:, : n_ch // 2].copy()) + bench("np.mean(data, axis=1) (a cheap real op)", lambda: np.mean(struct_msg.data, axis=1)) + + print("\n-- verdict (added cost over the O(1) baseline) --") + for name, cost in ( + ("hash(tobytes()), whole", naive), + ("crc32, whole", crc), + ("concat _fingerprint x2", concat * 2), + ): + delta = cost - base + print( + f" {name:<26} +{delta:6.3f} us/msg" + f" = {delta / guarded * 100:6.1f}% of the guarded copy" + # delta us/msg * 1000 msg/s = delta ms/s = delta/10 percent of a core. + f" {delta / 10:5.2f}% of a core @ 1 kHz" + ) + + +def section_mechanics(struct_msg: AxisArray) -> None: + """Why crc32, and why the dtype goes in as an object rather than a string.""" + struct = struct_msg.axes["ch"].data + n_ch = struct_msg.data.shape[1] + + print("\n== mechanics ==") + print(f"\nstruct itemsize {struct.dtype.itemsize} B, total {struct.nbytes} B") + for name in CHANNELMAP_DTYPE.names: + view = struct[name] + print( + f" field {name:<10} itemsize {view.dtype.itemsize:>3} B " + f"packed {view.dtype.itemsize * n_ch:>7} B contiguous={view.flags['C_CONTIGUOUS']}" + ) + two = struct[["array", "bank"]] + print(f" fields ('array','bank') -> view itemsize {two.dtype.itemsize} B, tobytes {len(two.tobytes())} B") + print(" ^ the multi-field view keeps the full itemsize: asking for 2 of 8") + print(" fields materializes all 8. Digest fields one at a time.") + + print("\n-- copy vs checksum: which dominates? --") + bench("struct.tobytes() (copy alone)", lambda: struct.tobytes(), nbytes=struct.nbytes) + bench("hash(struct.tobytes()) (copy + siphash)", lambda: hash(struct.tobytes()), nbytes=struct.nbytes) + bench("zlib.crc32(struct) (no intermediate copy)", lambda: zlib.crc32(struct), nbytes=struct.nbytes) + bench("zlib.adler32(struct)", lambda: zlib.adler32(struct), nbytes=struct.nbytes) + + print("\n-- carrying dtype alongside the checksum --") + bench("str(struct.dtype) structured repr, built field by field", lambda: str(struct.dtype), number=20_000) + bench("hash(struct.dtype) the object, hashable and value-equal", lambda: hash(struct.dtype)) + print(" ^ 'free' metadata is not free: str() of an 8-field dtype costs more") + print(" than the checksum it annotates. util.channels stores the object.") + + print("\n-- object-dtype axes --") + a = np.array([f"ch{i}" for i in range(n_ch)], dtype=object) + b = np.array(["".join(("ch", str(i))) for i in range(n_ch)], dtype=object) + print(f" equal object arrays, raw buffer: same crc32? {zlib.crc32(a.tobytes()) == zlib.crc32(b.tobytes())}") + print(" ^ the buffer is pointers, so a naive digest resets state every message.") + print(f" after .astype('U'): same crc32? {zlib.crc32(a.astype('U')) == zlib.crc32(b.astype('U'))}") + bench("a.astype('U') then crc32 (the widening path)", lambda: zlib.crc32(np.ascontiguousarray(a.astype("U")))) + + print("\n-- identity fast path, if a source reuses one axis object --") + cached_obj, cached = struct, zlib.crc32(struct) + other = channelmap_axis(n_ch).data + bench( + "hit: same array object -> reuse cached digest", + lambda: cached if struct is cached_obj else zlib.crc32(struct), + ) + bench("miss: fresh array -> full crc32", lambda: cached if other is cached_obj else zlib.crc32(other)) + print(" ^ ~50x cheaper on a hit, but only a transformer holding state can") + print(" do this; coord_value_fingerprint is a pure function.") + + +def section_shipped(struct_msg: AxisArray, plain_msg: AxisArray) -> None: + """The function as shipped -- these are the numbers in its docstring.""" + print("\n== shipped: util.channels.coord_value_fingerprint ==\n") + bench("cvf(struct, None) whole ChannelMap", lambda: coord_value_fingerprint(struct_msg, "ch", None)) + bench( + "cvf(struct, ('bank',)) U2, 7% of bytes", + lambda: coord_value_fingerprint(struct_msg, "ch", ("bank",)), + ) + bench( + "cvf(struct, ('array','bank')) 11% of bytes", + lambda: coord_value_fingerprint(struct_msg, "ch", ("array", "bank")), + ) + bench( + "cvf(struct, ('label',)) U16, 59% of bytes", + lambda: coord_value_fingerprint(struct_msg, "ch", ("label",)), + ) + bench("cvf(plain, None) plain label axis", lambda: coord_value_fingerprint(plain_msg, "ch", None)) + print("\n Restricting to a field is about invalidation correctness (an unread") + print(" x/y field churning must not reset the state), and is only sometimes") + print(" also a speedup -- a wide field costs more than the whole axis.") + + +def section_scaling(n_times: int) -> None: + print("\n== scaling with channel count ==\n") + for n in (32, 64, 128, 256, 512, 1024, 2048): + msg = make_msg(channelmap_axis(n), n, n_times) + arr = msg.axes["ch"].data + bench(f"n_ch={n:<5} cvf(struct, None)", lambda m=msg: coord_value_fingerprint(m, "ch", None), number=20_000) + bench(f"n_ch={n:<5} hash(tobytes())", lambda a=arr: hash(a.tobytes()), number=20_000, nbytes=arr.nbytes) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--n-ch", type=int, default=256, help="channels on the coordinate axis (default: 256)") + parser.add_argument("--n-times", type=int, default=30, help="samples per message (default: 30)") + parser.add_argument("--sections", default=",".join(SECTIONS), help=f"comma-separated subset of {SECTIONS}") + args = parser.parse_args() + + requested = [s.strip() for s in args.sections.split(",") if s.strip()] + unknown = [s for s in requested if s not in SECTIONS] + if unknown: + parser.error(f"unknown section(s) {unknown}; choose from {list(SECTIONS)}") + + struct_msg = make_msg(channelmap_axis(args.n_ch), args.n_ch, args.n_times) + plain_msg = make_msg(label_axis(args.n_ch), args.n_ch, args.n_times) + + print(f"{args.n_ch} channels, {args.n_times} samples/message, float32 data") + print( + f" coordinate axis: struct {struct_msg.axes['ch'].data.nbytes} B, " + f"plain labels {plain_msg.axes['ch'].data.nbytes} B" + ) + + if "strategies" in requested: + section_strategies(struct_msg, plain_msg) + if "mechanics" in requested: + section_mechanics(struct_msg) + if "shipped" in requested: + section_shipped(struct_msg, plain_msg) + if "scaling" in requested: + section_scaling(args.n_times) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_concat_fingerprint.py b/benchmarks/benchmark_concat_fingerprint.py new file mode 100644 index 00000000..c83ff9c2 --- /dev/null +++ b/benchmarks/benchmark_concat_fingerprint.py @@ -0,0 +1,224 @@ +"""What Concat spends deciding whether its axis cache is still valid. + +:class:`~ezmsg.sigproc.concat.ConcatProcessor` caches its merged output axes and +rebuilds them only when a *fingerprint* of each input changes. The fingerprint +runs on every message and the concatenate it guards does not: at 128 channels +with full ChannelMap metadata, fingerprinting used to be ~63% of ``_concat``'s +total time against ~10% for the ``xp.concat`` itself. + +Two things make it cheap, both measured here: + +* **Identity before content.** ``replace()`` carries ``axes``, the axis objects, + their ``.data`` arrays and ``attrs`` by reference, so in-process these are the + *same objects* message after message. An ``is`` check settles it in ~0.02 µs + where a checksum costs ~1 µs. A miss (what a cross-process hop produces) just + falls through to the digest, so it is a pure fast path. +* **crc32 over siphash.** See ``benchmark_axis_fingerprint.py``: the copy is not + the bottleneck, ``hash(bytes)`` is, and ``zlib.crc32`` reads the array buffer + directly at ~5x the throughput. + +The legacy fingerprint is reproduced verbatim below so the comparison stays +honest as the real one changes. Note it is not merely slower -- it digests only +the *concat* axis, while the cache holds every coordinate axis, and it serves +``LinearAxis`` objects from that cache too. The ``correctness`` section shows +what that costs you. + +Run from the repository root:: + + uv run python benchmarks/benchmark_concat_fingerprint.py + uv run python benchmarks/benchmark_concat_fingerprint.py --n-ch 512 + uv run python benchmarks/benchmark_concat_fingerprint.py --sections correctness +""" + +from __future__ import annotations + +import argparse +import timeit + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.sigproc.concat import ConcatProcessor, ConcatSettings + +SECTIONS = ("throughput", "breakdown", "correctness") +POOL = 64 + +CHANNELMAP_DTYPE = np.dtype( + [ + ("label", "U16"), + ("x", " CoordinateAxis: + d = np.zeros(n_ch, dtype=CHANNELMAP_DTYPE) + d["label"] = [f"elec{i + offset:04d}" for i in range(n_ch)] + d["x"] = np.arange(n_ch, dtype=float) + d["bank"] = "A" + d["elec"] = np.arange(n_ch, dtype=np.int32) + return CoordinateAxis(data=d, dims=["ch"]) + + +def make_msg(n_ch, n_times, key, ch, attrs, t=0.0) -> AxisArray: + return AxisArray( + np.zeros((n_times, n_ch, 2), dtype=np.float32), + dims=["time", "ch", "feature"], + axes={"time": AxisArray.TimeAxis(fs=30000.0, offset=t), "ch": ch, "feature": FEATURE_AXIS}, + key=key, + attrs=attrs, + ) + + +def legacy_fingerprint(proc: ConcatProcessor, msg: AxisArray) -> tuple: + """``ConcatProcessor._fingerprint`` as it stood before the rewrite.""" + ax = msg.axes.get(proc.settings.axis) + ax_hash = hash(ax.data.tobytes()) if ax is not None and hasattr(ax, "data") else None + attrs_fp = frozenset((k, type(v).__name__, repr(v)) for k, v in (msg.attrs or {}).items()) + return (tuple(msg.dims), msg.data.shape, ax_hash, attrs_fp) + + +def bench(label: str, fn, number: int = 5_000) -> float: + fn() + us = timeit.timeit(fn, number=number) / number * 1e6 + print(f" {label:<54} {us:8.3f} us") + return us + + +def _pools(n_ch, n_times, reuse_axis, n_attrs): + attrs = {f"k{i}": f"v{i}" for i in range(n_attrs)} + if reuse_axis: + sa, sb = ch_axis(n_ch, 0), ch_axis(n_ch, 10_000) + return ( + [make_msg(n_ch, n_times, "a", sa, attrs, i * 0.001) for i in range(POOL)], + [make_msg(n_ch, n_times, "b", sb, attrs, i * 0.001) for i in range(POOL)], + ) + return ( + [make_msg(n_ch, n_times, "a", ch_axis(n_ch, 0), dict(attrs), i * 0.001) for i in range(POOL)], + [make_msg(n_ch, n_times, "b", ch_axis(n_ch, 10_000), dict(attrs), i * 0.001) for i in range(POOL)], + ) + + +def section_throughput(n_ch: int, n_times: int) -> None: + print("\n== throughput: full _concat, per message ==") + for label, reuse, n_attrs in ( + ("axis objects reused (in-process steady state)", True, 2), + ("axis rebuilt per message (post-transport)", False, 2), + ("axis reused, 12 scalar attrs", True, 12), + ): + print(f"\n-- {label} --") + pa, pb = _pools(n_ch, n_times, reuse, n_attrs) + proc = ConcatProcessor(ConcatSettings(axis="ch")) + proc._concat(pa[0], pb[0]) + i = [0] + + def step(): + i[0] = (i[0] + 1) % POOL + return proc._concat(pa[i[0]], pb[i[0]]) + + total = bench("_concat, today", step) + + # Cycle the same pool the concat does: in the rebuilt case every message + # is a fresh object, so this has to pay the memo misses too. Measuring a + # single fixed message would hit the memo every time and flatter it. + j = [0] + + def cur_fp(): + j[0] = (j[0] + 1) % POOL + return (proc._fingerprint(pa[j[0]], proc.state.memo_a), proc._fingerprint(pb[j[0]], proc.state.memo_b)) + + def old_fp(): + j[0] = (j[0] + 1) % POOL + return (legacy_fingerprint(proc, pa[j[0]]), legacy_fingerprint(proc, pb[j[0]])) + + cur = bench(" its fingerprint (both inputs)", cur_fp) + old = bench(" legacy fingerprint (both inputs)", old_fp) + print( + f" => fingerprint {cur / total * 100:.0f}% of _concat; " + f"{old / cur:.1f}x cheaper than legacy ({old:.2f} -> {cur:.2f} us)" + ) + + +def section_breakdown(n_ch: int, n_times: int) -> None: + print("\n== breakdown: one input, memo hit vs miss ==\n") + pa, pb = _pools(n_ch, n_times, True, 2) + proc = ConcatProcessor(ConcatSettings(axis="ch")) + proc._concat(pa[0], pb[0]) + m = pa[0] + bench("memo hit (same axes + attrs objects)", lambda: proc._fingerprint(m, proc.state.memo_a)) + bench("memo bypassed (full content digest)", lambda: proc._fingerprint(m, None)) + bench("legacy (concat axis only, siphash + repr)", lambda: legacy_fingerprint(proc, m)) + print("\n the work being guarded, for scale:") + bench("np.concat([a.data, b.data], axis=1)", lambda: np.concat([pa[0].data, pb[0].data], axis=1)) + + +def section_correctness() -> None: + """The legacy fingerprint's blind spots, as running code.""" + print("\n== correctness ==\n") + + def msg(band, t, fill): + return AxisArray( + np.full((4, 2, 3), fill, dtype=np.float32), + dims=["time", "ch", "band"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0, offset=t), + "ch": CoordinateAxis(data=np.array(["c0", "c1"]), dims=["ch"]), + "band": CoordinateAxis(data=np.array(band), dims=["band"]), + }, + key="dev", + ) + + proc = ConcatProcessor(ConcatSettings(axis="ch", relabel_axis=False)) + proc._concat(msg(["alpha", "beta", "gamma"], 0.0, 1.0), msg(["alpha", "beta", "gamma"], 0.0, 2.0)) + out = proc._concat(msg(["delta", "theta", "mu"], 0.04, 1.0), msg(["delta", "theta", "mu"], 0.04, 2.0)) + print( + f" non-concat axis relabelled -> band {[str(x) for x in out.axes['band'].data]} " + f"(want ['delta', 'theta', 'mu'])" + ) + print(f" time advanced -> offset {out.axes['time'].offset} (want 0.04)") + + a = msg(["alpha", "beta", "gamma"], 0.0, 1.0) + b = msg(["delta", "theta", "mu"], 0.0, 1.0) + print( + f"\n legacy fingerprint tells those two messages apart? " + f"{legacy_fingerprint(proc, a) != legacy_fingerprint(proc, b)}" + ) + print( + f" current fingerprint does? " + f"{proc._fingerprint(a, None) != proc._fingerprint(b, None)}" + ) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--n-ch", type=int, default=128, help="channels per input (default: 128)") + p.add_argument("--n-times", type=int, default=30, help="samples per message (default: 30)") + p.add_argument("--sections", default=",".join(SECTIONS), help=f"subset of {SECTIONS}") + args = p.parse_args() + + requested = [s.strip() for s in args.sections.split(",") if s.strip()] + unknown = [s for s in requested if s not in SECTIONS] + if unknown: + p.error(f"unknown section(s) {unknown}; choose from {list(SECTIONS)}") + + print( + f"concat: 2 x ({args.n_times} x {args.n_ch} x 2) f32, ChannelMap ch axis " + f"({args.n_ch * CHANNELMAP_DTYPE.itemsize} B)" + ) + if "throughput" in requested: + section_throughput(args.n_ch, args.n_times) + if "breakdown" in requested: + section_breakdown(args.n_ch, args.n_times) + if "correctness" in requested: + section_correctness() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_hash_overhead.py b/benchmarks/benchmark_hash_overhead.py new file mode 100644 index 00000000..1bd0f120 --- /dev/null +++ b/benchmarks/benchmark_hash_overhead.py @@ -0,0 +1,512 @@ +"""What the safer default state hash costs, and what fingerprint caching gives back. + +Every stateful processor decides, once per message, whether its cached state is +still valid. That decision used to be nearly free and frequently wrong: the base +class returned a constant, so a processor without its own ``_hash_message`` +never reset, and a source that renamed its channels under a fixed key and +channel count kept being filtered through the previous channels' history. + +The default now folds in the message key, the dims, the length of every +dimension except the chunk dimension, the *values* on the coordinate axes and +the gain and offset of any linear axis among them. That is strictly more work +per message. This script measures how much more, and how much of it +``CoordinateAxis.fingerprint`` hands back by computing the expensive part -- a +crc32 over the axis data -- once per axis object rather than once per consumer. + +Three arms, so the two effects can be told apart: + +``before`` + The processors as they were: constant-hash default, hand-written overrides + on the few processors that had them. Cheap, and wrong in the ways above. +``after`` + What ships now. Correct everywhere, and the coordinate checksum is computed + on first access and cached on the axis, so a fan-out of N consumers pays for + it once. +``naive`` + The same correctness without the cache -- ``fingerprint`` recomputes on + every access. This is what the fix would have cost implemented the obvious + way; the gap between it and ``after`` is what the caching is worth. + +Fingerprints are cached *on the axis object*, so a benchmark that reuses one +message list across timed runs measures a warm cache and reports the checksum as +free. Every timed run below is preceded by an untimed pass that strips those +caches, which is what ``--arm naive`` and the cold/warm split are there to make +visible. + +``after`` and ``naive`` run against the working tree. ``before`` needs the +pre-sweep sources of *both* packages on ``PYTHONPATH``, which git worktrees +supply without disturbing either checkout:: + + git worktree add --detach /tmp/sigproc-pre + cp src/ezmsg/sigproc/__version__.py /tmp/sigproc-pre/src/ezmsg/sigproc/ + git -C ../ezmsg-baseproc worktree add --detach /tmp/baseproc-pre + cp ../ezmsg-baseproc/src/ezmsg/baseproc/__version__.py /tmp/baseproc-pre/src/ezmsg/baseproc/ + +(``__version__.py`` is generated at build time, so a fresh worktree lacks it.) +Then, from the repository root:: + + uv run python benchmarks/benchmark_hash_overhead.py --arm after + uv run python benchmarks/benchmark_hash_overhead.py --arm naive + PYTHONPATH=/tmp/sigproc-pre/src:/tmp/baseproc-pre/src \\ + uv run python benchmarks/benchmark_hash_overhead.py --arm before + +Pass ``--json`` to emit a machine-readable record for diffing arms. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import gc +import inspect +import json +import time +import timeit +import typing +import warnings + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +warnings.filterwarnings("ignore") + +HAS_CHUNK_DIM = "chunk_dim" in AxisArray.__dataclass_fields__ +DECLARE_CHUNK_DIM = HAS_CHUNK_DIM +"""Whether the simulated source declares its chunk dimension. + +The ``before`` arm must not: nothing set the field then, and the pre-sweep +``Spectrum`` does not clear it when it consumes ``time``, so a declared source +trips the validation that ships with it.""" + +# ezmsg-blackrock's ChannelMap ``ch`` axis: the "full metadata" case, 108 B per +# channel. Kept byte-identical to benchmark_axis_fingerprint.py so the numbers +# there line up with the ones here. +CHANNELMAP_DTYPE = np.dtype( + [ + ("label", "U16"), + ("x", " np.ndarray: + data = np.zeros(n_ch, dtype=CHANNELMAP_DTYPE) + data["label"] = [f"elec{i:04d}" for i in range(n_ch)] + data["x"] = np.arange(n_ch, dtype=np.float64) + data["y"] = np.arange(n_ch, dtype=np.float64) + data["size"] = 1.0 + data["array"] = np.arange(n_ch) // 128 + data["bank"] = np.array([("A", "B", "C", "D")[i // 64 % 4] for i in range(n_ch)]) + data["elec"] = np.arange(n_ch, dtype=np.int32) + 1 + data["headstage"] = np.arange(n_ch) // 32 + return data + + +def label_data(n_ch: int) -> np.ndarray: + """The plain label axis most sources emit.""" + return np.array([f"ch{i:04d}" for i in range(n_ch)]) + + +def make_message(data: np.ndarray, ch_data: np.ndarray, fs: float, key: str = "dev") -> AxisArray: + """A fresh message, with a *fresh* coordinate axis: a live source builds new + axis objects per message, so the fingerprint cache starts cold.""" + kwargs = {"chunk_dim": "time"} if DECLARE_CHUNK_DIM else {} + return AxisArray( + data, + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=fs), "ch": CoordinateAxis(data=ch_data, dims=["ch"])}, + key=key, + **kwargs, + ) + + +def chill(message: typing.Any) -> typing.Any: + """Drop any cached fingerprints, returning the message to its as-received state.""" + axes = getattr(message, "axes", None) + if axes: + for axis in axes.values(): + axis.__dict__.pop("_fingerprint", None) + return message + + +def uncache_fingerprint() -> bool: + """Make ``fingerprint`` recompute on every access (the ``naive`` arm).""" + if not hasattr(CoordinateAxis, "_compute_fingerprint"): + return False + CoordinateAxis.fingerprint = property(lambda self: self._compute_fingerprint()) + return True + + +# --------------------------------------------------------------------------- # +# Section 1: what one _hash_message call costs, per processor +# --------------------------------------------------------------------------- # + + +def discover_processors() -> dict[str, type]: + import importlib + import pkgutil + + from ezmsg.baseproc.stateful import Stateful + + import ezmsg.sigproc + + mods = [] + for m in pkgutil.walk_packages(ezmsg.sigproc.__path__, "ezmsg.sigproc."): + try: + mods.append(importlib.import_module(m.name)) + except Exception: + pass + names = {m.__name__ for m in mods} + found: dict[str, type] = {} + for mod in mods: + for obj in vars(mod).values(): + if ( + inspect.isclass(obj) + and issubclass(obj, Stateful) + and obj is not Stateful + and obj.__module__ in names + and not inspect.isabstract(obj) + ): + found[f"{obj.__module__.replace('ezmsg.sigproc.', '')}.{obj.__qualname__}"] = obj + return found + + +def _settings_type(cls: type) -> type | None: + for klass in cls.__mro__: + for base in getattr(klass, "__orig_bases__", ()): + for arg in typing.get_args(base): + if dataclasses.is_dataclass(arg): + return arg + return None + + +def _bare_instance(cls: type) -> typing.Any: + """``_hash_message`` reads the message and the settings, never the state, so + an uninitialised instance carrying default settings is enough to time it.""" + inst = object.__new__(cls) + settings_t = _settings_type(cls) + inst.settings = settings_t() if settings_t is not None else None + return inst + + +TARGET_S = 0.02 +"""Wall time to aim each timing loop at. ``timeit.autorange`` targets 0.2 s, +which is 40x more than sub-microsecond calls need to resolve and turns this +script into a ten-minute run across three arms.""" + + +def _time(fn: typing.Callable[[], typing.Any], repeats: int) -> tuple[float, int]: + timer = timeit.Timer(fn) + probe = timer.timeit(64) / 64 + n = max(64, min(int(TARGET_S / max(probe, 1e-9)), 200_000)) + return min(timer.repeat(repeats, n)) / n, n + + +def hash_cost_per_processor(msg: AxisArray, repeats: int) -> dict[str, dict[str, float]]: + """Cold and warm cost of one ``_hash_message`` call, per processor. + + *Warm* is what every consumer after the first pays: the fingerprint is + already on the axis. *Cold* includes computing it, and is what the first + consumer to touch a freshly built axis pays. The cache-clearing itself is + timed separately and subtracted, so cold is comparable to warm. + """ + baseline, _ = _time(lambda: chill(msg), repeats) + out: dict[str, dict[str, float]] = {} + for qualname, cls in sorted(discover_processors().items()): + try: + inst = _bare_instance(cls) + inst._hash_message(msg) + except Exception: + continue + warm, _ = _time(lambda: inst._hash_message(msg), repeats) + cold, _ = _time(lambda: inst._hash_message(chill(msg)), repeats) + out[qualname] = {"warm_us": warm * 1e6, "cold_us": max(cold - baseline, 0.0) * 1e6} + chill(msg) + return out + + +# --------------------------------------------------------------------------- # +# Section 2: a whole chain, end to end +# --------------------------------------------------------------------------- # + + +def build_chain(fs: float) -> list: + """An intracranial-features-shaped chain: rereference, band-limit, window, + spectrum, band-average, flatten to a feature vector.""" + from ezmsg.sigproc.affinetransform import CommonRereferenceSettings, CommonRereferenceTransformer + from ezmsg.sigproc.aggregate import AggregateSettings, AggregateTransformer, AggregationFunction + from ezmsg.sigproc.butterworthfilter import ButterworthFilterSettings, ButterworthFilterTransformer + from ezmsg.sigproc.flatten import FlattenSettings, FlattenTransformer + from ezmsg.sigproc.spectrum import SpectrumSettings, SpectrumTransformer + from ezmsg.sigproc.window import WindowSettings, WindowTransformer + + return [ + CommonRereferenceTransformer(CommonRereferenceSettings(mode="mean", axis="ch")), + ButterworthFilterTransformer(ButterworthFilterSettings(axis="time", order=4, cuton=70.0, cutoff=150.0)), + WindowTransformer(WindowSettings(axis="time", newaxis="win", window_dur=0.1, window_shift=0.02)), + SpectrumTransformer(SpectrumSettings(axis="time")), + AggregateTransformer(AggregateSettings(axis="freq", operation=AggregationFunction.MEAN)), + FlattenTransformer(FlattenSettings(preserve_axis="win", sample_axis="win", flatten_axes=("ch",))), + ] + + +CHAIN_LABELS = ("rereference", "butterworth", "window", "spectrum", "aggregate", "flatten") + + +def _empty(out: typing.Any) -> bool: + return out is None or (hasattr(out, "data") and out.data.size == 0) + + +def run_chain(chain: list, messages: list[AxisArray]) -> tuple[float, int]: + n_out = 0 + for msg in messages: + chill(msg) + gc.collect() + t0 = time.perf_counter() + for msg in messages: + out: typing.Any = msg + for stage in chain: + out = stage(out) + if _empty(out): + break + else: + n_out += 1 + return time.perf_counter() - t0, n_out + + +def freeze_hashes(chain: list) -> int: + """Neutralise ``_hash_message`` on every processor a stage owns, so the same + chain measures pure compute. Returns how many were reached.""" + from ezmsg.baseproc.stateful import Stateful + + seen: set[int] = set() + frozen = 0 + + def visit(obj) -> None: + nonlocal frozen + if id(obj) in seen or isinstance(obj, (np.ndarray, AxisArray, type)): + return + seen.add(id(obj)) + if isinstance(obj, Stateful): + obj._hash_message = lambda message: 0 + frozen += 1 + for value in list(getattr(obj, "__dict__", {}).values()): + if isinstance(value, (list, tuple)): + for item in value: + visit(item) + elif hasattr(value, "__dict__"): + visit(value) + + for stage in chain: + visit(stage) + return frozen + + +def chain_throughput(messages: list[AxisArray], fs: float, n_rounds: int) -> dict[str, typing.Any]: + """Time the chain with hashing live and with it neutralised. + + The difference is what hashing costs end to end -- but read it with the + per-stage numbers next to it. Hashing is a couple of microseconds against a + chain that spends hundreds on filtering and FFTs, so the difference sits + well inside run-to-run drift and can come out negative. The order of the two + runs alternates so that drift at least does not accumulate in one direction; + ``per_stage_hash_cost`` is the measurement that actually resolves it. + """ + live_times, frozen_times = [], [] + n_out = 0 + n_frozen = 0 + + def live_run() -> None: + nonlocal n_out + # A fresh chain each round, warmed on a few messages so the timed run + # measures steady state rather than every stage's first-message setup. + chain = build_chain(fs) + run_chain(chain, messages[:8]) + dt, n_out = run_chain(chain, messages) + live_times.append(dt) + + def frozen_run() -> None: + nonlocal n_frozen + chain = build_chain(fs) + run_chain(chain, messages[:8]) + n_frozen = freeze_hashes(chain) + dt, _ = run_chain(chain, messages) + frozen_times.append(dt) + + for round_ix in range(n_rounds): + for run in (live_run, frozen_run) if round_ix % 2 == 0 else (frozen_run, live_run): + run() + + per_msg = lambda ts: np.array(ts) / len(messages) * 1e6 # noqa: E731 + live_us, frozen_us = per_msg(live_times), per_msg(frozen_times) + live = float(live_us.min()) + frozen = float(frozen_us.min()) + # The spread across identical rounds. If it exceeds the live-minus-frozen + # difference -- which it does on this chain -- the difference is noise and + # the per-stage table is the number to quote. + spread = float(np.percentile(live_us, 90) - live_us.min()) + return { + "per_message_us": live, + "per_message_us_no_hashing": frozen, + "round_spread_us": spread, + "hash_overhead_us": live - frozen, + "hash_overhead_resolved": abs(live - frozen) > spread, + "msgs_per_sec": 1e6 / live, + "n_outputs": n_out, + "n_processors_frozen": n_frozen, + } + + +def per_stage_hash_cost(messages: list[AxisArray], fs: float, repeats: int) -> list[dict[str, typing.Any]]: + """Time each stage's ``_hash_message`` on the message that stage really sees. + + The intermediate messages matter: after ``Window`` the chunk dimension is + ``win``, after ``Spectrum`` the coordinate axes are different ones, and a + stage that hashes cheaply on the source message may not downstream. + """ + chain = build_chain(fs) + run_chain(chain, messages[:8]) + + inputs: list[typing.Any] = [] + out: typing.Any = messages[8] + for stage in chain: + inputs.append(out) + out = stage(out) + if _empty(out): + break + + rows = [] + for label, stage, msg in zip(CHAIN_LABELS, chain, inputs): + if not hasattr(stage, "_hash_message"): + # Stateless stages hold nothing to invalidate and never hash. + rows.append({"stage": label, "cls": type(stage).__name__, "dims": list(getattr(msg, "dims", []))}) + continue + baseline, _ = _time(lambda: chill(msg), repeats) + warm, _ = _time(lambda: stage._hash_message(msg), repeats) + cold, _ = _time(lambda: stage._hash_message(chill(msg)), repeats) + rows.append( + { + "stage": label, + "cls": type(stage).__name__, + "dims": list(getattr(msg, "dims", [])), + "warm_us": warm * 1e6, + "cold_us": max(cold - baseline, 0.0) * 1e6, + } + ) + return rows + + +# --------------------------------------------------------------------------- # +# Section 3: how many checksums the chain actually computes +# --------------------------------------------------------------------------- # + + +def count_checksums(messages: list[AxisArray], fs: float) -> dict[str, float]: + """Count crc32 calls over a run. This is the work ``fingerprint`` caches.""" + import zlib + + calls = 0 + real_crc32 = zlib.crc32 + + def counting(*args, **kwargs): + nonlocal calls + calls += 1 + return real_crc32(*args, **kwargs) + + chain = build_chain(fs) + run_chain(chain, messages[:8]) + for msg in messages: + chill(msg) + zlib.crc32 = counting + try: + for msg in messages: + out: typing.Any = msg + for stage in chain: + out = stage(out) + if _empty(out): + break + finally: + zlib.crc32 = real_crc32 + return {"crc32_calls": float(calls), "crc32_per_message": calls / len(messages)} + + +# --------------------------------------------------------------------------- # + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--arm", choices=("before", "after", "naive"), default="after") + p.add_argument("--n-ch", type=int, default=256) + p.add_argument("--fs", type=float, default=2000.0) + p.add_argument("--chunk-ms", type=float, default=20.0) + p.add_argument("--n-messages", type=int, default=200) + p.add_argument("--n-rounds", type=int, default=5) + p.add_argument("--repeats", type=int, default=7) + p.add_argument("--json", action="store_true") + args = p.parse_args() + + global DECLARE_CHUNK_DIM + if args.arm == "before": + DECLARE_CHUNK_DIM = False + if args.arm == "naive" and not uncache_fingerprint(): + raise SystemExit("--arm naive needs a build of ezmsg that has CoordinateAxis.fingerprint") + + n_time = int(round(args.fs * args.chunk_ms / 1000.0)) + ch_full = channelmap_data(args.n_ch) + ch_plain = label_data(args.n_ch) + signal = np.random.default_rng(0).standard_normal((n_time, args.n_ch)).astype(np.float32) + messages = [make_message(signal, ch_full, args.fs) for _ in range(args.n_messages)] + + record: dict[str, typing.Any] = { + "arm": args.arm, + "n_ch": args.n_ch, + "fs": args.fs, + "chunk_ms": args.chunk_ms, + "n_time": n_time, + "n_messages": args.n_messages, + "declares_chunk_dim": DECLARE_CHUNK_DIM, + "hash_us_full_metadata": hash_cost_per_processor(make_message(signal, ch_full, args.fs), args.repeats), + "hash_us_plain_labels": hash_cost_per_processor(make_message(signal, ch_plain, args.fs), args.repeats), + "per_stage": per_stage_hash_cost(messages, args.fs, args.repeats), + "chain": chain_throughput(messages, args.fs, args.n_rounds), + "checksums": count_checksums(messages, args.fs), + } + + if args.json: + print(json.dumps(record)) + return + + costs = record["hash_us_full_metadata"] + warm = [v["warm_us"] for v in costs.values()] + cold = [v["cold_us"] for v in costs.values()] + print(f"arm={args.arm} {args.n_ch} ch @ {args.fs:g} Hz, {n_time}-sample chunks\n") + print(f"_hash_message across {len(costs)} processors, 256-ch ChannelMap (us/call):") + print(f" warm (fingerprint cached) median {np.median(warm):6.3f} sum {sum(warm):7.2f}") + print(f" cold (fingerprint computed) median {np.median(cold):6.3f} sum {sum(cold):7.2f}\n") + print("per stage, on the message that stage really receives:") + hashing = 0.0 + for ix, row in enumerate(record["per_stage"]): + if "warm_us" not in row: + print(f" {row['stage']:<13}{str(row['dims']):<26}stateless, never hashes") + continue + # Only the first consumer of a freshly built axis pays the checksum. + hashing += row["cold_us"] if ix == 0 else row["warm_us"] + print(f" {row['stage']:<13}{str(row['dims']):<26}warm {row['warm_us']:6.3f} cold {row['cold_us']:6.3f}") + c = record["chain"] + print(f"\n6-stage chain ({c['n_processors_frozen']} stateful processors, {c['n_outputs']} outputs):") + print(f" per message {c['per_message_us']:8.2f} us ({c['msgs_per_sec']:.0f} msg/s)") + print(f" hashing, from the table {hashing:8.2f} us ({hashing / c['per_message_us'] * 100:.2f}% of the chain)") + print(f" crc32 calls per message {record['checksums']['crc32_per_message']:8.2f}") + print(f"\n cross-check by neutralising hashing: {c['hash_overhead_us']:+.2f} us") + print(f" round-to-round spread on this chain: {c['round_spread_us']:.2f} us", end=" ") + print("-- resolved" if c["hash_overhead_resolved"] else "-- below the noise floor, use the table above") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_hash_witness.py b/benchmarks/benchmark_hash_witness.py new file mode 100644 index 00000000..987d7aa0 --- /dev/null +++ b/benchmarks/benchmark_hash_witness.py @@ -0,0 +1,218 @@ +"""What the state-hash witness is worth in a graph that keeps growing. + +``Stateful._message_hash`` walks the dims, reaches into the axes and builds a +tuple to hash, once per message per processor. In a steady stream the answer is +the same every time, and the work to prove it is the same every time too. The +witness records the objects the last answer was derived from and, when none of +them has changed identity, returns the cached answer. + +The precondition is producer-side and every ezmsg source already satisfies it: +build the per-stream axes once, and replace only the chunk axis per message +(``replace(template, data=..., axes={**template.axes, "time": new_time_ax})``). +That hands every consumer the *same* coordinate axis object for the life of the +stream, and identity settles the question in one pointer comparison. + +Identity is not available everywhere: unpickling hands out a new axis object per +message, so every processor downstream of a process boundary sees fresh objects +carrying identical values. That arm is measured here as ``rebuilt``, and is why +the witness falls back to comparing the axis *value* -- the fingerprint, which +rides along already computed -- rather than giving up when identity fails. + +Run against the working tree for the witness arm, and against a checkout without +it for the baseline:: + + git -C ../ezmsg-baseproc worktree add --detach /tmp/baseproc-nowitness + cp ../ezmsg-baseproc/src/ezmsg/baseproc/__version__.py /tmp/baseproc-nowitness/src/ezmsg/baseproc/ + + uv run python benchmarks/benchmark_hash_witness.py --label witness + PYTHONPATH=/tmp/baseproc-nowitness/src uv run python benchmarks/benchmark_hash_witness.py --label baseline + +Pass ``--json`` to emit a record for diffing arms. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import time +import timeit +import typing + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis +from ezmsg.util.messages.util import replace + +from ezmsg.sigproc.ewma import EWMASettings, EWMATransformer + +TARGET_S = 0.02 + + +def _time(fn: typing.Callable[[], typing.Any], repeats: int = 7) -> float: + timer = timeit.Timer(fn) + probe = timer.timeit(64) / 64 + n = max(64, min(int(TARGET_S / max(probe, 1e-9)), 200_000)) + return min(timer.repeat(repeats, n)) / n * 1e6 + + +def make_template(n_ch: int, fs: float, n_time: int, key: str = "dev") -> AxisArray: + return AxisArray( + np.zeros((n_time, n_ch), np.float32), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array([f"ch{i:03d}" for i in range(n_ch)]), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +def stream(template: AxisArray, n: int, n_time: int, fs: float, *, hoist: bool) -> list[AxisArray]: + """``hoist=True`` is the template idiom every ezmsg source uses, and gives + every consumer in the producing process the same axis object. ``False`` is + what the far side of a process boundary sees: a new object per message, + carrying the same values and the same precomputed fingerprint.""" + out = [] + for i in range(n): + axes = {**template.axes, "time": replace(template.axes["time"], offset=i * n_time / fs)} + if not hoist: + axes["ch"] = CoordinateAxis(data=template.axes["ch"].data, dims=["ch"]) + msg = replace(template, data=np.full((n_time, template.data.shape[1]), float(i), np.float32), axes=axes) + for axis in msg.axes.values(): + getattr(axis, "fingerprint", None) # a source that primes its fingerprints + out.append(msg) + return out + + +def build_chain(n_nodes: int) -> list: + return [EWMATransformer(EWMASettings(axis="time", time_constant=0.5)) for _ in range(n_nodes)] + + +def run_chain(chain: list, messages: list[AxisArray]) -> float: + gc.collect() + t0 = time.perf_counter() + for msg in messages: + out: typing.Any = msg + for stage in chain: + out = stage(out) + return time.perf_counter() - t0 + + +def hit_rate(n_nodes: int, messages: list[AxisArray]) -> tuple[int, int]: + """How often the fast path answered, by counting the validator's verdicts. + + Patched before the chain is built, since a witness compiles its validator at + construction and a chain warmed beforehand would carry unpatched ones. + """ + from ezmsg.baseproc import stateful as st + + if not hasattr(st, "_build_witness"): + return (0, 0) + hits = total = 0 + real = st._build_witness + + def counting_build(*args, **kwargs): + witness = real(*args, **kwargs) + if witness is None: + return None + validator = witness[0] + + def counted(msg, _v=validator): + nonlocal hits, total + out = _v(msg) + total += 1 + hits += out + return out + + return (counted,) + witness[1:] + + st._build_witness = counting_build + try: + chain = build_chain(n_nodes) + run_chain(chain, messages[:4]) + hits = total = 0 + for msg in messages: + out: typing.Any = msg + for stage in chain: + out = stage(out) + finally: + st._build_witness = real + return hits, total + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--label", default="witness") + p.add_argument("--n-ch", type=int, default=16) + p.add_argument("--n-time", type=int, default=8) + p.add_argument("--fs", type=float, default=1000.0) + p.add_argument("--nodes", type=int, nargs="+", default=[1, 10, 30, 100]) + p.add_argument("--n-messages", type=int, default=200) + p.add_argument("--rounds", type=int, default=7) + p.add_argument("--json", action="store_true") + args = p.parse_args() + + template = make_template(args.n_ch, args.fs, args.n_time) + hoisted = stream(template, args.n_messages, args.n_time, args.fs, hoist=True) + rebuilt = stream(template, args.n_messages, args.n_time, args.fs, hoist=False) + + record: dict[str, typing.Any] = { + "label": args.label, + "n_ch": args.n_ch, + "n_time": args.n_time, + "n_messages": args.n_messages, + "per_call": {}, + "chains": {}, + } + + # --- one _hash_message call, in isolation + for name, msgs in (("hoisted", hoisted), ("rebuilt", rebuilt)): + proc = EWMATransformer(EWMASettings(axis="time", time_constant=0.5)) + for m in msgs[:4]: + proc(m) + cyc = iter(msgs * 100000) + overhead = _time(lambda: next(iter([None]))) + record["per_call"][name] = _time(lambda: proc._hash_message(next(cyc))) - overhead + + # --- whole chains, at several depths + for n_nodes in args.nodes: + for name, msgs in (("hoisted", hoisted), ("rebuilt", rebuilt)): + times = [] + for _ in range(args.rounds): + chain = build_chain(n_nodes) + run_chain(chain, msgs[:4]) + times.append(run_chain(chain, msgs)) + per_msg = min(times) / len(msgs) * 1e6 + hits, total = hit_rate(n_nodes, msgs) + record["chains"].setdefault(str(n_nodes), {})[name] = { + "us_per_message": per_msg, + "us_per_node": per_msg / n_nodes, + "msgs_per_sec": 1e6 / per_msg, + "witness_hit_pct": (100.0 * hits / total) if total else None, + } + + if args.json: + print(json.dumps(record)) + return + + print(f"{args.label}: {args.n_ch} ch x {args.n_time} samples, {args.n_messages} messages\n") + print("one _hash_message call:") + for name, us in record["per_call"].items(): + print(f" source {name:<9} {us:7.3f} us") + print( + f"\n{'nodes':>6} {'hoisted us/msg':>15}{'/node':>8}{'hit%':>7} {'rebuilt us/msg':>15}{'/node':>8}{'hit%':>7}" + ) + for n_nodes in args.nodes: + h = record["chains"][str(n_nodes)]["hoisted"] + r = record["chains"][str(n_nodes)]["rebuilt"] + hp = f"{h['witness_hit_pct']:.0f}" if h["witness_hit_pct"] is not None else "-" + rp = f"{r['witness_hit_pct']:.0f}" if r["witness_hit_pct"] is not None else "-" + print( + f"{n_nodes:>6} {h['us_per_message']:15.2f}{h['us_per_node']:8.3f}{hp:>7} " + f"{r['us_per_message']:15.2f}{r['us_per_node']:8.3f}{rp:>7}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_state_resets.py b/benchmarks/benchmark_state_resets.py new file mode 100644 index 00000000..b0240b02 --- /dev/null +++ b/benchmarks/benchmark_state_resets.py @@ -0,0 +1,267 @@ +"""How often does the base-class state hash actually reset a live graph? + +Most processors no longer implement ``_hash_message``: the base class decides, +folding in the message key, the dims, the length of every dimension except the +chunk dimension, the coordinate values on those dimensions, and the gain and +offset of any linear axis. This runs a real ``ez.run`` graph over simulated +256-channel data shaped like the intracranial feature pipeline and counts how +often each node rebuilds its state. + +What to look for: a stream whose configuration never changes should reset each +node exactly *once*, no matter how much the chunk size jitters. Anything more +means the hash is folding in something per-message, and every extra reset is a +filter redesigned or a buffer reallocated mid-stream. + +The source can rebuild its axis objects per message or reuse them, and the +graph can be split across a process boundary; neither should change the reset +counts, only the cost of arriving at them. + +Run from the repository root:: + + uv run python benchmarks/benchmark_memo_hit_rate.py + uv run python benchmarks/benchmark_memo_hit_rate.py --n-messages 500 --n-ch 512 +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import pathlib +import tempfile + +import ezmsg.core as ez +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.sigproc.affinetransform import AffineTransform, AffineTransformSettings, AffineTransformTransformer +from ezmsg.sigproc.flatten import Flatten, FlattenSettings, FlattenTransformer + +# Count state rebuilds per processor class, without touching the library. +# Patched on each concrete class rather than on Stateful: every processor +# overrides _reset_state, so a base-class patch would never be reached. +_RESETS: dict[str, int] = {} + + +def _count_resets(cls: type) -> None: + original = cls._reset_state + + def counting(self, *args, **kwargs): + _RESETS[type(self).__name__] = _RESETS.get(type(self).__name__, 0) + 1 + return original(self, *args, **kwargs) + + cls._reset_state = counting + + +for _cls in (AffineTransformTransformer, FlattenTransformer): + _count_resets(_cls) + +CHANNELMAP_DTYPE = np.dtype( + [ + ("label", "U16"), + ("x", " CoordinateAxis: + d = np.zeros(n_ch, dtype=CHANNELMAP_DTYPE) + d["label"] = [f"elec{i:04d}" for i in range(n_ch)] + d["x"] = np.arange(n_ch, dtype=float) + d["array"] = np.arange(n_ch) // 64 + d["bank"] = "A" + d["elec"] = np.arange(n_ch, dtype=np.int32) + return CoordinateAxis(data=d, dims=["ch"]) + + +def feature_axis() -> CoordinateAxis: + return CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]) + + +class SourceSettings(ez.Settings): + n_ch: int = 256 + n_times: int = 30 + n_messages: int = 200 + fs: float = 30000.0 + rebuild_axes: bool = False + """Build NEW axis objects per message, as a source that re-reads its channel + metadata each chunk would. False reuses them, which is what a source that + builds its channel map once at startup does.""" + + +class SourceState(ez.State): + ch_axis: CoordinateAxis | None = None + feat_axis: CoordinateAxis | None = None + + +class SimSource(ez.Unit): + """Emit AxisArrays shaped like one hub's broadband stream.""" + + SETTINGS = SourceSettings + STATE = SourceState + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + async def initialize(self) -> None: + self.STATE.ch_axis = channelmap(self.SETTINGS.n_ch) + self.STATE.feat_axis = feature_axis() + + @ez.publisher(OUTPUT_SIGNAL) + async def pub(self): + n_ch, n_t = self.SETTINGS.n_ch, self.SETTINGS.n_times + for i in range(self.SETTINGS.n_messages): + if self.SETTINGS.rebuild_axes: + ch, feat = channelmap(n_ch), feature_axis() + else: + ch, feat = self.STATE.ch_axis, self.STATE.feat_axis + yield ( + self.OUTPUT_SIGNAL, + AxisArray( + np.zeros((n_t, n_ch, 2), dtype=np.float32), + dims=["time", "ch", "feature"], + axes={ + "time": AxisArray.TimeAxis(fs=self.SETTINGS.fs, offset=i * n_t / self.SETTINGS.fs), + "ch": ch, + "feature": feat, + }, + key="sim", + ), + ) + await asyncio.sleep(0.001) + await asyncio.sleep(0.5) + raise ez.NormalTermination + + +class ReporterSettings(ez.Settings): + out_dir: str = "" + tag: str = "" + + +class Reporter(ez.Unit): + """Dump this process's fingerprint counters on shutdown.""" + + SETTINGS = ReporterSettings + INPUT_SIGNAL = ez.InputStream(AxisArray) + + @ez.subscriber(INPUT_SIGNAL) + async def sink(self, _: AxisArray) -> None: + pass + + async def shutdown(self) -> None: + path = pathlib.Path(self.SETTINGS.out_dir) / f"{self.SETTINGS.tag}-{os.getpid()}.json" + path.write_text(json.dumps({"pid": os.getpid(), "tag": self.SETTINGS.tag, "stats": dict(_RESETS)})) + + +class DownstreamSettings(ez.Settings): + flatten: FlattenSettings + reporter: ReporterSettings + + +class Downstream(ez.Collection): + """Flatten + its reporter, so both land in the same worker process.""" + + SETTINGS = DownstreamSettings + INPUT_SIGNAL = ez.InputStream(AxisArray) + + FLATTEN = Flatten() + REPORT = Reporter() + + def configure(self) -> None: + self.FLATTEN.apply_settings(self.SETTINGS.flatten) + self.REPORT.apply_settings(self.SETTINGS.reporter) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.INPUT_SIGNAL, self.FLATTEN.INPUT_SIGNAL), + (self.FLATTEN.OUTPUT_SIGNAL, self.REPORT.INPUT_SIGNAL), + ) + + +def build_and_run(n_ch, n_times, n_messages, rebuild_axes, out_dir, split) -> None: + """SRC -> LRR(non-square) -> LRR(square) -> {reporter, Flatten -> reporter}.""" + _RESETS.clear() # counters are process-global; isolate the runs + n_out = n_ch // 2 + weights = np.zeros((n_ch, n_out)) + weights[np.arange(n_out) * 2, np.arange(n_out)] = 1.0 # keep every other channel + + downstream = Downstream( + DownstreamSettings( + flatten=FlattenSettings(preserve_axis="time", flatten_axes=("ch", "feature"), output_axis="ch"), + reporter=ReporterSettings(out_dir=out_dir, tag="downstream"), + ) + ) + comps = { + "SRC": SimSource(SourceSettings(n_ch=n_ch, n_times=n_times, n_messages=n_messages, rebuild_axes=rebuild_axes)), + "LRR1": AffineTransform(AffineTransformSettings(weights=weights, axis="ch")), + "LRR2": AffineTransform(AffineTransformSettings(weights=np.eye(n_out), axis="ch")), + "UP_STATS": Reporter(ReporterSettings(out_dir=out_dir, tag="upstream")), + "DOWN": downstream, + } + conns = ( + (comps["SRC"].OUTPUT_SIGNAL, comps["LRR1"].INPUT_SIGNAL), + (comps["LRR1"].OUTPUT_SIGNAL, comps["LRR2"].INPUT_SIGNAL), + (comps["LRR2"].OUTPUT_SIGNAL, comps["UP_STATS"].INPUT_SIGNAL), + (comps["LRR2"].OUTPUT_SIGNAL, downstream.INPUT_SIGNAL), + ) + ez.run( + components=comps, + connections=conns, + process_components=(downstream,) if split else (), + force_single_process=not split, + ) + + +# How many instances of each class the graph below contains; a correctly +# behaving stream rebuilds each exactly once, on its first message. +INSTANCES = {"AffineTransformTransformer": 2, "FlattenTransformer": 1} + + +def report(out_dir: str, label: str) -> None: + print(f"\n=== {label} ===") + recs = [json.loads(f.read_text()) for f in sorted(pathlib.Path(out_dir).glob("*.json"))] + if not recs: + print(" (no counters written)") + return + print(f" pids reporting: {sorted({r['pid'] for r in recs})}") + for rec in sorted(recs, key=lambda r: r["tag"], reverse=True): + where = "source process" if rec["tag"] == "upstream" else "downstream of the boundary" + if not rec["stats"]: + print(f" {where} (pid {rec['pid']}): no stateful nodes here") + continue + print(f" {where} (pid {rec['pid']}):") + for name, n_resets in sorted(rec["stats"].items()): + # Counters are keyed by class, so a class used twice in the graph + # should show two rebuilds -- one each, on its first message. + expected = INSTANCES.get(name, 1) + flag = ( + "" if n_resets <= expected else f" <-- expected {expected}; the hash is seeing something per-message" + ) + print(f" {name:<40} state rebuilds: {n_resets} (one per instance, {expected} expected){flag}") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--n-ch", type=int, default=256) + p.add_argument("--n-times", type=int, default=30) + p.add_argument("--n-messages", type=int, default=200) + args = p.parse_args() + + print(f"{args.n_messages} messages, {args.n_ch} ch x {args.n_times} samples, ChannelMap ch axis") + for label, rebuild, split in ( + ("single process, source reuses its axis objects", False, False), + ("single process, source rebuilds axes per message", True, False), + ("split across processes, source reuses its axis objects", False, True), + ): + with tempfile.TemporaryDirectory() as d: + build_and_run(args.n_ch, args.n_times, args.n_messages, rebuild, d, split) + report(d, label) + + +if __name__ == "__main__": + main() diff --git a/docs/source/conf.py b/docs/source/conf.py index 7251f62e..30875690 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -33,6 +33,7 @@ # "sphinx_autodoc_typehints", # Disabled due to compatibility issue "sphinx_copybutton", "myst_parser", # For markdown files + "sphinxcontrib.mermaid", # For .. mermaid:: diagrams ] templates_path = ["_templates"] diff --git a/docs/source/guides/sigproc/axis_fingerprint.rst b/docs/source/guides/sigproc/axis_fingerprint.rst new file mode 100644 index 00000000..25ff697e --- /dev/null +++ b/docs/source/guides/sigproc/axis_fingerprint.rst @@ -0,0 +1,142 @@ +When a coordinate axis is fingerprinted +======================================== + +:attr:`~ezmsg.util.messages.axisarray.CoordinateAxis.fingerprint` is a small +hashable stand-in for an axis's *contents*. It is derived from the data rather +than assigned, computed on first access, and cached on the axis object -- so +nothing has to remember to bump it, and the cost is paid once per axis rather +than once per consumer per message. + +Transformers that cache anything resolved from coordinate *values* -- channel +labels into array indices (:obj:`~ezmsg.sigproc.slicer.Slicer`), or into output +labels (:obj:`~ezmsg.sigproc.flatten.Flatten`, +:obj:`~ezmsg.sigproc.affinetransform.AffineTransform`) -- fold it into their +state hash. Without it, a source that renames or reorders channels at a fixed +channel count keeps getting the previously resolved answer, and the operation +silently emits one channel's samples under another channel's label. + +The diagram below traces a serial graph split across two processes. The middle +unit on the left is a filter that needs the channel *count* to size its state +but never looks at the axis values, so it never triggers a fingerprint. + +.. mermaid:: + + sequenceDiagram + autonumber + box transparent Process A + participant SRC as Source
builds ch axis + participant BW as Butterworth
reads shape only + participant SL as Slicer
rewrites ch axis + end + box transparent Process B + participant RR as CommonRereference
reads ch values + participant FL as Flatten
reads ch values + end + + Note over SRC: build ch axis A once, keep as template
A._fingerprint: absent + + rect rgba(128, 128, 128, 0.12) + Note over SRC,FL: message 1 + SRC->>BW: msg(ch: A) + Note over BW: _hash_message = (key, sample_shape)
never touches axes, so never reads a fingerprint + BW->>SL: msg(ch: A) — replace(msg, data=...)
keeps the same axes dict + Note over SL: _hash_message reads A.fingerprint
COMPUTE, approx 1.06 us + Note over SRC,SL: A._fingerprint is now cached. A *is* Source's
template object, so the cache lands upstream. + Note over SL: _reset_state builds B = replace(A, data=A.data[sel])
fast_replace drops _fingerprint, so B is cold + SL->>RR: serialize, then msg(ch: B) + Note over RR: B' deserializes fresh and arrives cold
reads B'.fingerprint, COMPUTE approx 1.06 us + RR->>FL: msg(ch: B') in-process, same object + Note over FL: reads B'.fingerprint, cached, approx 0.05 us + end + + rect rgba(128, 128, 128, 0.12) + Note over SRC,FL: messages 2..N, steady state + SRC->>BW: msg(ch: A) + BW->>SL: msg(ch: A) + Note over SL: A.fingerprint cached, approx 0.05 us
hash unchanged, no reset, re-emits the same B + SL->>RR: serialize, then msg(ch: B) + Note over RR: B is still cold in Process A, so every message
deserializes cold and COMPUTEs again + RR->>FL: msg(ch: B'') + Note over FL: cached on B'', approx 0.05 us + end + +What the diagram is there to show +--------------------------------- + +**A downstream read mutates the upstream object.** ``Slicer`` reading +``A.fingerprint`` populates the cache on the very object ``Source`` holds as its +template, because ``replace(msg, data=...)`` passes axes along by reference. +That is the intended sharing: every later consumer of ``A`` in this process gets +the answer for free. + +**A newly built axis crosses a process boundary cold.** ``Slicer`` *creates* +``B`` and only ever reads ``A``'s fingerprint, so ``B`` is serialized without +one and each message deserializes cold in Process B -- one digest per message +there, shared between its consumers but not free. + +Touching the fingerprint once in whichever unit builds the axis fixes that. +Because the axis is a reused template, every subsequent serialization then +carries the cached value and the downstream process pays nothing:: + + def _reset_state(self, message: AxisArray) -> None: + ... + self._state.new_axis = replace(message.axes[axis], data=out_data) + _ = self._state.new_axis.fingerprint # so it rides the wire precomputed + +Forgetting it costs a microsecond, not correctness -- which is the difference +between this and a hand-maintained generation counter. + +Is it worth pinning axes across the boundary? +---------------------------------------------- + +Measured on a 30x256x2 float32 message with a 256-channel ChannelMap axis +(27 kB) and a feature axis: + +.. list-table:: + :header-rows: 1 + :widths: 55 15 15 15 + + * - per message off the boundary + - cost + - vs. the hop + - core @ 1 kHz + * - the hop itself (serialize + deserialize) + - 27.06 us + - -- + - -- + * - re-fingerprint every message + - 1.77 us + - 6.6% + - 0.18% + * - sender warms the template (one line, above) + - 0.31 us + - 1.1% + - 0.03% + * - staging area with pinned template axes + - 0.97 us + - 3.6% + - 0.10% + +Re-fingerprinting costs 6.6% of a boundary crossing that already costs 27 us, +so the do-nothing case is affordable. A receive-side staging area that compared +each arriving axis against a pinned template cannot beat the one-line sender +warm, because it still has to *read* the fingerprints in order to compare them; +it only adds back object identity, worth about 0.03 us per consumer. + +Why ``fast_replace`` drops the cache +------------------------------------- + +``fast_replace`` is ``arr.__class__(**{**arr.__dict__, **kwargs})``. It is +called on *axes*, not just on messages -- ``replace(message.axes[axis], +data=...)`` appears in :mod:`~ezmsg.sigproc.slicer`, +:mod:`~ezmsg.sigproc.affinetransform`, +:mod:`~ezmsg.sigproc.butterworthzerophase` and :mod:`~ezmsg.sigproc.window`. +Once ``_fingerprint`` is in ``__dict__``, that call becomes +``CoordinateAxis(data=..., dims=..., unit=..., _fingerprint=...)`` and raises +``TypeError: unexpected keyword argument '_fingerprint'``. + +So the drop is first of all what keeps ``replace()`` working, and only secondly +a correctness measure -- forwarding a digest of the *old* values onto a copy +that changes them would be silently wrong. ``replace()`` on an +:obj:`~ezmsg.util.messages.axisarray.AxisArray` is unaffected, since only +``CoordinateAxis`` ever gains the attribute. diff --git a/docs/source/guides/sigproc/content-sigproc.rst b/docs/source/guides/sigproc/content-sigproc.rst index 334c14f3..54a4bf5a 100644 --- a/docs/source/guides/sigproc/content-sigproc.rst +++ b/docs/source/guides/sigproc/content-sigproc.rst @@ -5,5 +5,6 @@ ezmsg-sigproc :maxdepth: 1 architecture + axis_fingerprint ../explanations/sigproc ../explanations/array_api diff --git a/pyproject.toml b/pyproject.toml index 88bb870d..171a0b17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,8 @@ requires-python = ">=3.10" dynamic = ["version"] dependencies = [ "array-api-compat>=1.11.1", - "ezmsg[axisarray]>=3.9.0", - "ezmsg-baseproc>=1.11.0", + "ezmsg[axisarray]>=3.10.0b2", + "ezmsg-baseproc>=1.12.0", "mlx>=0.18.0; sys_platform == 'darwin' and platform_machine == 'arm64'", "numba>=0.61.0", "numpy>=1.26.0", @@ -51,6 +51,7 @@ docs = [ "sphinx_autodoc_typehints>=3.0.0", "sphinx_copybutton", "myst_parser", + "sphinxcontrib-mermaid", ] profile = [ "snakeviz>=2.2.2", diff --git a/src/ezmsg/sigproc/adaptive_lattice_notch.py b/src/ezmsg/sigproc/adaptive_lattice_notch.py index 1462300f..68142210 100644 --- a/src/ezmsg/sigproc/adaptive_lattice_notch.py +++ b/src/ezmsg/sigproc/adaptive_lattice_notch.py @@ -70,11 +70,6 @@ class AdaptiveLatticeNotchFilterTransformer( # `_process`; `axis` and `init_notch_freq` seed cached state and shape. NONRESET_SETTINGS_FIELDS = frozenset({"gamma", "mu", "eta", "chunkwise"}) - def _hash_message(self, message: AxisArray) -> int: - ax_idx = message.get_axis_idx(self.settings.axis) - sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - return hash((message.key, message.axes[self.settings.axis].gain, sample_shape)) - def _reset_state(self, message: AxisArray) -> None: ax_idx = message.get_axis_idx(self.settings.axis) sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] diff --git a/src/ezmsg/sigproc/adaptive_lnc.py b/src/ezmsg/sigproc/adaptive_lnc.py index b2f4761e..cc56abc7 100644 --- a/src/ezmsg/sigproc/adaptive_lnc.py +++ b/src/ezmsg/sigproc/adaptive_lnc.py @@ -287,11 +287,6 @@ class AdaptiveLNCTransformer( # num_harmonics shape the state, so those do force a reset. NONRESET_SETTINGS_FIELDS = frozenset({"adapt_time_constant", "freq_time_constant"}) - def _hash_message(self, message: AxisArray) -> int: - ax_idx = message.get_axis_idx(self.settings.axis) - sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - return hash((message.key, message.axes[self.settings.axis].gain, sample_shape)) - def _reset_state(self, message: AxisArray) -> None: ax_idx = message.get_axis_idx(self.settings.axis) sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] diff --git a/src/ezmsg/sigproc/affinetransform.py b/src/ezmsg/sigproc/affinetransform.py index 9e0126c6..38da4d51 100644 --- a/src/ezmsg/sigproc/affinetransform.py +++ b/src/ezmsg/sigproc/affinetransform.py @@ -35,7 +35,8 @@ from ezmsg.sigproc.util.array import array_device, is_float_dtype, xp_asarray, xp_copy, xp_create, xp_empty from ezmsg.sigproc.util.blockdiag import plan_block_matmul -from ezmsg.sigproc.util.channels import ChannelGroupSpec, group_spec_fingerprint, resolve_channel_groups +from ezmsg.sigproc.util.channels import ChannelGroupSpec, resolve_channel_groups +from ezmsg.sigproc.util.message import with_fingerprint from ezmsg.sigproc.util.rereference import RereferenceKind, rereference_matrix KERNELS = ("auto", "dense", "blocks") @@ -237,14 +238,6 @@ def __call__(self, message: AxisArray) -> AxisArray: return message return super().__call__(message) - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[-1] - axis_idx = message.get_axis_idx(axis) - return hash( - (message.key, message.data.shape[axis_idx]) - + group_spec_fingerprint(message, axis, self.settings.channel_groups) - ) - def _reset_state(self, message: AxisArray) -> None: if self.settings.kernel not in KERNELS: raise ValueError(f"kernel must be one of {KERNELS}, got {self.settings.kernel!r}") @@ -310,7 +303,7 @@ def _reset_state(self, message: AxisArray) -> None: elif np.all(b_filled_outputs): new_labels = np.array(in_labels)[b_used_inputs] - self._state.new_axis = replace(message.axes[axis], data=np.array(new_labels)) + self._state.new_axis = with_fingerprint(replace(message.axes[axis], data=np.array(new_labels))) # Convert to match message.data namespace and device for _process. # Weights are numpy float64 up to here; some devices (e.g. MPS) don't @@ -592,14 +585,6 @@ class CommonRereferenceTransformer( downstream stage.) """ - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[-1] - axis_idx = message.get_axis_idx(axis) - return hash( - (message.key, message.data.shape[axis_idx]) - + group_spec_fingerprint(message, axis, self.settings.channel_groups) - ) - def _reset_state(self, message: AxisArray) -> None: xp = get_namespace(message.data) dev = array_device(message.data) diff --git a/src/ezmsg/sigproc/aggregate.py b/src/ezmsg/sigproc/aggregate.py index 0cafb0ff..997a00c7 100644 --- a/src/ezmsg/sigproc/aggregate.py +++ b/src/ezmsg/sigproc/aggregate.py @@ -29,6 +29,7 @@ ) from .spectral import OptionsEnum +from .util.message import with_fingerprint class AggregationFunction(OptionsEnum): @@ -258,17 +259,6 @@ async def __acall__(self, message: AxisArray) -> AxisArray: return message return await super().__acall__(message) - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[0] - target_axis = message.get_axis(axis) - - hash_components = (message.key,) - if hasattr(target_axis, "data"): - hash_components += (len(target_axis.data),) - elif isinstance(target_axis, AxisArray.LinearAxis): - hash_components += (target_axis.gain, target_axis.offset) - return hash(hash_components) - def _reset_state(self, message: AxisArray) -> None: axis = self.settings.axis or message.dims[0] target_axis = message.get_axis(axis) @@ -294,10 +284,12 @@ def _reset_state(self, message: AxisArray) -> None: ax_dat.append(sl_dat) self._state.slices = slices - self._state.out_axis = AxisArray.CoordinateAxis( - data=np.array(ax_dat), - dims=[axis], - unit=target_axis.unit, + self._state.out_axis = with_fingerprint( + AxisArray.CoordinateAxis( + data=np.array(ax_dat), + dims=[axis], + unit=target_axis.unit, + ) ) def _process(self, message: AxisArray) -> AxisArray: @@ -405,6 +397,9 @@ def _process(self, message: AxisArray) -> AxisArray: data=agg_data, dims=new_dims, axes=new_axes, + # Reducing over the dimension messages appended along leaves nothing + # to append along: each output is one aggregate. + chunk_dim=message.chunk_dim if message.chunk_dim in new_dims else None, ) diff --git a/src/ezmsg/sigproc/align.py b/src/ezmsg/sigproc/align.py index 36aed655..e41e7c90 100644 --- a/src/ezmsg/sigproc/align.py +++ b/src/ezmsg/sigproc/align.py @@ -63,6 +63,21 @@ class AlignAlongAxisProcessor( # -- Helpers ------------------------------------------------------------- + def _hash_message(self, message: AxisArray) -> int: + """Deliberately narrow: only the align axis's sample spacing. + + This processor is fed from two streams in alternation, and they differ + in key, in channel count and in channel labels by construction -- B is a + different device. The default hash would therefore see a change on every + single message and reset the buffers each time. The only property that + must force a reset here is a change in sample spacing, since that is + what the alignment arithmetic depends on. + + `_message_hash` cannot express this: for a coordinate-valued align axis + the spacing is derived from the values rather than read off a `gain`. + """ + return hash(self._extract_gain(message)) + def _extract_gain(self, message: AxisArray) -> float | None: align_name = self.settings.axis or message.dims[0] ax = message.axes.get(align_name) @@ -105,9 +120,6 @@ def _reset_b_state(self) -> None: # -- BaseStatefulTransformer interface ------------------------------------ - def _hash_message(self, message: AxisArray) -> int: - return hash(self._extract_gain(message)) - def _request_reset(self) -> None: # update_settings() calls this (via the base class) when a reset-relevant # setting such as buffer_dur or axis changes, and it invalidates _hash to diff --git a/src/ezmsg/sigproc/binned_aggregate.py b/src/ezmsg/sigproc/binned_aggregate.py index fdf6038f..baeddcae 100644 --- a/src/ezmsg/sigproc/binned_aggregate.py +++ b/src/ezmsg/sigproc/binned_aggregate.py @@ -56,7 +56,7 @@ from .aggregate import AggregationFunction, aggregate_slices, needs_coordinates from .util.array import xp_copy from .util.binning import BinSchedule, BinStep -from .util.message import is_empty_along +from .util.message import is_empty_along, with_fingerprint class BinnedAggregateSettings(ez.Settings): @@ -166,9 +166,6 @@ async def __acall__(self, message: AxisArray) -> AxisArray: return message return await super().__acall__(message) - def _hash_message(self, message: AxisArray) -> int: - return hash((message.axes[self.settings.axis].gain, message.key)) - def _reset_state(self, message: AxisArray) -> None: axis_info = message.get_axis(self.settings.axis) schedule = BinSchedule( @@ -178,9 +175,11 @@ def _reset_state(self, message: AxisArray) -> None: schedule.reset(1.0 / axis_info.gain) self._state.schedule = schedule self._state.metric_axis = ( - AxisArray.CoordinateAxis( - data=np.array([op.value for op in self._operations]), - dims=[self.settings.newaxis], + with_fingerprint( + AxisArray.CoordinateAxis( + data=np.array([op.value for op in self._operations]), + dims=[self.settings.newaxis], + ) ) if self._multi else None diff --git a/src/ezmsg/sigproc/concat.py b/src/ezmsg/sigproc/concat.py index 1d6125d9..22dd719b 100644 --- a/src/ezmsg/sigproc/concat.py +++ b/src/ezmsg/sigproc/concat.py @@ -14,6 +14,9 @@ from ezmsg.util.messages.axisarray import AxisArray, AxisBase, CoordinateAxis from ezmsg.util.messages.util import replace +from ezmsg.sigproc.util.channels import AxisFingerprintMemo +from ezmsg.sigproc.util.message import with_fingerprint + logger = logging.getLogger(__name__) # Sentinel for "attr key was missing on this side". Distinct from any user value. @@ -333,15 +336,23 @@ def _build_cached_axes( align_dim: str | None, merged_concat_axis: CoordinateAxis | None, ) -> dict[str, AxisBase]: - """Build an owned output-axis cache (everything except the alignment axis). + """Build an owned output-axis cache of the *coordinate* axes. Input axes may be views into an ezmsg transport buffer whose lifetime ends after the current subscriber callback. Axes kept across calls must therefore be copied into processor-owned memory. + + Only axes carrying ``.data`` are cached. A ``LinearAxis`` describes itself + with ``gain``/``offset`` scalars, and ``offset`` advances on *every* message + -- caching one freezes the output's time base at whatever the first message + said, which is what ``align_dim`` has been quietly working around (it is the + one axis excluded from the cache, so it alone is re-read from the live + message). Excluding every such axis fixes that generally, costs nothing to + carry live, and saves a deepcopy per rebuild. """ axes: dict[str, AxisBase] = {} for name, ax in a.axes.items(): - if name == align_dim: + if name == align_dim or getattr(ax, "data", None) is None: continue if name == concat_dim and merged_concat_axis is not None: axes[name] = merged_concat_axis @@ -406,6 +417,24 @@ class ConcatSettings(ez.Settings): always an upstream bug and silent coercion hides device<->host copies.""" +@dataclass +class _FingerprintMemo: + """Last-seen objects and their digests, for one input side. + + The axis half is :class:`~ezmsg.sigproc.util.channels.AxisFingerprintMemo`, + shared with the other transformers that cache axis-derived state; ``attrs`` + gets the same identity treatment here because concat is the only consumer + that fingerprints them. + + The cost of *not* doing this: fingerprinting was 63% of ``_concat``'s total + per-message time, against 10% for the concatenate it guards. + """ + + axes: AxisFingerprintMemo = field(default_factory=lambda: AxisFingerprintMemo(label="ConcatProcessor")) + attrs_obj: object = None + attrs_fp: frozenset | None = None + + @dataclass class ConcatState: queue_a: "asyncio.Queue[AxisArray]" = field(default_factory=asyncio.Queue) @@ -416,6 +445,8 @@ class ConcatState: # Fingerprints for cache invalidation. a_fingerprint: tuple | None = None b_fingerprint: tuple | None = None + memo_a: _FingerprintMemo = field(default_factory=_FingerprintMemo) + memo_b: _FingerprintMemo = field(default_factory=_FingerprintMemo) class ConcatProcessor: @@ -452,8 +483,8 @@ async def __acall__(self) -> AxisArray: def _concat(self, a: AxisArray, b: AxisArray) -> AxisArray: """Concatenate *a* and *b* along the configured axis.""" concat_dim = self.settings.axis - fp_a = self._fingerprint(a) - fp_b = self._fingerprint(b) + fp_a = self._fingerprint(a, self._state.memo_a) + fp_b = self._fingerprint(b, self._state.memo_b) if fp_a != self._state.a_fingerprint or fp_b != self._state.b_fingerprint: self._rebuild_cache(a, b) self._state.a_fingerprint = fp_a @@ -483,23 +514,63 @@ def _concat(self, a: AxisArray, b: AxisArray) -> AxisArray: concat_idx = a.dims.index(concat_dim) data = xp.concat([a.data, b.data], axis=concat_idx) - # Build axes: use cached axes + live alignment axis from a. - axes = dict(self._state.cached_axes) if self._state.cached_axes is not None else dict(a.axes) - # Re-insert any axis that changes per-message (e.g. time offset). - for name, ax in a.axes.items(): - if name not in axes: - axes[name] = ax + # Build axes: owned coordinate axes from the cache, everything else + # (alignment axis, LinearAxes) live from a, in a's original order. + cached = self._state.cached_axes + if cached is None: + axes = dict(a.axes) + else: + axes = {name: cached.get(name, ax) for name, ax in a.axes.items()} + # A concat axis created for a *new* dimension is not in a.axes. + for name, ax in cached.items(): + if name not in axes: + axes[name] = ax key = self.settings.new_key if self.settings.new_key is not None else a.key attrs = dict(self._state.merged_attrs) if self._state.merged_attrs else {} - return AxisArray(data, dims=list(a.dims), axes=axes, key=key, attrs=attrs) - - def _fingerprint(self, msg: AxisArray) -> tuple: - concat_dim = self.settings.axis - ax = msg.axes.get(concat_dim) - ax_hash = hash(ax.data.tobytes()) if ax is not None and hasattr(ax, "data") else None - attrs_fp = frozenset((k, type(v).__name__, repr(v)) for k, v in (msg.attrs or {}).items()) - return (tuple(msg.dims), msg.data.shape, ax_hash, attrs_fp) + # Built fresh rather than by replace(), so the layout has to be carried + # over explicitly. A concat along a *new* dimension leaves A's chunk + # dimension intact; concatenating along the chunk dimension itself would + # not, hence the membership check. + chunk_dim = a.chunk_dim if a.chunk_dim in a.dims else None + return AxisArray(data, dims=list(a.dims), axes=axes, key=key, attrs=attrs, chunk_dim=chunk_dim) + + def _fingerprint(self, msg: AxisArray, memo: _FingerprintMemo | None = None) -> tuple: + """Summarize everything ``_rebuild_cache`` reads, so the cache invalidates. + + Every coordinate axis, not just the concat axis: ``_build_cached_axes`` + copies all of them into ``cached_axes``, so a *different* axis changing + its values (e.g. a band axis relabelled mid-stream) leaves the output + carrying the first message's copy. Axes without ``.data`` are read live + rather than cached, so they contribute nothing here. + + *memo* short-circuits the content digests on object identity; see + :class:`_FingerprintMemo`. Passing ``None`` computes everything from + scratch, which is what the tests compare against. + """ + exclude = () if self.settings.align_axis is None else (self.settings.align_axis,) + axes_memo = memo.axes if memo is not None else AxisFingerprintMemo() + axes_fp = axes_memo.fingerprint(msg, exclude=exclude) + + attrs = msg.attrs + if memo is not None and attrs is memo.attrs_obj: + attrs_fp = memo.attrs_fp + else: + # _check_attr_type restricts merged attrs to str/int/float/bool, all + # hashable, so the value goes in as itself -- repr() would be both + # slower and, for anything numpy summarizes past 1000 elements, + # unable to tell two different values apart. But this runs *before* + # that validation, so an unsupported value must not raise here or it + # would mask _check_attr_type's much clearer error. + items = (attrs or {}).items() + try: + attrs_fp = frozenset((k, type(v).__name__, v) for k, v in items) + except TypeError: + attrs_fp = frozenset((k, type(v).__name__, repr(v)) for k, v in items) + if memo is not None: + memo.attrs_obj, memo.attrs_fp = attrs, attrs_fp + + return (tuple(msg.dims), msg.data.shape, axes_fp, attrs_fp) def _rebuild_cache(self, a: AxisArray, b: AxisArray) -> None: concat_dim = self.settings.axis @@ -566,6 +637,9 @@ def _rebuild_cache(self, a: AxisArray, b: AxisArray) -> None: ) self._state.merged_attrs = equal_attrs + if self._state.merged_concat_axis is not None: + with_fingerprint(self._state.merged_concat_axis) + self._state.cached_axes = _build_cached_axes( a, concat_dim, diff --git a/src/ezmsg/sigproc/coordinatespaces.py b/src/ezmsg/sigproc/coordinatespaces.py index c82c9298..57597dec 100644 --- a/src/ezmsg/sigproc/coordinatespaces.py +++ b/src/ezmsg/sigproc/coordinatespaces.py @@ -22,6 +22,8 @@ ) from ezmsg.util.messages.axisarray import AxisArray, replace +from .util.message import with_fingerprint + # -- Utility functions for coordinate transformations -- @@ -142,7 +144,7 @@ def _process(self, message: AxisArray) -> AxisArray: new_labels = np.array(["r", "theta"]) else: new_labels = np.array(["x", "y"]) - axes = {**axes, axis: replace(axes[axis], data=new_labels)} + axes = {**axes, axis: with_fingerprint(replace(axes[axis], data=new_labels))} return replace(message, data=result, axes=axes) diff --git a/src/ezmsg/sigproc/denormalize.py b/src/ezmsg/sigproc/denormalize.py index 2bc943df..68fdbc0d 100644 --- a/src/ezmsg/sigproc/denormalize.py +++ b/src/ezmsg/sigproc/denormalize.py @@ -22,6 +22,11 @@ class DenormalizeSettings(ez.Settings): distribution: str = "uniform" """Distribution to sample rates from. Options are 'uniform', 'normal', or 'constant'.""" + seed: int | None = None + """Seed for the per-channel draw. ``None`` draws fresh values on every state + reset, so a run is not reproducible and a mid-stream channel change silently + re-rolls every channel's rate. Set it for a repeatable simulation.""" + @processor_state class DenormalizeState: @@ -42,10 +47,15 @@ def _reset_state(self, message: AxisArray) -> None: ax_ix = message.get_axis_idx("ch") nch = message.data.shape[ax_ix] arr_size = (nch, 1) if ax_ix == 0 else (1, nch) + # Seeded per reset rather than per instance, so the same channel layout + # always denormalizes the same way -- a reset is triggered by a genuine + # change in the stream, and re-rolling every channel's rate on one would + # put a step discontinuity in the simulated signal. + rng = np.random.default_rng(self.settings.seed) if self.settings.distribution == "uniform": - self.state.offsets = np.random.uniform(2.0, 40.0, size=arr_size) + self.state.offsets = rng.uniform(2.0, 40.0, size=arr_size) elif self.settings.distribution == "normal": - self.state.offsets = np.random.normal( + self.state.offsets = rng.normal( loc=(self.settings.low_rate + self.settings.high_rate) / 2.0, scale=(self.settings.high_rate - self.settings.low_rate) / 6.0, size=arr_size, diff --git a/src/ezmsg/sigproc/diff.py b/src/ezmsg/sigproc/diff.py index 320c9134..609c0d1c 100644 --- a/src/ezmsg/sigproc/diff.py +++ b/src/ezmsg/sigproc/diff.py @@ -45,11 +45,6 @@ async def __acall__(self, message: AxisArray) -> AxisArray: return message return await super().__acall__(message) - def _hash_message(self, message: AxisArray) -> int: - ax_idx = message.get_axis_idx(self.settings.axis) - sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - return hash((sample_shape, message.key)) - def _reset_state(self, message) -> None: ax_idx = message.get_axis_idx(self.settings.axis) # Copied for the same reason as in `_process`: state must never alias the diff --git a/src/ezmsg/sigproc/ewma.py b/src/ezmsg/sigproc/ewma.py index b5257e2a..f33e1cea 100644 --- a/src/ezmsg/sigproc/ewma.py +++ b/src/ezmsg/sigproc/ewma.py @@ -273,12 +273,6 @@ async def __acall__(self, message: AxisArray) -> AxisArray: return message return await super().__acall__(message) - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[0] - axis_idx = message.get_axis_idx(axis) - sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((sample_shape, message.axes[axis].gain, message.key)) - def _reset_state(self, message: AxisArray) -> None: axis = self.settings.axis or message.dims[0] axis_idx = message.get_axis_idx(axis) diff --git a/src/ezmsg/sigproc/fbcca.py b/src/ezmsg/sigproc/fbcca.py index 1188c88d..05f86673 100644 --- a/src/ezmsg/sigproc/fbcca.py +++ b/src/ezmsg/sigproc/fbcca.py @@ -22,6 +22,7 @@ ) from .kaiser import KaiserFilterSettings from .sampler import SampleTriggerMessage +from .util.message import with_fingerprint from .window import WindowSettings, WindowTransformer @@ -158,8 +159,8 @@ def _process(self, message: AxisArray) -> AxisArray: if axis_name not in rm_dims and not (isinstance(axis, AxisArray.CoordinateAxis) and any(d in rm_dims for d in axis.dims)) } - out_axes[self.settings.target_freq_dim] = AxisArray.CoordinateAxis( - np.array(test_freqs), [self.settings.target_freq_dim] + out_axes[self.settings.target_freq_dim] = with_fingerprint( + AxisArray.CoordinateAxis(np.array(test_freqs), [self.settings.target_freq_dim]) ) if message.data.size == 0: diff --git a/src/ezmsg/sigproc/filter.py b/src/ezmsg/sigproc/filter.py index 93c85fc1..e0c5d20e 100644 --- a/src/ezmsg/sigproc/filter.py +++ b/src/ezmsg/sigproc/filter.py @@ -405,12 +405,6 @@ def __call__(self, message: AxisArray) -> AxisArray: self._hash = self._hash_message(message) return super().__call__(message) - def _hash_message(self, message: AxisArray) -> int: - axis = message.dims[0] if self.settings.axis is None else self.settings.axis - axis_idx = message.get_axis_idx(axis) - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((message.key, samp_shape)) - def _build_fir_taps(self, b: npt.NDArray, work_dtype: typing.Any, data: typing.Any, axis_idx: int) -> None: """Cache the converted taps that both dedicated FIR paths read.""" xp = get_namespace(data) @@ -842,13 +836,6 @@ def __call__(self, message: AxisArray) -> AxisArray: return super().__call__(message) - def _hash_message(self, message: AxisArray) -> int: - axis = message.dims[0] if self.settings.axis is None else self.settings.axis - gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1 - axis_idx = message.get_axis_idx(axis) - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((message.key, samp_shape, gain)) - def _reset_state(self, message: AxisArray) -> None: design_fun = self.get_design_function() axis = message.dims[0] if self.settings.axis is None else self.settings.axis diff --git a/src/ezmsg/sigproc/filterbank.py b/src/ezmsg/sigproc/filterbank.py index 6162cde7..5bf439e2 100644 --- a/src/ezmsg/sigproc/filterbank.py +++ b/src/ezmsg/sigproc/filterbank.py @@ -89,19 +89,13 @@ class FilterbankState: class FilterbankTransformer(BaseStatefulTransformer[FilterbankSettings, AxisArray, AxisArray, FilterbankState]): def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[0] - gain = message.axes[axis].gain if axis in message.axes else 1.0 - targ_ax_ix = message.get_axis_idx(axis) - in_shape = message.data.shape[:targ_ax_ix] + message.data.shape[targ_ax_ix + 1 :] + """Extend the default with the input dtype. - return hash( - ( - message.key, - gain if self.settings.mode in [FilterbankMode.FFT, FilterbankMode.AUTO] else None, - message.data.dtype.kind, - in_shape, - ) - ) + The kernels are cast to the message's dtype at reset, so a real input + followed by a complex one needs a rebuild even though nothing about the + stream's shape or channels changed. The default cannot see dtype. + """ + return self._message_hash(message, extra=(message.data.dtype.kind,)) def _reset_state(self, message: AxisArray) -> None: axis = self.settings.axis or message.dims[0] diff --git a/src/ezmsg/sigproc/filterbankdesign.py b/src/ezmsg/sigproc/filterbankdesign.py index 949e4d90..8aa59be5 100644 --- a/src/ezmsg/sigproc/filterbankdesign.py +++ b/src/ezmsg/sigproc/filterbankdesign.py @@ -107,11 +107,11 @@ def __call__(self, message: AxisArray) -> AxisArray: return super().__call__(message) def _hash_message(self, message: AxisArray) -> int: - axis = message.dims[0] if self.settings.axis is None else self.settings.axis - gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1 - axis_idx = message.get_axis_idx(axis) - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((message.key, samp_shape, gain)) + # The only state is a FilterbankTransformer whose kernels are a function + # of the sample rate. That inner transformer keeps its own hash and + # rebuilds itself when the channels change, so folding the channel + # fingerprint in here would only redesign kernels that came out the same. + return hash((message.key, getattr(message.axes.get(self.settings.axis), "gain", None))) def _reset_state(self, message: AxisArray) -> None: axis_obj = message.axes[self.settings.axis] diff --git a/src/ezmsg/sigproc/fir_hilbert.py b/src/ezmsg/sigproc/fir_hilbert.py index 66ba1a12..c0956503 100644 --- a/src/ezmsg/sigproc/fir_hilbert.py +++ b/src/ezmsg/sigproc/fir_hilbert.py @@ -257,13 +257,6 @@ class FIRHilbertEnvelopeTransformer( """ - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[0] - gain = getattr(self._state.filter, "gain", 0.0) - axis_idx = message.get_axis_idx(axis) - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((message.key, samp_shape, gain)) - def _reset_state(self, message: AxisArray) -> None: self._state.filter = FIRHilbertFilterTransformer(settings=self.settings) self._state.delay_buf = None diff --git a/src/ezmsg/sigproc/flatten.py b/src/ezmsg/sigproc/flatten.py index 4fff796c..5eac8349 100644 --- a/src/ezmsg/sigproc/flatten.py +++ b/src/ezmsg/sigproc/flatten.py @@ -42,6 +42,8 @@ ) from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis, replace +from .util.message import with_fingerprint + def normalize_axis_label(label): """Return a hashable string-or-tuple representation of a coord label. @@ -240,11 +242,6 @@ def _expand(arr: np.ndarray, axis_idx: int) -> np.ndarray: class FlattenTransformer(BaseStatefulTransformer[FlattenSettings, AxisArray, AxisArray, FlattenState]): - def _hash_message(self, message: AxisArray) -> int: - preserve_axis = self.settings.preserve_axis or message.dims[0] - non_preserve_shape = tuple(size for dim, size in zip(message.dims, message.data.shape) if dim != preserve_axis) - return hash((tuple(message.dims), non_preserve_shape)) - def _reset_state(self, message: AxisArray) -> None: preserve_axis = self.settings.preserve_axis or message.dims[0] if preserve_axis not in message.dims: @@ -279,12 +276,14 @@ def _reset_state(self, message: AxisArray) -> None: rest_shape = permuted_shape[1 + len(flatten_axes) :] target_inner_shape = (n_flat, *rest_shape) - output_axis_obj = _build_merged_axis( - message, - flatten_axes, - flatten_sizes, - output_axis, - self.settings.label_separator, + output_axis_obj = with_fingerprint( + _build_merged_axis( + message, + flatten_axes, + flatten_sizes, + output_axis, + self.settings.label_separator, + ) ) st = self._state @@ -322,7 +321,14 @@ def _process(self, message: AxisArray) -> AxisArray: if entry is not None: axes[ax] = entry - return replace(message, data=data, dims=list(st.output_dims), axes=axes) + # The preserved dimension may be renamed on the way out, and any + # dimension folded into the merged axis no longer exists to append along. + chunk_dim = message.chunk_dim + if chunk_dim == st.preserve_axis: + chunk_dim = st.sample_axis + elif chunk_dim not in st.output_dims: + chunk_dim = None + return replace(message, data=data, dims=list(st.output_dims), axes=axes, chunk_dim=chunk_dim) class Flatten(BaseTransformerUnit[FlattenSettings, AxisArray, AxisArray, FlattenTransformer]): diff --git a/src/ezmsg/sigproc/linear.py b/src/ezmsg/sigproc/linear.py index b92e566e..9375cc11 100644 --- a/src/ezmsg/sigproc/linear.py +++ b/src/ezmsg/sigproc/linear.py @@ -62,14 +62,6 @@ class LinearTransformTransformer( ... )) """ - def _hash_message(self, message: AxisArray) -> int: - """Hash based on shape and axis to detect when broadcast shapes need recalculation.""" - axis = self.settings.axis - if axis is not None: - axis_idx = message.get_axis_idx(axis) - return hash((message.data.ndim, axis_idx, message.data.shape[axis_idx])) - return hash(message.data.ndim) - def _reset_state(self, message: AxisArray) -> None: """Prepare scale/offset arrays with proper broadcast shapes.""" xp = get_namespace(message.data) diff --git a/src/ezmsg/sigproc/resample.py b/src/ezmsg/sigproc/resample.py index 30a7f6da..feb2576a 100644 --- a/src/ezmsg/sigproc/resample.py +++ b/src/ezmsg/sigproc/resample.py @@ -164,13 +164,6 @@ class ResampleProcessor(BaseStatefulProcessor[ResampleSettings, AxisArray, AxisA # `resample_rate` / `buffer_duration` / `axis` all size cached buffers. NONRESET_SETTINGS_FIELDS = frozenset({"max_chunk_delay", "fill_value", "reference_reset_after_chunks"}) - def _hash_message(self, message: AxisArray) -> int: - ax_idx: int = message.get_axis_idx(self.settings.axis) - sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - ax = message.axes[self.settings.axis] - gain = ax.gain if hasattr(ax, "gain") else None - return hash((message.key, gain) + sample_shape) - def _reset_state(self, message: AxisArray) -> None: """ Reset the internal state based on the incoming message. diff --git a/src/ezmsg/sigproc/rollingscaler.py b/src/ezmsg/sigproc/rollingscaler.py index dd0489b0..9e27f45d 100644 --- a/src/ezmsg/sigproc/rollingscaler.py +++ b/src/ezmsg/sigproc/rollingscaler.py @@ -108,13 +108,6 @@ class RollingScalerProcessor(BaseAdaptiveTransformer[RollingScalerSettings, Axis # are cached during `_reset_state`. NONRESET_SETTINGS_FIELDS = frozenset({"update_with_signal", "artifact_z_thresh", "clip"}) - def _hash_message(self, message: AxisArray) -> int: - axis = message.dims[0] if self.settings.axis is None else self.settings.axis - gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1 - axis_idx = message.get_axis_idx(axis) - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((message.key, samp_shape, gain)) - def _reset_state(self, message: AxisArray) -> None: xp = get_namespace(message.data) ch = message.data.shape[-1] diff --git a/src/ezmsg/sigproc/sampler.py b/src/ezmsg/sigproc/sampler.py index 04a2bc50..69a33607 100644 --- a/src/ezmsg/sigproc/sampler.py +++ b/src/ezmsg/sigproc/sampler.py @@ -87,13 +87,6 @@ def __call__(self, message: AxisArray | SampleTriggerMessage) -> list[AxisArray] else: return self.push_trigger(message) - def _hash_message(self, message: AxisArray) -> int: - # Compute hash based on message properties that require state reset - axis = self.settings.axis or message.dims[0] - axis_idx = message.get_axis_idx(axis) - sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - return hash((sample_shape, message.key)) - def _reset_state(self, message: AxisArray) -> None: self._state.buffer = HybridAxisArrayBuffer( duration=self.settings.buffer_dur, diff --git a/src/ezmsg/sigproc/scaler.py b/src/ezmsg/sigproc/scaler.py index 77e6654e..19ddb98a 100644 --- a/src/ezmsg/sigproc/scaler.py +++ b/src/ezmsg/sigproc/scaler.py @@ -138,6 +138,14 @@ def update_settings(self, new_settings: AdaptiveStandardScalerSettings) -> None: self.accumulate = new_settings.accumulate super().update_settings(new_settings) + def _hash_message(self, message: AxisArray) -> int: + # This transformer owns no array state of its own -- only the two child + # EWMATransformers, which hash the message themselves and rebuild on + # exactly the changes that matter. Hashing here would duplicate that work + # and, on a reset, discard the children wholesale rather than letting + # each rebuild the part of its state that actually went stale. + return 0 + def _reset_state(self, message: AxisArray) -> None: self._state.samps_ewma = EWMATransformer( time_constant=self.settings.time_constant, diff --git a/src/ezmsg/sigproc/signalinjector.py b/src/ezmsg/sigproc/signalinjector.py index ec660f2b..4083e25c 100644 --- a/src/ezmsg/sigproc/signalinjector.py +++ b/src/ezmsg/sigproc/signalinjector.py @@ -32,11 +32,6 @@ class SignalInjectorState: class SignalInjectorTransformer( BaseAsyncTransformer[SignalInjectorSettings, AxisArray, AxisArray, SignalInjectorState] ): - def _hash_message(self, message: AxisArray) -> int: - time_ax_idx = message.get_axis_idx(self.settings.time_dim) - sample_shape = message.data.shape[:time_ax_idx] + message.data.shape[time_ax_idx + 1 :] - return hash((message.key,) + sample_shape) - def _reset_state(self, message: AxisArray) -> None: if self._state.cur_frequency is None: self._state.cur_frequency = self.settings.frequency diff --git a/src/ezmsg/sigproc/slicer.py b/src/ezmsg/sigproc/slicer.py index 70592f66..2080df90 100644 --- a/src/ezmsg/sigproc/slicer.py +++ b/src/ezmsg/sigproc/slicer.py @@ -17,6 +17,8 @@ slice_along_axis, ) +from .util.message import with_fingerprint + """ Slicer:Select a subset of data along a particular axis. """ @@ -191,11 +193,6 @@ class SlicerState: class SlicerTransformer(BaseStatefulTransformer[SlicerSettings, AxisArray, AxisArray, SlicerState]): - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[-1] - axis_idx = message.get_axis_idx(axis) - return hash((message.key, message.data.shape[axis_idx])) - def _selects_positional_int(self, axinfo: AxisArray.CoordinateAxis | None) -> bool: """True iff the selection is a single bare-integer token that parse_slice resolved positionally (its int() path) rather than via a label match.""" @@ -296,7 +293,7 @@ def _reset_state(self, message: AxisArray) -> None: out_data = in_data[self._state.slice_ : self._state.slice_ + 1] else: out_data = in_data[self._state.slice_] - self._state.new_axis = replace(message.axes[axis], data=out_data) + self._state.new_axis = with_fingerprint(replace(message.axes[axis], data=out_data)) def _process(self, message: AxisArray) -> AxisArray: axis = self.settings.axis or message.dims[-1] diff --git a/src/ezmsg/sigproc/spectrum.py b/src/ezmsg/sigproc/spectrum.py index b7272452..93a28610 100644 --- a/src/ezmsg/sigproc/spectrum.py +++ b/src/ezmsg/sigproc/spectrum.py @@ -134,11 +134,23 @@ class SpectrumState: class SpectrumTransformer(BaseStatefulTransformer[SpectrumSettings, AxisArray, AxisArray, SpectrumState]): def _hash_message(self, message: AxisArray) -> int: + """Extend the default with the two things it cannot know about. + + The FFT is sized by the length of the axis being transformed, which is + normally the chunk dimension -- the one length the default deliberately + ignores because it changes with every message. Spectrum is the exception: + a different transform length is a different plan, so it has to be folded + back in. The dtype matters because a complex input takes a different + branch and produces a different frequency axis. + """ axis = self.settings.axis or message.dims[0] - ax_idx = message.get_axis_idx(axis) - ax_info = message.axes[axis] - targ_len = message.data.shape[ax_idx] - return hash((targ_len, message.data.ndim, is_complex_dtype(message.data.dtype), ax_idx, ax_info.gain)) + return self._message_hash( + message, + extra=( + message.data.shape[message.get_axis_idx(axis)], + is_complex_dtype(message.data.dtype), + ), + ) def _reset_state(self, message: AxisArray) -> None: axis = self.settings.axis or message.dims[0] @@ -261,7 +273,13 @@ def _process(self, message: AxisArray) -> AxisArray: spec = self.state.f_transform(spec) spec = slice_along_axis(spec, self.state.f_sl, message.get_axis_idx(axis)) - msg_out = replace(message, data=spec, dims=self.state.new_dims, axes=new_axes) + # The transformed axis is consumed: `time` becomes `freq`. If that was + # the dimension successive messages appended along, the output has none + # -- each message is one spectrum, and they stack rather than + # concatenate. Windowed input keeps its `win` dimension and so keeps + # appending along it. + out_chunk_dim = None if message.chunk_dim == axis else message.chunk_dim + msg_out = replace(message, data=spec, dims=self.state.new_dims, axes=new_axes, chunk_dim=out_chunk_dim) return msg_out diff --git a/src/ezmsg/sigproc/transpose.py b/src/ezmsg/sigproc/transpose.py index b69f851b..f8b9e16e 100644 --- a/src/ezmsg/sigproc/transpose.py +++ b/src/ezmsg/sigproc/transpose.py @@ -52,9 +52,6 @@ class TransposeTransformer(BaseStatefulTransformer[TransposeSettings, AxisArray, # `axes` drives the cached permutation in `_reset_state`. NONRESET_SETTINGS_FIELDS = frozenset({"order"}) - def _hash_message(self, message: AxisArray) -> int: - return hash(tuple(message.dims)) - def _reset_state(self, message: AxisArray) -> None: if self.settings.axes is None: self._state.axes_ints = None diff --git a/src/ezmsg/sigproc/util/channels.py b/src/ezmsg/sigproc/util/channels.py index 6a86c187..5ce6f442 100644 --- a/src/ezmsg/sigproc/util/channels.py +++ b/src/ezmsg/sigproc/util/channels.py @@ -22,16 +22,55 @@ Groups are returned in first-appearance order along the channel axis so a resolved grouping is reproducible and readable against the channel table. + +This module also owns the two ways a transformer can fold a channel axis into +its per-message state hash: :func:`group_spec_fingerprint` (O(1), notices only +that the metadata *field* appeared or vanished) and +:func:`coord_value_fingerprint` (O(bytes), notices the *values* changing). Their +docstrings explain which failure mode each one is for. """ from __future__ import annotations +import os +import zlib from collections.abc import Callable, Sequence from typing import Union import numpy as np from ezmsg.util.messages.axisarray import AxisArray +# Whether AxisFingerprintMemo counts its hits and misses. Off unless the env var +# is set, because the answer it gives -- do axis objects survive, or is every +# message a fresh deserialization? -- is a property of a *deployed* graph's +# process layout, not something a unit test can tell you. See +# benchmarks/benchmark_memo_hit_rate.py. +_STATS_ENABLED = bool(os.environ.get("EZMSG_SIGPROC_FINGERPRINT_STATS")) +_STATS: dict[str, list[int]] = {} + + +def fingerprint_stats() -> dict[str, tuple[int, int, int]]: + """``{label: (calls, digests_computed, mapping_hits)}`` for this process. + + Empty unless ``EZMSG_SIGPROC_FINGERPRINT_STATS=1``. Counters are + per-process, so a multi-process graph has to collect them from each worker. + + ``digests_computed`` is the number that matters -- how often the memo failed + to save the O(bytes) work. ``mapping_hits`` separates the two shortcuts: the + whole-``axes``-mapping check, which only fires when an upstream node passed + the mapping through untouched, from the per-array check that does the real + work (most nodes rebuild the mapping via ``replace(..., axes={...})`` even + when the axis objects inside it are unchanged). + """ + return {label: (c[0], c[1], c[2]) for label, c in _STATS.items()} + + +def reset_fingerprint_stats() -> None: + """Zero every counter in this process.""" + for counts in _STATS.values(): + counts[0] = counts[1] = counts[2] = 0 + + ChannelGroupSpec = Union[ str, Sequence[str], @@ -181,6 +220,15 @@ def group_spec_fingerprint( fixed key and channel count is deliberately not detected — a genuine remap arrives with a new key or channel count. + That concession is safe for a *grouping* (a stale grouping is arithmetic on + the wrong partition, which a changed key or channel count would have caught) + but not for every consumer of channel metadata. An operation whose cached + state is a set of resolved *indices* -- where a stale answer emits one + channel's samples under another channel's label -- wants + :func:`coord_value_fingerprint` instead, which costs O(bytes) but actually + tracks the values. The two are a deliberate pair; pick by whether a silent + stale answer is recoverable downstream. + The two common specs -- ``None`` and a single field name -- are classified inline rather than through :func:`group_spec_fields`, because at this call rate the function call itself is a measurable share of the cost. Both @@ -195,3 +243,200 @@ def group_spec_fingerprint( ax = message.axes.get(axis or message.dims[-1]) names = getattr(getattr(getattr(ax, "data", None), "dtype", None), "names", None) return (bool(names) and all(field in names for field in fields),) + + +def array_value_fingerprint(arr: np.ndarray) -> tuple: + """Content digest of one array: ``(dtype, shape, checksum)``. + + The shared primitive under :func:`coord_value_fingerprint`, also used + directly by consumers that already hold the array (e.g. + :class:`~ezmsg.sigproc.concat.ConcatProcessor`, fingerprinting each axis it + caches). + + ``zlib.crc32`` rather than ``hash(arr.tobytes())`` because the bottleneck is + the hash, not the copy. Measured on a 256-channel ChannelMap axis (27.6 kB, + Apple M-series): the ``tobytes()`` copy is 0.30 µs (94 GB/s) while CPython's + siphash over the result is 4.7 µs (5.5 GB/s); ``crc32`` reads the array's + buffer directly at 29 GB/s for 0.94 µs total -- 5.3x cheaper. + + The tradeoff is a 32-bit checksum, so a collision means a missed state + reset. ``dtype`` and ``shape`` ride along both because they are nearly free + and because they carry most of the structural change a checksum could alias. + + The dtype goes in as the ``np.dtype`` object, not ``str(dtype)``: numpy + builds a structured dtype's repr field by field, which costs 9.8 µs for the + eight-field ChannelMap above -- ten times the checksum it was annotating. + The object is hashable and compares by value, so it does the same job for + 0.02 µs. + + ``crc32`` needs a C-contiguous buffer, which a struct-array *field* view + never is, so the gather is explicit here rather than left to fail. + """ + arr = np.ascontiguousarray(arr) + if arr.dtype.hasobject: + # An object array's buffer is pointers: two equal arrays built from + # distinct string objects have different bytes, so checksumming it + # would reset the state on every message. Widen to a real dtype first. + try: + arr = np.ascontiguousarray(arr.astype("U")) + except (TypeError, ValueError): + # Elements with no string form -- vanishingly rare on a coordinate + # axis. Correctness over speed: repr is content-based and stable. + return (arr.dtype, arr.shape, repr(arr.tolist())) + return (arr.dtype, arr.shape, zlib.crc32(arr)) + + +def coord_value_fingerprint( + message: AxisArray, + axis: str | None, + fields: Sequence[str] | None = None, +) -> tuple: + """Digest of the coordinate *values* on *axis*, restricted to *fields*. + + The value-sensitive counterpart to :func:`group_spec_fingerprint`, for + transformers that cache indices resolved against coordinate values (labels, + regex matches, field matches). Folding this into ``_hash_message`` makes such + a cache re-resolve when a source renames, reorders or swaps out channels + without changing its key or channel count. + + Args: + message: The message whose axis is being fingerprinted. + axis: Coordinate axis name. ``None`` defaults to the last dimension. + fields: Struct-array fields the consumer actually matches against. + ``None`` digests the whole coordinate array, which is what an + unstructured (plain label) axis needs. A named field absent from the + dtype contributes ``None``, so gaining or losing it still registers. + + Returns: + A hashable tuple, empty when the axis carries no coordinate data. + + Restricting to *fields* is about **invalidation correctness, not speed**: a + source that recomputes float ``x``/``y`` positions each message would + otherwise reset the state continuously, even for a selection that only ever + reads ``label``. It is sometimes also cheaper and sometimes not. Measured on + a 256-channel ChannelMap (eight fields, 108 B itemsize, 27.6 kB total): + + ============================== ========== ========================== + ``fields`` cost vs. whole axis (1.15 µs) + ============================== ========== ========================== + ``('bank',)`` (U2, 7% of bytes) 0.95 µs cheaper + ``('array', 'bank')`` (11%) 1.39 µs *more expensive* + ``('label',)`` (U16, 59%) 2.65 µs *more expensive* + ============================== ========== ========================== + + A wide field loses because extracting it is a strided gather (7-20 GB/s) + while the whole axis is one contiguous read (29 GB/s). Fields are digested + one at a time for the same reason numpy makes multi-field indexing a trap: + ``arr[['array', 'bank']]`` returns a view that keeps the *original* + itemsize, so its ``tobytes()`` is the entire 27.6 kB -- asking for two of + eight fields would otherwise cost more than asking for all of them. + """ + ax = message.axes.get(axis or message.dims[-1]) + data = getattr(ax, "data", None) + if data is None: + return () + names = getattr(getattr(data, "dtype", None), "names", None) + if not fields or names is None: + return array_value_fingerprint(data) + return tuple(array_value_fingerprint(data[f]) if f in names else None for f in fields) + + +class AxisFingerprintMemo: + """Per-consumer, identity-first fingerprints of a message's coordinate axes. + + A transformer that caches anything derived from axis *values* -- resolved + indices, output labels -- has to notice when those values change under a + fixed key and shape, and the honest check is O(bytes). This makes it O(1) + in the case that actually occurs. + + ``replace()`` carries ``axes``, the axis objects and their ``.data`` arrays + by reference, so a message threaded through a chain of transformers presents + the *same objects* every time. Two ``is`` checks -- first the whole ``axes`` + mapping, then each array -- settle it without touching the bytes. A miss + just computes the digest, so this is a pure fast path: it can make the check + cheaper, never wrong. + + Measured across a 20-node graph checking one 256-channel ChannelMap axis: + 99.7 µs to digest per node per message, 0.6 µs with this. After a + cross-process hop every object is fresh, so it degrades to ~20 µs -- see + ``benchmarks/benchmark_axis_fingerprint.py``. + + **The contract this assumes**: a coordinate array is never mutated in place. + Messages fan out to multiple graph branches, so mutating one is already + unsafe; this turns that into a requirement. + + One memo belongs to one consumer, which must pass the same *names* and + *exclude* on every call -- the whole-mapping shortcut caches a single answer + per ``axes`` object and cannot tell that the question changed. + """ + + __slots__ = ("_axes_obj", "_axes_fp", "_per_axis", "_counts") + + def __init__(self, label: str | None = None) -> None: + self._axes_obj: object = None + self._axes_fp: tuple = () + self._per_axis: dict[str, tuple] = {} + # None unless stats are enabled, so the hot path is one `is not None`. + # [calls, digests_computed, mapping_hits]; None unless stats are on, + # so the hot path costs one `is not None`. + self._counts: list[int] | None = _STATS.setdefault(label, [0, 0, 0]) if _STATS_ENABLED and label else None + + def fingerprint( + self, + message: AxisArray, + names: Sequence[str] | None = None, + exclude: Sequence[str] = (), + ) -> tuple: + """Digest the coordinate axes' values, as ``((name, digest), ...)``. + + Args: + message: Message whose axes are being fingerprinted. + names: Restrict to these axes. ``None`` covers every coordinate + axis, which is the safe default when the consumer's own axis + selection is not known until state reset. + exclude: Axes to skip even when *names* is ``None`` -- for an axis + deliberately read live rather than cached. + + Returns: + A hashable tuple; empty when no axis carries coordinate data. Axes + without ``.data`` contribute nothing, since a ``LinearAxis`` + compares by value for free. + """ + axes = message.axes + if self._counts is not None: + self._counts[0] += 1 + if axes is self._axes_obj: + if self._counts is not None: + self._counts[2] += 1 + return self._axes_fp + + computed = 0 + parts = [] + for name, ax in axes.items(): + if name in exclude or (names is not None and name not in names): + continue + data = getattr(ax, "data", None) + if data is None: + continue + seen = self._per_axis.get(name) + if seen is not None and seen[0] is data: + parts.append((name, seen[1])) + continue + # A CoordinateAxis from ezmsg >= 3.10 derives and caches its own + # fingerprint, which rides the pickle across a process boundary -- + # the one place this memo can never hit, since every message + # deserializes fresh objects. Ask before digesting. Older ezmsg has + # no such attribute and falls through, so this stays version- + # agnostic; within one process the answer is the same for every + # axis, so the tuple never mixes the two forms. + fp = getattr(ax, "fingerprint", None) + if fp is None: + fp = array_value_fingerprint(data) + computed += 1 + self._per_axis[name] = (data, fp) + parts.append((name, fp)) + + if self._counts is not None and computed: + self._counts[1] += 1 + self._axes_obj, self._axes_fp = axes, tuple(parts) + return self._axes_fp diff --git a/src/ezmsg/sigproc/util/message.py b/src/ezmsg/sigproc/util/message.py index fc8eb031..5011b634 100644 --- a/src/ezmsg/sigproc/util/message.py +++ b/src/ezmsg/sigproc/util/message.py @@ -20,9 +20,34 @@ "has_samples_along", "is_empty_along", "is_sample_message", + "with_fingerprint", ] +def with_fingerprint(axis: AxisArray.CoordinateAxis) -> AxisArray.CoordinateAxis: + """Compute *axis*'s fingerprint now, and return the axis. + + Every stateful consumer reads the fingerprint of the coordinate axes that + describe a stream's configuration, and the value is cached on the instance + and pickled with it. Computing it at the point of construction therefore + pays the checksum once, for everybody: + + * In this process, the axis object is reused for the life of the stream, so + one call covers every message and every consumer downstream of it. + * Across a process boundary it is better than that. Unpickling hands out a + *new* axis object per message, so a cold axis is re-checksummed by the + first consumer in every receiving process, on every message, forever. + A primed one arrives with the answer already attached. + + Apply it to axes that describe the stream -- channel labels, frequency + labels, feature labels -- not to per-message coordinates along the chunk + dimension, whose fingerprint no consumer reads and whose data is new every + message anyway. + """ + axis.fingerprint # noqa: B018 -- evaluated for the caching side effect + return axis + + def is_empty_along(message: AxisArray, dims: typing.Iterable[str]) -> bool: """True iff any of the named dims is present in ``message`` with zero length. diff --git a/src/ezmsg/sigproc/wavelets.py b/src/ezmsg/sigproc/wavelets.py index 71b87385..3890b714 100644 --- a/src/ezmsg/sigproc/wavelets.py +++ b/src/ezmsg/sigproc/wavelets.py @@ -16,6 +16,7 @@ from ezmsg.util.messages.util import replace from .filterbank import FilterbankMode, MinPhaseMode, filterbank +from .util.message import with_fingerprint class CWTSettings(ez.Settings): @@ -42,16 +43,8 @@ class CWTState: class CWTTransformer(BaseStatefulTransformer[CWTSettings, AxisArray, AxisArray, CWTState]): def _hash_message(self, message: AxisArray) -> int: - ax_idx = message.get_axis_idx(self.settings.axis) - in_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] - return hash( - ( - message.data.dtype.kind, - message.axes[self.settings.axis].gain, - in_shape, - message.key, - ) - ) + """Extend the default with the input dtype, which sizes the output buffer.""" + return self._message_hash(message, extra=(message.data.dtype.kind,)) def _reset_state(self, message: AxisArray) -> None: if "freq" in message.dims: @@ -121,7 +114,7 @@ def _reset_state(self, message: AxisArray) -> None: dims=message.dims[:ax_idx] + message.dims[ax_idx + 1 :] + ["freq", self.settings.axis], axes={ **{k: deepcopy(v) for k, v in message.axes.items()}, - "freq": AxisArray.CoordinateAxis(unit="Hz", data=freqs, dims=["freq"]), + "freq": with_fingerprint(AxisArray.CoordinateAxis(unit="Hz", data=freqs, dims=["freq"])), }, key=message.key, ) diff --git a/src/ezmsg/sigproc/window.py b/src/ezmsg/sigproc/window.py index a4df0c33..ba0b213b 100644 --- a/src/ezmsg/sigproc/window.py +++ b/src/ezmsg/sigproc/window.py @@ -127,6 +127,7 @@ class WindowState: """Target axis re-anchored per ``anchor``; constant for the life of the state.""" out_dims: list[str] | None = None + out_chunk_dim: str | None = None empty_out: npt.NDArray | sparse.SparseArray | None = None """Cached zero-window output, returned unchanged whenever no window is due.""" @@ -251,15 +252,6 @@ def is_batcher(self) -> bool: and self.settings.window_shift == self.settings.window_dur ) - def _hash_message(self, message: AxisArray) -> int: - axis = self.settings.axis or message.dims[0] - axis_idx = message.get_axis_idx(axis) - axis_info = message.get_axis(axis) - fs = 1.0 / axis_info.gain - samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :] - - return hash(samp_shape + (fs, message.key)) - def _reset_state(self, message: AxisArray) -> None: _newaxis = self.settings.newaxis or "win" if not self._state.newaxis_warned and _newaxis in message.dims: @@ -316,9 +308,11 @@ def _reset_state(self, message: AxisArray) -> None: self._state.out_axis = None self._state.empty_out = None if self.is_batcher: - # Windows tile the target axis, so they need no axis of their own. + # Windows tile the target axis, so they need no axis of their own, + # and the stream still grows along whichever dim it did before. self._state.out_dims = list(message.dims) self._state.out_newaxis = None + self._state.out_chunk_dim = message.chunk_dim else: self._state.out_dims = list(message.dims[:axis_idx]) + [_newaxis] + list(message.dims[axis_idx:]) self._state.out_newaxis = replace( @@ -326,6 +320,12 @@ def _reset_state(self, message: AxisArray) -> None: gain=0.0 if self.settings.window_shift is None else axis_info.gain * self._state.window_shift_samples, offset=0.0, # offset modified per-msg below ) + # Successive messages now append along `newaxis`: its length is the + # number of windows this message happened to yield, while the target + # axis has become a fixed-length within-window axis. Declaring it + # spares every downstream consumer from having to guess, and gets + # the guess right where a "time" convention would not. + self._state.out_chunk_dim = _newaxis def __call__(self, message: AxisArray) -> AxisArray: if self.settings.window_dur is None: @@ -409,7 +409,13 @@ def _process(self, message: AxisArray) -> AxisArray: out_dat = self._state.empty_out out_offset = axis_info.offset out_axes[axis] = replace(axis_info, offset=out_offset) - return replace(message, data=out_dat, dims=self._state.out_dims, axes=out_axes) + return replace( + message, + data=out_dat, + dims=self._state.out_dims, + axes=out_axes, + chunk_dim=self._state.out_chunk_dim, + ) # Update targeted (windowed) axis so that its offset is relative to the new axis. # The result depends only on the axis gain and the settings, both fixed for @@ -474,6 +480,7 @@ def _process(self, message: AxisArray) -> AxisArray: data=out_dat, dims=self._state.out_dims, axes={**out_axes, _newaxis: self._state.out_newaxis}, + chunk_dim=self._state.out_chunk_dim, ) return msg_out @@ -523,6 +530,10 @@ async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator: data=slice_along_axis(ret.data, msg_ix, axis_idx), dims=ret.dims[:axis_idx] + ret.dims[axis_idx + 1 :], axes=_out_axes, + # Unbundling drops `win`, so the published stream is + # back to appending along the target axis: one + # message per window, each carrying its own offset. + chunk_dim=axis, ) yield self.OUTPUT_SIGNAL, _ret diff --git a/tests/helpers/recycled_shm.py b/tests/helpers/recycled_shm.py index 253fc171..bd7d7bcf 100644 --- a/tests/helpers/recycled_shm.py +++ b/tests/helpers/recycled_shm.py @@ -72,10 +72,22 @@ def _detach(result): The output of a transformer may itself alias the message it came from, which is legitimate -- it is handed straight downstream and not retained -- but it means the collected outputs have to be copied before the next publish. + + Coordinate axes are copied too, not just ``data``. A passed-through ``ch`` + axis is as much a view onto the slot as the samples are, so leaving it + attached made every collected output's labels turn to garbage on the next + publish -- which the comparison could not see while + ``CoordinateAxis.__eq__`` compared only ``unit`` (ezmsg-org/ezmsg#258 + stack). With that fixed, failing to copy here reports every transformer as + retaining its axes. """ if result is None: return None - return replace(result, data=np.array(result.data)) + axes = { + name: (replace(ax, data=np.array(ax.data)) if getattr(ax, "data", None) is not None else ax) + for name, ax in result.axes.items() + } + return replace(result, data=np.array(result.data), axes=axes) def run_recycled(proc, messages: typing.Sequence[AxisArray], *, slot_bytes: int = SLOT_BYTES) -> list: diff --git a/tests/unit/test_affine_transform.py b/tests/unit/test_affine_transform.py index 4159c461..13fe38b1 100644 --- a/tests/unit/test_affine_transform.py +++ b/tests/unit/test_affine_transform.py @@ -389,12 +389,16 @@ def test_common_rereference_all_singleton_groups_exclude_current(): assert np.array_equal(xformer(AxisArray(in_dat, dims=["time", "ch"])).data, in_dat) -def test_common_rereference_field_values_change_is_not_detected(): - """Intentional concession: a live bank remap at fixed key + channel count is - NOT re-derived. _hash_message folds only an O(1) "field present" boolean, not - the field's bytes, to keep the per-message hash from scaling with channel - count. A genuine remap on real hardware arrives with a new key or channel - count (see the escape-hatch assertion below).""" +def test_common_rereference_field_values_change_is_detected(): + """A live bank remap at a fixed key and channel count re-derives the groups. + + This used to be a documented concession: the hash folded only an O(1) + "is the field present" boolean, so a remap that kept the same key and + channel count went unnoticed and the channels were referenced against the + wrong neighbours. The base-class hash now folds in the channel axis's + content fingerprint, which is computed once per axis object and cached on + it, so detecting the remap no longer costs a per-message walk of the field. + """ n_times = 80 rng = np.random.default_rng(11) in_dat = rng.standard_normal((n_times, 4)) @@ -406,16 +410,16 @@ def test_common_rereference_field_values_change_is_not_detected(): xformer(msg1) assert [list(g) for g in xformer._state.groups] == [[0, 1], [2, 3]] - # Same key and channel count, different bank assignment -> hash unchanged, - # so the cached groups are (deliberately) NOT re-derived. + # Same key, same channel count, different bank assignment -> re-derived. msg2 = AxisArray(in_dat, dims=["time", "ch"], axes={"ch": _banked_ch_axis(["A", "B", "A", "B"])}, key="dev") xformer(msg2) - assert [list(g) for g in xformer._state.groups] == [[0, 1], [2, 3]] + assert [list(g) for g in xformer._state.groups] == [[0, 2], [1, 3]] - # Escape hatch: a new key (as a real remap would carry) forces re-derivation. - msg3 = AxisArray(in_dat, dims=["time", "ch"], axes={"ch": _banked_ch_axis(["A", "B", "A", "B"])}, key="dev2") + # An unchanged layout must not re-derive, or every message would pay for it. + before = xformer._hash + msg3 = AxisArray(in_dat, dims=["time", "ch"], axes={"ch": _banked_ch_axis(["A", "B", "A", "B"])}, key="dev") xformer(msg3) - assert [list(g) for g in xformer._state.groups] == [[0, 2], [1, 3]] + assert xformer._hash == before def test_common_rereference_explicit_groups_beat_field(): @@ -1228,3 +1232,47 @@ def test_stacked_bias_repeat_processing_is_stable(): first = np.asarray(proc(msg).data).copy() second = np.asarray(proc(msg).data) assert np.array_equal(first, second) + + +def test_affine_transform_non_square_follows_a_relabel(): + """A non-square transform caches an output axis selected from the input + labels, so a relabel at a fixed channel count has to re-derive it.""" + weights = np.array([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]]) # drops channel 3 + + def msg(labels): + return AxisArray( + np.zeros((4, 3), np.float32), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": AxisArray.CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key="dev", + ) + + xformer = AffineTransformTransformer(AffineTransformSettings(weights=weights, axis="ch")) + assert [str(x) for x in xformer(msg(["A", "B", "C"])).axes["ch"].data] == ["A", "B"] + assert [str(x) for x in xformer(msg(["X", "Y", "Z"])).axes["ch"].data] == ["X", "Y"] + + +def test_affine_transform_unchanged_labels_do_not_reset(): + """An equal axis rebuilt per message must not look like a change, or the + weights would be re-derived at the sample rate.""" + weights = np.eye(3) + + def msg(): + return AxisArray( + np.zeros((4, 3), np.float32), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": AxisArray.CoordinateAxis(data=np.array(["A", "B", "C"]), dims=["ch"]), + }, + key="dev", + ) + + xformer = AffineTransformTransformer(AffineTransformSettings(weights=weights, axis="ch")) + xformer(msg()) + first_hash = xformer._hash + xformer(msg()) + assert xformer._hash == first_hash diff --git a/tests/unit/test_axis_fingerprint_priming.py b/tests/unit/test_axis_fingerprint_priming.py new file mode 100644 index 00000000..5bb0968b --- /dev/null +++ b/tests/unit/test_axis_fingerprint_priming.py @@ -0,0 +1,182 @@ +"""A transformer that builds a coordinate axis should hand it over ready to use. + +``CoordinateAxis.fingerprint`` is what every stateful consumer keys its state on. +It is computed on first access and cached on the instance, and the cache pickles +with the axis -- so whoever touches it first pays, and everyone after gets it +free. + +In-process that first toucher is usually the next stateful node, which primes the +axis as a side effect of hashing it. The gap is the boundary: unpickling builds a +*new* axis object per message, so an axis that left its producing process cold is +re-checksummed by the first consumer in every receiving process, on every +message, forever. :func:`~ezmsg.sigproc.util.message.with_fingerprint` closes +that by priming at the point of construction, and these tests pin it, since +nothing else would notice it stopping. + +Deliberately not primed: coordinate axes along the chunk dimension. Their data is +new every message and no consumer reads their fingerprint -- the default hash +takes only ``gain`` from the chunk axis. +""" + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.sigproc.affinetransform import AffineTransformSettings, AffineTransformTransformer +from ezmsg.sigproc.aggregate import ( + AggregationFunction, + RangedAggregateSettings, + RangedAggregateTransformer, +) +from ezmsg.sigproc.binned_aggregate import BinnedAggregateSettings, BinnedAggregateTransformer +from ezmsg.sigproc.concat import ConcatProcessor, ConcatSettings +from ezmsg.sigproc.coordinatespaces import ( + CoordinateMode, + CoordinateSpacesSettings, + CoordinateSpacesTransformer, +) +from ezmsg.sigproc.flatten import FlattenSettings, FlattenTransformer +from ezmsg.sigproc.slicer import SlicerSettings, SlicerTransformer +from ezmsg.sigproc.wavelets import CWTSettings, CWTTransformer + + +def signal(n_time: int = 64, labels: list[str] | None = None, fs: float = 100.0, key: str = "dev") -> AxisArray: + labels = labels or [f"c{i}" for i in range(4)] + return AxisArray( + np.random.default_rng(0).standard_normal((n_time, len(labels))).astype(np.float32), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +def spectrum(n_win: int = 8, n_freq: int = 16, n_ch: int = 4) -> AxisArray: + return AxisArray( + np.random.default_rng(0).standard_normal((n_win, n_freq, n_ch)).astype(np.float32), + dims=["win", "freq", "ch"], + axes={ + "win": AxisArray.TimeAxis(fs=10.0), + "freq": CoordinateAxis(data=np.arange(n_freq, dtype=float), dims=["freq"], unit="Hz"), + "ch": CoordinateAxis(data=np.array([f"c{i}" for i in range(n_ch)]), dims=["ch"]), + }, + key="dev", + chunk_dim="win", + ) + + +def freshly_built(source: AxisArray, result: AxisArray) -> dict[str, CoordinateAxis]: + """The coordinate axes on *result* that are not objects *source* handed in. + + Identity, not equality: an axis that merely passed through was primed by + whichever consumer hashed it, which would mask a producer that never primes + anything. + """ + incoming = {id(axis) for axis in source.axes.values()} + return { + dim: axis for dim, axis in result.axes.items() if isinstance(axis, CoordinateAxis) and id(axis) not in incoming + } + + +def assert_primed(source: AxisArray, result: AxisArray, expected: set[str]) -> None: + built = freshly_built(source, result) + assert set(built) >= expected, f"expected new axes {expected}, got {set(built)}" + cold = sorted(dim for dim, axis in built.items() if "_fingerprint" not in axis.__dict__) + assert not cold, f"axes handed downstream without a fingerprint: {cold}" + + +class TestCreatedAxesArePrimed: + def test_affine_transform_output_labels(self): + msg = signal() + proc = AffineTransformTransformer(AffineTransformSettings(weights=np.ones((4, 2)), axis="ch")) + assert_primed(msg, proc(msg), {"ch"}) + + def test_slicer_selection(self): + msg = signal() + proc = SlicerTransformer(SlicerSettings(selection="0:2", axis="ch")) + assert_primed(msg, proc(msg), {"ch"}) + + def test_flatten_merged_axis(self): + msg = signal() + proc = FlattenTransformer(FlattenSettings(preserve_axis="time", sample_axis="time", flatten_axes=("ch",))) + assert_primed(msg, proc(msg), {"ch"}) + + def test_ranged_aggregate_band_axis(self): + msg = spectrum() + proc = RangedAggregateTransformer( + RangedAggregateSettings(axis="freq", bands=[(0.0, 4.0), (8.0, 12.0)], operation=AggregationFunction.MEAN) + ) + assert_primed(msg, proc(msg), {"freq"}) + + def test_binned_aggregate_metric_axis(self): + msg = signal() + proc = BinnedAggregateTransformer( + BinnedAggregateSettings( + axis="time", + bin_duration=0.1, + operation=(AggregationFunction.MEAN, AggregationFunction.MAX), + newaxis="metric", + ) + ) + assert_primed(msg, proc(msg), {"metric"}) + + def test_coordinate_spaces_relabel(self): + msg = signal(labels=["x", "y"]) + proc = CoordinateSpacesTransformer(CoordinateSpacesSettings(axis="ch", mode=CoordinateMode.CART2POL)) + assert_primed(msg, proc(msg), {"ch"}) + + def test_wavelet_frequency_axis(self): + msg = signal(n_time=128) + proc = CWTTransformer(CWTSettings(frequencies=(8.0, 12.0, 20.0), wavelet="morl", axis="time")) + assert_primed(msg, proc(msg), {"freq"}) + + def test_concat_merged_axis(self): + a = signal(labels=["a0", "a1"], key="A") + b = signal(labels=["b0", "b1"], key="B") + proc = ConcatProcessor(ConcatSettings(axis="ch")) + assert_primed(a, proc._concat(a, b), {"ch"}) + + +class TestTheChunkAxisIsLeftAlone: + """Priming a per-message chunk coordinate would be pure cost: its data is + new every message and the default hash reads only ``gain`` from it.""" + + def test_an_irregular_time_axis_is_not_primed(self): + n_time = 32 + msg = AxisArray( + np.zeros((n_time, 2), np.float32), + dims=["time", "ch"], + axes={ + "time": CoordinateAxis(data=np.arange(n_time, dtype=float), dims=["time"], unit="s"), + "ch": CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"]), + }, + key="dev", + chunk_dim="time", + ) + SlicerTransformer(SlicerSettings(selection="0:1", axis="ch"))(msg) + assert "_fingerprint" not in msg.axes["time"].__dict__ + + +def test_priming_is_idempotent_and_returns_the_axis(): + from ezmsg.sigproc.util.message import with_fingerprint + + axis = CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"]) + assert with_fingerprint(axis) is axis + first = axis.__dict__["_fingerprint"] + assert with_fingerprint(axis).__dict__["_fingerprint"] is first + + +@pytest.mark.parametrize("dtype", ["U8", "f8", "i4"]) +def test_priming_survives_a_pickle_round_trip(dtype): + """The whole point: the far side gets the answer without recomputing.""" + import pickle + + from ezmsg.sigproc.util.message import with_fingerprint + + axis = with_fingerprint(CoordinateAxis(data=np.arange(8).astype(dtype), dims=["ch"])) + landed = pickle.loads(pickle.dumps(axis)) + assert "_fingerprint" in landed.__dict__ + assert landed.__dict__["_fingerprint"] == axis.fingerprint diff --git a/tests/unit/test_buffer_recycling.py b/tests/unit/test_buffer_recycling.py index 56ca1475..8419ebd5 100644 --- a/tests/unit/test_buffer_recycling.py +++ b/tests/unit/test_buffer_recycling.py @@ -190,3 +190,87 @@ def test_binned_aggregate_multi_op_does_not_retain(): ), _msgs(_equal_blocks(4, 12, seed=0)), ) + + +def _axis_equality_is_content_based() -> bool: + """Whether the installed ezmsg compares coordinate axes by value. + + Before ezmsg-org/ezmsg#258's stack, ``CoordinateAxis.__eq__`` resolved + through the MRO to ``AxisBase.__eq__``, which compares only ``unit`` -- so + two axes with different labels compared equal and this suite could not see a + retained coordinate axis at all. + """ + ch = AxisArray.CoordinateAxis + return ch(data=np.array(["A"]), dims=["ch"]) != ch(data=np.array(["B"]), dims=["ch"]) + + +requires_content_axis_equality = pytest.mark.skipif( + not _axis_equality_is_content_based(), + reason="ezmsg's CoordinateAxis.__eq__ compares only `unit`; axis retention is undetectable", +) + + +class _RetainsChAxis: + """Caches the first message's ch axis and re-emits it -- a view into the slot.""" + + def __init__(self) -> None: + self.cached = None + + def __call__(self, msg: AxisArray) -> AxisArray: + if self.cached is None: + self.cached = msg.axes["ch"] + return replace(msg, axes={**msg.axes, "ch": self.cached}) + + +class _CopiesChAxis(_RetainsChAxis): + """The same thing done correctly: the cached axis owns its memory.""" + + def __call__(self, msg: AxisArray) -> AxisArray: + if self.cached is None: + axis = msg.axes["ch"] + self.cached = replace(axis, data=np.array(axis.data)) + return replace(msg, axes={**msg.axes, "ch": self.cached}) + + +def _msgs_with_distinct_ch_labels(n: int = 3, n_time: int = 8) -> list[AxisArray]: + """Messages whose ch labels differ, so a stale axis view is observable. + + With identical labels on every message a retained view reads the *same* + bytes back out of the recycled slot and the corruption is invisible. + """ + out, offset = [], 0.0 + for i in range(n): + out.append( + AxisArray( + data=np.random.default_rng(i).standard_normal((n_time, N_CH)), + dims=["time", "ch"], + axes=frozendict( + { + "time": AxisArray.TimeAxis(fs=FS, offset=offset), + "ch": AxisArray.CoordinateAxis(data=np.array([f"m{i}c{c}" for c in range(N_CH)]), dims=["ch"]), + } + ), + key="test_buffer_recycling", + ) + ) + offset += n_time / FS + return out + + +@requires_content_axis_equality +def test_harness_detects_a_retained_coordinate_axis(): + """The suite's own check: a retained *axis* must fail, not just retained data. + + Coordinate axes are views onto the transport slot exactly as the samples + are, so a transformer that caches one across calls reads recycled bytes. + Without this, a helper that forgot to compare axes would leave every + transformer here untested for axis retention and nothing would say so. + """ + with pytest.raises(AssertionError, match="axis 'ch' differs"): + assert_survives_buffer_recycling(_RetainsChAxis, _msgs_with_distinct_ch_labels()) + + +@requires_content_axis_equality +def test_harness_accepts_a_copied_coordinate_axis(): + """...and does not cry wolf when the transformer copies it properly.""" + assert_survives_buffer_recycling(_CopiesChAxis, _msgs_with_distinct_ch_labels()) diff --git a/tests/unit/test_concat.py b/tests/unit/test_concat.py index bfa7b18b..69853a4c 100644 --- a/tests/unit/test_concat.py +++ b/tests/unit/test_concat.py @@ -502,6 +502,140 @@ def test_cache_reused(self): proc._concat(msg_a, msg_b) assert proc.state.cached_axes is first_cache # Same dict object. + def test_time_offset_is_taken_from_the_live_message(self): + """A LinearAxis must never be served from the cache. + + ``offset`` advances every message, so caching the time axis pins every + output to the first message's time base. ``align_axis`` used to be the + only thing that saved you from this, because the alignment axis is the + one axis excluded from the cache — but it defaults to None. + """ + proc = ConcatProcessor(ConcatSettings(axis="ch")) + offsets = [0.0, 0.04, 0.08] + out = [ + proc._concat( + _make_msg(np.ones((4, 2)), offset=t, ch_labels=["A0", "A1"]), + _make_msg(np.ones((4, 2)), offset=t, ch_labels=["B0", "B1"]), + ) + for t in offsets + ] + assert [o.axes["time"].offset for o in out] == offsets + # The cache is not rebuilt to achieve it: a changing offset is not a + # configuration change. + assert proc.state.cached_axes is not None + assert "time" not in proc.state.cached_axes + + def test_non_concat_coordinate_axis_change_invalidates_cache(self): + """Every cached axis has to be fingerprinted, not just the concat axis. + + ``_build_cached_axes`` copies all coordinate axes into the cache, so a + relabelled *band* axis would otherwise keep emitting the first + message's labels — the concat-axis hash cannot see it. + """ + + def msg(band: list[str], fill: float) -> AxisArray: + return AxisArray( + np.full((4, 2, 3), fill), + dims=["time", "ch", "band"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": CoordinateAxis(data=np.array(["c0", "c1"]), dims=["ch"]), + "band": CoordinateAxis(data=np.array(band), dims=["band"]), + }, + ) + + proc = ConcatProcessor(ConcatSettings(axis="ch", relabel_axis=False)) + first = proc._concat(msg(["alpha", "beta", "gamma"], 1.0), msg(["alpha", "beta", "gamma"], 2.0)) + assert list(first.axes["band"].data) == ["alpha", "beta", "gamma"] + + second = proc._concat(msg(["delta", "theta", "mu"], 1.0), msg(["delta", "theta", "mu"], 2.0)) + assert list(second.axes["band"].data) == ["delta", "theta", "mu"] + + def test_unchanged_coordinate_axes_do_not_invalidate_cache(self): + """The converse: equal axes rebuilt per message must not look like a + change, or the cache would be rebuilt at the sample rate.""" + proc = ConcatProcessor(ConcatSettings(axis="ch")) + + def pair(offset: float): + return ( + _make_msg(np.ones((4, 2)), offset=offset, ch_labels=["A0", "A1"]), + _make_msg(np.ones((4, 2)), offset=offset, ch_labels=["B0", "B1"]), + ) + + proc._concat(*pair(0.0)) + first_cache = proc.state.cached_axes + for i in range(1, 4): + proc._concat(*pair(i * 0.04)) + assert proc.state.cached_axes is first_cache + + def test_cached_axes_are_still_owned(self): + """Coordinate axes are copied, not aliased: input axes may be views into + a transport buffer whose lifetime ends with the callback.""" + proc = ConcatProcessor(ConcatSettings(axis="ch", relabel_axis=False)) + band = CoordinateAxis(data=np.array(["b0", "b1", "b2"]), dims=["band"]) + + def msg(fill: float) -> AxisArray: + return AxisArray( + np.full((4, 2, 3), fill), + dims=["time", "ch", "band"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": CoordinateAxis(data=np.array(["c0", "c1"]), dims=["ch"]), + "band": band, + }, + ) + + a = msg(1.0) + out = proc._concat(a, msg(2.0)) + assert out.axes["band"] is not a.axes["band"] + assert out.axes["band"].data is not a.axes["band"].data + assert list(out.axes["band"].data) == ["b0", "b1", "b2"] + + def test_output_axes_follow_input_dim_order(self): + """Serving some axes from the cache and some live must not reorder them.""" + proc = ConcatProcessor(ConcatSettings(axis="ch", relabel_axis=False)) + + def msg(fill: float) -> AxisArray: + return AxisArray( + np.full((4, 2, 3), fill), + dims=["time", "ch", "band"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": CoordinateAxis(data=np.array(["c0", "c1"]), dims=["ch"]), + "band": CoordinateAxis(data=np.array(["b0", "b1", "b2"]), dims=["band"]), + }, + ) + + out = proc._concat(msg(1.0), msg(2.0)) + assert list(out.axes.keys()) == ["time", "ch", "band"] + + def test_memoized_fingerprint_matches_the_uncached_one(self): + """The identity fast path must be invisible: same answer, every message. + + Covers both regimes — a source reusing one axis object (memo hits) and + one rebuilding an equal axis per message (memo misses). + """ + proc = ConcatProcessor(ConcatSettings(axis="ch")) + shared = CoordinateAxis(data=np.array(["A0", "A1"]), dims=["ch"], unit="label") + for reuse in (True, False): + for i in range(4): + msg = _make_msg( + np.ones((4, 2)), + offset=i * 0.04, + ch_axis=shared if reuse else CoordinateAxis(data=np.array(["A0", "A1"]), dims=["ch"], unit="label"), + ) + memoized = proc._fingerprint(msg, proc.state.memo_a) + assert memoized == proc._fingerprint(msg, None) + + def test_memo_notices_a_new_axis_object_with_new_values(self): + """A memo miss must fall through to the content digest, not reuse.""" + proc = ConcatProcessor(ConcatSettings(axis="ch")) + first = proc._fingerprint(_make_msg(np.ones((4, 2)), ch_labels=["A0", "A1"]), proc.state.memo_a) + same = proc._fingerprint(_make_msg(np.ones((4, 2)), ch_labels=["A0", "A1"]), proc.state.memo_a) + other = proc._fingerprint(_make_msg(np.ones((4, 2)), ch_labels=["Z0", "Z1"]), proc.state.memo_a) + assert first == same + assert first != other + def test_cache_invalidated_on_shape_change(self): settings = ConcatSettings(axis="ch", relabel_axis=False) proc = ConcatProcessor(settings) diff --git a/tests/unit/test_flatten.py b/tests/unit/test_flatten.py index 55a03895..22b17a93 100644 --- a/tests/unit/test_flatten.py +++ b/tests/unit/test_flatten.py @@ -330,3 +330,73 @@ def test_state_caching_across_messages_mlx(self): ) # Same cached output_axis_obj is reused across messages. assert out1.axes["ch"] is out2.axes["ch"] + + +class TestRelabelledInput: + """The merged output axis is built from the input labels, so it has to + follow them — a relabel at a fixed channel count is invisible to shape.""" + + @staticmethod + def _msg(ch_labels): + return AxisArray( + np.zeros((4, len(ch_labels), 2), np.float32), + dims=["time", "ch", "feature"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": CoordinateAxis(data=np.array(ch_labels), dims=["ch"]), + "feature": CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]), + }, + key="dev", + ) + + def _labels(self, out): + return [str(row["label"]) for row in out.axes["ch"].data] + + def test_relabelled_channels_are_followed(self): + proc = FlattenTransformer( + FlattenSettings(preserve_axis="time", flatten_axes=("ch", "feature"), output_axis="ch") + ) + assert self._labels(proc(self._msg(["A", "B"]))) == ["A/spk", "A/sbp", "B/spk", "B/sbp"] + assert self._labels(proc(self._msg(["X", "Y"]))) == ["X/spk", "X/sbp", "Y/spk", "Y/sbp"] + + def test_unchanged_labels_reuse_the_cached_axis(self): + """The converse: an equal axis rebuilt per message must not reset, or + the merged axis would be rebuilt at the sample rate.""" + proc = FlattenTransformer( + FlattenSettings(preserve_axis="time", flatten_axes=("ch", "feature"), output_axis="ch") + ) + first = proc(self._msg(["A", "B"])) + again = proc(self._msg(["A", "B"])) + assert first.axes["ch"] is again.axes["ch"] + + def test_labels_outside_flatten_axes_also_reset(self): + """A coordinate axis the output labels do not depend on still resets. + + Flatten used to narrow its hash to the flattened axes, so a change to an + unrelated `ch` axis left the cached merged axis in place. It now takes + the base-class hash, which fingerprints every coordinate axis, so this + rebuilds. That is conservative rather than wrong -- the rebuilt axis is + identical -- and it buys a single rule for every processor instead of a + per-processor exception. + """ + proc = FlattenTransformer(FlattenSettings(preserve_axis="time", flatten_axes=("feature",), output_axis="feat")) + + def msg(ch_labels): + return AxisArray( + np.zeros((4, 2, 2), np.float32), + dims=["time", "ch", "feature"], + axes={ + "time": AxisArray.TimeAxis(fs=100.0), + "ch": CoordinateAxis(data=np.array(ch_labels), dims=["ch"]), + "feature": CoordinateAxis(data=np.array(["spk", "sbp"]), dims=["feature"]), + }, + key="dev", + ) + + first = proc(msg(["A", "B"])) + rebuilt = proc(msg(["X", "Y"])) + assert rebuilt.axes["feat"] is not first.axes["feat"] + # ...and the rebuilt labels are unchanged, since ch does not feed them. + assert [str(r["label"]) for r in rebuilt.axes["feat"].data] == [ + str(r["label"]) for r in first.axes["feat"].data + ] diff --git a/tests/unit/test_state_reset_semantics.py b/tests/unit/test_state_reset_semantics.py new file mode 100644 index 00000000..dbc8f6b2 --- /dev/null +++ b/tests/unit/test_state_reset_semantics.py @@ -0,0 +1,167 @@ +"""When a stateful processor resets, now that the base class decides by default. + +Most processors no longer implement ``_hash_message`` at all: the base class +folds in the message key, the dims, the length of every dimension except the one +the stream is chunked along, the coordinate values on those dimensions and the +gain and offset of any linear axis among them. These tests pin the behaviour +that fell out of removing those overrides, and the ``chunk_dim`` bookkeeping the +default depends on. +""" + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.sigproc.aggregate import AggregateSettings, AggregateTransformer, AggregationFunction +from ezmsg.sigproc.butterworthfilter import ButterworthFilterSettings, ButterworthFilterTransformer +from ezmsg.sigproc.flatten import FlattenSettings, FlattenTransformer +from ezmsg.sigproc.spectrum import SpectrumSettings, SpectrumTransformer +from ezmsg.sigproc.window import WindowSettings, WindowTransformer + +FS = 100.0 + + +def _msg(data, labels, fs=FS, key="dev", chunk_dim="time"): + return AxisArray( + data, + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim=chunk_dim, + ) + + +class TestPerChannelStateFollowsTheChannels: + """A filter's `zi` belongs to the channels it was built for. + + This is what the sweep was for. A device reconfigured mid-session sends the + same number of channels under different labels; keying the state on shape + alone left the new channels being filtered through the old ones' history. + """ + + @staticmethod + def _filter(): + return ButterworthFilterTransformer(ButterworthFilterSettings(axis="time", order=2, cuton=1.0, cutoff=20.0)) + + def test_relabel_at_fixed_channel_count_resets(self): + rng = np.random.default_rng(0) + warmup = rng.standard_normal((50, 2)) * 10.0 + fresh = rng.standard_normal((20, 2)) * 0.01 + + carried = self._filter() + carried(_msg(warmup, ["armA-1", "armA-2"])) + got = carried(_msg(fresh, ["armB-1", "armB-2"])) + + # What a correctly reset filter produces for the same input. + want = self._filter()(_msg(fresh, ["armB-1", "armB-2"])) + np.testing.assert_allclose(got.data, want.data) + + def test_chunk_size_jitter_does_not_reset(self): + """The filter must carry state across chunks, or it would ring forever.""" + rng = np.random.default_rng(1) + proc = self._filter() + proc(_msg(rng.standard_normal((50, 2)), ["a", "b"])) + # FilterByDesign delegates to an inner FilterTransformer that owns `zi`. + state_before = proc._state.filter._state.zi.copy() + proc(_msg(rng.standard_normal((17, 2)), ["a", "b"])) + assert not np.array_equal(proc._state.filter._state.zi, state_before), ( + "the filter should have advanced its state across the chunk, not reset it" + ) + + def test_channel_count_change_resets(self): + rng = np.random.default_rng(2) + proc = self._filter() + proc(_msg(rng.standard_normal((50, 2)), ["a", "b"])) + out = proc(_msg(rng.standard_normal((50, 3)), ["a", "b", "c"])) + assert out.data.shape[1] == 3 + + +class TestChunkDimBookkeeping: + """Every operation that renames, consumes or invents the chunked dimension + has to say so, or the base class excludes the wrong one.""" + + @staticmethod + def _stream(n_time=64, n_ch=2): + rng = np.random.default_rng(3) + return _msg(rng.standard_normal((n_time, n_ch)), [f"c{i}" for i in range(n_ch)]) + + def test_window_declares_the_new_axis(self): + """After windowing, messages append along `win`, not `time`.""" + proc = WindowTransformer(WindowSettings(axis="time", newaxis="win", window_dur=0.2, window_shift=0.1)) + out = proc(self._stream()) + assert out.dims[:2] == ["win", "time"] + assert out.chunk_dim == "win" + + def test_batcher_mode_keeps_the_incoming_chunk_dim(self): + """Batcher mode adds no axis: windows tile the target axis.""" + proc = WindowTransformer( + WindowSettings(axis="time", newaxis=None, window_dur=0.2, window_shift=0.2, batch_windows=True) + ) + out = proc(self._stream()) + assert "win" not in out.dims + assert out.chunk_dim == "time" + + def test_spectrum_clears_it_when_it_consumes_it(self): + """`time` becomes `freq`: each output is one spectrum, nothing appends.""" + out = SpectrumTransformer(SpectrumSettings(axis="time"))(self._stream()) + assert "time" not in out.dims + assert out.chunk_dim is None + + def test_spectrum_keeps_a_windowed_chunk_dim(self): + """Windowed input still appends along `win` after the transform.""" + win = WindowTransformer(WindowSettings(axis="time", newaxis="win", window_dur=0.2, window_shift=0.1))( + self._stream() + ) + out = SpectrumTransformer(SpectrumSettings(axis="time"))(win) + assert out.chunk_dim == "win" + assert "win" in out.dims + + def test_aggregate_clears_it_when_reducing_it(self): + out = AggregateTransformer(AggregateSettings(axis="time", operation=AggregationFunction.MEAN))(self._stream()) + assert "time" not in out.dims + assert out.chunk_dim is None + + def test_flatten_follows_a_renamed_preserve_axis(self): + proc = FlattenTransformer( + FlattenSettings(preserve_axis="time", sample_axis="sample", flatten_axes=("ch",), output_axis="ch") + ) + out = proc(self._stream()) + assert out.chunk_dim == "sample" + + @pytest.mark.parametrize("declared", ["time", None]) + def test_a_declaration_is_optional(self, declared): + """Undeclared messages still work -- the base class falls back.""" + rng = np.random.default_rng(4) + msg = _msg(rng.standard_normal((32, 2)), ["a", "b"], chunk_dim=declared) + out = ButterworthFilterTransformer(ButterworthFilterSettings(axis="time", order=2, cuton=1.0, cutoff=20.0))(msg) + assert out.chunk_dim == declared + + +class TestOverridesThatRemain: + """The three processors that still need something the default cannot know.""" + + @staticmethod + def _msg(dtype): + rng = np.random.default_rng(5) + return _msg(rng.standard_normal((64, 2)).astype(dtype), ["a", "b"]) + + def test_spectrum_reacts_to_a_dtype_change(self): + """A complex input takes a different branch and a different freq axis.""" + proc = SpectrumTransformer(SpectrumSettings(axis="time")) + proc(self._msg(np.float32)) + before = proc._hash + proc(self._msg(np.complex64)) + assert proc._hash != before + + def test_spectrum_reacts_to_a_transform_length_change(self): + """The FFT is sized by the chunk dimension -- the one length the default + ignores -- so Spectrum has to fold it back in.""" + proc = SpectrumTransformer(SpectrumSettings(axis="time")) + rng = np.random.default_rng(6) + proc(_msg(rng.standard_normal((64, 2)), ["a", "b"])) + before = proc._hash + proc(_msg(rng.standard_normal((128, 2)), ["a", "b"])) + assert proc._hash != before diff --git a/tests/unit/test_util_channels.py b/tests/unit/test_util_channels.py index ef9b22e8..013f2132 100644 --- a/tests/unit/test_util_channels.py +++ b/tests/unit/test_util_channels.py @@ -4,6 +4,7 @@ from ezmsg.sigproc.util.channels import ( channel_groups_from_field, + coord_value_fingerprint, group_spec_fields, group_spec_fingerprint, resolve_channel_groups, @@ -232,3 +233,139 @@ def test_field_values_are_not_folded(self): assert group_spec_fingerprint(_msg(["A", "B"]), "ch", "bank") == group_spec_fingerprint( _msg(["B", "A"]), "ch", "bank" ) + + +def _multifield_msg(labels: list[str], banks: list[str], xs: list[float] | None = None) -> AxisArray: + """Message with a ChannelMap-shaped structured ch axis (label/bank/x).""" + dt = np.dtype([("label", "U8"), ("bank", "U2"), ("x", " AxisArray: + return AxisArray( + data=np.zeros((3, len(labels))), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=np.array(labels), dims=["ch"])}, + ) + + +class TestCoordValueFingerprint: + def test_no_coordinate_data_is_empty(self): + """Nothing to fingerprint, so it contributes nothing to the hash.""" + bare = AxisArray(data=np.zeros((3, 2)), dims=["time", "ch"]) + assert coord_value_fingerprint(bare, "ch", None) == () + assert coord_value_fingerprint(bare, "ch", ["bank"]) == () + + def test_equal_axes_rebuilt_per_message_agree(self): + """Sources commonly rebuild an identical axis every message; that must + not read as a change, or the state would reset at the sample rate.""" + assert coord_value_fingerprint(_plain_msg(["a", "b"]), "ch") == coord_value_fingerprint( + _plain_msg(["a", "b"]), "ch" + ) + + def test_reorder_and_rename_are_detected(self): + base = coord_value_fingerprint(_plain_msg(["Ch5", "Ch7", "Ch10"]), "ch") + assert coord_value_fingerprint(_plain_msg(["Ch7", "Ch5", "Ch10"]), "ch") != base + assert coord_value_fingerprint(_plain_msg(["Ch1", "Ch2", "Ch3"]), "ch") != base + + def test_axis_defaults_to_last_dim(self): + msg = _plain_msg(["a", "b"]) + assert coord_value_fingerprint(msg, None) == coord_value_fingerprint(msg, "ch") + + def test_only_requested_fields_matter(self): + """The whole point of the restriction: an unrelated field churning must + not invalidate a selection that never reads it.""" + a = _multifield_msg(["e1", "e2"], ["A", "B"], xs=[0.0, 1.0]) + b = _multifield_msg(["e1", "e2"], ["A", "B"], xs=[9.9, 8.8]) + assert coord_value_fingerprint(a, "ch", ["label"]) == coord_value_fingerprint(b, "ch", ["label"]) + assert coord_value_fingerprint(a, "ch", ["x"]) != coord_value_fingerprint(b, "ch", ["x"]) + # ...and with no restriction, it does invalidate. + assert coord_value_fingerprint(a, "ch", None) != coord_value_fingerprint(b, "ch", None) + + def test_requested_field_values_are_detected(self): + a = _multifield_msg(["e1", "e2"], ["A", "B"]) + b = _multifield_msg(["e1", "e2"], ["B", "A"]) + assert coord_value_fingerprint(a, "ch", ["bank"]) != coord_value_fingerprint(b, "ch", ["bank"]) + assert coord_value_fingerprint(a, "ch", ["label"]) == coord_value_fingerprint(b, "ch", ["label"]) + + def test_missing_field_is_distinct_from_present(self): + """Gaining or losing the field still registers, as it does for + group_spec_fingerprint.""" + present = _multifield_msg(["e1", "e2"], ["A", "B"]) + plain = _plain_msg(["e1", "e2"]) + assert coord_value_fingerprint(present, "ch", ["nosuch"]) == (None,) + # A plain axis has no fields at all, so it digests whole rather than + # reporting every requested field as absent. + assert coord_value_fingerprint(plain, "ch", ["bank"]) == coord_value_fingerprint(plain, "ch", None) + + def test_multiple_fields_are_order_sensitive_and_independent(self): + msg = _multifield_msg(["e1", "e2"], ["A", "B"]) + both = coord_value_fingerprint(msg, "ch", ["label", "bank"]) + assert both == coord_value_fingerprint(msg, "ch", ["label"]) + coord_value_fingerprint(msg, "ch", ["bank"]) + assert both != coord_value_fingerprint(msg, "ch", ["bank", "label"]) + + def test_result_is_hashable(self): + """Callers fold it into hash((key, n_ch) + fingerprint).""" + msg = _multifield_msg(["e1", "e2"], ["A", "B"]) + for fields in (None, ["label"], ["label", "bank"], ["nosuch"]): + hash(("key", 2) + coord_value_fingerprint(msg, "ch", fields)) + + def test_object_dtype_axis_is_content_based(self): + """An object array's buffer is pointers, so checksumming it directly + would make two equal axes disagree and reset the state every message.""" + a = AxisArray( + data=np.zeros((3, 3)), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=np.array(["c1", "c2", "c3"], dtype=object), dims=["ch"])}, + ) + b = AxisArray( + data=np.zeros((3, 3)), + dims=["time", "ch"], + axes={ + "ch": AxisArray.CoordinateAxis( + data=np.array(["".join(["c", str(i)]) for i in (1, 2, 3)], dtype=object), dims=["ch"] + ) + }, + ) + assert coord_value_fingerprint(a, "ch") == coord_value_fingerprint(b, "ch") + c = AxisArray( + data=np.zeros((3, 3)), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=np.array(["c1", "c9", "c3"], dtype=object), dims=["ch"])}, + ) + assert coord_value_fingerprint(a, "ch") != coord_value_fingerprint(c, "ch") + + def test_non_contiguous_coordinate_data(self): + """A sliced/strided coordinate array has no C-contiguous buffer; the + digest must gather rather than raise.""" + labels = np.array(["a", "X", "b", "X", "c", "X"])[::2] + assert not labels.flags["C_CONTIGUOUS"] + msg = AxisArray( + data=np.zeros((3, 3)), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=labels, dims=["ch"])}, + ) + assert coord_value_fingerprint(msg, "ch") == coord_value_fingerprint(_plain_msg(["a", "b", "c"]), "ch") + + def test_dtype_change_alone_is_detected(self): + """Same values, different width -- shape and dtype ride along so this + does not depend on the checksum alone.""" + wide = AxisArray( + data=np.zeros((3, 2)), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=np.array([1, 2], dtype=np.int64), dims=["ch"])}, + ) + narrow = AxisArray( + data=np.zeros((3, 2)), + dims=["time", "ch"], + axes={"ch": AxisArray.CoordinateAxis(data=np.array([1, 2], dtype=np.int32), dims=["ch"])}, + ) + assert coord_value_fingerprint(wide, "ch") != coord_value_fingerprint(narrow, "ch")