Fix chunk indexing for streams with differing start times

src/ezmsg/nwb/iterator.py

@@ -77,9 +77,17 @@ def _build_chunk_messages_static(
     out: list[AxisArray] = []
     for strm_name, strm_dict in streams.items():
         info = strm_dict["info"]
-        start_idx = strm_dict["chunk_offsets"][chunk_ix]
-        if chunk_ix + 1 < len(strm_dict["chunk_offsets"]):
-            stop_idx = strm_dict["chunk_offsets"][chunk_ix + 1]
+        chunk_offsets = strm_dict["chunk_offsets"]
+        # Defensive: offset tables are built one entry per global chunk, so this
+        # never trips in normal operation. It guards against a caller handing us
+        # a stream whose table is shorter than ``n_chunks`` — index out of bounds
+        # would otherwise crash the whole chunk instead of just dropping that
+        # stream for this index.
+        if chunk_ix >= len(chunk_offsets):
+            continue
+        start_idx = chunk_offsets[chunk_ix]
+        if chunk_ix + 1 < len(chunk_offsets):
+            stop_idx = chunk_offsets[chunk_ix + 1]
         else:
             stop_idx = info.dset.shape[0]
         template = info.template
@@ -107,9 +115,14 @@ def _build_chunk_messages_static(
         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]
             else:
-                chunk_t0 = template.axes["time"].gain * start_idx
+                # 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,
@@ -253,11 +266,28 @@ def _preload(self):
                 chunk_boundaries = start_time + np.arange(n_chunks) * self.settings.chunk_dur - slicer.ts_off
                 chunk_ix_offsets = np.searchsorted(timestamps, chunk_boundaries, side="left").astype(int)
             else:
-                samps_per_chunk = self.settings.chunk_dur / template.axes["time"].gain
-                t0_abs = float(info.t0) + float(slicer.ts_off)
-                first_chunk = max(0, int((t0_abs - float(start_time)) // self.settings.chunk_dur))
-                chunk_ix_offsets = np.arange(n_chunks - first_chunk) * samps_per_chunk
-                chunk_ix_offsets = chunk_ix_offsets.astype(int)
+                # Sample index at each GLOBAL chunk boundary, computed from the
+                # stream's own start time (``info.t0``) and nominal gain. Building
+                # offsets on the shared global grid — rather than from each
+                # stream's own first sample — keeps streams that start at
+                # different times mutually aligned: chunk ``j`` covers the same
+                # wall-clock window for every stream. Boundaries before this
+                # stream begins go negative and boundaries past its end overshoot
+                # ``n_samples``; clamping turns both into empty slices, so a
+                # late-starting / early-ending stream simply contributes nothing
+                # to the chunks outside its span instead of being shifted.
+                gain = template.axes["time"].gain
+                chunk_boundaries = start_time + np.arange(n_chunks) * self.settings.chunk_dur - slicer.ts_off
+                # First sample at/after each boundary — ``searchsorted(side="left")``
+                # semantics on a regular grid, matching the event branch above.
+                # Use ceil, not round: round assigns a boundary to the nearest
+                # sample, which can pull a pre-boundary sample into the chunk when
+                # chunk_dur isn't an integer multiple of the sample period and
+                # disagree with the timestamped/event paths. The epsilon absorbs
+                # floating-point drift so an exact boundary isn't bumped up a sample.
+                rel = (chunk_boundaries - float(info.t0)) / gain
+                chunk_ix_offsets = np.ceil(rel - 1e-6).astype(int)
+                chunk_ix_offsets = np.clip(chunk_ix_offsets, 0, info.dset.shape[0])
 
             self._state.streams[name] = {
                 "info": info,

src/ezmsg/nwb/slicer.py

@@ -393,9 +393,15 @@ def read_by_index(self, stream_key: str, start_idx: int, stop_idx: int) -> AxisA
         template = info.template
 
         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:
-            chunk_t0 = template.axes["time"].gain * start_idx
+            # Rate-only: ``start_idx`` is 0-based from the stream's own start, so
+            # its absolute time is ``info.t0`` plus the within-stream offset.
+            # Omitting ``info.t0`` would mis-time a stream whose ``starting_time``
+            # is non-zero and disagree with the clock-driven producer's
+            # ``file_t = t0 + idx/fs`` bookkeeping.
+            chunk_t0 = float(info.t0) + template.axes["time"].gain * start_idx
 
         return replace(
             template,

tests/test_iterator.py

@@ -404,6 +404,148 @@ def test_multi_stream_interleaving(test_nwb_path):
     assert abs(first_binned - first_raw) <= 1
 
 
+def test_ragged_stream_lengths_do_not_crash(test_nwb_path):
+    """Defensive guard: a stream whose chunk_offsets table is shorter than the
+    file-wide ``n_chunks`` must not raise an IndexError; that stream simply
+    stops contributing once its table runs out. Offset tables are normally
+    built one-entry-per-chunk (see ``test_late_starting_stream_stays_aligned``),
+    so this exercises the bounds guard against a short table directly.
+    """
+    it = NWBAxisArrayIterator(
+        NWBIteratorSettings(
+            filepath=test_nwb_path,
+            chunk_dur=1.0,
+            reference_clock=ReferenceClockType.UNKNOWN,
+            stream_keys=["BinnedSpikes", "RawAnalog"],
+        )
+    )
+    # Simulate a ragged stream: truncate RawAnalog's offset table so it has
+    # fewer chunks than the file-wide n_chunks (3). Without the guard in
+    # _build_chunk_messages_static this raises IndexError at the last chunk.
+    short = it._state.streams["RawAnalog"]
+    short["chunk_offsets"] = short["chunk_offsets"][:-1]
+
+    keys = [m.key for m in it]  # full iteration must not raise
+
+    # The full-length stream still produces messages for every chunk...
+    assert keys.count("BinnedSpikes") == it._state.n_chunks
+    # ...while the truncated stream contributes one fewer.
+    assert keys.count("RawAnalog") == it._state.n_chunks - 1
+
+
+def test_late_starting_stream_stays_aligned(tmp_path):
+    """A stream that starts partway into the recording must be emitted in the
+    chunks that match its real wall-clock time — not shifted to chunk 0 — and
+    its message time offsets must reflect its true start time. Regression for
+    streams with different ``starting_time`` (e.g. CereLink Hub2 vs NPLAY).
+    """
+    import datetime
+
+    from pynwb import NWBHDF5IO, NWBFile, TimeSeries
+
+    path = tmp_path / "late_start.nwb"
+    nwb = NWBFile(
+        session_description="m",
+        identifier="m",
+        session_start_time=datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc),
+    )
+    rate = 100.0
+    # Stream A: t=0..5s. Stream B: t=2..5s (starts 2 chunks / 2 s later).
+    nwb.add_acquisition(
+        TimeSeries(
+            name="A",
+            data=np.arange(int(5 * rate), dtype=np.float32)[:, None],
+            unit="V",
+            rate=rate,
+            starting_time=0.0,
+        )
+    )
+    nwb.add_acquisition(
+        TimeSeries(
+            name="B",
+            data=(1000 + np.arange(int(3 * rate), dtype=np.float32))[:, None],
+            unit="V",
+            rate=rate,
+            starting_time=2.0,
+        )
+    )
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwb)
+
+    it = NWBAxisArrayIterator(
+        NWBIteratorSettings(
+            filepath=path,
+            chunk_dur=1.0,
+            reference_clock=ReferenceClockType.UNKNOWN,
+        )
+    )
+    # Offset tables are built one entry per global chunk for every stream.
+    n_chunks = it._state.n_chunks
+    assert len(it._state.streams["A"]["chunk_offsets"]) == n_chunks
+    assert len(it._state.streams["B"]["chunk_offsets"]) == n_chunks
+
+    # First non-empty message for each stream: when does each first appear and
+    # at what time offset?
+    first_offset = {}
+    for m in it:
+        if m.data.size and m.key not in first_offset:
+            first_offset[m.key] = (m.axes["time"].offset, float(m.data.flat[0]))
+
+    # A begins at t=0; B begins at t=2.0 with its own first sample (1000) —
+    # NOT shifted to t=0.
+    assert first_offset["A"][0] == pytest.approx(0.0)
+    assert first_offset["B"][0] == pytest.approx(2.0)
+    assert first_offset["B"][1] == 1000.0
+
+
+def test_chunk_offsets_match_searchsorted_for_noninteger_period(tmp_path):
+    """When chunk_dur is not an integer multiple of the sample period, chunk
+    offsets must be the first sample at/after each boundary (searchsorted
+    side='left' / ceil), not the nearest sample (round) — otherwise a
+    pre-boundary sample leaks into the next chunk and disagrees with the
+    event/timestamped paths.
+    """
+    import datetime
+
+    from pynwb import NWBHDF5IO, NWBFile, TimeSeries
+
+    path = tmp_path / "noninteger.nwb"
+    nwb = NWBFile(
+        session_description="m",
+        identifier="m",
+        session_start_time=datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc),
+    )
+    rate = 256.0  # 256 samples/s -> 25.6 samples per 0.1 s chunk (non-integer)
+    nwb.add_acquisition(
+        TimeSeries(
+            name="S",
+            data=np.arange(int(3 * rate), dtype=np.float32)[:, None],
+            unit="V",
+            rate=rate,
+            starting_time=0.0,
+        )
+    )
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwb)
+
+    chunk_dur = 0.1
+    it = NWBAxisArrayIterator(
+        NWBIteratorSettings(
+            filepath=path,
+            chunk_dur=chunk_dur,
+            reference_clock=ReferenceClockType.UNKNOWN,
+        )
+    )
+    offsets = np.asarray(it._state.streams["S"]["chunk_offsets"])
+    n_chunks = it._state.n_chunks
+    n_samples = int(3 * rate)
+
+    boundaries = np.arange(n_chunks) * chunk_dur  # start_time=0, ts_off=0, t0=0
+    sample_times = np.arange(n_samples) / rate
+    expected = np.clip(np.searchsorted(sample_times, boundaries, side="left"), 0, n_samples)
+    np.testing.assert_array_equal(offsets, expected)
+
+
 # --- Channel axis preserved ---
 
 

tests/test_slicer.py

@@ -277,6 +277,47 @@ def test_read_by_index_has_linear_axis(slicer):
     assert hasattr(msg.axes["time"], "gain")  # LinearAxis
 
 
+def test_read_by_index_offset_includes_starting_time(tmp_path):
+    """For a rate-only stream with non-zero ``starting_time``, the emitted time
+    offset must include that start time (sample indices are 0-based from the
+    stream start). Regression: previously the offset was ``ts_off + gain*idx``,
+    omitting ``t0``, which mis-timed late-starting streams.
+    """
+    import datetime
+
+    from pynwb import NWBHDF5IO, NWBFile, TimeSeries
+
+    path = tmp_path / "late_start.nwb"
+    nwb = NWBFile(
+        session_description="m",
+        identifier="m",
+        session_start_time=datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc),
+    )
+    rate = 100.0
+    nwb.add_acquisition(
+        TimeSeries(
+            name="Late",
+            data=np.arange(int(3 * rate), dtype=np.float32)[:, None],
+            unit="V",
+            rate=rate,
+            starting_time=2.0,
+        )
+    )
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwb)
+
+    s = NWBSlicer(filepath=path, reference_clock=ReferenceClockType.UNKNOWN)
+    try:
+        gain = 1.0 / rate
+        # Sample 0 is at the stream's start time (2.0 s); sample 50 is 0.5 s later.
+        msg0 = s.read_by_index("Late", 0, 10)
+        msg50 = s.read_by_index("Late", 50, 60)
+        assert msg0.axes["time"].offset == pytest.approx(2.0)
+        assert msg50.axes["time"].offset == pytest.approx(2.0 + 50 * gain)
+    finally:
+        s.close()
+
+
 # --- Timestamped continuous slicing ---