Coverage for src/evutils/io/buffer.py: 97%

1"""Event buffering module. 

2 

3Provides `EventAccumulator` for buffering and rotating event structures 

4in memory efficiently. 

5""" 

6 

7import numpy as np 

8 

9from ..types import EventArray, TriggerArray 

10from ._native_core import EventSoABuffers, TriggerSoABuffers 

11 

12class EventAccumulator(): 

13 """Reused struct-of-arrays staging buffer that native decoders fill in place. 

14 

15 The whole point is to avoid a copy on the read path: instead of a decoder 

16 parsing into its own buffer and the reader copying that into a ring, the 

17 decoder's parser writes *directly* into this accumulator's storage (via 

18 :meth:`prepare` + ``decoder.parse_step``). Consumed events at the front are 

19 reclaimed by :meth:`_rotate` (moving only the small unconsumed remainder), so 

20 the backing arrays are allocated once and never re-faulted. Only the final 

21 window handed to the caller is copied (:meth:`slice_copy`). 

22 

23 Timestamps are stored as ``uint64`` (matching the native ``timestamp64_t``) 

24 and exposed as ``int64`` via a zero-copy ``.view`` (values are positive and 

25 in range). 

26 """ 

27 

28 def __init__(self, capacity: int): 

29 """Initialize the event accumulator. 

30  

31 Parameters 

32 ---------- 

33 capacity : int 

34 The initial capacity for the event buffer. 

35 

36 """ 

37 self._capacity = int(capacity) 

38 self._buf = EventSoABuffers(self._capacity) 

