ezmsg-sigproc 2.2.0__py3-none-any.whl → 2.4.0__py3-none-any.whl

ezmsg/sigproc/sampler.py

@@ -1,18 +1,19 @@
 import asyncio
 from collections import deque
+import copy
 import traceback
 import typing
 
 import numpy as np
-import numpy.typing as npt
 import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import (
     AxisArray,
-    slice_along_axis,
 )
 from ezmsg.util.messages.util import replace
 
 from .util.profile import profile_subpub
+from .util.axisarray_buffer import HybridAxisArrayBuffer
+from .util.buffer import UpdateStrategy
 from .util.message import SampleMessage, SampleTriggerMessage
 from .base import (
     BaseStatefulTransformer,
@@ -43,6 +44,7 @@ class SamplerSettings(ez.Settings):
         None (default) will choose the first axis in the first input.
         Note: (for now) the axis must exist in the msg .axes and be of type AxisArray.LinearAxis
     """
+
     period: tuple[float, float] | None = None
     """Optional default period (in seconds) if unspecified in SampleTriggerMessage."""
 
@@ -51,20 +53,25 @@ class SamplerSettings(ez.Settings):
 
     estimate_alignment: bool = True
     """
-    If true, use message timestamp fields and reported sampling rate to estimate sample-accurate alignment for samples.
+    If true, use message timestamp fields and reported sampling rate to estimate
+     sample-accurate alignment for samples.
     If false, sampling will be limited to incoming message rate -- "Block timing"
     NOTE: For faster-than-realtime playback --  Incoming timestamps must reflect
     "realtime" operation for estimate_alignment to operate correctly.
     """
 
+    buffer_update_strategy: UpdateStrategy = "immediate"
+    """
+    The buffer update strategy. See :obj:`ezmsg.sigproc.util.buffer.UpdateStrategy`.
+    If you expect to push data much more frequently than triggers, then "on_demand"
+    might be more efficient. For most other scenarios, "immediate" is best.
+    """
+
 
 @processor_state
 class SamplerState:
-    fs: float = 0.0
-    offset: float | None = None
-    buffer: npt.NDArray | None = None
+    buffer: HybridAxisArrayBuffer | None = None
     triggers: deque[SampleTriggerMessage] | None = None
-    n_samples: int = 0
 
 
 class SamplerTransformer(
@@ -73,6 +80,16 @@ class SamplerTransformer(
     def __call__(
         self, message: AxisArray | SampleTriggerMessage
     ) -> list[SampleMessage]:
+        # TODO: Currently we have a single entry point that accepts both
+        #  data and trigger messages and we choose a code path based on
+        #  the message type. However, in the future we will likely replace
+        #  SampleTriggerMessage with an agumented form of AxisArray,
+        #  leveraging its attrs field, which makes this a bit harder.
+        #  We should probably force callers of this object to explicitly
+        #  call `push_trigger` for trigger messages. This will also
+        #  simplify typing somewhat because `push_trigger` should not
+        #  return anything yet we currently have it returning an empty
+        #  list just to be compatible with __call__.
         if isinstance(message, AxisArray):
             return super().__call__(message)
         else:
@@ -82,102 +99,75 @@ class SamplerTransformer(
         # Compute hash based on message properties that require state reset
         axis = self.settings.axis or message.dims[0]
         axis_idx = message.get_axis_idx(axis)
-        fs = 1.0 / message.get_axis(axis).gain
         sample_shape = (
             message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
         )
-        return hash((fs, sample_shape, axis_idx, message.key))
+        return hash((sample_shape, message.key))
 
     def _reset_state(self, message: AxisArray) -> None:
-        axis = self.settings.axis or message.dims[0]
-        axis_idx = message.get_axis_idx(axis)
-        axis_info = message.get_axis(axis)
-        self._state.fs = 1.0 / axis_info.gain
-        self._state.buffer = None
+        self._state.buffer = HybridAxisArrayBuffer(
+            duration=self.settings.buffer_dur,
+            axis=self.settings.axis or message.dims[0],
+            update_strategy=self.settings.buffer_update_strategy,
+            overflow_strategy="warn-overwrite",  # True circular buffer
+        )
         if self._state.triggers is None:
             self._state.triggers = deque()
         self._state.triggers.clear()
-        self._state.n_samples = message.data.shape[axis_idx]
 
     def _process(self, message: AxisArray) -> list[SampleMessage]:
-        axis = self.settings.axis or message.dims[0]
-        axis_idx = message.get_axis_idx(axis)
-        axis_info = message.get_axis(axis)
-        self._state.offset = axis_info.offset
-
-        # Update buffer
-        self._state.buffer = (
-            message.data
-            if self._state.buffer is None
-            else np.concatenate((self._state.buffer, message.data), axis=axis_idx)
-        )
+        self._state.buffer.write(message)
 
-        # Calculate timestamps associated with buffer.
-        buffer_offset = np.arange(self._state.buffer.shape[axis_idx], dtype=float)
-        buffer_offset -= buffer_offset[-message.data.shape[axis_idx]]
-        buffer_offset *= axis_info.gain
-        buffer_offset += axis_info.offset
+        # How much data in the buffer?
+        buff_t_range = (
+            self._state.buffer.axis_first_value,
+            self._state.buffer.axis_final_value,
+        )
 
-        # ... for each trigger, collect the message (if possible) and append to msg_out
-        msg_out: list[SampleMessage] = []
-        for trig in list(self._state.triggers):
+        # Process in reverse order so that we can remove triggers safely as we iterate.
+        msgs_out: list[SampleMessage] = []
+        for trig_ix in range(len(self._state.triggers) - 1, -1, -1):
+            trig = self._state.triggers[trig_ix]
             if trig.period is None:
-                # This trigger was malformed; drop it.
-                self._state.triggers.remove(trig)
+                ez.logger.warning("Sampling failed: trigger period not specified")
+                del self._state.triggers[trig_ix]
+                continue
+
+            trig_range = trig.timestamp + np.array(trig.period)
 
             # If the previous iteration had insufficient data for the trigger timestamp + period,
             #  and buffer-management removed data required for the trigger, then we will never be able
             #  to accommodate this trigger. Discard it. An increase in buffer_dur is recommended.
-            if (trig.timestamp + trig.period[0]) < buffer_offset[0]:
+            if trig_range[0] < buff_t_range[0]:
                 ez.logger.warning(
-                    f"Sampling failed: Buffer span {buffer_offset[0]} is beyond the "
-                    f"requested sample period start: {trig.timestamp + trig.period[0]}"
+                    f"Sampling failed: Buffer span {buff_t_range} begins beyond the "
+                    f"requested sample period start: {trig_range[0]}"
                 )
-                self._state.triggers.remove(trig)
+                del self._state.triggers[trig_ix]
+                continue
 
-            t_start = trig.timestamp + trig.period[0]
-            if t_start >= buffer_offset[0]:
-                start = np.searchsorted(buffer_offset, t_start)
-                stop = start + int(
-                    np.round(self._state.fs * (trig.period[1] - trig.period[0]))
-                )
-                if self._state.buffer.shape[axis_idx] > stop:
-                    # Trigger period fully enclosed in buffer.
-                    msg_out.append(
-                        SampleMessage(
-                            trigger=trig,
-                            sample=replace(
-                                message,
-                                data=slice_along_axis(
-                                    self._state.buffer, slice(start, stop), axis_idx
-                                ),
-                                axes={
-                                    **message.axes,
-                                    axis: replace(
-                                        axis_info, offset=buffer_offset[start]
-                                    ),
-                                },
-                            ),
-                        )
-                    )
-                    self._state.triggers.remove(trig)
-
-        # Trim buffer
-        buf_len = int(self.settings.buffer_dur * self._state.fs)
-        self._state.buffer = slice_along_axis(
-            self._state.buffer, np.s_[-buf_len:], axis_idx
-        )
+            if trig_range[1] > buff_t_range[1]:
+                # We don't *yet* have enough data to satisfy this trigger.
+                continue
+
+            # We know we have enough data in the buffer to satisfy this trigger.
+            buff_idx = self._state.buffer.axis_searchsorted(trig_range, side="right")
+            self._state.buffer.seek(buff_idx[0])  # FFWD to starting position.
+            buff_axarr = self._state.buffer.peek(buff_idx[1] - buff_idx[0])
+            self._state.buffer.seek(-buff_idx[0])  # Rewind it back.
+            # Note: buffer will trim itself as needed based on buffer_dur.
+
+            # Prepare output and drop trigger
+            msgs_out.append(SampleMessage(trigger=copy.copy(trig), sample=buff_axarr))
+            del self._state.triggers[trig_ix]
 
-        return msg_out
+        msgs_out.reverse()  # in-place
+        return msgs_out
 
     def push_trigger(self, message: SampleTriggerMessage) -> list[SampleMessage]:
         # Input is a trigger message that we will use to sample the buffer.
 
-        if (
-            self._state.buffer is None
-            or not self._state.fs
-            or self._state.offset is None
-        ):
+        if self._state.buffer is None:
             # We've yet to see any data; drop the trigger.
             return []
 
@@ -194,11 +184,9 @@ class SamplerTransformer(
             return []
 
         # Check that period is compatible with buffer duration.
-        max_buf_len = int(np.round(self.settings.buffer_dur * self._state.fs))
-        req_buf_len = int(np.round((_period[1] - _period[0]) * self._state.fs))
-        if req_buf_len >= max_buf_len:
+        if (_period[1] - _period[0]) > self.settings.buffer_dur:
             ez.logger.warning(
-                f"Sampling failed: {_period=} >= {self.settings.buffer_dur=}"
+                f"Sampling failed: trigger period {_period=} >= buffer capacity {self.settings.buffer_dur=}"
             )
             return []
 
@@ -206,7 +194,7 @@ class SamplerTransformer(
         if not self.settings.estimate_alignment:
             # Override the trigger timestamp with the next sample's likely timestamp.
             trigger_ts = (
-                self._state.offset + (self.state.n_samples + 1) / self._state.fs
+                self._state.buffer.axis_final_value + self._state.buffer.axis_gain
             )
 
         new_trig_msg = replace(

ezmsg/sigproc/util/axisarray_buffer.py

@@ -0,0 +1,379 @@
+import math
+import typing
+
+from array_api_compat import get_namespace
+import numpy as np
+from ezmsg.util.messages.axisarray import AxisArray, LinearAxis, CoordinateAxis
+from ezmsg.util.messages.util import replace
+
+from .buffer import HybridBuffer
+
+
+Array = typing.TypeVar("Array")
+
+
+class HybridAxisBuffer:
+    """
+    A buffer that intelligently handles ezmsg.util.messages.AxisArray _axes_ objects.
+     LinearAxis is maintained internally by tracking its offset, gain, and the number
+     of samples that have passed through.
+     CoordinateAxis has its data values maintained in a `HybridBuffer`.
+
+    Args:
+        duration: The desired duration of the buffer in seconds. This is non-limiting
+         when managing a LinearAxis.
+        **kwargs: Additional keyword arguments to pass to the underlying HybridBuffer
+            (e.g., `update_strategy`, `threshold`, `overflow_strategy`, `max_size`).
+    """
+
+    _coords_buffer: HybridBuffer | None
+    _coords_template: CoordinateAxis | None
+    _coords_gain_estimate: float | None = None
+    _linear_axis: LinearAxis | None
+    _linear_n_available: int
+
+    def __init__(self, duration: float, **kwargs):
+        self.duration = duration
+        self.buffer_kwargs = kwargs
+        # Delay initialization until the first message arrives
+        self._coords_buffer = None
+        self._coords_template = None
+        self._linear_axis = None
+        self._linear_n_available = 0
+
+    @property
+    def capacity(self) -> int:
+        """The maximum number of samples that can be stored in the buffer."""
+        if self._coords_buffer is not None:
+            return self._coords_buffer.capacity
+        elif self._linear_axis is not None:
+            return int(math.ceil(self.duration / self._linear_axis.gain))
+        else:
+            return 0
+
+    def available(self) -> int:
+        if self._coords_buffer is None:
+            return self._linear_n_available
+        return self._coords_buffer.available()
+
+    def is_empty(self) -> bool:
+        return self.available() == 0
+
+    def is_full(self) -> bool:
+        if self._coords_buffer is not None:
+            return self._coords_buffer.is_full()
+        return 0 < self.capacity == self.available()
+
+    def _initialize(self, first_axis: LinearAxis | CoordinateAxis) -> None:
+        if hasattr(first_axis, "data"):
+            # Initialize a CoordinateAxis buffer
+            if len(first_axis.data) > 1:
+                _axis_gain = (first_axis.data[-1] - first_axis.data[0]) / (
+                    len(first_axis.data) - 1
+                )
+            else:
+                _axis_gain = 1.0
+            self._coords_gain_estimate = _axis_gain
+            capacity = int(self.duration / _axis_gain)
+            self._coords_buffer = HybridBuffer(
+                get_namespace(first_axis.data),
+                capacity,
+                other_shape=(),
+                dtype=first_axis.data.dtype,
+                **self.buffer_kwargs,
+            )
+            self._coords_template = replace(first_axis, data=first_axis.data[:0].copy())
+        else:
+            # Initialize a LinearAxis buffer
+            self._linear_axis = replace(first_axis, offset=first_axis.offset)
+            self._linear_n_available = 0
+
+    def write(self, axis: LinearAxis | CoordinateAxis, n_samples: int) -> None:
+        if self._linear_axis is None and self._coords_buffer is None:
+            self._initialize(axis)
+
+        if self._coords_buffer is not None:
+            if axis.__class__ is not self._coords_template.__class__:
+                raise TypeError(
+                    f"Buffer initialized with {self._coords_template.__class__.__name__}, "
+                    f"but received {axis.__class__.__name__}."
+                )
+            self._coords_buffer.write(axis.data)
+        else:
+            if axis.__class__ is not self._linear_axis.__class__:
+                raise TypeError(
+                    f"Buffer initialized with {self._linear_axis.__class__.__name__}, "
+                    f"but received {axis.__class__.__name__}."
+                )
+            if axis.gain != self._linear_axis.gain:
+                raise ValueError(
+                    f"Buffer initialized with gain={self._linear_axis.gain}, "
+                    f"but received gain={axis.gain}."
+                )
+            if self._linear_n_available + n_samples > self.capacity:
+                # Simulate overflow by advancing the offset and decreasing
+                # the number of available samples.
+                n_to_discard = self._linear_n_available + n_samples - self.capacity
+                self.seek(n_to_discard)
+            # Update the offset corresponding to the oldest sample in the buffer
+            #  by anchoring on the new offset and accounting for the samples already available.
+            self._linear_axis.offset = (
+                axis.offset - self._linear_n_available * axis.gain
+            )
+            self._linear_n_available += n_samples
+
+    def peek(self, n_samples: int | None = None) -> LinearAxis | CoordinateAxis:
+        if self._coords_buffer is not None:
+            return replace(
+                self._coords_template, data=self._coords_buffer.peek(n_samples)
+            )
+        else:
+            # Return a shallow copy.
+            return replace(self._linear_axis, offset=self._linear_axis.offset)
+
+    def seek(self, n_samples: int) -> int:
+        if self._coords_buffer is not None:
+            return self._coords_buffer.seek(n_samples)
+        else:
+            n_to_seek = min(n_samples, self._linear_n_available)
+            self._linear_n_available -= n_to_seek
+            self._linear_axis.offset += n_to_seek * self._linear_axis.gain
+            return n_to_seek
+
+    def prune(self, n_samples: int) -> int:
+        """Discards all but the last n_samples from the buffer."""
+        n_to_discard = self.available() - n_samples
+        if n_to_discard <= 0:
+            return 0
+        return self.seek(n_to_discard)
+
+    @property
+    def final_value(self) -> float | None:
+        """
+        The axis-value (timestamp, typically) of the last sample in the buffer.
+        This does not advance the read head.
+        """
+        if self._coords_buffer is not None:
+            return self._coords_buffer.peek_last()[0]
+        elif self._linear_axis is not None:
+            return self._linear_axis.value(self._linear_n_available - 1)
+        else:
+            return None
+
+    @property
+    def first_value(self) -> float | None:
+        """
+        The axis-value (timestamp, typically) of the first sample in the buffer.
+        This does not advance the read head.
+        """
+        if self.available() == 0:
+            return None
+        if self._coords_buffer is not None:
+            return self._coords_buffer.peek_at(0)[0]
+        elif self._linear_axis is not None:
+            return self._linear_axis.value(0)
+        else:
+            return None
+
+    @property
+    def gain(self) -> float | None:
+        if self._coords_buffer is not None:
+            return self._coords_gain_estimate
+        elif self._linear_axis is not None:
+            return self._linear_axis.gain
+        else:
+            return None
+
+    def searchsorted(
+        self, values: typing.Union[float, Array], side: str = "left"
+    ) -> typing.Union[int, Array]:
+        if self._coords_buffer is not None:
+            return self._coords_buffer.xp.searchsorted(
+                self._coords_buffer.peek(self.available()), values, side=side
+            )
+        else:
+            if self.available() == 0:
+                if isinstance(values, float):
+                    return 0
+                else:
+                    _xp = get_namespace(values)
+                    return _xp.zeros_like(values, dtype=int)
+
+            f_inds = (values - self._linear_axis.offset) / self._linear_axis.gain
+            res = np.ceil(f_inds)
+            if side == "right":
+                res[np.isclose(f_inds, res)] += 1
+            return res.astype(int)
+
+
+class HybridAxisArrayBuffer:
+    """A buffer that intelligently handles ezmsg.util.messages.AxisArray objects.
+
+    This buffer defers its own initialization until the first message arrives,
+    allowing it to automatically configure its size, shape, dtype, and array backend
+    (e.g., NumPy, CuPy) based on the message content and a desired buffer duration.
+
+    Args:
+        duration: The desired duration of the buffer in seconds.
+        axis: The name of the axis to buffer along.
+        **kwargs: Additional keyword arguments to pass to the underlying HybridBuffer
+            (e.g., `update_strategy`, `threshold`, `overflow_strategy`, `max_size`).
+    """
+
+    _data_buffer: HybridBuffer | None
+    _axis_buffer: HybridAxisBuffer
+    _template_msg: AxisArray | None
+
+    def __init__(self, duration: float, axis: str = "time", **kwargs):
+        self.duration = duration
+        self._axis = axis
+        self.buffer_kwargs = kwargs
+        self._axis_buffer = HybridAxisBuffer(duration=duration, **kwargs)
+        # Delay initialization until the first message arrives
+        self._data_buffer = None
+        self._template_msg = None
+
+    def available(self) -> int:
+        """The total number of unread samples currently available in the buffer."""
+        if self._data_buffer is None:
+            return 0
+        return self._data_buffer.available()
+
+    def is_empty(self) -> bool:
+        return self.available() == 0
+
+    def is_full(self) -> bool:
+        return 0 < self._data_buffer.capacity == self.available()
+
+    @property
+    def axis_first_value(self) -> float | None:
+        """The axis-value (timestamp, typically) of the first sample in the buffer."""
+        return self._axis_buffer.first_value
+
+    @property
+    def axis_final_value(self) -> float | None:
+        """The axis-value (timestamp, typically) of the last sample in the buffer."""
+        return self._axis_buffer.final_value
+
+    def _initialize(self, first_msg: AxisArray) -> None:
+        # Create a template message that has everything except the data are length 0
+        #  and the target axis is missing.
+        self._template_msg = replace(
+            first_msg,
+            data=first_msg.data[:0],
+            axes={k: v for k, v in first_msg.axes.items() if k != self._axis},
+        )
+
+        in_axis = first_msg.axes[self._axis]
+        self._axis_buffer._initialize(in_axis)
+
+        capacity = int(self.duration / self._axis_buffer.gain)
+        self._data_buffer = HybridBuffer(
+            get_namespace(first_msg.data),
+            capacity,
+            other_shape=first_msg.data.shape[1:],
+            dtype=first_msg.data.dtype,
+            **self.buffer_kwargs,
+        )
+
+    def write(self, msg: AxisArray) -> None:
+        """Adds an AxisArray message to the buffer, initializing on the first call."""
+        in_axis_idx = msg.get_axis_idx(self._axis)
+        if in_axis_idx > 0:
+            # This class assumes that the target axis is the first axis.
+            # If it is not, we move it to the front.
+            dims = list(msg.dims)
+            dims.insert(0, dims.pop(in_axis_idx))
+            _xp = get_namespace(msg.data)
+            msg = replace(msg, data=_xp.moveaxis(msg.data, in_axis_idx, 0), dims=dims)
+
+        if self._data_buffer is None:
+            self._initialize(msg)
+
+        self._data_buffer.write(msg.data)
+        self._axis_buffer.write(msg.axes[self._axis], msg.shape[0])
+
+    def peek(self, n_samples: int | None = None) -> AxisArray | None:
+        """Retrieves the oldest unread data as a new AxisArray without advancing the read head."""
+
+        if self._data_buffer is None:
+            return None
+
+        data_array = self._data_buffer.peek(n_samples)
+
+        if data_array is None:
+            return None
+
+        out_axis = self._axis_buffer.peek(n_samples)
+
+        return replace(
+            self._template_msg,
+            data=data_array,
+            axes={**self._template_msg.axes, self._axis: out_axis},
+        )
+
+    def peek_axis(
+        self, n_samples: int | None = None
+    ) -> LinearAxis | CoordinateAxis | None:
+        """Retrieves the axis data without advancing the read head."""
+        if self._data_buffer is None:
+            return None
+
+        out_axis = self._axis_buffer.peek(n_samples)
+
+        if out_axis is None:
+            return None
+
+        return out_axis
+
+    def seek(self, n_samples: int) -> int:
+        """Advances the read pointer by n_samples."""
+        if self._data_buffer is None:
+            return 0
+
+        skipped_data_count = self._data_buffer.seek(n_samples)
+        axis_skipped = self._axis_buffer.seek(skipped_data_count)
+        assert (
+            axis_skipped == skipped_data_count
+        ), f"Axis buffer skipped {axis_skipped} samples, but data buffer skipped {skipped_data_count}."
+
+        return skipped_data_count
+
+    def read(self, n_samples: int | None = None) -> AxisArray | None:
+        """Retrieves the oldest unread data as a new AxisArray and advances the read head."""
+        retrieved_axis_array = self.peek(n_samples)
+
+        if retrieved_axis_array is None or retrieved_axis_array.shape[0] == 0:
+            return None
+
+        self.seek(retrieved_axis_array.shape[0])
+
+        return retrieved_axis_array
+
+    def prune(self, n_samples: int) -> int:
+        """Discards all but the last n_samples from the buffer."""
+        if self._data_buffer is None:
+            return 0
+
+        n_to_discard = self.available() - n_samples
+        if n_to_discard <= 0:
+            return 0
+
+        return self.seek(n_to_discard)
+
+    @property
+    def axis_gain(self) -> float | None:
+        """
+        The gain of the target axis, which is the time step between samples.
+        This is typically the sampling rate (e.g., 1 / fs).
+        """
+        return self._axis_buffer.gain
+
+    def axis_searchsorted(
+        self, values: typing.Union[float, Array], side: str = "left"
+    ) -> typing.Union[int, Array]:
+        """
+        Find the indices into which the given values would be inserted
+        into the target axis data to maintain order.
+        """
+        return self._axis_buffer.searchsorted(values, side=side)