improve benchmarking code

bench.py

@@ -4,46 +4,16 @@ from time import time
 import matplotlib.pyplot as plt
 import math
 
-nthr = 1
-
 
 def _l2error(a, b):
     return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
 
 
-def prime_factorize(n):
-    factors = []
-    number = math.fabs(n)
-
-    while number > 1:
-        factor = get_next_prime_factor(number)
-        factors.append(factor)
-        number /= factor
-
-    if n < -1:  # If we'd check for < 0, -1 would give us trouble
-        factors[0] = -factors[0]
-
-    return tuple(factors)
-
-
-def get_next_prime_factor(n):
-    if n % 2 == 0:
-        return 2
-
-    # Not 'good' [also] checking non-prime numbers I guess?
-    # But the alternative, creating a list of prime numbers,
-    # wouldn't it be more demanding? Process of creating it.
-    for x in range(3, int(math.ceil(math.sqrt(n)) + 1), 2):
-        if n % x == 0:
-            return x
-    return int(n)
-
-
-def measure_fftw(a, nrepeat):
+def measure_fftw(a, nrepeat, nthr, flags=('FFTW_MEASURE',)):
     import pyfftw
     f1 = pyfftw.empty_aligned(a.shape, dtype=a.dtype)
     f2 = pyfftw.empty_aligned(a.shape, dtype=a.dtype)
-    fftw = pyfftw.FFTW(f1, f2, flags=('FFTW_MEASURE',), axes=range(a.ndim), threads=nthr)
+    fftw = pyfftw.FFTW(f1, f2, flags=flags, axes=range(a.ndim), threads=nthr)
     f1[()] = a
     tmin = 1e38
     for i in range(nrepeat):
@@ -51,10 +21,14 @@ def measure_fftw(a, nrepeat):
         fftw()
         t1 = time()
         tmin = min(tmin, t1-t0)
-    return tmin
+    return tmin, f2
+
 
+def measure_fftw_est(a, nrepeat, nthr):
+    return measure_fftw(a, nrepeat, nthr, flags=('FFTW_ESTIMATE',))
 
-def measure_fftw_np_interface(a, nrepeat):
+
+def measure_fftw_np_interface(a, nrepeat, nthr):
     import pyfftw
     pyfftw.interfaces.cache.enable()
     tmin = 1e38
@@ -63,10 +37,10 @@ def measure_fftw_np_interface(a, nrepeat):
         b = pyfftw.interfaces.numpy_fft.fftn(a)
         t1 = time()
         tmin = min(tmin, t1-t0)
-    return tmin
+    return tmin, b
 
 
-def measure_pypocketfft(a, nrepeat):
+def measure_pypocketfft(a, nrepeat, nthr):
     import pypocketfft as ppf
     tmin = 1e38
     for i in range(nrepeat):
@@ -74,30 +48,59 @@ def measure_pypocketfft(a, nrepeat):
         b = ppf.c2c(a, forward=True, nthreads=nthr)
         t1 = time()
         tmin = min(tmin, t1-t0)
-    return tmin
+    return tmin, b
 
 
-def measure_scipy_fftpack(a, nrepeat):
+def measure_scipy_fftpack(a, nrepeat, nthr):
     import scipy.fftpack
     tmin = 1e38
+    if nthr!=1:
+        raise NotImplementedError("scipy.fftpack does not support multiple threads")
     for i in range(nrepeat):
         t0 = time()
         b = scipy.fftpack.fftn(a)
         t1 = time()
         tmin = min(tmin, t1-t0)
-    return tmin
+    return tmin, b
 
 
-def bench_nd(ndim, nmax, ntry, tp, funcs, nrepeat, ttl="", filename=""):
+def measure_scipy_fft(a, nrepeat, nthr):
+    import scipy.fft
+    tmin = 1e38
+    for i in range(nrepeat):
+        t0 = time()
+        b = scipy.fft.fftn(a, workers=nthr)
+        t1 = time()
+        tmin = min(tmin, t1-t0)
+    return tmin, b
+
+
+def measure_numpy_fft(a, nrepeat, nthr):
+    import scipy.fft
+    tmin = 1e38
+    if nthr!=1:
+        raise NotImplementedError("numpy.fft does not support multiple threads")
+    for i in range(nrepeat):
+        t0 = time()
+        b = np.fft.fftn(a)
+        t1 = time()
+        tmin = min(tmin, t1-t0)
+    return tmin, b
+
+
+def bench_nd(ndim, nmax, nthr, ntry, tp, funcs, nrepeat, ttl="", filename=""):
     print("{}D, type {}, max extent is {}:".format(ndim, tp, nmax))
     results = [[] for i in range(len(funcs))]
     for n in range(ntry):
         shp = np.random.randint(nmax//3, nmax+1, ndim)
         print("  {0:4d}/{1}: shape={2} ...".format(n, ntry, shp), end=" ", flush=True)
-        a = (np.random.rand(*shp) + 1j*np.random.rand(*shp)).astype(tp)
+        a = (np.random.rand(*shp)-0.5 + 1j*(np.random.rand(*shp)-0.5)).astype(tp)
+        output=[]
         for func, res in zip(funcs, results):
-            res.append(func(a, nrepeat))
-        print("{0:5.2e}/{1:5.2e} = {2:5.2f}".format(results[0][n],results[1][n],results[0][n]/results[1][n]))
+            tmp = func(a, nrepeat, nthr)
+            res.append(tmp[0])
+            output.append(tmp[1])
+        print("{0:5.2e}/{1:5.2e} = {2:5.2f}  L2 error={3}".format(results[0][n],results[1][n],results[0][n]/results[1][n],_l2error(output[0],output[1])))
     results = np.array(results)
     plt.title("{}: {}D, {}, max_extent={}".format(
         ttl, ndim, str(tp), nmax))
@@ -111,9 +114,10 @@ def bench_nd(ndim, nmax, ntry, tp, funcs, nrepeat, ttl="", filename=""):
 
 funcs = (measure_pypocketfft, measure_fftw)
 ttl = "pypocketfft/FFTW()"
-bench_nd(1, 8192, 100, "c16", funcs, 10, ttl, "1d.png")
-bench_nd(2, 2048, 100, "c16", funcs, 2, ttl, "2d.png")
-bench_nd(3, 256, 100, "c16", funcs, 2, ttl, "3d.png")
-bench_nd(1, 8192, 100, "c8", funcs, 10, ttl, "1d_single.png")
-bench_nd(2, 2048, 100, "c8", funcs, 2, ttl, "2d_single.png")
-bench_nd(3, 256, 100, "c8", funcs, 2, ttl, "3d_single.png")
+nthr = 1
+bench_nd(1, 8192, nthr, 100,  "c16", funcs, 10, ttl, "1d.png")
+bench_nd(2, 2048, nthr, 100, "c16", funcs, 2, ttl, "2d.png")
+bench_nd(3, 256, nthr, 100, "c16", funcs, 2, ttl, "3d.png")
+bench_nd(1, 8192, nthr, 100, "c8", funcs, 10, ttl, "1d_single.png")
+bench_nd(2, 2048, nthr, 100, "c8", funcs, 2, ttl, "2d_single.png")
+bench_nd(3, 256, nthr, 100, "c8", funcs, 2, ttl, "3d_single.png")