Coverage for src/evutils/random.py: 98%

1"""Generation and perturbation of synthetic events. 

2 

3Create random event arrays for testing and benchmarking 

4(``random_events``, ``random_events_generator``) and add random timestamp 

5jitter to existing events (``event_jitter``, ``event_jitter_n``). 

6""" 

7 

8from typing import Generator 

9from .types import Event_dtype 

10import numpy as np 

11 

12def random_events(n_events: int, width: int = 1280, height: int = 720, start_ts: int = 0, end_ts: int = 10_000_000) -> np.ndarray: 

13 """Generates n_events random events with x and y coordinates in the range [0, width) and [0, height) respectively. 

14 

15 Examples 

16 -------- 

17 >>> import numpy as np 

18 >>> from evutils.random import random_events 

19 >>> events = random_events(10, width=640, height=480, start_ts=0, end_ts=1000) 

20 >>> len(events) 

21 10 

22 >>> bool(events["x"].max() < 640) 

23 True 

24 """ 

25 events = np.empty(n_events, dtype=Event_dtype) 

26 events["x"] = np.random.randint(0, width, n_events) 

27 events["y"] = np.random.randint(0, height, n_events) 

28 events["p"] = np.random.randint(0, 2, n_events) 

29 events["t"] = np.random.randint(start_ts, end_ts, n_events) 

30 

31 # Sort the timestamps 

32 events["t"].sort() 

33 

34 return events 

35 

36def random_events_generator(n_events: int, width: int = 1280, height: int = 720, start_ts: int = 0, end_ts: int = 10_000_000, chunk_size: int = 10000) -> Generator[np.ndarray, None, None]: 

37 """Generates n_events random events with x and y coordinates in the range [0, width) and [0, height) respectively. 

38 

39 Examples 

40 -------- 

41 >>> from evutils.random import random_events_generator 

42 >>> gen = random_events_generator(25000, chunk_size=10000) 

43 >>> for chunk in gen: 

44 ... print(len(chunk)) 

45 10000 

46 10000 

47 5000 

48 """ 

49 if n_events == 0: 

50 return 

51 

52 n_chunks = int(np.ceil(n_events / chunk_size)) 

53 

54 chunk_ts_len = (end_ts - start_ts) // n_chunks 

55 if chunk_ts_len == 0: 

56 chunk_ts_len = 1 

57 

58 chunk_end_ts = start_ts + chunk_ts_len 

59 

60 for chunk in range(n_chunks): 

61 if chunk >= n_chunks - 1: 

62 chunk_end_ts = end_ts 

63 chunk_size = n_events - (chunk * chunk_size) 

64 

65 

66 events = random_events(chunk_size, width, height, start_ts, chunk_end_ts) 

67 

68 start_ts += chunk_ts_len 

69 chunk_end_ts += chunk_ts_len 

70 

71 yield events 

72 

73def event_jitter_n(events: np.ndarray, mean: float = 0.0, std: float = 1.0, sort: bool = True, in_place: bool = False) -> np.ndarray: 

74 """Adds a random jitter to the timestamps of the events. 

75 

76 Examples 

77 -------- 

78 >>> import numpy as np 

79 >>> from evutils.random import random_events, event_jitter_n 

80 >>> events = random_events(5) 

81 >>> jittered = event_jitter_n(events.copy(), mean=0.0, std=5.0) 

82 >>> len(jittered) == len(events) 

83 True 

84 """ 

85 if not in_place: 

86 events = events.copy() 

87 events["t"] += np.round(np.random.normal(mean, std, len(events))).astype(np.int64) 

88 

89 if sort: 

90 events = np.sort(events, order="t") 

91 

92 return events 

93 

94def event_jitter(events: np.ndarray, jitter: int = 1, sort: bool = True, in_place: bool = False) -> np.ndarray: 

95 """Adds a random jitter to the timestamps of the events. 

96 

97 Examples 

98 -------- 

99 >>> import numpy as np 

100 >>> from evutils.random import random_events, event_jitter 

101 >>> events = random_events(10) 

102 >>> jittered = event_jitter(events.copy(), jitter=5, sort=True) 

103 >>> len(jittered) == len(events) 

104 True 

105 """ 

106 jitter = int(jitter) 

107 

108 if not in_place: 

109 events = events.copy() 

110 # +1: randint's high bound is exclusive -- without it the jitter is 

111 # asymmetric ([-j, j-1]) and jitter=0 raises ValueError. 

112 events["t"] += np.random.randint(-jitter, jitter + 1, len(events)) 

113 

114 if sort: 

115 events = np.sort(events, order="t") 

116 

117 return events 

118