diff --git a/src/ezmsg/tools/plot/shmem_sweep.py b/src/ezmsg/tools/plot/shmem_sweep.py index 943ecb7..8dd4ec5 100644 --- a/src/ezmsg/tools/plot/shmem_sweep.py +++ b/src/ezmsg/tools/plot/shmem_sweep.py @@ -208,7 +208,7 @@ def _close_figure(self) -> None: logger.exception("closing the sweep figure raised; continuing teardown") def _on_tick(self) -> None: - samples, _overflow = self._mirror.auto_view() + samples, times, _overflow = self._mirror.timed_view() shape = describe_mirror(self._mirror, label_fields=self._label_fields) if shape is None or shape.srate <= 0: @@ -232,7 +232,9 @@ def _on_tick(self) -> None: self._apply_shape(shape) if samples is not None and samples.size: - self._sweep.push_data(np.ascontiguousarray(flatten_for_plot(samples, shape), dtype=np.float32)) + self._sweep.push_data( + np.ascontiguousarray(flatten_for_plot(samples, shape), dtype=np.float32), timestamps=times + ) self.on_frame(shape) diff --git a/src/ezmsg/tools/shmem/shmem.py b/src/ezmsg/tools/shmem/shmem.py index 3c30191..a47a03e 100644 --- a/src/ezmsg/tools/shmem/shmem.py +++ b/src/ezmsg/tools/shmem/shmem.py @@ -20,6 +20,10 @@ The other half must monitor the metadata shared memory to see if it changes, and if it does then it must recreate the data shared memory buffer reader at the new location. +The data segment holds two arrays: first one float64 timestamp per frame (the buffered axis's value for that frame), +then the frames themselves. Both are written before `write_index` advances, so a reader that takes frames from the +ring takes their times from the same indices and the two always agree. See `ring_views`. + Finally, there is a third piece of shared memory carrying everything about the AxisArray that does not fit in the fixed-size metadata header: the non-buffered coordinate axes (e.g. a `ch` axis holding per-channel bank/elec/label), axis units, and the message `attrs`. It lives at shorten_shmem_name(f"{shmem_name}/meta{meta_generation}") and is @@ -83,7 +87,10 @@ def shorten_shmem_name(long_name: str) -> str: # from a header written by a build old enough to predate this check. SHMEM_META_MAGIC = 0x455A4D53 -# Bumped on any change to ShmemArrMeta._fields_ or to the .aux_meta wire format. +# Bumped on any change to ShmemArrMeta._fields_, to the data segment's layout +# (see ring_views), or to the .aux_meta wire format. +# +# 2: the data segment leads with one float64 timestamp per frame. # # The two halves of a shmem link must be the same version -- there is no # compatibility shim, by choice: the layouts are an internal detail between two @@ -91,7 +98,7 @@ def shorten_shmem_name(long_name: str) -> str: # cost more than it is worth. What we do owe is a loud failure rather than a # quiet one, so the reader validates the magic and version up front and raises # instead of misreading a header it does not understand. -SHMEM_META_STRUCT_VERSION = 1 +SHMEM_META_STRUCT_VERSION = 2 class ShmemVersionError(RuntimeError): @@ -153,6 +160,23 @@ def key(self, value: str) -> None: ctypes.memmove(self._key_bytes, key_bytes[: self._key_len], self._key_len) +def ring_views(buf: memoryview, meta: ShmemArrMeta) -> typing.Tuple[npt.NDArray, npt.NDArray]: + """The data segment's two arrays: per-frame timestamps (float64), then the frames. + + Times lead so both arrays are aligned for any frame dtype. Writer and reader + both call this, so the layout is stated once. + """ + n_frames = meta.shape[0] + times = np.ndarray((n_frames,), dtype=np.float64, buffer=buf[: n_frames * 8]) + frames = np.ndarray(meta.shape[: meta.ndim], dtype=np.dtype(meta.dtype), buffer=buf[n_frames * 8 :]) + return times, frames + + +def ring_nbytes(n_frames: int, frame_shape: typing.Tuple[int, ...], itemsize: int) -> int: + """Size of the data segment described by :func:`ring_views`.""" + return int(n_frames * 8 + n_frames * np.prod(frame_shape) * itemsize) + + class ShMemCircBuffSettings(ez.Settings): shmem_name: typing.Optional[str] buf_dur: float @@ -180,6 +204,8 @@ class ShMemCircBuffState(ez.State): meta_struct: typing.Optional[ShmemArrMeta] = None buffer_shmem: typing.Optional[SharedMemory] = None buffer_arr: typing.Optional[npt.NDArray] = None + # One timestamp per frame of buffer_arr, at the same ring indices. + times_arr: typing.Optional[npt.NDArray] = None meta_hash: int = -1 # The dimension currently buffered along; see ShMemCircBuffSettings.axis. buff_axis: typing.Optional[str] = None @@ -331,6 +357,7 @@ def _cleanup_buffer(self): # Destroy the buffer self.STATE.buffer_arr = None + self.STATE.times_arr = None # Destroy the shared memory object if self.STATE.buffer_shmem is not None: @@ -547,7 +574,7 @@ def _reset_buffer(self, msg: AxisArray) -> None: """ self.STATE.meta_struct.buffer_generation += 1 msg_dtype, msg_srate, n_frames, frame_shape = self._get_msg_meta(msg) - buff_size = int(n_frames * np.prod(frame_shape) * msg.data.itemsize) + buff_size = ring_nbytes(n_frames, frame_shape, msg.data.itemsize) buff_shm_name = self.SETTINGS.shmem_name + "/buffer" + str(self.STATE.meta_struct.buffer_generation) short_name = shorten_shmem_name(buff_shm_name) self.STATE.buffer_shmem = _persist_create_shmem( @@ -556,11 +583,7 @@ def _reset_buffer(self, msg: AxisArray) -> None: purpose=f"data ring gen {self.STATE.meta_struct.buffer_generation} " f"({'x'.join(str(d) for d in (n_frames,) + frame_shape)})", ) - self.STATE.buffer_arr = np.ndarray( - self.STATE.meta_struct.shape[: self.STATE.meta_struct.ndim], - dtype=np.dtype(self.STATE.meta_struct.dtype.decode("utf8")), - buffer=self.STATE.buffer_shmem.buf[:], - ) + self.STATE.times_arr, self.STATE.buffer_arr = ring_views(self.STATE.buffer_shmem.buf, self.STATE.meta_struct) self.STATE.meta_struct.write_index = 0 self.STATE.meta_struct.wrap_counter = 0 self.STATE.meta_struct.bvalid = True @@ -612,14 +635,23 @@ async def on_message(self, msg: AxisArray): self._update_aux_if_needed(msg) n_samples = data.shape[0] - write_stop = self.STATE.meta_struct.write_index + n_samples + axis = msg.axes[buff_axis] + times = axis.data if hasattr(axis, "data") else axis.offset + np.arange(n_samples) * axis.gain + write_index = self.STATE.meta_struct.write_index + write_stop = write_index + n_samples + # Frames and times both land before write_index moves, so a reader + # never sees one without the other. if write_stop > self.STATE.buffer_arr.shape[0]: overflow = write_stop - self.STATE.buffer_arr.shape[0] - self.STATE.buffer_arr[self.STATE.meta_struct.write_index :] = data[: n_samples - overflow] - self.STATE.buffer_arr[:overflow] = data[n_samples - overflow :] + split = n_samples - overflow + self.STATE.buffer_arr[write_index:] = data[:split] + self.STATE.buffer_arr[:overflow] = data[split:] + self.STATE.times_arr[write_index:] = times[:split] + self.STATE.times_arr[:overflow] = times[split:] self.STATE.meta_struct.write_index = overflow self.STATE.meta_struct.wrap_counter += 1 else: - self.STATE.buffer_arr[self.STATE.meta_struct.write_index : write_stop] = data[:] + self.STATE.buffer_arr[write_index:write_stop] = data[:] + self.STATE.times_arr[write_index:write_stop] = times self.STATE.meta_struct.write_index = write_stop diff --git a/src/ezmsg/tools/shmem/shmem_mirror.py b/src/ezmsg/tools/shmem/shmem_mirror.py index 8e90056..9f7c14f 100644 --- a/src/ezmsg/tools/shmem/shmem_mirror.py +++ b/src/ezmsg/tools/shmem/shmem_mirror.py @@ -24,6 +24,7 @@ ShmemArrMeta, ShMemCircBuffState, ShmemVersionError, + ring_views, shorten_shmem_name, ) @@ -231,6 +232,7 @@ def _cleanup_buffer(self): if self._mirror_state.buffer_arr is not None: del self._mirror_state.buffer_arr self._mirror_state.buffer_arr = None + self._mirror_state.times_arr = None if self._mirror_state.buffer_shmem is not None: # Note: Uncommenting the following does not eliminate the resource_tracker warnings. @@ -308,10 +310,8 @@ def _reset_buffer(self) -> bool: buff_name = self._shmem_name + "/buffer" + str(self._mirror_state.meta_struct.buffer_generation) short_name = shorten_shmem_name(buff_name) self._mirror_state.buffer_shmem = SharedMemory(short_name, create=False) - self._mirror_state.buffer_arr = np.ndarray( - self._mirror_state.meta_struct.shape[: self._mirror_state.meta_struct.ndim], - dtype=np.dtype(self._mirror_state.meta_struct.dtype), - buffer=self._mirror_state.buffer_shmem.buf[:], + self._mirror_state.times_arr, self._mirror_state.buffer_arr = ring_views( + self._mirror_state.buffer_shmem.buf, self._mirror_state.meta_struct ) self._last_meta = self.meta # Copy if self._change_callback is not None: @@ -319,10 +319,12 @@ def _reset_buffer(self) -> bool: return True except FileNotFoundError: self._mirror_state.buffer_arr = None + self._mirror_state.times_arr = None self._mirror_state.buffer_shmem = None except TypeError: # buffer is too small for requested array self._mirror_state.buffer_arr = None + self._mirror_state.times_arr = None self._mirror_state.buffer_shmem = None print("DEBUG!") return False @@ -351,6 +353,17 @@ def connect(self, name: str) -> None: self._last_connect_try = time.time() def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, bool]: + """The frames written since the last read, and whether the ring overran them.""" + frames, _times, overflow = self.timed_view(n) + return frames, overflow + + def timed_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, npt.NDArray, bool]: + """Like :meth:`auto_view`, plus each returned frame's timestamp. + + The times are the buffered axis's values for those frames, as the + writer received them: the time axis's clock for a time-buffered stream. + """ + no_times = np.empty(0, dtype=np.float64) if self._mirror_state.meta_struct is None: self.connect(self._shmem_name) @@ -361,7 +374,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, if self._mirror_state.meta_struct is None or not self._mirror_state.meta_struct.bvalid: # Still not connected # or we are connected but the buffer data is invalid. - return np.array([[]]), False + return np.array([[]]), no_times, False b_connected = True # Determine if we need to reset the buffer @@ -374,7 +387,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, if not b_connected: # We STILL aren't connected. - return np.array([[]]), False + return np.array([[]]), no_times, False # -- From here, we should know we have a good connection to a valid buffer -- # @@ -405,7 +418,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, if n_available <= 1 or (n is not None and n_available < n): # Not enough samples available. # Return a null-slice of the buffer. This provides correct dimensions. - return self._mirror_state.buffer_arr[:0], b_overflow + return self._mirror_state.buffer_arr[:0], no_times, b_overflow # We have enough samples. if n is None: @@ -415,6 +428,7 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, # Return a contiguous chunk t_slice = np.s_[max(0, self._read_index) : self._read_index + n] result = self._mirror_state.buffer_arr[t_slice, :] + times = self._mirror_state.times_arr[t_slice].copy() else: # Split read into two chunks n_after_wrap = n - (self._mirror_state.meta_struct.shape[0] - self._read_index) @@ -425,8 +439,11 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, ), axis=0, ) + times = np.concatenate( + (self._mirror_state.times_arr[self._read_index :], self._mirror_state.times_arr[:n_after_wrap]) + ) self._read_index = (self._read_index + n) % self._mirror_state.meta_struct.shape[0] self._last_meta.wrap_counter = self._mirror_state.meta_struct.wrap_counter - return result, b_overflow + return result, times, b_overflow diff --git a/tests/test_shmem_mirror.py b/tests/test_shmem_mirror.py index e0d591d..e7a6ad8 100644 --- a/tests/test_shmem_mirror.py +++ b/tests/test_shmem_mirror.py @@ -182,3 +182,40 @@ def test_shmem_mirror_switch_buffer(): print(f"\tFrom {dt0} to {switch_dt} after {np.sum(chunk_lens[:switch_idx])} samples") assert dt0 == np.float64 assert switch_dt == np.float16 + + +def test_each_frame_carries_its_timestamp_across_a_wrap(): + """The reader gets each frame's time from the same ring slots as the frame. + + Driven in-process (no ez.run) so it is deterministic: 100 Hz into a 1 s + ring, in chunks that straddle the wrap, read back after each push. + """ + import asyncio + + from ezmsg.tools.shmem.shmem import ShMemCircBuffSettings + + name = f"times{os.getpid()}" + sink = ShMemCircBuff(ShMemCircBuffSettings(shmem_name=name, buf_dur=1.0, axis="time")) + sink.STATE = ShMemCircBuff.STATE() + sink._reset_meta() + mirror = EZShmMirror(name) + try: + got_frames, got_times = [], [] + t0, n = 1000.0, 0 + for chunk in (40, 40, 40, 40): # 160 frames through a 100-frame ring + data = np.arange(n, n + chunk, dtype=np.float32)[:, None] * np.ones((1, 2), np.float32) + msg = AxisArray(data, dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=100.0, offset=t0 + n / 100)}) + asyncio.run(sink.on_message(msg)) + n += chunk + frames, times, overflow = mirror.timed_view() + assert not overflow + got_frames.append(frames[:, 0].copy()) + got_times.append(times) + frames, times = np.concatenate(got_frames), np.concatenate(got_times) + # Whatever frames came back, each one's time is the time it was sent with. + assert len(frames) > 100 + np.testing.assert_allclose(times, t0 + frames / 100) + finally: + mirror.disconnect() + sink._cleanup_buffer() + sink._cleanup_meta()