39 self._tr_capacity = max(self._capacity // 16, 1) 

40 self._tr = TriggerSoABuffers(self._tr_capacity) 

41 self._start = 0 # events before this index are consumed 

42 self._tr_start = 0 

43 

44 def __len__(self) -> int: 

45 return self._buf.size - self._start 

46 

47 def t_window(self) -> np.ndarray: 

48 """int64 view of the currently unconsumed timestamps (zero-copy). 

49  

50 Returns 

51 ------- 

52 np.ndarray 

53 The unconsumed timestamps as an int64 array. 

54 

55 """ 

56 s, e = self._start, self._buf.size 

57 return self._buf.t[s:e].view(np.int64) 

58 

59 def t_window_tr(self) -> np.ndarray: 

60 """int64 view of the currently unconsumed trigger timestamps (zero-copy). 

61 

62 Returns 

63 ------- 

64 np.ndarray 

65 The unconsumed trigger timestamps as an int64 array. 

66 

67 """ 

68 s, e = self._tr_start, self._tr.size 

69 return self._tr.t[s:e].view(np.int64) 

70 

71 def _rotate(self) -> None: 

72 """Move the unconsumed remainder to the front, freeing consumed space.""" 

73 s = self._start 

74 if s > 0: 

75 b = self._buf 

76 n = b.size - s 

77 if n > 0: 

78 import ctypes 

79 ctypes.memmove(b.t.ctypes.data, b.t.ctypes.data + s * b.t.itemsize, n * b.t.itemsize) 

80 ctypes.memmove(b.x.ctypes.data, b.x.ctypes.data + s * b.x.itemsize, n * b.x.itemsize) 

81 ctypes.memmove(b.y.ctypes.data, b.y.ctypes.data + s * b.y.itemsize, n * b.y.itemsize) 

82 ctypes.memmove(b.p.ctypes.data, b.p.ctypes.data + s * b.p.itemsize, n * b.p.itemsize) 

83 b.c.size = n 

84 self._start = 0 

85 

86 s_tr = self._tr_start 

87 if s_tr > 0: 

88 t = self._tr 

89 n_tr = t.size - s_tr 

90 if n_tr > 0: 

91 import ctypes 

92 ctypes.memmove(t.t.ctypes.data, t.t.ctypes.data + s_tr * t.t.itemsize, n_tr * t.t.itemsize) 

93 ctypes.memmove(t.id.ctypes.data, t.id.ctypes.data + s_tr * t.id.itemsize, n_tr * t.id.itemsize) 

94 ctypes.memmove(t.p.ctypes.data, t.p.ctypes.data + s_tr * t.p.itemsize, n_tr * t.p.itemsize) 

95 t.c.size = n_tr 

96 self._tr_start = 0 

97 

98 def _ensure_events(self, headroom: int) -> None: 

99 """Guarantee ``headroom`` free event slots past the current size: 

100 reclaim consumed front first, then grow the backing arrays if that is 

101 still not enough (an oversized window, or drain-to-EOF mode). Growing is 

102 geometric so a stream of large windows costs at most a few reallocs.""" 

103 b = self._buf 

104 if self._capacity - b.size < headroom: 

105 self._rotate() 

106 if self._capacity - b.size < headroom: 

107 new_cap = max(b.size + headroom, int(self._capacity * 1.5) + 1) 

108 b.grow(new_cap) 

109 self._capacity = new_cap 

110 

111 def prepare(self, step: int) -> tuple[EventSoABuffers, TriggerSoABuffers]: 

112 """Ready the buffer for the decoder to append up to ``step`` events, and 

113 return ``(events_soa, triggers_soa)`` for ``decoder.parse_step``. 

114 

115 Reclaims consumed events (and grows the buffer if a window outgrows the 

116 modest initial capacity), then caps the SoA capacity the parser sees to 

117 ``size + step`` so a single step does not overshoot the requested window 

118 by more than one step's worth. 

119 

120 Parameters 

121 ---------- 

122 step : int 

123 The number of events the decoder is expected to append. 

124 

125 Returns 

126 ------- 

127 tuple 

128 A tuple of (events_soa, triggers_soa) buffers. 

129 

130 """ 

131 b = self._buf 

132 t = self._tr 

133 self._ensure_events(step) 

134 if t.capacity - t.size < step // 16: 

135 self._rotate() 

136 if t.capacity - t.size < step // 16: 

137 t.grow(t.size + max(step // 16, 1)) 

138 b.c.capacity = min(self._capacity, b.size + step) 

139 t.c.capacity = t.capacity 

140 return b, t 

141 

142 def append(self, data: EventArray, triggers: TriggerArray | None = None) -> None: 

143 """Copy an EventArray in (fallback path for non-native decoders). 

144 

145 Parameters 

146 ---------- 

147 data : EventArray 

148 The events to append. 

149 triggers : TriggerArray, optional 

150 The triggers to append. 

151 

152 """ 

153 n = len(data) 

154 if n > 0: 

155 b = self._buf 

156 self._ensure_events(n) # rotate + grow so a large chunk always fits 

157 e = b.size 

158 b.t[e:e + n] = data.t 

159 b.x[e:e + n] = data.x 

160 b.y[e:e + n] = data.y 

161 b.p[e:e + n] = data.p 

162 b.c.size = e + n 

163 

164 if triggers is not None and len(triggers) > 0: 

165 n_tr = len(triggers) 

166 t = self._tr 

167 if t.capacity - t.size < n_tr: 

168 self._rotate() 

169 if t.capacity - t.size < n_tr: 

170 t.grow(t.size + n_tr) 

171 e_tr = t.size 

172 t.t[e_tr:e_tr + n_tr] = triggers.t 

173 t.id[e_tr:e_tr + n_tr] = triggers.id 

174 t.p[e_tr:e_tr + n_tr] = triggers.p 

175 t.c.size = e_tr + n_tr 

176 

177 def slice_copy(self, k: int, tr_k: int = 0) -> tuple[EventArray, TriggerArray]: 

178 """Return an independent copy of the first ``k`` unconsumed events and 

179 advance past them. 

180  

181 Parameters 

182 ---------- 

183 k : int 

184 The number of unconsumed events to slice and copy. 

185 tr_k : int, optional 

186 The number of unconsumed triggers to slice and copy. 

187 

188 Returns 

189 ------- 

190 tuple 

191 A tuple of (EventArray, TriggerArray) containing the copied slices. 

192 

193 """ 

194 s = self._buf 

195 i = self._start 

196 out = EventArray( 

197 s.t[i:i + k].view(np.int64), s.x[i:i + k], s.y[i:i + k], s.p[i:i + k] 

198 ).copy() 

199 self._start += k 

200 

201 t = self._tr 

202 j = self._tr_start 

203 out_tr = TriggerArray( 

204 t.t[j:j + tr_k].view(np.int64), t.p[j:j + tr_k], t.id[j:j + tr_k] 

205 ).copy() 

206 self._tr_start += tr_k 

207 

208 return out, out_tr 

209 

210 def reset(self) -> None: 

211 """Reset the buffer and trigger size to 0 and clear consumed offsets.""" 

212 self._buf.c.size = 0 

213 self._tr.c.size = 0 

214 self._start = 0 

215 self._tr_start = 0 

216