Coverage for src/evutils/io/_native_aer.py: 100%

1"""ctypes bindings for the native AER parser. 

2 

3AER records carry no timestamps; :class:`AerParser` configures how ``t`` is 

4generated (``AER_TS_ZERO`` or ``AER_TS_SEQUENTIAL``, carried across chunks). 

5Shared SoA buffers and parse helpers live in :mod:`evutils.io._native_core`. 

6""" 

7from __future__ import annotations 

8import ctypes 

9from ctypes import POINTER, c_uint32, cast as c_cast, c_char, c_void_p, c_uint64, byref 

10from typing import cast 

11import numpy as np 

12from ._native_core import register_bindings, NativeError, EventSoABuffers, TriggerSoABuffers, ParserResult, lib 

13 

14class AerInputBuffer(ctypes.Structure): 

15 _fields_ = [("begin", POINTER(c_uint32)), ("end", POINTER(c_uint32))] 

16 

17class AerInput: 

18 __slots__ = ("arr", "c") 

19 def __init__(self, words: np.ndarray): 

20 if words.dtype != np.uint32 or not words.flags["C_CONTIGUOUS"]: raise NativeError("AerInput needs a C-contiguous uint32 array") 

21 self.arr = words 

22 base = words.ctypes.data 

23 self.c = AerInputBuffer() 

24 self.c.begin = c_cast(base, POINTER(c_uint32)) 

25 self.c.end = c_cast(base + words.nbytes, POINTER(c_uint32)) 

26 def consumed(self, result: ParserResult) -> int: 

27 cur = c_cast(result.current, c_void_p).value or self.arr.ctypes.data 

28 return (cur - self.arr.ctypes.data) // 4 

29 

30AER_TS_ZERO = 0 

31AER_TS_SEQUENTIAL = 1 

32 

33def _bind_aer(handle: ctypes.CDLL) -> None: 

34 from ._native_core import EventBufferSOA, TriggerBufferSOA, ParserResult 

35 if hasattr(handle, "AER_state_size"): 

36 handle.AER_state_size.argtypes = [] 

37 handle.AER_state_size.restype = ctypes.c_size_t 

38 if hasattr(handle, "AER_state_configure"): 

39 handle.AER_state_configure.argtypes = [c_void_p, ctypes.c_int32, c_uint64, c_uint64] 

40 handle.AER_state_configure.restype = None 

41 if hasattr(handle, "AER_parse_chunk_soa"): 

42 handle.AER_parse_chunk_soa.argtypes = [c_void_p, POINTER(AerInputBuffer), POINTER(EventBufferSOA), POINTER(TriggerBufferSOA)] 

43 handle.AER_parse_chunk_soa.restype = ParserResult 

44 

45register_bindings(_bind_aer) 

46 

47class AerParser: 

48 __slots__ = ("_state", "_buf", "_mode", "_t_start", "_t_step") 

49 def __init__(self, mode: int = AER_TS_ZERO, t_start: int = 0, t_step: int = 1): 

50 self._buf = (c_char * int(lib().AER_state_size()))() 

51 self._state = c_cast(self._buf, c_void_p) 

52 self._mode = mode 

53 self._t_start = t_start 

54 self._t_step = t_step 

55 lib().AER_state_configure(self._state, mode, t_start, t_step) 

56 def reset(self) -> None: 

57 lib().AER_state_configure(self._state, self._mode, self._t_start, self._t_step) 

58 def parse_chunk_soa(self, inp: AerInput, events: EventSoABuffers, triggers: TriggerSoABuffers) -> ParserResult: 

59 return cast(ParserResult, lib().AER_parse_chunk_soa(self._state, byref(inp.c), byref(events.c), byref(triggers.c))) 

60 def __enter__(self) -> "AerParser": return self