iterator gap fix

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
 

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,

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,19 +497,43 @@ 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,
                 data=out_data,
                 axes={
                     **template.axes,
                     "time": replace(
-                        template.axes["time"],
+                        time_axis,
                         offset=self._ts_off + chunk_t0,
                     ),
                 },