Coverage for src/evutils/transforms/functional/_drop.py: 66%

1import math 

2 

3import numpy as np 

4from evutils.jit import lazy_njit 

5 

6from ._common import apply_kernel, sample_range 

7 

8@lazy_njit 

9def _drop_random_events_jit(t: np.ndarray, x: np.ndarray, y: np.ndarray, p: np.ndarray, drop_rate: float, seed: int): 

10 """ 

11 Drops a percentage of events randomly using slicing, compiled via Numba. 

12  

13 Parameters 

14 ---------- 

15 t, x, y, p : np.ndarray 

16 Constituent event arrays. 

17 drop_rate : float 

18 Percentage of events to drop (0 to 1). 

19 seed : int 

20 Seed for the random number generator. If -1, no seed is set. 

21  

22 Returns 

23 ------- 

24 tuple 

25 (new_t, new_x, new_y, new_p) 

26 """ 

27 if seed != -1: 

28 np.random.seed(seed) 

29 # Using random.rand to generate a boolean mask is fully supported and highly optimized in Numba 

30 mask = np.random.rand(len(t)) >= drop_rate 

31 return t[mask], x[mask], y[mask], p[mask] 

32 

33def drop_random_events(events, drop_rate: float = 0.1, seed: int | None = None): 

34 """Drops a percentage of events randomly. 

35  

36 Parameters 

37 ---------- 

38 events : np.ndarray or EventArray 

39 Array of events to drop from. 

40 drop_rate : float, optional 

41 Percentage of events to drop, by default 0.1 (10%). 

42  

43 Returns 

44 ------- 

45 np.ndarray or EventArray 

46 Array of events with the specified percentage dropped. 

47 """ 

48 import math 

49 if math.isnan(drop_rate) or drop_rate <= 0 or drop_rate >= 1: 

50 raise ValueError("drop_rate must be between 0 and 1") 

51 

52 from evutils.transforms.compose import unwrap_events, repack_events 

53 if len(events) == 0: 

54 return events 

55 

56 t, x, y, p = unwrap_events(events) 

57 t, x, y, p = _drop_random_events_jit(t, x, y, p, drop_rate, seed if seed is not None else -1) 

58 return repack_events(events, t, x, y, p) 

59 

60def drop_event(events, p=0.1, seed: int | None = None): 

61 """Randomly drop each event independently with probability ``p``. 

62 

63 The tonic-compatible name for :func:`drop_random_events`. Unlike tonic's 

64 ``drop_event_numpy`` (which drops exactly ``round(p * n)`` events via 

65 shuffle), this drops *independently*, which is O(N) rather than O(N log N). 

66 (sampling without replacement) this uses an independent Bernoulli mask per 

67 event, which is what keeps the kernel JIT-friendly. The expected number of 

68 dropped events is the same; the exact count is binomially distributed. 

69 

70 Parameters 

71 ---------- 

72 events : np.ndarray or EventArray 

73 Events to drop from. 

74 p : float or tuple of float, optional 

75 Drop probability in ``[0, 1)``. A ``(lo, hi)`` tuple is sampled uniformly 

76 per call. Defaults to ``0.1``. 

77 seed : int, optional 

78 Seed for the kernel's random number generator (numba's per-thread 

79 legacy ``np.random`` state). ``None`` leaves the state unseeded. 

80 

81 Returns 

82 ------- 

83 np.ndarray or EventArray 

84 Events that survived the drop, in their original container type. 

85 """ 

86 p = sample_range(p) 

87 if math.isnan(p) or p < 0 or p >= 1: 

88 raise ValueError("p must be in [0, 1)") 

89 if p == 0: 

90 return events 

91 return apply_kernel(events, _drop_random_events_jit, p, 

92 seed if seed is not None else -1) 

93 

94@lazy_njit 

95def _drop_by_time_jit(t, x, y, p, duration_ratio: float, seed: int): 

96 """Drop a single contiguous time window covering ``duration_ratio`` of the span.""" 

97 if seed != -1: 

98 np.random.seed(seed) 

99 if len(t) == 0: 

100 return t, x, y, p 

101 t_start = t.min() 

102 t_end = t.max() 

103 span = t_end - t_start 

104 drop_duration = span * duration_ratio 

105 hi = t_end - drop_duration 

106 # np.random.uniform requires low < high; clamp the degenerate case. 

107 drop_start = np.random.uniform(t_start, hi) if hi > t_start else t_start 

108 keep = (t < drop_start) | (t > drop_start + drop_duration) 

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

110 

111def drop_by_time(events, duration_ratio=0.2, seed: int | None = None): 

112 """Drop every event inside one randomly-placed time window. 

113 

114 The window length is ``duration_ratio`` of the recording span (``[t.min(), t.max()]``, 

115 following tonic), positioned uniformly at random within it. 

116 

117 Parameters 

118 ---------- 

119 events : np.ndarray or EventArray 

120 Events to drop from. 

121 duration_ratio : float or tuple of float, optional 

122 Window length as a fraction of the span, in ``[0, 1)``. A ``(lo, hi)`` 

123 tuple is sampled uniformly per call. Defaults to ``0.2``. 

124 seed : int, optional 

125 Seed for the random number generator. Defaults to None. 

126 

127 Returns 

128 ------- 

129 np.ndarray or EventArray 

130 Events outside the dropped window, in their original container type. 

131 """ 

132 if seed is not None: 

133 np.random.seed(seed) 

134 ratio = sample_range(duration_ratio) 

135 if math.isnan(ratio) or ratio < 0 or ratio >= 1: 

136 raise ValueError("duration_ratio must be in [0, 1)") 

137 if ratio == 0: 

138 return events 

139 return apply_kernel(events, _drop_by_time_jit, ratio, seed if seed is not None else -1)