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

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 

8 

9from typing import Generator 

10from .types import Event_dtype 

11import numpy as np 

12 

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

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

15 

16 Examples 

17 -------- 

18 >>> import numpy as np 

19 >>> from evutils.random import random_events 

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

21 >>> len(events) 

22 10 

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

24 True 

25 """ 

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

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

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

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

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

31 

32 # Sort the timestamps 

33 events["t"].sort() 

34 

35 return events 

36 

37 

38def 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]: 

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

40 

41 Examples 

42 -------- 

43 >>> from evutils.random import random_events_generator 

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

45 >>> for chunk in gen: 

46 ... print(len(chunk)) 

47 10000 

48 10000 

49 5000 

50 """ 

51 if n_events == 0: 

52 return 

53 

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

55 

56 chunk_ts_len = (end_ts - start_ts) // n_chunks 

57 if chunk_ts_len == 0: 

58 chunk_ts_len = 1 

59 

60 chunk_end_ts = start_ts + chunk_ts_len 

61 

62 

63 

64 for chunk in range(n_chunks): 

65 if chunk >= n_chunks - 1: 

66 chunk_end_ts = end_ts 

67 chunk_size = n_events - (chunk * chunk_size) 

68 

69 

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

71 

72 start_ts += chunk_ts_len 

73 chunk_end_ts += chunk_ts_len 

74 

75 yield events 

76 

77 

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

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

80 

81 Examples 

82 -------- 

83 >>> import numpy as np 

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

85 >>> events = random_events(5) 

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

87 >>> len(jittered) == len(events) 

88 True 

89 """ 

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

91 

92 if sort: 

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

94 

95 return events 

96 

97 

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

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

100 

101 Examples 

102 -------- 

103 >>> import numpy as np 

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

105 >>> events = random_events(10) 

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

107 >>> len(jittered) == len(events) 

108 True 

109 """ 

110 jitter = int(jitter) 

111 

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

113 

114 if sort: 

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

116 

117 return events 

118 

119