Coverage for src/evutils/types.py: 96%

1"""Core data types for event streams. 

2 

3Defines the structured NumPy dtypes used throughout evutils — ``Events`` 

4(timestamp, x, y, polarity) and ``Triggers`` (timestamp, polarity, id) — 

5together with small helpers for checking event arrays. 

6""" 

7 

8 

9import ctypes 

10from typing import Any, TypeVar 

11 

12import numpy as np 

13 

14__all__ = ['Event_dtype', 'Trigger_dtype', 'Event', 'EventArray', 'TriggerArray', 'is_monotonically_increasing'] 

15 

16 

17#: A structured numpy dtype for event data. 

18#: 

19#: Fields: 

20#: 

21#: - `t` (np.int64): Timestamp of the event (us). 

22#: - `x` (np.uint16): X-coordinate. 

23#: - `y` (np.uint16): Y-coordinate. 

24#: - `p` (np.uint8): Polarity (0: off, 1: on). 

25Event_dtype = np.dtype([('t', np.int64), ('x', np.uint16), ('y', np.uint16), ('p', np.uint8)]) 

26 

27 

28#: A structured numpy dtype for trigger data. 

29#: 

30#: Fields: 

31#: 

32#: - `t` (np.int64): Timestamp of the event (us). 

33#: - `p` (np.uint8): Polarity (0: off, 1: on). 

34#: - `id` (np.uint8): Identifier. 

35Trigger_dtype = np.dtype([('t', np.int64), ('p', np.uint8), ('id', np.uint8)]) 

36 

37 

38class Event(ctypes.Structure): 

39 """Ctypes structure representing an event. 

40 

41 Fields 

42 ------ 

43 t : ctypes.c_int64 

44 Timestamp of the event (us). 

45 x : ctypes.c_uint16 

46 X-coordinate. 

47 y : ctypes.c_uint16 

48 Y-coordinate. 

49 p : ctypes.c_uint8 

50 Polarity (0: off, 1: on). 

51 """ 

52 

53 _fields_ = [("t", ctypes.c_int64), 

54 ("x", ctypes.c_uint16), 

55 ("y", ctypes.c_uint16), 

56 ("p", ctypes.c_uint8)] 

57 

58 

59 

60def is_monotonically_increasing(events: np.ndarray) -> bool: 

61 """Checks if the event ts is monotonically increasing. 

62 

63 Parameters 

64 ---------- 

65 events : np.ndarray 

66 Array of events with a 't' field for timestamps. 

67 

68 Returns 

69 ------- 

70 bool 

71 True if timestamps are monotonically increasing, False otherwise. 

72 

73 """ 

74 return bool(np.all(np.diff(events['t']) >= 0)) 

75 

76 

77 

78_S = TypeVar("_S", bound="SoaArray") 

79 

80 

81class SoaArray: 

82 """Abstract base class for struct-of-arrays (SoA) layout. 

83 

84 Examples 

85 -------- 

86 >>> import numpy as np 

87 >>> from evutils.types import EventArray 

88 >>> events = EventArray(t=[1, 2], x=[10, 20], y=[30, 40], p=[1, 0]) 

89 >>> float(np.mean(events.x)) # SoA layout allows fast operations on single columns 

90 15.0 

91 >>> events.to_numpy() # doctest: +SKIP 

92 array([(1, 10, 30, 1), (2, 20, 40, 0)], 

93 dtype=[('t', '<i8'), ('x', '<u2'), ('y', '<u2'), ('p', 'u1')]) 

94 """ 

95 

96 __slots__ = () 

97 _aos_dtype: np.dtype 

98 _fields: tuple[str, ...] 

99 

100 def __getitem__(self, key: Any) -> Any: 

101 if isinstance(key, str): 

102 return getattr(self, key) 

103 

104 # When indexing a single element, return a NumPy void record to match AoS behaviour exactly. 

105 if isinstance(key, (int, np.integer)): 

106 record = np.empty((), dtype=self._aos_dtype) 

107 for f in self._fields: 

108 record[f] = getattr(self, f)[key] 

109 return record[()] # Returns a scalar np.void 

110 

111 # Otherwise, slice all columns and return a new SoA array 

112 sliced_args = {f: getattr(self, f)[key] for f in self._fields} 

113 return self.__class__(**sliced_args) 

114 

115 def __len__(self) -> int: 

116 return len(getattr(self, self._fields[0])) 

117 

118 def __repr__(self) -> str: 

119 n = len(self) 

120 name = self.__class__.__name__ 

121 if n == 0: 

122 return f"{name}(empty)" 

123 

124 if n <= 10: 

125 return f"{name}(n={n}):\n{self.to_aos()}" 

126 

127 # Slice before AoS conversion for speed 

128 head_str = str(self[:3].to_aos()).rstrip(']') 

129 tail_str = str(self[-3:].to_aos()).lstrip('[') 

