ezmsg-sigproc 2.1.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

ezmsg/sigproc/util/buffer.py

@@ -0,0 +1,470 @@
+import collections
+import math
+import typing
+import warnings
+
+Array = typing.TypeVar("Array")
+ArrayNamespace = typing.Any
+DType = typing.Any
+UpdateStrategy = typing.Literal["immediate", "threshold", "on_demand"]
+OverflowStrategy = typing.Literal["grow", "raise", "drop", "warn-overwrite"]
+
+
+class HybridBuffer:
+    """A stateful, FIFO buffer that combines a deque for fast appends with a
+    contiguous circular buffer for efficient, advancing reads.
+
+    This buffer is designed to be agnostic to the array library used (e.g., NumPy,
+    CuPy, PyTorch) via the Python Array API standard.
+
+    Args:
+        array_namespace: The array library (e.g., numpy, cupy) that conforms to the Array API.
+        capacity: The current maximum number of samples to store in the circular buffer.
+        other_shape: A tuple defining the shape of the non-sample dimensions.
+        dtype: The data type of the samples, belonging to the provided array_namespace.
+        update_strategy: The strategy for synchronizing the deque to the circular buffer (flushing).
+        threshold: The number of samples to accumulate in the deque before flushing.
+          Ignored if update_strategy is "immediate" or "on_demand".
+        overflow_strategy: The strategy for handling overflow when the buffer is full.
+            Options are "grow", "raise", "drop", or "warn-overwrite". If "grow" (default), the buffer will
+            increase its capacity to accommodate new samples up to max_size. If "raise", an error will be
+            raised when the buffer is full. If "drop", the overflowing samples will be ignored.
+            If "warn-overwrite", a warning will be logged then the overflowing samples will
+            overwrite previously-unread samples.
+        max_size: The maximum size of the buffer in bytes.
+            If the buffer exceeds this size, it will raise an error.
+        warn_once: If True, will only warn once on overflow when using "warn-overwrite" strategy.
+    """
+
+    def __init__(
+        self,
+        array_namespace: ArrayNamespace,
+        capacity: int,
+        other_shape: tuple[int, ...],
+        dtype: DType,
+        update_strategy: UpdateStrategy = "on_demand",
+        threshold: int = 0,
+        overflow_strategy: OverflowStrategy = "grow",
+        max_size: int = 1024**3,  # 1 GB default max size
+        warn_once: bool = True,
+    ):
+        self.xp = array_namespace
+        self._capacity = capacity
+        self._deque = collections.deque()
+        self._update_strategy = update_strategy
+        self._threshold = threshold
+        self._overflow_strategy = overflow_strategy
+        self._max_size = max_size
+        self._warn_once = warn_once
+
+        self._buffer = self.xp.empty((capacity, *other_shape), dtype=dtype)
+        self._head = 0  # Write pointer
+        self._tail = 0  # Read pointer
+        self._buff_unread = 0  # Number of unread samples in the circular buffer
+        self._buff_read = 0  # Tracks samples read and still in buffer
+        self._deque_len = 0  # Number of unread samples in the deque
+        self._last_overflow = (
+            0  # Tracks the last overflow count, overwritten or skipped
+        )
+        self._warned = False  # Tracks if we've warned already (for warn_once)
+
+    @property
+    def capacity(self) -> int:
+        """The maximum number of samples that can be stored in the buffer."""
+        return self._capacity
+
+    def available(self) -> int:
+        """The total number of unread samples available (in buffer and deque)."""
+        return self._buff_unread + self._deque_len
+
+    def is_empty(self) -> bool:
+        """Returns True if there are no unread samples in the buffer or deque."""
+        return self.available() == 0
+
+    def is_full(self) -> bool:
+        """Returns True if the buffer is full and cannot _flush_ more samples without overwriting."""
+        return self._buff_unread == self._capacity
+
+    def tell(self) -> int:
+        """Returns the number of samples that have been read and are still in the buffer."""
+        return self._buff_read
+
+    def write(self, block: Array):
+        """Appends a new block (an array of samples) to the internal deque."""
+        other_shape = self._buffer.shape[1:]
+        if other_shape == (1,) and block.ndim == 1:
+            block = block[:, self.xp.newaxis]
+
+        if block.shape[1:] != other_shape:
+            raise ValueError(
+                f"Block shape {block.shape[1:]} does not match buffer's other_shape {other_shape}"
+            )
+
+        # Most overflow strategies are handled during flush, but there are a couple
+        # scenarios that can be evaluated on write to give immediate feedback.
+        new_len = self._deque_len + block.shape[0]
+        if new_len > self._capacity and self._overflow_strategy == "raise":
+            raise OverflowError(
+                f"Buffer overflow: {new_len} samples awaiting in deque exceeds buffer capacity {self._capacity}."
+            )
+        elif new_len * block.dtype.itemsize > self._max_size:
+            raise OverflowError(
+                f"deque contents would exceed max_size ({self._max_size}) on subsequent flush."
+                "Are you reading samples frequently enough?"
+            )
+
+        self._deque.append(block)
+        self._deque_len += block.shape[0]
+
+        if self._update_strategy == "immediate" or (
+            self._update_strategy == "threshold"
+            and (0 < self._threshold <= self._deque_len)
+        ):
+            self.flush()
+
+    def _estimate_overflow(self, n_samples: int) -> int:
+        """
+        Estimates the number of samples that would overflow we requested n_samples
+        from the buffer.
+        """
+        if n_samples > self.available():
+            raise ValueError(
+                f"Requested {n_samples} samples, but only {self.available()} are available."
+            )
+        n_overflow = 0
+        if self._deque and (n_samples > self._buff_unread):
+            # We would cause a flush, but would that cause an overflow?
+            n_free = self._capacity - self._buff_unread
+            n_overflow = max(0, self._deque_len - n_free)
+        return n_overflow
+
+    def read(
+        self,
+        n_samples: int | None = None,
+    ) -> Array:
+        """
+        Retrieves the oldest unread samples from the buffer with padding
+        and advances the read head.
+
+        Args:
+            n_samples: The number of samples to retrieve. If None, returns all
+                unread samples.
+
+        Returns:
+            An array containing the requested samples. This may be a view or a copy.
+            Note: The result may have more samples than the buffer.capacity as it
+            may include samples from the deque in the output.
+        """
+        n_samples = n_samples if n_samples is not None else self.available()
+        data = None
+        offset = 0
+        n_overflow = self._estimate_overflow(n_samples)
+        if n_overflow > 0:
+            first_read = self._buff_unread
+            if (n_overflow - first_read) < self.capacity or (
+                self._overflow_strategy == "drop"
+            ):
+                # We can prevent the overflow (or at least *some* if using "drop"
+                # strategy) by reading the samples in the buffer first to make room.
+                data = self.xp.empty(
+                    (n_samples, *self._buffer.shape[1:]), dtype=self._buffer.dtype
+                )
+                self.peek(first_read, out=data[:first_read])
+                offset += first_read
+                self.seek(first_read)
+                n_samples -= first_read
+        if data is None:
+            data = self.peek(n_samples)
+            self.seek(data.shape[0])
+        else:
+            d2 = self.peek(n_samples, out=data[offset:])
+            self.seek(d2.shape[0])
+
+        return data
+
+    def peek(self, n_samples: int | None = None, out: Array | None = None) -> Array:
+        """
+        Retrieves the oldest unread samples from the buffer with padding without
+        advancing the read head.
+
+        Args:
+            n_samples: The number of samples to retrieve. If None, returns all
+                unread samples.
+            out: Optionally, a destination array to store the samples.
+                If provided, must have shape (n_samples, *other_shape) where
+                other_shape matches the shape of the samples in the buffer.
+                If `out` is provided then the data will always be copied into it,
+                even if they are contiguous in the buffer.
+
+        Returns:
+            An array containing the requested samples. This may be a view or a copy.
+            Note: The result may have more samples than the buffer.capacity as it
+            may include samples from the deque in the output.
+        """
+        if n_samples is None:
+            n_samples = self.available()
+        elif n_samples > self.available():
+            raise ValueError(
+                f"Requested to peek {n_samples} samples, but only {self.available()} are available."
+            )
+        if out is not None and out.shape[0] < n_samples:
+            raise ValueError(
+                f"Output array shape {out.shape} is smaller than requested {n_samples} samples."
+            )
+
+        if n_samples == 0:
+            return self._buffer[:0]
+
+        self._flush_if_needed(n_samples=n_samples)
+
+        if self._tail + n_samples > self._capacity:
+            # discontiguous read (wraps around)
+            part1_len = self._capacity - self._tail
+            part2_len = n_samples - part1_len
+            out = (
+                out
+                if out is not None
+                else self.xp.empty(
+                    (n_samples, *self._buffer.shape[1:]), dtype=self._buffer.dtype
+                )
+            )
+            out[:part1_len] = self._buffer[self._tail :]
+            out[part1_len:] = self._buffer[:part2_len]
+        else:
+            if out is not None:
+                out[:] = self._buffer[self._tail : self._tail + n_samples]
+            else:
+                # No output array provided, just return a view
+                out = self._buffer[self._tail : self._tail + n_samples]
+
+        return out
+
+    def peek_at(self, idx: int, allow_flush: bool = False) -> Array:
+        """
+        Retrieves a specific sample from the buffer without advancing the read head.
+
+        Args:
+            idx: The index of the sample to retrieve, relative to the read head.
+            allow_flush: If True, allows flushing the deque to the buffer if the
+                requested sample is not in the buffer. If False and the sample is
+                in the deque, the sample will be retrieved from the deque (slow!).
+
+        Returns:
+            An array containing the requested sample. This may be a view or a copy.
+        """
+        if idx < 0 or idx >= self.available():
+            raise IndexError(f"Index {idx} out of bounds for unread samples.")
+
+        if not allow_flush and idx >= self._buff_unread:
+            # The requested sample is in the deque.
+            idx -= self._buff_unread
+            deq_splits = self.xp.cumsum(
+                [0] + [_.shape[0] for _ in self._deque], dtype=int
+            )
+            arr_idx = self.xp.searchsorted(deq_splits, idx, side="right") - 1
+            idx -= deq_splits[arr_idx]
+            return self._deque[arr_idx][idx : idx + 1]
+
+        self._flush_if_needed(n_samples=idx + 1)
+
+        # The requested sample is within the unread samples in the buffer.
+        idx = (self._tail + idx) % self._capacity
+        return self._buffer[idx : idx + 1]
+
+    def peek_last(self) -> Array:
+        """
+        Retrieves the last sample in the buffer without advancing the read head.
+        """
+        if self._deque:
+            return self._deque[-1][-1:]
+        elif self._buff_unread > 0:
+            idx = (self._head - 1 + self._capacity) % self._capacity
+            return self._buffer[idx : idx + 1]
+        else:
+            raise IndexError("Cannot peek last from an empty buffer.")
+
+    def seek(self, n_samples: int) -> int:
+        """
+        Advances the read head by n_samples.
+
+        Args:
+            n_samples: The number of samples to seek.
+            Will seek forward if positive or backward if negative.
+
+        Returns:
+            The number of samples actually skipped.
+        """
+        self._flush_if_needed(n_samples=n_samples)
+
+        n_to_seek = max(min(n_samples, self._buff_unread), -self._buff_read)
+
+        if n_to_seek == 0:
+            return 0
+
+        self._tail = (self._tail + n_to_seek) % self._capacity
+        self._buff_unread -= n_to_seek
+        self._buff_read += n_to_seek
+
+        return n_to_seek
+
+    def _flush_if_needed(self, n_samples: int | None = None):
+        if (
+            self._update_strategy == "on_demand"
+            and self._deque
+            and (n_samples is None or n_samples > self._buff_unread)
+        ):
+            self.flush()
+
+    def flush(self):
+        """
+        Transfers all data from the deque to the circular buffer.
+        Note: This may overwrite data depending on the overflow strategy,
+            which will invalidate previous state variables.
+        """
+        if not self._deque:
+            return
+
+        n_new = self._deque_len
+        n_free = self._capacity - self._buff_unread
+        n_overflow = max(0, n_new - n_free)
+
+        # If new data is larger than buffer and overflow strategy is "warn-overwrite",
+        #  then we can take a shortcut and replace the entire buffer.
+        if n_new >= self._capacity and self._overflow_strategy == "warn-overwrite":
+            if n_overflow > 0 and (not self._warn_once or not self._warned):
+                self._warned = True
+                warnings.warn(
+                    f"Buffer overflow: {n_new} samples received, but only {self._capacity - self._buff_unread} available. "
+                    f"Overwriting {n_overflow} previous samples.",
+                    RuntimeWarning,
+                )
+
+            # We need to grab the last `self._capacity` samples from the deque
+            samples_to_copy = self._capacity
+            copied_samples = 0
+            for block in reversed(self._deque):
+                if copied_samples >= samples_to_copy:
+                    break
+                n_to_copy = min(block.shape[0], samples_to_copy - copied_samples)
+                start_idx = block.shape[0] - n_to_copy
+                self._buffer[
+                    samples_to_copy - copied_samples - n_to_copy : samples_to_copy
+                    - copied_samples
+                ] = block[start_idx:]
+                copied_samples += n_to_copy
+
+            self._head = 0
+            self._tail = 0
+            self._buff_unread = self._capacity
+            self._buff_read = 0
+            self._last_overflow = n_overflow
+
+        else:
+            if n_overflow > 0:
+                if self._overflow_strategy == "raise":
+                    raise OverflowError(
+                        f"Buffer overflow: {n_new} samples received, but only {n_free} available."
+                    )
+                elif self._overflow_strategy == "warn-overwrite":
+                    if not self._warn_once or not self._warned:
+                        self._warned = True
+                        warnings.warn(
+                            f"Buffer overflow: {n_new} samples received, but only {n_free} available. "
+                            f"Overwriting {n_overflow} previous samples.",
+                            RuntimeWarning,
+                        )
+                    # Move the tail forward to make room for the new data.
+                    self.seek(n_overflow)
+                    # Adjust the read pointer to account for the overflow. Should always be 0.
+                    self._buff_read = max(0, self._buff_read - n_overflow)
+                    self._last_overflow = n_overflow
+                elif self._overflow_strategy == "drop":
+                    # Drop the overflow samples from the deque
+                    samples_to_drop = n_overflow
+                    while samples_to_drop > 0 and self._deque:
+                        block = self._deque[-1]
+                        if samples_to_drop >= block.shape[0]:
+                            samples_to_drop -= block.shape[0]
+                            self._deque.pop()
+                        else:
+                            block = self._deque.pop()
+                            self._deque.append(block[:-samples_to_drop])
+                            samples_to_drop = 0
+                    n_new -= n_overflow
+                    self._last_overflow = n_overflow
+
+                elif self._overflow_strategy == "grow":
+                    self._grow_buffer(self._capacity + n_new)
+                    self._last_overflow = 0
+
+            # Copy data to buffer by iterating over the deque
+            for block in self._deque:
+                n_block = block.shape[0]
+                space_til_end = self._capacity - self._head
+                if n_block > space_til_end:
+                    # Two-part copy (wraps around)
+                    part1_len = space_til_end
+                    part2_len = n_block - part1_len
+                    self._buffer[self._head :] = block[:part1_len]
+                    self._buffer[:part2_len] = block[part1_len:]
+                else:
+                    # Single-part copy
+                    self._buffer[self._head : self._head + n_block] = block
+                self._head = (self._head + n_block) % self._capacity
+
+            self._buff_unread += n_new
+            if (self._buff_read > self._tail) or (self._tail > self._head):
+                # We have wrapped around the buffer; our count of read samples
+                #  is simply the buffer capacity minus the count of unread samples.
+                self._buff_read = self._capacity - self._buff_unread
+
+        self._deque.clear()
+        self._deque_len = 0
+
+    def _grow_buffer(self, min_capacity: int):
+        """
+        Grows the buffer to at least min_capacity.
+        This is a helper method for the overflow strategy "grow".
+        """
+        if self._capacity >= min_capacity:
+            return
+
+        other_shape = self._buffer.shape[1:]
+        max_capacity = self._max_size / (
+            self._buffer.dtype.itemsize * math.prod(other_shape)
+        )
+        if min_capacity > max_capacity:
+            raise OverflowError(
+                f"Cannot grow buffer to {min_capacity} samples, "
+                f"maximum capacity is {max_capacity} samples ({self._max_size} bytes)."
+            )
+
+        new_capacity = min(max_capacity, max(self._capacity * 2, min_capacity))
+        new_buffer = self.xp.empty(
+            (new_capacity, *other_shape), dtype=self._buffer.dtype
+        )
+
+        # Copy existing data to new buffer
+        total_samples = self._buff_read + self._buff_unread
+        if total_samples > 0:
+            start_idx = (self._tail - self._buff_read) % self._capacity
+            stop_idx = (self._tail + self._buff_unread) % self._capacity
+            if stop_idx > start_idx:
+                # Data is contiguous
+                new_buffer[:total_samples] = self._buffer[start_idx:stop_idx]
+            else:
+                # Data wraps around. We write it in 2 parts.
+                part1_len = self._capacity - start_idx
+                part2_len = stop_idx
+                new_buffer[:part1_len] = self._buffer[start_idx:]
+                new_buffer[part1_len : part1_len + part2_len] = self._buffer[:stop_idx]
+            # self._buff_read stays the same
+            self._tail = self._buff_read
+            # self._buff_unread stays the same
+            self._head = self._tail + self._buff_unread
+        else:
+            self._tail = 0
+            self._head = 0
+
+        self._buffer = new_buffer
+        self._capacity = new_capacity

