Stop doing full-array work at open to answer approximate questions

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)
 

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
@@ -66,7 +66,7 @@ def infer_nominal_rate(dts: np.ndarray) -> float:
     """
     if dts.size == 0:
         return 0.0
-    median = float(np.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)]
@@ -436,7 +436,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 * interval_median(dts):
                         rate = infer_nominal_rate(dts)
                     else:
                         rate = 0.0
@@ -475,16 +479,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``

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))

tests/test_slicer_perf.py

@@ -0,0 +1,102 @@
+"""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 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:
+    """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 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 interval_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"
+
+
+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