BENCH: Modify benchmarks for radix sort.

1 files changed, 32 insertions, 13 deletions

diff --git a/benchmarks/benchmarks/bench_function_base.py b/benchmarks/benchmarks/bench_function_base.py
index 7ea8c39b0..f730bf6ba 100644
--- a/benchmarks/benchmarks/bench_function_base.py
+++ b/benchmarks/benchmarks/bench_function_base.py
@@ -95,6 +95,17 @@ class Select(Benchmark):
         np.select(self.cond_large, ([self.d, self.e] * 10))
 
 
+def memoize(f):
+    _memoized = {}
+    def wrapped(*args):
+        if args not in _memoized:
+            _memoized[args] = f(*args)
+        
+        return _memoized[args].copy()
+
+    return f
+
+
 class SortGenerator(object):
     # The size of the unsorted area in the "random unsorted area"
     # benchmarks
@@ -103,6 +114,7 @@ class SortGenerator(object):
     BUBBLE_SIZE = 100
 
     @staticmethod
+    @memoize
     def random(size, dtype):
         """
         Returns a randomly-shuffled array.
@@ -112,6 +124,7 @@ class SortGenerator(object):
         return arr
     
     @staticmethod
+    @memoize
     def ordered(size, dtype):
         """
         Returns an ordered array.
@@ -119,6 +132,7 @@ class SortGenerator(object):
         return np.arange(size, dtype=dtype)
 
     @staticmethod
+    @memoize
     def reversed(size, dtype):
         """
         Returns an array that's in descending order.
@@ -126,6 +140,7 @@ class SortGenerator(object):
         return np.arange(size-1, -1, -1, dtype=dtype)
 
     @staticmethod
+    @memoize
     def uniform(size, dtype):
         """
         Returns an array that has the same value everywhere.
@@ -133,6 +148,7 @@ class SortGenerator(object):
         return np.ones(size, dtype=dtype)
 
     @staticmethod
+    @memoize
     def swapped_pair(size, dtype, swap_frac):
         """
         Returns an ordered array, but one that has ``swap_frac * size``
@@ -145,6 +161,7 @@ class SortGenerator(object):
         return a
 
     @staticmethod
+    @memoize
     def sorted_block(size, dtype, block_size):
         """
         Returns an array with blocks that are all sorted.
@@ -159,6 +176,7 @@ class SortGenerator(object):
         return np.array(b)
 
     @classmethod
+    @memoize
     def random_unsorted_area(cls, size, dtype, frac, area_size=None):
         """
         This type of array has random unsorted areas such that they
@@ -176,6 +194,7 @@ class SortGenerator(object):
         return a
 
     @classmethod
+    @memoize
     def random_bubble(cls, size, dtype, bubble_num, bubble_size=None):
         """
         This type of array has ``bubble_num`` random unsorted areas.
@@ -197,7 +216,7 @@ class Sort(Benchmark):
         # In NumPy 1.17 and newer, 'merge' can be one of several
         # stable sorts, it isn't necessarily merge sort.
         ['quick', 'merge', 'heap'],
-        ['float64', 'int64', 'uint64'],
+        ['float64', 'int64', 'int16'],
         [
             ('random',),
             ('ordered',),
@@ -206,15 +225,15 @@ class Sort(Benchmark):
             ('sorted_block', 10),
             ('sorted_block', 100),
             ('sorted_block', 1000),
-            ('swapped_pair', 0.01),
-            ('swapped_pair', 0.1),
-            ('swapped_pair', 0.5),
-            ('random_unsorted_area', 0.5),
-            ('random_unsorted_area', 0.1),
-            ('random_unsorted_area', 0.01),
-            ('random_bubble', 1),
-            ('random_bubble', 5),
-            ('random_bubble', 10),
+            # ('swapped_pair', 0.01),
+            # ('swapped_pair', 0.1),
+            # ('swapped_pair', 0.5),
+            # ('random_unsorted_area', 0.5),
+            # ('random_unsorted_area', 0.1),
+            # ('random_unsorted_area', 0.01),
+            # ('random_bubble', 1),
+            # ('random_bubble', 5),
+            # ('random_bubble', 10),
         ],
     ]
     param_names = ['kind', 'dtype', 'array_type']
@@ -227,10 +246,10 @@ class Sort(Benchmark):
         array_class = array_type[0]
         self.arr = getattr(SortGenerator, array_class)(self.ARRAY_SIZE, dtype, *array_type[1:])
 
-    def time_sort_inplace(self, kind, dtype, array_type):
-        self.arr.sort(kind=kind)
-
     def time_sort(self, kind, dtype, array_type):
+        # Using np.sort(...) instead of arr.sort(...) because it makes a copy.
+        # This is important because the data is prepared once per benchmark, but
+        # used across multiple runs.
         np.sort(self.arr, kind=kind)
 
     def time_argsort(self, kind, dtype, array_type):