130 

131 return f"{name}(n={n}):\n{head_str}\n ...\n {tail_str}]" 

132 

133 def copy(self: _S) -> _S: 

134 """Return a deep copy with independent column arrays.""" 

135 copied_args = {f: getattr(self, f).copy() for f in self._fields} 

136 return self.__class__(**copied_args) 

137 

138 @classmethod 

139 def empty(cls: 'type[_S]') -> _S: 

140 """Return an empty SoA array with correctly-typed (zero-length) columns.""" 

141 args = {f: np.empty(0, dtype=cls._aos_dtype[f]) for f in cls._fields} 

142 return cls(**args) 

143 

144 @classmethod 

145 def from_aos(cls: 'type[_S]', aos_array: np.ndarray) -> _S: 

146 """Constructs a SoA array from an array of structures (AoS) numpy array.""" 

147 args = {f: np.ascontiguousarray(aos_array[f]) for f in cls._fields} 

148 return cls(**args) 

149 

150 def to_aos(self) -> np.ndarray: 

151 """Converts the SoA array to an array of structures (AoS) numpy array.""" 

152 aos_array = np.empty(len(self), dtype=self._aos_dtype) 

153 for f in self._fields: 

154 aos_array[f] = getattr(self, f) 

155 return aos_array 

156 

157 def to_numpy(self) -> np.ndarray: 

158 """Converts the SoA array to a structured numpy array. Alias for to_aos.""" 

159 return self.to_aos() 

160 

161 def __array__(self, dtype: Any = None, copy: Any = None) -> np.ndarray: 

162 """Numpy interop: ``np.asarray(arr)`` returns the AoS structured array.""" 

163 aos = self.to_aos() 

164 if dtype is not None: 

165 return aos.astype(dtype) 

166 return aos 

167 

168 

169class EventArray(SoaArray): 

170 """A container for storing event data in a struct-of-arrays (SoA) layout. 

171 

172 The four fields ``t``, ``x``, ``y`` and ``p`` are kept as separate 

173 contiguous numpy arrays. This is the native layout of the C parser and 

174 avoids the padding of the packed :data:`Event_dtype` struct. 

175 

176 Both attribute access (``events.t``) and key access (``events['t']``) return 

177 the underlying column, so most code written for structured arrays keeps 

178 working. ``np.asarray(events)`` yields the array-of-structures form (see 

179 :meth:`__array__`), which lets EventArray flow into code that still expects 

180 :data:`Event_dtype`. 

181 

182 Examples 

183 -------- 

184 >>> from evutils.types import EventArray 

185 >>> events = EventArray(t=[100, 150], x=[10, 20], y=[30, 40], p=[1, 0]) 

186 >>> events.t 

187 array([100, 150]) 

188 >>> events.x 

189 array([10, 20], dtype=uint16) 

190 >>> events.y 

191 array([30, 40], dtype=uint16) 

192 >>> events.p 

193 array([1, 0], dtype=uint8) 

194 >>> events[:10] 

195 EventArray(n=2): 

196 [(100, 10, 30, 1) (150, 20, 40, 0)] 

197 >>> events[0] # doctest: +SKIP 

198 np.void((100, 10, 30, 1), dtype=[('t', '<i8'), ('x', '<u2'), ('y', '<u2'), ('p', 'u1')]) 

199 """ 

200 

201 __slots__ = ['t', 'x', 'y', 'p'] 

202 _aos_dtype = Event_dtype 

203 _fields = ('t', 'x', 'y', 'p') 

204 

205 def __init__(self, t: Any, x: Any, y: Any, p: Any) -> None: 

206 self.t = np.asarray(t, dtype=np.int64) 

207 self.x = np.asarray(x, dtype=np.uint16) 

208 self.y = np.asarray(y, dtype=np.uint16) 

209 self.p = np.asarray(p, dtype=np.uint8) 

210 

211 

212class TriggerArray(SoaArray): 

213 """A container for storing trigger data in a struct-of-arrays (SoA) layout. 

214 

215 Examples 

216 -------- 

217 >>> from evutils.types import TriggerArray 

218 >>> triggers = TriggerArray(t=[1000, 2000], p=[1, 0], id=[0, 1]) 

219 >>> triggers.t 

220 array([1000, 2000]) 

221 >>> triggers.p 

222 array([1, 0], dtype=uint8) 

223 >>> triggers.id 

224 array([0, 1], dtype=uint8) 

225 """ 

226 

227 __slots__ = ['t', 'p', 'id'] 

228 _aos_dtype = Trigger_dtype 

229 _fields = ('t', 'p', 'id') 

230 

231 def __init__(self, t: Any, p: Any, id: Any) -> None: 

232 self.t = np.asarray(t, dtype=np.int64) 

233 self.p = np.asarray(p, dtype=np.uint8) 

234 self.id = np.asarray(id, dtype=np.uint8)