1import time 

2 

3from pystencils.integer_functions import modulo_ceil 

4 

5 

6class TimeLoop: 

7 def __init__(self, steps=2): 

8 self._call_data = [] 

9 self._fixed_steps = steps 

10 self._pre_run_functions = [] 

11 self._post_run_functions = [] 

12 self._single_step_functions = [] 

13 self.time_steps_run = 0 

14 

15 @property 

16 def fixed_steps(self): 

17 return self._fixed_steps 

18 

19 def add_pre_run_function(self, f): 

20 self._pre_run_functions.append(f) 

21 

22 def add_post_run_function(self, f): 

23 self._post_run_functions.append(f) 

24 

25 def add_single_step_function(self, f): 

26 self._single_step_functions.append(f) 

27 

28 def add_call(self, functor, argument_list): 

29 if hasattr(functor, 'kernel'): 

30 functor = functor.kernel 

31 if not isinstance(argument_list, list): 

32 argument_list = [argument_list] 

33 

34 for argument_dict in argument_list: 

35 self._call_data.append((functor, argument_dict)) 

36 

37 def pre_run(self): 

38 for f in self._pre_run_functions: 

39 f() 

40 

41 def post_run(self): 

42 for f in self._post_run_functions: 

43 f() 

44 

45 def run(self, time_steps=1): 

46 self.pre_run() 

47 fixed_steps = self._fixed_steps 

48 call_data = self._call_data 

49 main_iterations, rest_iterations = divmod(time_steps, fixed_steps) 

50 try: 

51 for _ in range(main_iterations): 

52 for func, kwargs in call_data: 

53 func(**kwargs) 

54 self.time_steps_run += fixed_steps 

55 for _ in range(rest_iterations): 

56 for func in self._single_step_functions: 

57 func() 

58 self.time_steps_run += 1 

59 except KeyboardInterrupt: 

60 pass 

61 self.post_run() 

62 

63 def benchmark_run(self, time_steps=0, init_time_steps=0): 

64 init_time_steps_rounded = modulo_ceil(init_time_steps, self._fixed_steps) 

65 time_steps_rounded = modulo_ceil(time_steps, self._fixed_steps) 

66 call_data = self._call_data 

67 

68 self.pre_run() 

69 for i in range(init_time_steps_rounded // self._fixed_steps): 

70 for func, kwargs in call_data: 

71 func(**kwargs) 

72 self.time_steps_run += init_time_steps_rounded 

73 

74 start = time.perf_counter() 

75 for i in range(time_steps_rounded // self._fixed_steps): 

76 for func, kwargs in call_data: 

77 func(**kwargs) 

78 end = time.perf_counter() 

79 self.time_steps_run += time_steps_rounded 

80 self.post_run() 

81 

82 time_for_one_iteration = (end - start) / time_steps 

83 return time_for_one_iteration 

84 

85 def run_time_span(self, seconds): 

86 iterations = 0 

87 self.pre_run() 

88 start = time.perf_counter() 

89 while time.perf_counter() < start + seconds: 

90 for func, kwargs in self._call_data: 

91 func(**kwargs) 

92 iterations += self._fixed_steps 

93 end = time.perf_counter() 

94 self.post_run() 

95 self.time_steps_run += iterations 

96 return iterations, end - start 

97 

98 def benchmark(self, time_for_benchmark=5, init_time_steps=2, number_of_time_steps_for_estimation='auto'): 

99 """Returns the time in seconds for one time step. 

100 

101 Args: 

102 time_for_benchmark: number of seconds benchmark should take 

103 init_time_steps: number of time steps run initially for warm up, to get arrays into cache etc 

104 number_of_time_steps_for_estimation: time steps run before real benchmarks, to determine number of time 

105 steps that approximately take 'time_for_benchmark' or 'auto' 

106 """ 

107 # Run a few time step to get first estimate 

108 if number_of_time_steps_for_estimation == 'auto': 

109 self.run(1) 

110 iterations, total_time = self.run_time_span(0.5) 

111 duration_of_time_step = total_time / iterations 

112 else: 

113 duration_of_time_step = self.benchmark_run(number_of_time_steps_for_estimation, init_time_steps) 

114 

115 # Run for approximately 'time_for_benchmark' seconds 

116 time_steps = int(time_for_benchmark / duration_of_time_step) 

117 time_steps = max(time_steps, 4) 

118 return self.benchmark_run(time_steps, init_time_steps)