Coverage for src/evutils/chunking.py: 90%
185 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""Splitting event streams into chunks.
3Slice a continuous event stream into fixed-size windows, either by event
4count or by time interval.
5"""
7import numpy as np
9import time
10import queue
11import threading
12from typing import Iterator
13from evutils.io.buffer import EventAccumulator
14from evutils.types import EventArray
16def window_delta_t(events: np.ndarray, delta_t: int = 10_000) -> Iterator[np.ndarray]:
17 """Returns a generator that chunks the events array into windows of size delta_t.
19 Parameters
20 ----------
21 events : np.ndarray
22 Array of events
23 delta_t : int, optional
24 Size of the window in microseconds, by default 10_000
26 Examples
27 --------
28 >>> from evutils.random import random_events
29 >>> from evutils.chunking import window_delta_t
30 >>> events = random_events(1000, start_ts=0, end_ts=30_000)
31 >>> chunks = list(window_delta_t(events, delta_t=10_000))
32 >>> len(chunks) > 0
33 True
34 """
35 if delta_t <= 0:
36 raise ValueError("delta_t must be positive")
37 if len(events) == 0:
38 return
40 index_start = 0
42 ts = events["t"]
43 current_ts = ts[0]
45 while index_start < len(events):
46 next_index = np.searchsorted(ts[index_start:], current_ts + delta_t)
48 window = events[index_start:index_start + next_index]
49 yield window
51 current_ts += delta_t
52 index_start += next_index
54def sliding_window(events: np.ndarray, delta_t: int = 10_000, window_size: int = 20_000, full_window: bool = False) -> Iterator[np.ndarray]:
55 """Returns a generator that chunks the events array into windows of size delta_t.
57 Parameters
58 ----------
59 events : np.ndarray
60 Array of events
61 delta_t : int, optional
62 Time delta between frames in microseconds, by default 10_000
63 window_size : int, optional
64 Size of the window in microseconds, by default 20_000
65 can overlap with the next frame
66 full_window : bool, optional
67 If True, the last window will be full, by default False
68 If False, the last window will be the remaining events
70 Examples
71 --------
72 >>> from evutils.random import random_events
73 >>> from evutils.chunking import sliding_window
74 >>> events = random_events(1000, start_ts=0, end_ts=50_000)
75 >>> chunks = list(sliding_window(events, delta_t=10_000, window_size=20_000))
76 >>> len(chunks) > 0
77 True
78 """
79 if delta_t <= 0 or window_size <= 0:
80 raise ValueError("delta_t and window_size must be positive")
81 if len(events) == 0:
82 return
84 index_start = 0
86 ts = events["t"]
87 current_ts = ts[0]
89 while index_start < len(events):
91 next_frame_index = np.searchsorted(ts[index_start:], current_ts + delta_t)
92 next_window_index = np.searchsorted(ts[index_start:], current_ts + window_size)
96 # Exit if the next window index is not full
97 if full_window and index_start + next_window_index >= len(events):
98 break
100 window = events[index_start:index_start + next_window_index]
101 yield window
103 current_ts += delta_t
104 index_start += next_frame_index
106def sort_events(events: np.ndarray) -> np.ndarray:
107 """Sorts the events array by timestamp.
109 Parameters
110 ----------
111 events : np.ndarray
112 Array of events
114 Examples
115 --------
116 >>> import numpy as np
117 >>> from evutils.chunking import sort_events
118 >>> from evutils.random import random_events
119 >>> events = random_events(10)
120 >>> events["t"] = np.arange(10, 0, -1)
121 >>> sorted_events = sort_events(events)
122 >>> int(sorted_events["t"][0])
123 1
124 """
125 return np.sort(events, order="t")
127def get_dt_events(events: np.ndarray, dt: int =10_000) -> np.ndarray:
128 """Returns the events that are within a time window of dt from the first event's timestamp.
130 Parameters
131 ----------
132 events : np.ndarray
133 Array of events
134 dt : int, optional
135 Time window in microseconds, by default 10_000
137 Examples
138 --------
139 >>> import numpy as np
140 >>> from evutils.random import random_events
141 >>> from evutils.chunking import get_dt_events
142 >>> events = random_events(100, start_ts=0, end_ts=50_000)
143 >>> sub_events = get_dt_events(events, dt=10_000)
144 >>> bool((sub_events["t"] <= events["t"][0] + 10_000).all())
145 True
146 """
147 if len(events) == 0:
148 return events
150 first_ts = events[0]['t']
151 last_ts = first_ts + dt
153 next_index = np.searchsorted(events['t'], last_ts)
155 return events[:next_index]
156def stream_delta_t(raw_stream: Iterator[np.ndarray | EventArray], delta_t: int) -> Iterator[np.ndarray | EventArray]:
157 """A pipeline generator that turns a raw stream into perfect delta_t chunks.
158 This maintains a small internal buffer for events that cross boundaries.
159 """
160 if delta_t <= 0:
161 raise ValueError("delta_t must be positive")
162 # Start small: the accumulator grows geometrically on demand, so a large
163 # up-front capacity only page-faults ~GBs of memory for nothing.
164 acc = EventAccumulator(capacity=1_000_000)
165 current_ts = None
166 has_triggers = False # set from the stream's items: tuple => (events, triggers)
168 for incoming in raw_stream:
169 # Handle unpacking depending on if the stream yields triggers or not
170 if isinstance(incoming, tuple):
171 has_triggers = True
172 ev, tr = incoming
173 acc.append(ev, tr)
174 else:
175 acc.append(incoming, None)
177 if len(acc) == 0 and (not has_triggers or (acc._tr.size - acc._tr_start) == 0):
178 continue
180 # Initialize our absolute time anchor from the very first event or trigger
181 if current_ts is None:
182 if len(acc) > 0:
183 current_ts = int(acc.t_window()[0])
184 elif has_triggers and (acc._tr.size - acc._tr_start) > 0:
185 current_ts = int(acc.t_window_tr()[0])
186 else:
187 continue
189 # Yield as many full windows as we have accumulated
190 while True:
191 end_ts = current_ts + delta_t
192 t = acc.t_window()
194 max_ts = -1
195 if len(t) > 0:
196 max_ts = t[-1]
197 if has_triggers:
198 tr_t = acc.t_window_tr()
199 if len(tr_t) > 0 and tr_t[-1] > max_ts:
200 max_ts = tr_t[-1]
202 if max_ts < end_ts:
203 # Not enough data for a full window yet; fetch more chunks
204 break
206 # Find boundary
207 idx = int(np.searchsorted(t, end_ts, side='left')) if len(t) > 0 else 0
209 # Slice and yield. Mirror the input shape: only a trigger-carrying
210 # stream ((events, triggers) tuples) yields tuples back.
211 if has_triggers:
212 tr_t = acc.t_window_tr()
213 tr_idx = int(np.searchsorted(tr_t, end_ts, side='left')) if len(tr_t) > 0 else 0
214 chunk_ev, chunk_tr = acc.slice_copy(idx, tr_idx)
215 yield chunk_ev, chunk_tr
216 else:
217 chunk_ev, _ = acc.slice_copy(idx, 0)
218 yield chunk_ev
220 current_ts += delta_t
222 # Stream finished! Yield whatever is leftover in the buffer
223 if len(acc) > 0 or (has_triggers and acc._tr.size - acc._tr_start > 0):
224 if has_triggers:
225 yield acc.slice_copy(len(acc), acc._tr.size - acc._tr_start)
226 else:
227 yield acc.slice_copy(len(acc), 0)[0]
229def stream_n_events(raw_stream: Iterator[np.ndarray | EventArray], n_events: int) -> Iterator[np.ndarray | EventArray]:
230 """Pipeline generator: chunks stream by event count."""
231 if n_events <= 0:
232 raise ValueError("n_events must be positive")
233 acc = EventAccumulator(capacity=max(1_000_000, n_events * 2))
234 has_triggers = False # set from the stream's items: tuple => (events, triggers)
235 for incoming in raw_stream:
236 if isinstance(incoming, tuple):
237 has_triggers = True
238 acc.append(incoming[0], incoming[1])
239 else:
240 acc.append(incoming, None)
242 while len(acc) >= n_events:
243 if has_triggers:
244 if len(acc) == n_events:
245 tr_idx = acc._tr.size - acc._tr_start
246 else:
247 tr_idx = int(np.searchsorted(acc.t_window_tr(), acc.t_window()[n_events], side='left'))
248 yield acc.slice_copy(n_events, tr_idx)
249 else:
250 yield acc.slice_copy(n_events, 0)[0]
252 if len(acc) > 0 or (has_triggers and acc._tr.size - acc._tr_start > 0):
253 if has_triggers:
254 yield acc.slice_copy(len(acc), acc._tr.size - acc._tr_start)
255 else:
256 yield acc.slice_copy(len(acc), 0)[0]
258def stream_skip_to_time(stream: Iterator[np.ndarray | EventArray], start_ts: int) -> Iterator[np.ndarray | EventArray]:
259 """Pipeline generator: drops events until start_ts is reached."""
260 skipping = True
261 for incoming in stream:
262 ev = incoming[0] if isinstance(incoming, tuple) else incoming
263 if skipping:
264 # ["t"] works for both EventArray and plain structured ndarrays
265 # (attribute access .t does not exist on the latter).
266 if len(ev) == 0 or ev["t"][-1] < start_ts:
267 continue # Drop whole chunk
269 # Found the boundary! Slice the chunk and stop skipping
270 idx = int(np.searchsorted(ev["t"], start_ts))
271 skipping = False
273 if isinstance(incoming, tuple):
274 tr_idx = int(np.searchsorted(incoming[1]["t"], start_ts))
275 yield incoming[0][idx:], incoming[1][tr_idx:]
276 else:
277 yield incoming[idx:]
278 else:
279 yield incoming
281def stream_skip_to_event(stream: Iterator[np.ndarray | EventArray], n: int) -> Iterator[np.ndarray | EventArray]:
282 """Pipeline generator: drops the first ``n`` events, then passes through.
284 The event-index counterpart of :func:`stream_skip_to_time` -- the linear
285 fallback for seeking by event index on a non-seekable stream.
286 """
287 seen = 0
288 skipping = True
289 for incoming in stream:
290 ev = incoming[0] if isinstance(incoming, tuple) else incoming
291 if skipping:
292 if seen + len(ev) <= n:
293 seen += len(ev)
294 continue # whole chunk is before the target
295 idx = n - seen
296 skipping = False
297 if isinstance(incoming, tuple):
298 # Slice triggers at the first kept event's timestamp, mirroring
299 # stream_skip_to_time (triggers before the target are dropped).
300 if len(incoming[1]) > 0 and idx < len(ev):
301 tr_idx = int(np.searchsorted(incoming[1]["t"], ev["t"][idx]))
302 else:
303 tr_idx = 0
304 yield incoming[0][idx:], incoming[1][tr_idx:]
305 else:
306 yield incoming[idx:]
307 else:
308 yield incoming
310def stream_async(stream: Iterator[np.ndarray | EventArray], maxsize: int = 5) -> Iterator[np.ndarray | EventArray]:
311 """Pipeline generator: runs upstream decoding in a background thread."""
312 q: queue.Queue[np.ndarray | EventArray] = queue.Queue(maxsize=maxsize)
314 def worker() -> None:
315 try:
316 for item in stream:
317 # IMPORTANT: C-parsers reuse internal buffers! We MUST copy the chunk
318 # before placing it in the queue to prevent the next read_chunk()
319 # from overwriting the memory of the chunk we just yielded!
320 if isinstance(item, tuple):
321 ev = item[0].copy() if item[0] is not None else None
322 tr = item[1].copy() if item[1] is not None else None
323 q.put((ev, tr))
324 else:
325 q.put(item.copy())
326 except Exception as e:
327 q.put(e)
328 finally:
329 q.put(None) # Sentinel
331 t = threading.Thread(target=worker, daemon=True)
332 t.start()
334 while True:
335 item = q.get()
336 if item is None:
337 break
338 if isinstance(item, Exception):
339 raise item
340 yield item
342def stream_paced_playback(stream: Iterator[np.ndarray | EventArray], playback_speed: float = 1.0) -> Iterator[np.ndarray | EventArray]:
343 """Pipeline generator: spaces out yielding chunks to match wall-clock real-time."""
344 start_wall = None
345 start_ts = None
347 for incoming in stream:
348 ev = incoming[0] if isinstance(incoming, tuple) else incoming
349 if len(ev) == 0:
350 yield incoming
351 continue
353 if start_ts is None:
354 start_ts = ev["t"][0]
355 start_wall = time.perf_counter()
357 # How far into the stream is this chunk's end?
358 stream_elapsed_us = ev["t"][-1] - start_ts
359 expected_wall_elapsed = (stream_elapsed_us / 1_000_000) / playback_speed
361 target_wall = start_wall + expected_wall_elapsed
362 now = time.perf_counter()
364 if target_wall > now:
365 time.sleep(target_wall - now)
367 yield incoming