Coverage for src/evutils/chunking.py: 97%

1"""Splitting event streams into chunks. 

2 

3Slice a continuous event stream into fixed-size windows, either by event 

4count or by time interval. 

5""" 

6 

7import numpy as np 

8 

9import time 

10import queue 

11import threading 

12from typing import Iterator, Any 

13from evutils.io.buffer import EventAccumulator 

14 

15 

16 

17 

18def window_delta_t(events: np.ndarray, delta_t: int = 10_000) -> Iterator[np.ndarray]: 

19 """Returns a generator that chunks the events array into windows of size delta_t. 

20  

21 Parameters 

22 ---------- 

23 events : np.ndarray 

24 Array of events 

25 delta_t : int, optional 

26 Size of the window in microseconds, by default 10_000 

27 

28 Examples 

29 -------- 

30 >>> from evutils.random import random_events 

31 >>> from evutils.chunking import window_delta_t 

32 >>> events = random_events(1000, start_ts=0, end_ts=30_000) 

33 >>> chunks = list(window_delta_t(events, delta_t=10_000)) 

34 >>> len(chunks) > 0 

35 True 

36 """ 

37 if len(events) == 0: 

38 return 

39 

40 index_start = 0 

41 

42 ts = events["t"] 

43 current_ts = ts[0] 

44 

45 while index_start < len(events): 

46 next_index = np.searchsorted(ts[index_start:], current_ts + delta_t) 

47 

48 window = events[index_start:index_start + next_index] 

49 yield window 

50 

51 current_ts += delta_t 

52 index_start += next_index 

53 

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. 

56  

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 

69 

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 len(events) == 0: 

80 return 

81 

82 index_start = 0 

83 

84 ts = events["t"] 

85 current_ts = ts[0] 

86 

87 while index_start < len(events): 

88 

89 

90 next_frame_index = np.searchsorted(ts[index_start:], current_ts + delta_t) 

91 next_window_index = np.searchsorted(ts[index_start:], current_ts + window_size) 

92 

93 

94 

95 window = events[index_start:index_start + next_window_index] 

96 yield window 

97 

98 

99 # Exit if the next window index is not full 

100 if full_window and index_start + next_window_index >= len(events): 

101 break 

102 

103 current_ts += delta_t 

104 index_start += next_frame_index 

105 

106 

107 

108 

109def sort_events(events: np.ndarray) -> np.ndarray: 

110 """Sorts the events array by timestamp. 

111  

112 Parameters 

113 ---------- 

114 events : np.ndarray 

115 Array of events 

116 

117 Examples 

118 -------- 

119 >>> import numpy as np 

120 >>> from evutils.chunking import sort_events 

121 >>> from evutils.random import random_events 

122 >>> events = random_events(10) 

123 >>> events["t"] = np.arange(10, 0, -1) 

124 >>> sorted_events = sort_events(events) 

125 >>> int(sorted_events["t"][0]) 

126 1 

127 """ 

128 return np.sort(events, order="t") 

129 

130def get_dt_events(events: np.ndarray, dt: int =10_000) -> np.ndarray: 

131 """Returns the events that are within a time window of dt from the first event's timestamp. 

132 

133 Parameters 

134 ---------- 

135 events : np.ndarray 

136 Array of events 

137 dt : int, optional 

138 Time window in microseconds, by default 10_000 

139 

140 Examples 

141 -------- 

142 >>> import numpy as np 

143 >>> from evutils.random import random_events 

144 >>> from evutils.chunking import get_dt_events 

145 >>> events = random_events(100, start_ts=0, end_ts=50_000) 

146 >>> sub_events = get_dt_events(events, dt=10_000) 

147 >>> bool((sub_events["t"] <= events["t"][0] + 10_000).all()) 

148 True 

149 """ 

150 if len(events) == 0: 

151 return events 

152 

153 first_ts = events[0]['t'] 

154 last_ts = first_ts + dt 

155 

156 next_index = np.searchsorted(events['t'], last_ts) 

157 

158 return events[:next_index] 

159def stream_delta_t(raw_stream: Iterator[Any], delta_t: int) -> Iterator[Any]: 

160 """A pipeline generator that turns a raw stream into perfect delta_t chunks. 

161 This maintains a small internal buffer for events that cross boundaries. 

162 """ 

163 acc = EventAccumulator(capacity=100_000_000) 

164 current_ts = None 

165 

166 for incoming in raw_stream: 

167 # Handle unpacking depending on if the stream yields triggers or not 

168 if isinstance(incoming, tuple): 

169 ev, tr = incoming 

170 acc.append(ev, tr) 

171 else: 

172 acc.append(incoming, None) 

173 

174 if len(acc) == 0: 

175 continue 

176 

177 # Initialize our absolute time anchor from the very first event 

178 if current_ts is None: 

179 current_ts = int(acc.t_window()[0]) 

180 

181 # Yield as many full windows as we have accumulated 

182 while True: 

183 end_ts = current_ts + delta_t 

184 t = acc.t_window() 

185 

186 if len(t) == 0 or t[-1] < end_ts: 

187 # Not enough data for a full window yet; fetch more chunks 

188 break 

189 

