Coverage for src/evutils/transforms/functional/_refractory.py: 90%

1"""Refractory-period functional transform.""" 

2import numpy as np 

3from evutils.jit import lazy_njit 

4 

5from ._common import apply_kernel 

6 

7@lazy_njit 

8def _refractory_period_jit(t, x, y, p, delta: int): 

9 """Discard events that fire within ``delta`` of the previous event at the same pixel. 

10 

11 An event survives when ``t - t_last > delta`` for its pixel. The per-pixel 

12 clock is updated on *every* event (including dropped ones), matching tonic's 

13 ``refractory_period_numpy``. This loop is exactly the kind of scalar-heavy 

14 code Numba turns into tight machine code. 

15 """ 

16 n = len(t) 

17 if n == 0: 

18 return t, x, y, p 

19 width = int(x.max()) + 1 

20 height = int(y.max()) + 1 

21 # Init to t[0] - delta - 1 so the first event at each pixel is always kept, even if t[0] is 0. 

22 last = np.full((width, height), t[0] - delta - 1, dtype=np.int64) 

23 keep = np.zeros(n, dtype=np.bool_) 

24 

25 for i in range(n): 

26 xi = x[i] 

27 yi = y[i] 

28 if t[i] - last[xi, yi] > delta: 

29 keep[i] = True 

30 last[xi, yi] = t[i] 

31 

32 return t[keep], x[keep], y[keep], p[keep] 

33 

34def refractory_period(events, delta): 

35 """Enforce a per-pixel refractory period. 

36 

37 Parameters 

38 ---------- 

39 events : np.ndarray or EventArray 

40 Events to filter. Must be sorted by timestamp. 

41 delta : int 

42 Refractory period in the same time unit as the timestamps. Events at a 

43 pixel that fire within ``delta`` of that pixel's previous event are 

44 dropped. 

45 

46 Returns 

47 ------- 

48 np.ndarray or EventArray 

49 Filtered events, in their original container type. 

50 """ 

51 return apply_kernel(events, _refractory_period_jit, int(delta))