diff --git a/quantecon/util/timing.py b/quantecon/util/timing.py index 5ff9c8dc..4b41277e 100644 --- a/quantecon/util/timing.py +++ b/quantecon/util/timing.py @@ -5,6 +5,7 @@ import time import numpy as np from ..timings.timings import get_default_precision +from numba import njit class __Timer__: @@ -161,8 +162,9 @@ def loop_timer(self, n, function, args=None, verbose=True, digits=2, print("Total run time: %d:%02d:%0d.%0*d" % (h, m, s, digits, (s % 1)*(10**digits))) - average_time = all_times.mean() - average_of_best = np.sort(all_times)[:best_of].mean() + average_time = _mean(all_times) + average_of_best = _mean(_sort_slice(all_times, best_of)) + if verbose: m, s = divmod(average_time, 60) @@ -456,6 +458,27 @@ def loop_timer(n, function, args=None, verbose=True, digits=2, best_of=3): return __timer__.loop_timer(n, function, args, verbose, digits, best_of) + +@njit(cache=True, fastmath=True) +def _mean(arr: np.ndarray) -> float: + total = 0.0 + n = arr.size + for i in range(n): + total += arr[i] + return total / n if n > 0 else 0.0 + +@njit(cache=True, fastmath=True) +def _sort_slice(arr: np.ndarray, best_of: int) -> np.ndarray: + arr_copy = np.empty(arr.size) + for i in range(arr.size): + arr_copy[i] = arr[i] + arr_sorted = np.sort(arr_copy) + result = np.empty(best_of) + for i in range(best_of): + result[i] = arr_sorted[i] + return result + + # Set docstring _names = ['tic', 'tac', 'toc', 'loop_timer'] _funcs = [eval(name) for name in _names]