diff --git a/src/ezmsg/nwb/clockdriven.py b/src/ezmsg/nwb/clockdriven.py index 3815165..27133d2 100644 --- a/src/ezmsg/nwb/clockdriven.py +++ b/src/ezmsg/nwb/clockdriven.py @@ -12,7 +12,7 @@ from ezmsg.baseproc.units import BaseClockDrivenUnit from ezmsg.util.messages.axisarray import AxisArray, LinearAxis -from .slicer import NWBSlicer +from .slicer import DEFAULT_GAP_TOL, NWBSlicer from .util import ReferenceClockType @@ -32,6 +32,18 @@ class NWBClockDrivenSettings(ClockDrivenSettings): meaningful when ``n_time`` is ``None`` (variable chunk size); with a fixed ``n_time`` the rate acts as a pause/unpause gate since chunk size is fixed. """ + gap_tol: float = DEFAULT_GAP_TOL + """Gap threshold for timestamped continuous streams, as a fraction of the + nominal sample period (forwarded to ``NWBSlicer.read_by_time``). + + The clock cadence is never altered: ``file_t`` always advances by ``dt`` per + tick so sibling producers on the same clock stay aligned, and windows that + fall inside a gap emit zero-length chunks. When a single window straddles a + gap (interval > ``(1 + gap_tol)`` nominal periods), that chunk's time axis + becomes a ``CoordinateAxis`` carrying the true per-sample timestamps rather + than a uniform ``LinearAxis`` that would misplace post-gap samples. Set very + large to keep a single LinearAxis across gaps. No effect on event streams. + """ @processor_state @@ -242,7 +254,7 @@ def _process_time_window(self, clock_tick: LinearAxis) -> AxisArray | None: t_start = self._state.file_t t_end = t_start + dt - output = slicer.read_by_time(self.settings.stream_key, t_start, t_end) + output = slicer.read_by_time(self.settings.stream_key, t_start, t_end, gap_tol=self.settings.gap_tol) self._state.file_t = t_end return output diff --git a/src/ezmsg/nwb/iterator.py b/src/ezmsg/nwb/iterator.py index cc3fd9f..5c8edcb 100644 --- a/src/ezmsg/nwb/iterator.py +++ b/src/ezmsg/nwb/iterator.py @@ -15,7 +15,7 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace -from .slicer import NWBSlicer +from .slicer import DEFAULT_GAP_TOL, NWBSlicer, find_gaps from .util import ReferenceClockType # Sentinel pushed to the prefetch queue to indicate end-of-stream. Identity-compared. @@ -45,6 +45,17 @@ class NWBIteratorSettings(ez.Settings): """HDF5 raw data chunk cache size in bytes (forwarded to NWBSlicer).""" rdcc_nslots: int = NWBSlicer.DEFAULT_RDCC_NSLOTS """HDF5 raw data chunk cache slot count (forwarded to NWBSlicer).""" + gap_tol: float = DEFAULT_GAP_TOL + """Gap threshold for timestamped continuous streams, as a fraction of the + nominal sample period: a gap is declared when an interval exceeds + ``(1 + gap_tol) / fs``. Chunks spanning a gap are split into gap-free + messages so the regular ``LinearAxis`` never misplaces post-gap samples. + + The default ``0.5`` (1.5x period) sits between neural-data jitter (<~1.05x) + and the smallest real gap (one dropped sample = ~2x), so it catches every + gap without splitting on jitter. Set very large to disable splitting. No + effect on rate-only streams or event tables. + """ @processor_state @@ -63,6 +74,7 @@ def _build_chunk_messages_static( slicer: NWBSlicer, streams: dict, chunk_ix: int, + gap_tol: float = DEFAULT_GAP_TOL, ) -> list[AxisArray]: """Build the messages for ``chunk_ix`` from explicit slicer/streams refs. @@ -114,29 +126,56 @@ def _build_chunk_messages_static( out.append(template) else: out_data = info.dset[start_idx:stop_idx] - if info.timestamps is not None and start_idx < len(info.timestamps): - # Explicit timestamps are already absolute (file-relative) times. - chunk_t0 = info.timestamps[start_idx] + time_axis = template.axes["time"] + + # Timestamped continuous stream on a regular LinearAxis: a single + # chunk that spans a gap in the explicit timestamps would emit a + # uniform time axis that misplaces every post-gap sample. Split it + # into gap-free runs, each anchored on its own first timestamp. + # Rate-only streams (no per-sample timestamps) and CoordinateAxis + # streams (no ``gain`` to compare against) keep the old single-chunk + # path. + if info.has_timestamps and info.timestamps is not None and hasattr(time_axis, "gain") and len(out_data): + ts_chunk = np.asarray(info.timestamps[start_idx:stop_idx]) + gap_after = find_gaps(ts_chunk, time_axis.gain, gap_tol) + # Run boundaries within the chunk: [0, gap1+1, gap2+1, ..., len]. + bounds = [0, *(gap_after + 1).tolist(), ts_chunk.shape[0]] + for b0, b1 in zip(bounds[:-1], bounds[1:]): + out.append( + replace( + template, + data=out_data[b0:b1], + axes={ + **template.axes, + "time": replace(time_axis, offset=ts_off + ts_chunk[b0]), + }, + key=strm_name, + ) + ) else: - # Rate-only: the absolute time of ``start_idx`` is the stream's - # own start (``info.t0``) plus the within-stream offset. Omitting - # ``info.t0`` would label a late-starting stream as if it began - # at the file origin, mis-timing it against other streams. - chunk_t0 = float(info.t0) + template.axes["time"].gain * start_idx - out.append( - replace( - template, - data=out_data, - axes={ - **template.axes, - "time": replace( - template.axes["time"], - offset=ts_off + chunk_t0, - ), - }, - key=strm_name, + if info.timestamps is not None and start_idx < len(info.timestamps): + # Explicit timestamps are already absolute (file-relative) times. + chunk_t0 = info.timestamps[start_idx] + else: + # Rate-only: the absolute time of ``start_idx`` is the stream's + # own start (``info.t0``) plus the within-stream offset. Omitting + # ``info.t0`` would label a late-starting stream as if it began + # at the file origin, mis-timing it against other streams. + chunk_t0 = float(info.t0) + time_axis.gain * start_idx + out.append( + replace( + template, + data=out_data, + axes={ + **template.axes, + "time": replace( + time_axis, + offset=ts_off + chunk_t0, + ), + }, + key=strm_name, + ) ) - ) return out @@ -146,6 +185,7 @@ def _prefetch_worker( n_chunks: int, q: queue.Queue, stop: threading.Event, + gap_tol: float = DEFAULT_GAP_TOL, ) -> None: """Prefetch worker target. Top-level function (no closure over the iterator) so the iterator can be garbage-collected as soon as the user @@ -156,7 +196,7 @@ def _prefetch_worker( for chunk_ix in range(n_chunks): if stop.is_set(): return - msgs = _build_chunk_messages_static(slicer, streams, chunk_ix) + msgs = _build_chunk_messages_static(slicer, streams, chunk_ix, gap_tol) # Block on a full queue, but wake periodically to honour stop. while not stop.is_set(): try: @@ -302,7 +342,7 @@ def _preload(self): def _build_chunk_messages(self, chunk_ix: int) -> list[AxisArray]: """Sync-side wrapper around :func:`_build_chunk_messages_static`.""" - return _build_chunk_messages_static(self._state.slicer, self._state.streams, chunk_ix) + return _build_chunk_messages_static(self._state.slicer, self._state.streams, chunk_ix, self.settings.gap_tol) def _chunk_step(self): """Sync path: build the next chunk and append to the deque.""" @@ -332,6 +372,7 @@ def _start_prefetch(self) -> None: self._state.n_chunks, self._state.prefetch_queue, self._state.prefetch_stop, + self.settings.gap_tol, ), name="NWBIterator-prefetch", daemon=True, diff --git a/src/ezmsg/nwb/slicer.py b/src/ezmsg/nwb/slicer.py index 47135b7..bcc5d32 100644 --- a/src/ezmsg/nwb/slicer.py +++ b/src/ezmsg/nwb/slicer.py @@ -20,6 +20,27 @@ from .util import ReferenceClockType +# 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 +# sample = ~2x), so it catches every gap without splitting on jitter. Shared as the +# default for the slicer, iterator, and clock-driven producer. +DEFAULT_GAP_TOL = 0.5 + + +def find_gaps(timestamps: np.ndarray, gain: float, gap_tol: float) -> np.ndarray: + """Indices ``i`` where a gap separates sample ``i`` from sample ``i + 1``. + + A gap is an inter-sample interval exceeding ``(1 + gap_tol) * gain`` — i.e. + the timestamps jump by more than ``(1 + gap_tol)`` nominal sample periods. + Returns the left-edge indices as an ``int`` array (empty when there are + fewer than two samples, no usable ``gain``, or no gaps). Shared by the + iterator (which splits chunks at gaps) and the slicer's ``read_by_time`` + (which switches to a CoordinateAxis when a window spans one). + """ + if gain <= 0.0 or timestamps.shape[0] < 2: + return np.empty(0, dtype=int) + return np.flatnonzero(np.diff(timestamps) > gain * (1.0 + gap_tol)) + @dataclass class StreamInfo: @@ -416,12 +437,22 @@ def read_by_index(self, stream_key: str, start_idx: int, stop_idx: int) -> AxisA key=stream_key, ) - def read_by_time(self, stream_key: str, t_start: float, t_end: float) -> AxisArray: + def read_by_time( + self, stream_key: str, t_start: float, t_end: float, gap_tol: float = DEFAULT_GAP_TOL + ) -> AxisArray: """Read data by time window [t_start, t_end). For timestamped continuous streams and event/interval tables. t_start and t_end are in the same reference frame as the stored timestamps (i.e., file-relative, before ts_off). + + ``gap_tol`` is the gap threshold as a fraction of the nominal sample + period (see :func:`find_gaps`). When a continuous window's samples are + regularly spaced the result carries a cheap ``LinearAxis``; when the + window straddles a gap, the time axis becomes a ``CoordinateAxis`` + carrying the true per-sample timestamps so the gap is represented + faithfully instead of silently flattening post-gap samples onto a + uniform axis. Ignored for event streams. """ info = self._streams[stream_key] template = info.template @@ -466,11 +497,35 @@ def read_by_time(self, stream_key: str, t_start: float, t_end: float) -> AxisArr stop_idx = int(np.searchsorted(ts_arr, t_end, side="left")) out_data = info.dset[start_idx:stop_idx] + time_axis = template.axes["time"] + ts_window = np.asarray(ts_arr[start_idx:stop_idx]) + has_gain = hasattr(time_axis, "gain") + + # A window straddling a gap (or any stream with no usable rate) + # cannot be described by a uniform LinearAxis without misplacing + # samples. Emit a CoordinateAxis carrying the true per-sample + # timestamps instead — the same representation events use. Gap-free + # windows keep the cheaper LinearAxis. + emit_coords = (not has_gain) or find_gaps(ts_window, time_axis.gain, gap_tol).size > 0 + if emit_coords: + return replace( + template, + data=out_data, + axes={ + **template.axes, + "time": AxisArray.CoordinateAxis( + data=self._ts_off + ts_window, + dims=["time"], + unit=getattr(time_axis, "unit", "s"), + ), + }, + key=stream_key, + ) if start_idx < len(ts_arr): chunk_t0 = ts_arr[start_idx] else: - chunk_t0 = template.axes["time"].gain * start_idx if hasattr(template.axes["time"], "gain") else 0.0 + chunk_t0 = time_axis.gain * start_idx return replace( template, @@ -478,7 +533,7 @@ def read_by_time(self, stream_key: str, t_start: float, t_end: float) -> AxisArr axes={ **template.axes, "time": replace( - template.axes["time"], + time_axis, offset=self._ts_off + chunk_t0, ), }, diff --git a/tests/conftest.py b/tests/conftest.py index 78fcfc2..cfa981b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,9 @@ """Shared fixtures for ezmsg-nwb tests.""" +import datetime from pathlib import Path +import numpy as np import pytest from create_test_nwb import create_test_nwb from filelock import FileLock @@ -22,3 +24,130 @@ def test_nwb_path(): if not path.exists(): create_test_nwb(path) return path + + +# --- Gappy-stream fixture --------------------------------------------------- +# +# A timestamped continuous stream whose explicit timestamps contain a single +# large gap, while still advertising a nominal ``rate`` (so the slicer picks a +# regular ``LinearAxis`` for it). This is the dangerous case for both Solution +# A (iterator) and Solution B (slicer / clock-driven): 100 Hz, 150 samples +# before a 1.0 s gap and 150 after. The per-sample data value equals its global +# sample index so callers can recover where each emitted sample came from. + +GAPPY_RATE = 100.0 +GAPPY_GAIN = 1.0 / GAPPY_RATE +GAPPY_N_PRE = 150 +GAPPY_N_POST = 150 +GAPPY_GAP = 1.0 + + +def gappy_timestamps() -> np.ndarray: + """True per-sample timestamps of the gappy stream (file-relative).""" + pre = np.arange(GAPPY_N_PRE) * GAPPY_GAIN # 0.00 .. 1.49 + post = pre[-1] + GAPPY_GAIN + GAPPY_GAP + np.arange(GAPPY_N_POST) * GAPPY_GAIN # 2.50 .. 3.99 + return np.concatenate([pre, post]) + + +def _index_data(n: int, n_ch: int = 3) -> np.ndarray: + """Data whose row ``i`` carries the constant value ``i`` across channels, so + a caller can recover each emitted sample's global index from its value.""" + return np.arange(n, dtype=np.float32)[:, None] + np.zeros((1, n_ch), dtype=np.float32) + + +def _write_nwb(path, series, rate_attrs): + """Write *series* (list of TimeSeries) to *path*; stamp ``rate`` attrs after + so the slicer treats those streams as regular (LinearAxis).""" + import h5py + from pynwb import NWBHDF5IO, NWBFile + + nwbfile = NWBFile( + session_description="synthetic", + identifier="synthetic001", + session_start_time=datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc), + ) + for s in series: + nwbfile.add_acquisition(s) + with NWBHDF5IO(str(path), "w") as io: + io.write(nwbfile) + with h5py.File(str(path), "a") as f: + for name, rate in rate_attrs.items(): + f[f"acquisition/{name}/timestamps"].attrs["rate"] = rate + return path + + +@pytest.fixture(scope="session") +def gappy_nwb_path(tmp_path_factory): + """Build a minimal NWB file with one gappy timestamped stream.""" + from pynwb import TimeSeries + + path = tmp_path_factory.mktemp("gappy") / "gappy.nwb" + n = GAPPY_N_PRE + GAPPY_N_POST + series = TimeSeries( + name="Gappy", + data=_index_data(n), + unit="V", + timestamps=gappy_timestamps(), + description="gappy timestamped stream", + ) + # The rate attr makes the slicer assign a regular LinearAxis, mirroring a + # recording with a known fs but dropped samples. + return _write_nwb(path, [series], {"Gappy": GAPPY_RATE}) + + +@pytest.fixture(scope="session") +def gappy_and_clean_nwb_path(tmp_path_factory): + """Two regular (rate-stamped) containers: ``Gappy`` (has the 1.0 s gap) and + ``Clean`` (400 contiguous samples, no gap), both 100 Hz spanning 0..3.99 s. + + Mirrors the CA001 layout — multiple acquisition containers in one file — so + a test can drive one clock-driven producer per stream and confirm a gap in + one container does not desync the other. + """ + from pynwb import TimeSeries + + path = tmp_path_factory.mktemp("gappy_clean") / "gappy_clean.nwb" + gappy = TimeSeries( + name="Gappy", data=_index_data(GAPPY_N_PRE + GAPPY_N_POST), unit="V", + timestamps=gappy_timestamps(), description="gappy stream", + ) + n_clean = 400 + clean = TimeSeries( + name="Clean", data=_index_data(n_clean), unit="V", + timestamps=np.arange(n_clean) * GAPPY_GAIN, description="contiguous stream", + ) + return _write_nwb(path, [gappy, clean], {"Gappy": GAPPY_RATE, "Clean": GAPPY_RATE}) + + +# --- Irregular-stream fixture ----------------------------------------------- +# +# A timestamped continuous stream with NO rate attr and high inter-sample +# variance, so the slicer's rate heuristic gives up (rate=0) and assigns a +# CoordinateAxis template (no ``gain``). Exercises the ``not has_gain`` branch +# of ``read_by_time``, which must still emit true per-sample timestamps. + +IRREGULAR_N = 200 + + +def irregular_timestamps() -> np.ndarray: + """Monotonic, highly-irregular per-sample timestamps (fixed seed).""" + rng = np.random.default_rng(7) + intervals = rng.uniform(0.01, 0.5, size=IRREGULAR_N) + return np.cumsum(intervals) + + +@pytest.fixture(scope="session") +def irregular_nwb_path(tmp_path_factory): + """Build an NWB file with one irregular (rate-0 → CoordinateAxis) stream.""" + from pynwb import TimeSeries + + path = tmp_path_factory.mktemp("irregular") / "irregular.nwb" + series = TimeSeries( + name="Irregular", + data=_index_data(IRREGULAR_N), + unit="V", + timestamps=irregular_timestamps(), + description="irregular timestamped stream (no rate)", + ) + # No rate attr -> slicer detects rate 0 -> CoordinateAxis template. + return _write_nwb(path, [series], {}) diff --git a/tests/test_clockdriven.py b/tests/test_clockdriven.py index 9ff6bf6..fdbd815 100644 --- a/tests/test_clockdriven.py +++ b/tests/test_clockdriven.py @@ -2,6 +2,9 @@ import threading +import numpy as np +import pytest +from conftest import GAPPY_N_POST, GAPPY_N_PRE, gappy_timestamps from ezmsg.util.messages.axisarray import AxisArray, LinearAxis from ezmsg.nwb.clockdriven import NWBClockDrivenProducer, NWBClockDrivenSettings @@ -393,6 +396,125 @@ def test_playback_rate_zero_pauses(test_nwb_path): assert producer._state.sample_idx == 0 +# --- Timestamp gaps (Solution B) --- + + +def _gappy_producer(path, **overrides): + settings = NWBClockDrivenSettings( + fs=0.0, + n_time=10, + filepath=path, + stream_key="Gappy", + reference_clock=ReferenceClockType.UNKNOWN, + **overrides, + ) + return NWBClockDrivenProducer(settings=settings) + + +def test_clockdriven_clean_window_linear_axis(gappy_nwb_path): + """A gap-free window emits a regular LinearAxis chunk.""" + p = _gappy_producer(gappy_nwb_path, start_offset=0.0) + r = p(AxisArray.LinearAxis(gain=1.0, offset=0.0)) # window [0, 1.0) + assert r is not None + assert r.data.shape[0] == 100 + assert hasattr(r.axes["time"], "gain") # LinearAxis + assert not hasattr(r.axes["time"], "data") + + +def test_clockdriven_window_inside_gap_emits_empty(gappy_nwb_path): + """During the gap the producer emits zero-length chunks (no fabricated + data); the clock cursor still advances so sibling streams stay aligned.""" + p = _gappy_producer(gappy_nwb_path, start_offset=1.6) + r = p(AxisArray.LinearAxis(gain=0.1, offset=0.0)) # window [1.6, 1.7) inside gap + assert r is not None + assert r.data.shape[0] == 0 + # Cursor advanced by dt (from t0 + start_offset = 1.6 to 1.7) despite + # emitting nothing, so sibling producers on the same clock stay aligned. + assert p._state.file_t == pytest.approx(1.7, abs=1e-9) + + +def test_clockdriven_window_spanning_gap_coordinate_axis(gappy_nwb_path): + """A window containing samples on both sides of the gap emits a + CoordinateAxis with the true timestamps, not a misleading LinearAxis.""" + p = _gappy_producer(gappy_nwb_path, start_offset=1.0) + r = p(AxisArray.LinearAxis(gain=2.0, offset=0.0)) # window [1.0, 3.0) + assert r is not None + assert r.data.shape[0] == 100 + assert hasattr(r.axes["time"], "data") # CoordinateAxis + ts = np.asarray(r.axes["time"].data) + assert ts[49] == pytest.approx(1.49, abs=1e-6) + assert ts[50] == pytest.approx(2.50, abs=1e-6) + + +def test_clockdriven_gap_tol_keeps_linear_axis(gappy_nwb_path): + """A large gap_tol keeps a gap-spanning window on a single LinearAxis.""" + p = _gappy_producer(gappy_nwb_path, start_offset=1.0, gap_tol=1e6) + r = p(AxisArray.LinearAxis(gain=2.0, offset=0.0)) + assert r is not None + assert hasattr(r.axes["time"], "gain") # LinearAxis despite the gap + assert r.data.shape[0] == 100 + + +def test_clockdriven_sweep_across_gap_preserves_all_samples(gappy_nwb_path): + """Sweeping the whole stream tick-by-tick recovers every sample in order: + no data lost in the gap, no duplication, cursor advances monotonically.""" + p = _gappy_producer(gappy_nwb_path, start_offset=0.0) + chunks = [] + # 40 ticks * 0.1 s = 4.0 s covers the full 0..3.99 s stream. + for _ in range(40): + r = p(AxisArray.LinearAxis(gain=0.1, offset=0.0)) + if r is not None and r.data.shape[0] > 0: + chunks.append(np.asarray(r.data)) + data = np.concatenate(chunks, axis=0) + n = len(gappy_timestamps()) + assert data.shape[0] == n + np.testing.assert_array_equal(data[:, 0], np.arange(n, dtype=np.float32)) + + +def test_clockdriven_two_producers_stay_synced_across_gap(gappy_and_clean_nwb_path): + """The core Solution-B guarantee: one producer per container, both on the + same clock. A gap in ``Gappy`` makes it emit empty chunks for a few ticks + while ``Clean`` keeps producing — yet their cursors never diverge, because + every producer advances ``file_t`` by the same ``dt`` regardless of gaps. + This is the multi-container (CA001-style) playback case. + """ + + def producer(key): + return NWBClockDrivenProducer( + settings=NWBClockDrivenSettings( + fs=0.0, n_time=10, filepath=gappy_and_clean_nwb_path, stream_key=key, + reference_clock=ReferenceClockType.UNKNOWN, start_offset=0.0, + ) + ) + + pg, pc = producer("Gappy"), producer("Clean") + g_total = c_total = 0 + gap_tick_seen = False # a tick where Gappy is empty but Clean still produces + g_chunks = [] + # 0.5 s windows; the 1.0 s gap (1.49..2.50) lands wholly inside ticks 3 & 4. + for _ in range(9): + tick = AxisArray.LinearAxis(gain=0.5, offset=0.0) + rg, rc = pg(tick), pc(tick) + # Cursors stay in lockstep on every single tick — the sync invariant. + assert pg._state.file_t == pytest.approx(pc._state.file_t, abs=1e-9) + gn = 0 if rg is None else rg.data.shape[0] + cn = 0 if rc is None else rc.data.shape[0] + if gn == 0 and cn > 0: + gap_tick_seen = True + if gn: + g_chunks.append(np.asarray(rg.data)) + g_total += gn + c_total += cn + + assert gap_tick_seen, "expected ticks where the gappy stream was empty while the clean stream produced" + # No samples lost or duplicated in either container. + assert g_total == GAPPY_N_PRE + GAPPY_N_POST + assert c_total == 400 + # Gappy data still in order across the gap. + g_data = np.concatenate(g_chunks, axis=0) + np.testing.assert_array_equal(g_data[:, 0], np.arange(g_total, dtype=np.float32)) + + # --- Unit class --- diff --git a/tests/test_iterator.py b/tests/test_iterator.py index af4c136..68437d2 100644 --- a/tests/test_iterator.py +++ b/tests/test_iterator.py @@ -7,9 +7,13 @@ import numpy as np import pytest +from conftest import GAPPY_N_POST, GAPPY_N_PRE from ezmsg.nwb import NWBAxisArrayIterator, NWBIteratorSettings, ReferenceClockType +# The gappy-stream fixture (``gappy_nwb_path``) and its GAPPY_* parameters live +# in conftest.py so the slicer and clock-driven tests can share them. + async def test_areset_state_runs_reset_in_worker_thread(test_nwb_path): """``_areset_state`` must offload sync ``_reset_state`` to a worker @@ -600,9 +604,9 @@ def test_prefetch_runs_in_worker_thread(test_nwb_path, monkeypatch): seen_tids: list[int] = [] real = iterator_mod._build_chunk_messages_static - def spy(slicer, streams, chunk_ix): + def spy(slicer, streams, chunk_ix, gap_tol=0.5): seen_tids.append(threading.get_ident()) - return real(slicer, streams, chunk_ix) + return real(slicer, streams, chunk_ix, gap_tol) monkeypatch.setattr(iterator_mod, "_build_chunk_messages_static", spy) @@ -818,3 +822,145 @@ def spy_file(*args, **kwargs): open_kwargs = seen_kwargs[0] assert open_kwargs["rdcc_nbytes"] == custom_nbytes assert open_kwargs["rdcc_nslots"] == custom_nslots + + +# --- Timestamp gaps (Solution A: split chunks at gaps) ---------------------- + + +def test_gappy_whole_stream_not_one_gap_spanning_chunk(gappy_nwb_path): + """A chunk big enough to cover the whole gappy stream must be split into + two gap-free messages, not emitted as one chunk that silently spans the + gap with a uniform LinearAxis. + """ + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=gappy_nwb_path, + chunk_dur=100.0, # whole stream in a single chunk + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + ) + msgs = [m for m in it if m.data.shape[0] > 0] + + assert len(msgs) == 2, "gappy stream was not split at the gap" + assert msgs[0].data.shape[0] == GAPPY_N_PRE + assert msgs[1].data.shape[0] == GAPPY_N_POST + # Offsets reflect the real first-timestamp of each gap-free run. + assert msgs[0].axes["time"].offset == pytest.approx(0.0, abs=1e-6) + assert msgs[1].axes["time"].offset == pytest.approx(2.50, abs=1e-6) + # Sample ordering and identity preserved across the split. + assert msgs[0].data[0, 0] == 0 + assert msgs[0].data[-1, 0] == GAPPY_N_PRE - 1 + assert msgs[1].data[0, 0] == GAPPY_N_PRE + assert msgs[1].data[-1, 0] == GAPPY_N_PRE + GAPPY_N_POST - 1 + + +def test_gappy_midchunk_split(gappy_nwb_path): + """When the gap falls in the middle of an index-based chunk, that chunk is + split into two messages while gap-free chunks pass through unchanged. + + chunk_dur=1.0 @ 100 Hz -> 100 samples/chunk. The gap sits between sample + 149 and 150, i.e. inside the second chunk (samples 100..199). + """ + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=gappy_nwb_path, + chunk_dur=1.0, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + ) + msgs = [m for m in it if m.data.shape[0] > 0] + + # chunk0: samples 0..99 (gap-free) -> 1 msg + # chunk1: samples 100..199 spans gap -> split 100..149 / 150..199 + # chunk2: samples 200..299 (gap-free) -> 1 msg + assert len(msgs) == 4 + sizes = [m.data.shape[0] for m in msgs] + assert sizes == [100, 50, 50, 100] + # The two halves of the split chunk sit on either side of the gap. + assert msgs[1].axes["time"].offset == pytest.approx(1.00, abs=1e-6) + assert msgs[2].axes["time"].offset == pytest.approx(2.50, abs=1e-6) + + +def test_gappy_segments_match_true_timestamps(gappy_nwb_path): + """Every emitted message's reconstructed time axis (offset + i*gain) must + match the file's true per-sample timestamps within a fraction of a sample. + """ + from ezmsg.nwb.slicer import NWBSlicer + + slicer = NWBSlicer( + filepath=gappy_nwb_path, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + true_ts = np.asarray(slicer.get_stream_info("Gappy").timestamps[:]) + slicer.close() + + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=gappy_nwb_path, + chunk_dur=1.0, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + ) + for m in it: + if m.data.shape[0] == 0: + continue + gain = m.axes["time"].gain + offset = m.axes["time"].offset + idx = m.data[:, 0].astype(int) # data value == global sample index + reconstructed = offset + np.arange(m.data.shape[0]) * gain + np.testing.assert_allclose(reconstructed, true_ts[idx], atol=gain * 0.5) + + +def test_gappy_total_samples_and_order_preserved(gappy_nwb_path): + """Splitting at gaps must not drop, duplicate, or reorder samples.""" + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=gappy_nwb_path, + chunk_dur=1.0, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + ) + data = np.concatenate([m.data for m in it if m.data.shape[0] > 0], axis=0) + n = GAPPY_N_PRE + GAPPY_N_POST + assert data.shape[0] == n + np.testing.assert_array_equal(data[:, 0], np.arange(n, dtype=np.float32)) + + +def test_gap_tol_disables_split(gappy_nwb_path): + """A large ``gap_tol`` widens the gap threshold enough that the stream is + emitted as a single (gap-spanning) chunk again — the knob works. + """ + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=gappy_nwb_path, + chunk_dur=100.0, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + gap_tol=1e6, + ) + ) + msgs = [m for m in it if m.data.shape[0] > 0] + assert len(msgs) == 1 + assert msgs[0].data.shape[0] == GAPPY_N_PRE + GAPPY_N_POST + + +def test_jittered_stream_not_oversplit(test_nwb_path): + """The lightly-jittered Broadband stream (no real gaps) must not be split: + sub-microsecond jitter stays well under the gap threshold. + """ + it = NWBAxisArrayIterator( + NWBIteratorSettings( + filepath=test_nwb_path, + chunk_dur=1.0, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Broadband"], + ) + ) + msgs = [m for m in it if m.data.shape[0] > 0] + # 3 s of 1 kHz data in 1 s chunks -> exactly 3 non-empty messages. + assert len(msgs) == 3 diff --git a/tests/test_slicer.py b/tests/test_slicer.py index 537350b..4b374fb 100644 --- a/tests/test_slicer.py +++ b/tests/test_slicer.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from conftest import gappy_timestamps from ezmsg.nwb.slicer import NWBSlicer from ezmsg.nwb.util import ReferenceClockType @@ -17,6 +18,28 @@ def slicer(test_nwb_path): s.close() +@pytest.fixture +def gappy_slicer(gappy_nwb_path): + s = NWBSlicer( + filepath=gappy_nwb_path, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Gappy"], + ) + yield s + s.close() + + +@pytest.fixture +def irregular_slicer(irregular_nwb_path): + s = NWBSlicer( + filepath=irregular_nwb_path, + reference_clock=ReferenceClockType.UNKNOWN, + stream_keys=["Irregular"], + ) + yield s + s.close() + + # --- Stream discovery --- @@ -349,6 +372,74 @@ def test_read_by_time_events_empty_window(slicer): assert msg.data.shape[0] == 0 +# --- Timestamp gaps (Solution B: CoordinateAxis on gappy windows) --------- + + +def test_read_by_time_clean_window_linear_axis(gappy_slicer): + """A gap-free window keeps the cheap regular LinearAxis representation.""" + msg = gappy_slicer.read_by_time("Gappy", 0.0, 1.0) + assert msg.data.shape[0] == 100 # samples 0..99 @ 100 Hz + assert hasattr(msg.axes["time"], "gain") # LinearAxis + assert not hasattr(msg.axes["time"], "data") + assert msg.axes["time"].offset == pytest.approx(0.0, abs=1e-6) + + +def test_read_by_time_gap_spanning_window_coordinate_axis(gappy_slicer): + """A window that straddles the gap must emit a CoordinateAxis carrying the + true per-sample timestamps rather than a uniform LinearAxis that would + misplace every post-gap sample. + """ + msg = gappy_slicer.read_by_time("Gappy", 1.0, 3.0) + # 50 pre-gap (t 1.00..1.49) + 50 post-gap (t 2.50..2.99). + assert msg.data.shape[0] == 100 + assert hasattr(msg.axes["time"], "data") # CoordinateAxis + assert not hasattr(msg.axes["time"], "gain") + ts = np.asarray(msg.axes["time"].data) + assert ts.shape[0] == 100 + # The jump across the gap is preserved, not flattened. + assert ts[49] == pytest.approx(1.49, abs=1e-6) + assert ts[50] == pytest.approx(2.50, abs=1e-6) + + +def test_read_by_time_coordinate_timestamps_match_true(gappy_slicer): + """CoordinateAxis timestamps equal the file's true timestamps per sample.""" + msg = gappy_slicer.read_by_time("Gappy", 1.0, 3.0) + true_ts = gappy_timestamps() + idx = msg.data[:, 0].astype(int) # data value == global sample index + np.testing.assert_allclose(np.asarray(msg.axes["time"].data), true_ts[idx], atol=1e-9) + + +def test_read_by_time_window_inside_gap_empty(gappy_slicer): + """A window entirely inside the gap returns a zero-length result.""" + msg = gappy_slicer.read_by_time("Gappy", 1.7, 2.3) + assert msg.data.shape[0] == 0 + + +def test_read_by_time_gap_tol_disables_coordinate_axis(gappy_slicer): + """A large gap_tol widens the threshold so even a gap-spanning window stays + a single LinearAxis chunk — the knob works (and loses no samples).""" + msg = gappy_slicer.read_by_time("Gappy", 1.0, 3.0, gap_tol=1e6) + assert hasattr(msg.axes["time"], "gain") # LinearAxis + assert msg.data.shape[0] == 100 + + +def test_read_by_time_irregular_stream_coordinate_axis(irregular_slicer): + """A stream with no usable rate (CoordinateAxis template, no ``gain``) can + never be described by a uniform LinearAxis, so every window emits a + CoordinateAxis carrying the true per-sample timestamps. Covers the + ``not has_gain`` branch of read_by_time. + """ + from conftest import irregular_timestamps + + true_ts = irregular_timestamps() + msg = irregular_slicer.read_by_time("Irregular", float(true_ts[50]), float(true_ts[150])) + assert msg.data.shape[0] == 100 # samples 50..149 + assert hasattr(msg.axes["time"], "data") # CoordinateAxis + assert not hasattr(msg.axes["time"], "gain") + idx = msg.data[:, 0].astype(int) # data value == global sample index + np.testing.assert_allclose(np.asarray(msg.axes["time"].data), true_ts[idx], atol=1e-9) + + # --- Lifecycle ---