ezmsg/sigproc/util/sparse.py

@@ -1,11 +1,13 @@
+import numpy as np
 import sparse
 
 
-def sliding_win_oneaxis(
+def sliding_win_oneaxis_old(
     s: sparse.SparseArray, nwin: int, axis: int, step: int = 1
 ) -> sparse.SparseArray:
     """
     Like `ezmsg.util.messages.axisarray.sliding_win_oneaxis` but for sparse arrays.
+    This approach is about 4x slower than the version that uses coordinate arithmetic below.
 
     Args:
         s: The input sparse array.
@@ -25,5 +27,97 @@ def sliding_win_oneaxis(
         full_slices[: axis + 1] + (sl,) + full_slices[axis + 2 :] for sl in targ_slices
     ]
     result = sparse.concatenate([s[_] for _ in full_slices], axis=axis)
-    # TODO: Profile this approach vs modifying coords only.
     return result
+
+
+def sliding_win_oneaxis(
+    s: sparse.SparseArray, nwin: int, axis: int, step: int = 1
+) -> sparse.SparseArray:
+    """
+    Generates a view-like sparse array using a sliding window of specified length along a specified axis.
+    Sparse analog of an optimized dense as_strided-based implementation with these properties:
+
+    - Accepts a single `nwin` and a single `axis`.
+    - Inserts a new 'win' axis immediately BEFORE the original target axis.
+      Output shape:
+        s.shape[:axis] + (W,) + (nwin,) + s.shape[axis+1:]
+      where W = s.shape[axis] - (nwin - 1).
+    - If `step > 1`, stepping is applied by slicing along the new windows axis (same observable behavior
+      as doing `slice_along_axis(result, slice(None, None, step), axis)` in the dense version).
+
+    Args:
+        s: Input sparse array (pydata/sparse COO-compatible).
+        nwin: Sliding window size (must be > 0).
+        axis: Axis of `s` along which the window slides (supports negative indexing).
+        step: Stride between windows. If > 1, applied by slicing the windows axis after construction.
+
+    Returns:
+        A sparse array with a new windows axis inserted before the original axis.
+
+    Notes:
+        - Mirrors the dense function’s known edge case: when nwin == shape[axis] + 1, W becomes 0 and
+          an empty windows axis is returned.
+        - Built by coordinate arithmetic; no per-window indexing or concatenation.
+    """
+    if -s.ndim <= axis < 0:
+        axis = s.ndim + axis
+    if not (0 <= axis < s.ndim):
+        raise ValueError(f"Invalid axis {axis} for array with {s.ndim} dimensions")
+    if nwin <= 0:
+        raise ValueError("nwin must be > 0")
+    dim = s.shape[axis]
+
+    last_win_start = dim - nwin
+    win_starts = list(range(0, last_win_start + 1, step))
+    n_win_out = len(win_starts)
+    if n_win_out <= 0:
+        # Return array with proper shape except empty along windows axis
+        return sparse.zeros(
+            s.shape[:axis] + (0,) + (nwin,) + s.shape[axis + 1 :], dtype=s.dtype
+        )
+
+    coo = s.asformat("coo")
+    coords = coo.coords  # shape: (ndim, nnz)
+    data = coo.data  # shape: (nnz,)
+    ia = coords[axis]  # indices along sliding axis, shape: (nnz,)
+
+    # We emit contributions for each offset o in [0, nwin-1].
+    # For a nonzero at index i, it contributes to window start w = i - o when 0 <= w < W.
+    out_coords_blocks = []
+    out_data_blocks = []
+
+    # Small speed/memory tweak: reuse dtypes and pre-allocate o-array once per loop.
+    idx_dtype = coords.dtype
+
+    for win_ix, win_start in enumerate(win_starts):
+        w = ia - win_start
+        # Valid window starts are those within [0, nwin]
+        mask = (w >= 0) & (w < nwin)
+        if not mask.any():
+            continue
+
+        sel = np.nonzero(mask)[0]
+        w_sel = w[sel]
+
+        # Build new coords with windows axis inserted at `axis` and the original axis
+        # becoming the next axis with fixed offset value `o`.
+        # Output ndim = s.ndim + 1
+        before = coords[:axis, sel]  # unchanged
+        after_other = coords[axis + 1 :, sel]  # dims after original axis
+        win_idx_row = np.full((1, sel.size), win_ix, dtype=idx_dtype)
+
+        new_coords = np.vstack([before, win_idx_row, w_sel[None, :], after_other])
+
+        out_coords_blocks.append(new_coords)
+        out_data_blocks.append(data[sel])
+
+    if not out_coords_blocks:
+        return sparse.zeros(
+            s.shape[:axis] + (n_win_out,) + (nwin,) + s.shape[axis + 1 :], dtype=s.dtype
+        )
+
+    out_coords = np.hstack(out_coords_blocks)
+    out_data = np.hstack(out_data_blocks)
+    out_shape = s.shape[:axis] + (n_win_out,) + (nwin,) + s.shape[axis + 1 :]
+
+    return sparse.COO(out_coords, out_data, shape=out_shape)

ezmsg/sigproc/window.py

@@ -209,13 +209,15 @@ class WindowTransformer(
         )
 
         # Create a vector of buffer timestamps to track axis `offset` in output(s)
-        buffer_tvec = xp.asarray(range(self._state.buffer.shape[axis_idx]), dtype=float)
+        buffer_t0 = 0.0
+        buffer_tlen = self._state.buffer.shape[axis_idx]
 
         # Adjust so first _new_ sample at index 0.
-        buffer_tvec -= buffer_tvec[-message.data.shape[axis_idx]]
+        buffer_t0 -= self._state.buffer.shape[axis_idx] - message.data.shape[axis_idx]
+
         # Convert form indices to 'units' (probably seconds).
-        buffer_tvec *= axis_info.gain
-        buffer_tvec += axis_info.offset
+        buffer_t0 *= axis_info.gain
+        buffer_t0 += axis_info.offset
 
         if self.settings.window_shift is not None and self._state.shift_deficit > 0:
             n_skip = min(self._state.buffer.shape[axis_idx], self._state.shift_deficit)
@@ -223,7 +225,8 @@ class WindowTransformer(
                 self._state.buffer = slice_along_axis(
                     self._state.buffer, slice(n_skip, None), axis_idx
                 )
-                buffer_tvec = buffer_tvec[n_skip:]
+                buffer_t0 += n_skip * axis_info.gain
+                buffer_tlen -= n_skip
                 self._state.shift_deficit -= n_skip
 
         # Generate outputs.
@@ -250,7 +253,9 @@ class WindowTransformer(
                 + (1,)
                 + self._state.buffer.shape[axis_idx:]
             )
-            win_offset = buffer_tvec[-self._state.window_samples]
+            win_offset = buffer_t0 + axis_info.gain * (
+                buffer_tlen - self._state.window_samples
+            )
         elif self._state.buffer.shape[axis_idx] >= self._state.window_samples:
             # Deterministic window shifts.
             sliding_win_fun = (
@@ -264,10 +269,7 @@ class WindowTransformer(
                 axis_idx,
                 step=self._state.window_shift_samples,
             )
-            offset_view = sliding_win_fun(buffer_tvec, self._state.window_samples, 0)[
-                :: self._state.window_shift_samples
-            ]
-            win_offset = offset_view[0, 0]
+            win_offset = buffer_t0
 
             # Drop expired beginning of buffer and update shift_deficit
             multi_shift = self._state.window_shift_samples * out_dat.shape[axis_idx]