190 # Find boundary 

191 idx = int(np.searchsorted(t, end_ts, side='left')) 

192 

193 # Slice and yield 

194 if acc._tr is not None: 

195 tr_t = acc.t_window_tr() 

196 tr_idx = int(np.searchsorted(tr_t, end_ts, side='left')) 

197 chunk_ev, chunk_tr = acc.slice_copy(idx, tr_idx) 

198 yield chunk_ev, chunk_tr 

199 else: 

200 chunk_ev, _ = acc.slice_copy(idx, 0) 

201 yield chunk_ev 

202 

203 current_ts += delta_t 

204 

205 # Stream finished! Yield whatever is leftover in the buffer 

206 if len(acc) > 0: 

207 if acc._tr is not None: 

208 yield acc.slice_copy(len(acc), acc._tr.size - acc._tr_start) 

209 else: 

210 yield acc.slice_copy(len(acc), 0)[0] 

211 

212 

213def stream_n_events(raw_stream: Iterator[Any], n_events: int) -> Iterator[Any]: 

214 """Pipeline generator: chunks stream by event count.""" 

215 acc = EventAccumulator(capacity=max(1_000_000, n_events * 2)) 

216 for incoming in raw_stream: 

217 if isinstance(incoming, tuple): 

218 acc.append(incoming[0], incoming[1]) 

219 else: 

220 acc.append(incoming, None) 

221 

222 while len(acc) >= n_events: 

223 if acc._tr is not None: 

224 if len(acc) == n_events: 

225 tr_idx = acc._tr.size - acc._tr_start 

226 else: 

227 tr_idx = int(np.searchsorted(acc.t_window_tr(), acc.t_window()[n_events], side='left')) 

228 yield acc.slice_copy(n_events, tr_idx) 

229 else: 

230 yield acc.slice_copy(n_events, 0)[0] 

231 

232 if len(acc) > 0: 

233 if acc._tr is not None: 

234 yield acc.slice_copy(len(acc), acc._tr.size - acc._tr_start) 

235 else: 

236 yield acc.slice_copy(len(acc), 0)[0] 

237 

238def stream_skip_to_time(stream: Iterator[Any], start_ts: int) -> Iterator[Any]: 

239 """Pipeline generator: drops events until start_ts is reached.""" 

240 skipping = True 

241 for incoming in stream: 

242 ev = incoming[0] if isinstance(incoming, tuple) else incoming 

243 if skipping: 

244 if len(ev) == 0 or ev.t[-1] < start_ts: 

245 continue # Drop whole chunk 

246 

247 # Found the boundary! Slice the chunk and stop skipping 

248 idx = int(np.searchsorted(ev.t, start_ts)) 

249 skipping = False 

250 

251 if isinstance(incoming, tuple): 

252 tr_idx = int(np.searchsorted(incoming[1].t, start_ts)) 

253 yield incoming[0][idx:], incoming[1][tr_idx:] 

254 else: 

255 yield incoming[idx:] 

256 else: 

257 yield incoming 

258 

259def stream_async(stream: Iterator[Any], maxsize: int = 5) -> Iterator[Any]: 

260 """Pipeline generator: runs upstream decoding in a background thread.""" 

261 q: queue.Queue[Any] = queue.Queue(maxsize=maxsize) 

262 

263 def worker() -> None: 

264 try: 

265 for item in stream: 

266 # IMPORTANT: C-parsers reuse internal buffers! We MUST copy the chunk 

267 # before placing it in the queue to prevent the next read_chunk()  

268 # from overwriting the memory of the chunk we just yielded! 

269 if isinstance(item, tuple): 

270 ev = item[0].copy() if item[0] is not None else None 

271 tr = item[1].copy() if item[1] is not None else None 

272 q.put((ev, tr)) 

273 else: 

274 q.put(item.copy()) 

275 except Exception as e: 

276 q.put(e) 

277 finally: 

278 q.put(None) # Sentinel 

279 

280 t = threading.Thread(target=worker, daemon=True) 

281 t.start() 

282 

283 while True: 

284 item = q.get() 

285 if item is None: 

286 break 

287 if isinstance(item, Exception): 

288 raise item 

289 yield item 

290 

291def stream_paced_playback(stream: Iterator[Any], playback_speed: float = 1.0) -> Iterator[Any]: 

292 """Pipeline generator: spaces out yielding chunks to match wall-clock real-time.""" 

293 start_wall = None 

294 start_ts = None 

295 

296 for incoming in stream: 

297 ev = incoming[0] if isinstance(incoming, tuple) else incoming 

298 if len(ev) == 0: 

299 yield incoming 

300 continue 

301 

302 if start_ts is None: 

303 start_ts = ev.t[0] 

304 start_wall = time.perf_counter() 

305 

306 # How far into the stream is this chunk's end? 

307 stream_elapsed_us = ev.t[-1] - start_ts 

308 expected_wall_elapsed = (stream_elapsed_us / 1_000_000) / playback_speed 

309 

310 target_wall = start_wall + expected_wall_elapsed 

311 now = time.perf_counter() 

312 

313 if target_wall > now: 

314 time.sleep(target_wall - now) 

315 

316 yield incoming