From 7c3eb5efd5ac252ffc2040a3fb5c0dfa309b0e47 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 29 Aug 2026 00:10:19 -0400 Subject: [PATCH 1/2] Stop doing full-array work at open to answer approximate questions NWBSlicer construction on a 5-stream 5.7 GB recording: 218 ms -> 119 ms with the dejitter cache warm, 250 -> 119 cold. Two thirds of what is left is inside pynwb's read(), building Python objects for the file, which is not ours to fix cheaply. np.median over every interval, 45 ms per stream. The median here is never the answer -- infer_nominal_rate returns a mean over the intervals that survive a trim window, and the median only says where to centre that window. Partitioning 7 million floats to place a window is precision nobody asked for. It is not even approximate in practice. Timestamps arrive on a quantisation grid, so the median snaps to a grid point that a stride cannot move: strides of 16 through 1024 over a real 7.09M-interval stream all give 3.3301000002e-05. Verified end to end that every stream in that recording reports a bit-identical rate with the subsample as with the full median, including the two _device_ts streams that actually reach the inference path at 30137.773554857125 and 30000.082189805245. Strided rather than head-of-array, so a stream whose intervals drift between its start and its end is represented across the estimate instead of by whichever end happened to be read first. np.var was also computed twice to evaluate one `or`. Short-circuiting meant the second pass only ran when the first comparison failed -- which is exactly the irregular-stream case, i.e. the slow one. electrodes.table.to_dataframe(), 12 ms per electrical series. It materializes every column the writer stored -- position, group, filtering -- into pandas so we can read one of them, and it drags in a reference resolution per row on the way. Indexing the label column directly is the same values for a thirtieth of the work. The fallback for a table with no label column now names channels from the id column rather than from a DataFrame index. That is not a change: the index to_dataframe() built was the id column. Making it explicit is what keeps it correct now that the DataFrame is gone, since positional indices coincide with ids only when a series references the whole table in order. Tests pin equivalence rather than speed: _fast_median against np.median across the threshold, inferred rates unchanged, and -- because the trim is the only reason a mean is usable here -- that a stream with 400 dropped-packet holes still reports 30 kHz rather than an average including the holes. --- src/ezmsg/nwb/slicer.py | 55 ++++++++++++++++++++++++++++----- tests/test_slicer_perf.py | 65 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 7 deletions(-) create mode 100644 tests/test_slicer_perf.py diff --git a/src/ezmsg/nwb/slicer.py b/src/ezmsg/nwb/slicer.py index 73f796b..22890d0 100644 --- a/src/ezmsg/nwb/slicer.py +++ b/src/ezmsg/nwb/slicer.py @@ -42,6 +42,34 @@ RATE_TRIM_HI = 1.5 +MEDIAN_SUBSAMPLE_MAX = 100_000 +"""Above this many intervals, estimate the median from a strided subsample. + +``np.median`` partitions the whole array: 45 ms on a recording-hour of 30 kHz +timestamps, and it runs once per stream at open. The median here is only ever a +*centre* for a trim window or a scale for a variance comparison -- never the +returned estimate, which is a mean over the surviving intervals -- so it needs +to be in the right place, not exact. + +And it lands in exactly the same place. Timestamps arrive on a quantisation +grid, so the median snaps to a grid point that a stride cannot move: on a real +7.09M-interval stream, strides of 16 through 1024 all give 3.3301000002e-05, +and :func:`infer_nominal_rate` returns 30147.9172 from every one of them. +""" + + +def _fast_median(values: np.ndarray) -> float: + """``np.median``, on a strided subsample once the array is large enough. + + Strided rather than random or head-of-array: a stride samples the whole + recording, so a stream whose intervals differ between its start and its end + is represented across the estimate rather than by whichever part was read. + """ + if values.size > MEDIAN_SUBSAMPLE_MAX: + values = values[:: -(-values.size // MEDIAN_SUBSAMPLE_MAX)] + return float(np.median(values)) + + def infer_nominal_rate(dts: np.ndarray) -> float: """Sample rate from per-sample intervals: the mean of the non-outlier ones. @@ -66,7 +94,7 @@ def infer_nominal_rate(dts: np.ndarray) -> float: """ if dts.size == 0: return 0.0 - median = float(np.median(dts)) + median = _fast_median(dts) if not np.isfinite(median) or median <= 0.0: return 0.0 core = dts[(dts > RATE_TRIM_LO * median) & (dts < RATE_TRIM_HI * median)] @@ -436,7 +464,11 @@ def _load(self) -> None: rate = child.timestamps.attrs["rate"] else: dts = np.diff(child.timestamps[:]) - if np.var(dts) < 1e-3 or np.var(dts) < 0.05 * np.median(dts): + variance = float(np.var(dts)) + # ``np.var`` twice was two full passes to evaluate one + # ``or``; the second only ran when the first failed, but it + # ran on every irregular stream, which is the slow case. + if variance < 1e-3 or variance < 0.05 * _fast_median(dts): rate = infer_nominal_rate(dts) else: rate = 0.0 @@ -475,16 +507,25 @@ def _load(self) -> None: # otherwise an ElectricalSeries that references a # strict subset of the electrodes table produces a # ch-axis whose length does not match data.shape[1]. + # One column, not the whole table. ``to_dataframe()`` materializes + # every electrode column into pandas -- position, group, filtering, + # whatever the writer stored -- to read one of them, and it costs + # ~12 ms per electrical series at open. Indexing the column + # directly is the same values for a thirtieth of the work. region_idx = np.asarray(child.electrodes.data) - full_df = child.electrodes.table.to_dataframe() - el_df = full_df.iloc[region_idx] - if "label" in el_df.columns: + table = child.electrodes.table + if "label" in table.colnames: # Decoded here, at the read boundary: these labels become the # ch-axis coordinates every downstream name-based channel # selection matches against. - ch_labels = as_text_array(el_df["label"].values) + ch_labels = as_text_array(np.asarray(table["label"].data[:])[region_idx]) else: - ch_labels = np.array([f"ch_{idx}" for idx in el_df.index.tolist()]) + # ``to_dataframe`` indexes by the table's ``id`` column, so the + # fallback names must come from ``id`` too -- not from the + # positional indices, which coincide with it only when the + # series references the whole table in order. + ids = np.asarray(table.id.data[:])[region_idx] + ch_labels = np.array([f"ch_{idx}" for idx in ids.tolist()]) axes["ch"] = AxisArray.CoordinateAxis(data=ch_labels, dims=["ch"]) # ``matched_key`` is the user-facing key — equal to ``child.name`` diff --git a/tests/test_slicer_perf.py b/tests/test_slicer_perf.py new file mode 100644 index 0000000..d287b34 --- /dev/null +++ b/tests/test_slicer_perf.py @@ -0,0 +1,65 @@ +"""Shortcuts taken at open, and the properties that make them safe. + +Both are pure speedups: they must not move a single value. What is pinned here +is equivalence with the straightforward computation, not the speed. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from ezmsg.nwb.slicer import MEDIAN_SUBSAMPLE_MAX, _fast_median, infer_nominal_rate + + +def quantised_intervals(n: int, period: float = 1 / 30000, grid: float = 4e-8, seed: int = 3) -> np.ndarray: + """Intervals as a real acquisition delivers them: jittered, then snapped to + the device's timestamp grid. The snapping is what makes a subsampled median + land on exactly the same value as a full one.""" + rng = np.random.default_rng(seed) + return np.round((period + rng.normal(0, period * 0.02, n)) / grid) * grid + + +@pytest.mark.parametrize("n", [10, MEDIAN_SUBSAMPLE_MAX, MEDIAN_SUBSAMPLE_MAX + 1, 3_000_000]) +def test_fast_median_matches_the_real_one(n): + dts = quantised_intervals(n) + assert _fast_median(dts) == float(np.median(dts)) + + +def test_small_arrays_are_not_subsampled(): + """Below the threshold there is nothing to gain, so take the exact answer.""" + dts = np.linspace(1.0, 2.0, 1001) # even-length medians interpolate; no grid + assert _fast_median(dts) == float(np.median(dts)) + + +def test_inferred_rate_is_unchanged_by_subsampling(): + """The median only centres the trim window; the returned estimate is a mean + over what survives it. So the median may be approximate, and here it is not + even that.""" + dts = quantised_intervals(3_000_000) + small = dts[:1000] + assert infer_nominal_rate(small) == pytest.approx(1 / np.mean(small), rel=1e-9) + assert infer_nominal_rate(dts) == pytest.approx(30000.0, rel=1e-3) + + +def test_gaps_still_do_not_drag_the_estimate(): + """The trim is the whole reason the mean is usable. Subsampling the median + must not weaken it: a stream with real gaps still reports its sample rate, + not an average that includes the holes.""" + dts = quantised_intervals(2_000_000) + dts[::5000] = 0.01 # 400 dropped-packet holes, 300x the sample period + assert infer_nominal_rate(dts) == pytest.approx(30000.0, rel=1e-3) + + +def test_ch_labels_come_from_the_electrodes_region(scaled_nwb_path): + """Reading the label column directly must still honour the region: a series + referencing a subset of the table gets that subset's labels, in order.""" + from ezmsg.nwb import NWBSlicer + + slicer = NWBSlicer(scaled_nwb_path, dejitter=False) + try: + labels = slicer.get_stream_info("Broadband").template.axes["ch"].data + finally: + slicer.close() + assert list(labels) == ["elec0", "elec1", "elec2", "elec3"] + assert labels.dtype.kind == "U" From 17ac7577d2e4574d83d88d82266cfe5d22d27f33 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sat, 29 Aug 2026 01:16:47 -0400 Subject: [PATCH 2/2] Share the interval-median shortcut with the clock model Cold open -- the first time a file is seen, which for batch analysis is every file -- spent ~0.67 s taking medians over full recordings inside clockmodel.py, which the earlier slicer-side fix did not reach. Reconstruction time for a 5-stream 5.7 GB recording: 1.02 s -> 0.57 s, measured by alternating the two implementations so drift hits both. _fast_median moves to util.py as interval_median, since clockmodel cannot import from slicer (slicer imports clockmodel) and util imports neither. The rename says what the shortcut is licensed by: these are medians of inter-sample intervals, and intervals inherit the device's quantisation grid, so the median snaps to a grid point a stride cannot move. Verified on a real recording that all five streams reconstruct bit-identical timestamps and the gap threshold matches to the last digit. _auto_gap_threshold takes the shortcut for its median and deliberately not for its percentile. Subsampling the whole function is 16x faster and moves the threshold 1.4%, because a 99.9th percentile over 200k values rests on their top 200 instead of on 7000. That threshold decides which jumps are real gaps, so moving it re-segments the stream and re-times its samples -- a behaviour change wearing a speedup's clothes. A test asserts the asymmetry so nobody 'finishes the job' later. through_c's median is left alone for the same reason: it averages fit residuals, which are continuous and have no grid to snap to. --- src/ezmsg/nwb/clockmodel.py | 12 +++++++++-- src/ezmsg/nwb/slicer.py | 34 +++-------------------------- src/ezmsg/nwb/util.py | 30 ++++++++++++++++++++++++++ tests/test_slicer_perf.py | 43 ++++++++++++++++++++++++++++++++++--- 4 files changed, 83 insertions(+), 36 deletions(-) diff --git a/src/ezmsg/nwb/clockmodel.py b/src/ezmsg/nwb/clockmodel.py index 97c2fd1..c957ca9 100644 --- a/src/ezmsg/nwb/clockmodel.py +++ b/src/ezmsg/nwb/clockmodel.py @@ -43,6 +43,8 @@ import numpy as np +from .util import interval_median + DEFAULT_DEJITTER_KNOTS = 40 """Number of knots for the piecewise-linear fit. ~40 over a multi-million-sample stream puts a knot every few seconds -- fine enough to track sub-millisecond @@ -214,7 +216,7 @@ def _nominal_period(times: np.ndarray) -> float: the median is unmoved by them).""" if times.shape[0] < 2: return 0.0 - return float(np.median(np.diff(times))) + return interval_median(np.diff(times)) def _auto_gap_threshold(times: np.ndarray) -> float: @@ -230,7 +232,13 @@ def _auto_gap_threshold(times: np.ndarray) -> float: dt = np.diff(np.asarray(times, dtype=float)) if dt.size == 0: return np.inf - period = float(np.median(dt)) + period = interval_median(dt) + # The percentile stays on the full array. It is the one statistic here that a + # subsample genuinely moves: the 99.9th percentile of 200k values rests on + # their top 200 rather than on 7000, which shifted this threshold by 1.4% on + # a real stream. The threshold decides which jumps count as real gaps, so a + # shift there re-segments the stream and re-times its samples -- a behaviour + # change, not a speedup. envelope = float(np.percentile(np.abs(dt - period), GAP_ENVELOPE_PCT)) return period + max(GAP_MIN_S, GAP_ENVELOPE_MULT * envelope) diff --git a/src/ezmsg/nwb/slicer.py b/src/ezmsg/nwb/slicer.py index 22890d0..853c048 100644 --- a/src/ezmsg/nwb/slicer.py +++ b/src/ezmsg/nwb/slicer.py @@ -27,7 +27,7 @@ reconstruct_group, ) from .scaling import SCALING_ATTR, StreamScaling, describe_stream_scaling -from .util import ReferenceClockType, as_text, as_text_array +from .util import ReferenceClockType, as_text, as_text_array, interval_median # Default gap threshold as a fraction of the nominal sample period (1.5x period). # Sits between neural-data jitter (<~1.05x) and the smallest real gap (one dropped @@ -42,34 +42,6 @@ RATE_TRIM_HI = 1.5 -MEDIAN_SUBSAMPLE_MAX = 100_000 -"""Above this many intervals, estimate the median from a strided subsample. - -``np.median`` partitions the whole array: 45 ms on a recording-hour of 30 kHz -timestamps, and it runs once per stream at open. The median here is only ever a -*centre* for a trim window or a scale for a variance comparison -- never the -returned estimate, which is a mean over the surviving intervals -- so it needs -to be in the right place, not exact. - -And it lands in exactly the same place. Timestamps arrive on a quantisation -grid, so the median snaps to a grid point that a stride cannot move: on a real -7.09M-interval stream, strides of 16 through 1024 all give 3.3301000002e-05, -and :func:`infer_nominal_rate` returns 30147.9172 from every one of them. -""" - - -def _fast_median(values: np.ndarray) -> float: - """``np.median``, on a strided subsample once the array is large enough. - - Strided rather than random or head-of-array: a stride samples the whole - recording, so a stream whose intervals differ between its start and its end - is represented across the estimate rather than by whichever part was read. - """ - if values.size > MEDIAN_SUBSAMPLE_MAX: - values = values[:: -(-values.size // MEDIAN_SUBSAMPLE_MAX)] - return float(np.median(values)) - - def infer_nominal_rate(dts: np.ndarray) -> float: """Sample rate from per-sample intervals: the mean of the non-outlier ones. @@ -94,7 +66,7 @@ def infer_nominal_rate(dts: np.ndarray) -> float: """ if dts.size == 0: return 0.0 - median = _fast_median(dts) + median = interval_median(dts) if not np.isfinite(median) or median <= 0.0: return 0.0 core = dts[(dts > RATE_TRIM_LO * median) & (dts < RATE_TRIM_HI * median)] @@ -468,7 +440,7 @@ def _load(self) -> None: # ``np.var`` twice was two full passes to evaluate one # ``or``; the second only ran when the first failed, but it # ran on every irregular stream, which is the slow case. - if variance < 1e-3 or variance < 0.05 * _fast_median(dts): + if variance < 1e-3 or variance < 0.05 * interval_median(dts): rate = infer_nominal_rate(dts) else: rate = 0.0 diff --git a/src/ezmsg/nwb/util.py b/src/ezmsg/nwb/util.py index b8d05d3..7adf8b8 100644 --- a/src/ezmsg/nwb/util.py +++ b/src/ezmsg/nwb/util.py @@ -43,3 +43,33 @@ def build_nwb_fname(metadata: DeepDict) -> str: ses = metadata["NWBFile"].get("session_id", metadata["NWBFile"]["session_start_time"].strftime("%Y%m%dT%H%M%S")) fname_str += f"_ses-{ses}" return f"{fname_str}_ephys.nwb" + + +MEDIAN_SUBSAMPLE_MAX = 100_000 +"""Above this many values, :func:`interval_median` uses a strided subsample. + +``np.median`` partitions the whole array -- 45 ms on a recording-hour of 30 kHz +inter-sample intervals -- and the open path does it several times per stream: +once to infer a nominal rate, once per stream to derive a real-gap threshold. +""" + + +def interval_median(values: np.ndarray) -> float: + """Median of a set of *inter-sample intervals*, cheaply. + + For intervals specifically, and not for residuals or other continuous + quantities: timestamps arrive on a device's quantisation grid, so their + differences cluster onto that grid and the median snaps to a grid point a + stride cannot move. On a real 7.09M-interval stream, strides of 16 through + 1024 all return 3.3301000002e-05. That is what makes the shortcut exact here + and not elsewhere -- a median of fit residuals, say, is a continuous + quantity with no grid to snap to, and subsampling it would genuinely change + the answer. + + Strided rather than random or head-of-array: a stride samples the whole + recording, so a stream whose intervals differ between its start and its end + is represented across the estimate rather than by whichever part was read. + """ + if values.size > MEDIAN_SUBSAMPLE_MAX: + values = values[:: -(-values.size // MEDIAN_SUBSAMPLE_MAX)] + return float(np.median(values)) diff --git a/tests/test_slicer_perf.py b/tests/test_slicer_perf.py index d287b34..b3ba832 100644 --- a/tests/test_slicer_perf.py +++ b/tests/test_slicer_perf.py @@ -9,7 +9,8 @@ import numpy as np import pytest -from ezmsg.nwb.slicer import MEDIAN_SUBSAMPLE_MAX, _fast_median, infer_nominal_rate +from ezmsg.nwb.slicer import infer_nominal_rate +from ezmsg.nwb.util import MEDIAN_SUBSAMPLE_MAX, interval_median def quantised_intervals(n: int, period: float = 1 / 30000, grid: float = 4e-8, seed: int = 3) -> np.ndarray: @@ -23,13 +24,13 @@ def quantised_intervals(n: int, period: float = 1 / 30000, grid: float = 4e-8, s @pytest.mark.parametrize("n", [10, MEDIAN_SUBSAMPLE_MAX, MEDIAN_SUBSAMPLE_MAX + 1, 3_000_000]) def test_fast_median_matches_the_real_one(n): dts = quantised_intervals(n) - assert _fast_median(dts) == float(np.median(dts)) + assert interval_median(dts) == float(np.median(dts)) def test_small_arrays_are_not_subsampled(): """Below the threshold there is nothing to gain, so take the exact answer.""" dts = np.linspace(1.0, 2.0, 1001) # even-length medians interpolate; no grid - assert _fast_median(dts) == float(np.median(dts)) + assert interval_median(dts) == float(np.median(dts)) def test_inferred_rate_is_unchanged_by_subsampling(): @@ -63,3 +64,39 @@ def test_ch_labels_come_from_the_electrodes_region(scaled_nwb_path): slicer.close() assert list(labels) == ["elec0", "elec1", "elec2", "elec3"] assert labels.dtype.kind == "U" + + +def test_clock_model_helpers_use_the_same_shortcut(): + """``_nominal_period`` and ``_auto_gap_threshold`` run over full recordings + at open, several times per stream. Both take medians of intervals, so both + get the subsample -- and both must be unmoved by it.""" + from ezmsg.nwb.clockmodel import _auto_gap_threshold, _nominal_period + + dts = quantised_intervals(2_000_000) + times = np.concatenate([[0.0], np.cumsum(dts)]) + + assert _nominal_period(times) == float(np.median(np.diff(times))) + + # And the threshold, against the same function computed with a full median. + import ezmsg.nwb.clockmodel as cm + + subsampled = _auto_gap_threshold(times) + original = cm.interval_median + cm.interval_median = lambda v: float(np.median(v)) + try: + reference = _auto_gap_threshold(times) + finally: + cm.interval_median = original + assert subsampled == reference + + +def test_gap_threshold_percentile_stays_on_the_full_array(): + """Guard the deliberate asymmetry: subsampling the percentile too would be + 16x faster and would move the threshold ~1.4%, re-segmenting streams.""" + import inspect + + from ezmsg.nwb.clockmodel import _auto_gap_threshold + + src = inspect.getsource(_auto_gap_threshold) + assert "interval_median(dt)" in src + assert "np.percentile(np.abs(dt - period)" in src