add some FFT demos

pypocketfft/demos/bench.py

+import numpy as np
+import pypocketfft
+from time import time
+import matplotlib.pyplot as plt
+
+
+def _l2error(a, b):
+    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
+
+
+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=flags, axes=range(a.ndim), threads=nthr)
+    f1[()] = a
+    tmin = 1e38
+    for i in range(nrepeat):
+        t0 = time()
+        fftw()
+        t1 = time()
+        tmin = min(tmin, t1-t0)
+    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, nthr):
+    import pyfftw
+    pyfftw.interfaces.cache.enable()
+    tmin = 1e38
+    for i in range(nrepeat):
+        t0 = time()
+        b = pyfftw.interfaces.numpy_fft.fftn(a)
+        t1 = time()
+        tmin = min(tmin, t1-t0)
+    return tmin, b
+
+
+def measure_pypocketfft(a, nrepeat, nthr):
+    import pypocketfft as ppf
+    tmin = 1e38
+    b = a.copy()
+    for i in range(nrepeat):
+        t0 = time()
+        b = ppf.c2c(a, out=b, forward=True, nthreads=nthr)
+        t1 = time()
+        tmin = min(tmin, t1-t0)
+    return tmin, b
+
+
+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, b
+
+
+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 measure_mkl_fft(a, nrepeat, nthr):
+    import os
+    os.environ['OMP_NUM_THREADS'] = str(nthr)
+    import mkl_fft
+    tmin = 1e38
+    for i in range(nrepeat):
+        t0 = time()
+        b = mkl_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="",
+             nice_sizes=True):
+    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)
+        if nice_sizes:
+            shp = np.array([pypocketfft.good_size(sz) for sz in shp])
+        print("  {0:4d}/{1}: shape={2} ...".format(n, ntry, shp), end=" ", flush=True)
+        a = (np.random.rand(*shp)-0.5 + 1j*(np.random.rand(*shp)-0.5)).astype(tp)
+        output=[]
+        for func, res in zip(funcs, results):
+            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))
+    plt.xlabel("time ratio")
+    plt.ylabel("counts")
+    plt.hist(results[0, :]/results[1, :], bins="auto")
+    if filename != "":
+        plt.savefig(filename)
+    plt.show()
+
+
+funcs = (measure_pypocketfft, measure_fftw)
+ttl = "pypocketfft/FFTW()"
+nthr = 1
+nice_sizes = True
+bench_nd(1, 8192, nthr, 100, "c16", funcs, 10, ttl, "1d.png", nice_sizes)
+bench_nd(2, 2048, nthr, 100, "c16", funcs, 2, ttl, "2d.png", nice_sizes)
+bench_nd(3, 256, nthr, 100, "c16", funcs, 2, ttl, "3d.png", nice_sizes)
+bench_nd(1, 8192, nthr, 100, "c8", funcs, 10, ttl, "1d_single.png", nice_sizes)
+bench_nd(2, 2048, nthr, 100, "c8", funcs, 2, ttl, "2d_single.png", nice_sizes)
+bench_nd(3, 256, nthr, 100, "c8", funcs, 2, ttl, "3d_single.png", nice_sizes)

pypocketfft/demos/stress.py

+import numpy as np
+import pypocketfft
+
+
+def _l2error(a, b, axes):
+    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))/np.log2(np.max([2,np.prod(np.take(a.shape,axes))]))
+
+
+def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
+    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
+                           out=out, nthreads=nthreads)
+
+
+def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
+    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
+                           out=out, nthreads=nthreads)
+
+
+def rfftn(a, axes=None, inorm=0, nthreads=1):
+    return pypocketfft.r2c(a, axes=axes, forward=True, inorm=inorm,
+                           nthreads=nthreads)
+
+
+def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
+    return pypocketfft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
+                           inorm=inorm, nthreads=nthreads)
+
+
+nthreads = 0
+
+
+def update_err(err, name, value):
+    if name in err and err[name] >= value:
+        return err
+    err[name] = value
+    for (nm, v) in err.items():
+        print("{}: {}".format(nm, v))
+    print()
+    return err
+
+
+def test(err):
+    ndim = np.random.randint(1, 5)
+    axlen = int((2**20)**(1./ndim))
+    shape = np.random.randint(1, axlen, ndim)
+    axes = np.arange(ndim)
+    np.random.shuffle(axes)
+    nax = np.random.randint(1, ndim+1)
+    axes = axes[:nax]
+    lastsize = shape[axes[-1]]
+    a = np.random.rand(*shape)-0.5 + 1j*np.random.rand(*shape)-0.5j
+    a_32 = a.astype(np.complex64)
+    b = ifftn(fftn(a, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
+              nthreads=nthreads)
+    err = update_err(err, "cmax", _l2error(a, b, axes))
+    b = ifftn(fftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
+              nthreads=nthreads)
+    err = update_err(err, "cmax", _l2error(a.real, b, axes))
+    b = fftn(ifftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
+             nthreads=nthreads)
+    err = update_err(err, "cmax", _l2error(a.real, b, axes))
+    b = ifftn(fftn(a.astype(np.complex64), axes=axes, nthreads=nthreads),
+              axes=axes, inorm=2, nthreads=nthreads)
+    err = update_err(err, "cmaxf", _l2error(a.astype(np.complex64), b, axes))
+    b = irfftn(rfftn(a.real, axes=axes, nthreads=nthreads), axes=axes, inorm=2,
+               lastsize=lastsize, nthreads=nthreads)
+    err = update_err(err, "rmax", _l2error(a.real, b, axes))
+    b = irfftn(rfftn(a.real.astype(np.float32), axes=axes, nthreads=nthreads),
+               axes=axes, inorm=2, lastsize=lastsize, nthreads=nthreads)
+    err = update_err(err, "rmaxf", _l2error(a.real.astype(np.float32), b, axes))
+    b = pypocketfft.separable_hartley(
+        pypocketfft.separable_hartley(a.real, axes=axes, nthreads=nthreads),
+        axes=axes, inorm=2, nthreads=nthreads)
+    err = update_err(err, "hmax", _l2error(a.real, b, axes))
+    b = pypocketfft.genuine_hartley(
+        pypocketfft.genuine_hartley(a.real, axes=axes, nthreads=nthreads),
+        axes=axes, inorm=2, nthreads=nthreads)
+    err = update_err(err, "hmax", _l2error(a.real, b, axes))
+    b = pypocketfft.separable_hartley(
+            pypocketfft.separable_hartley(
+                a.real.astype(np.float32), axes=axes, nthreads=nthreads),
+            axes=axes, inorm=2, nthreads=nthreads)
+    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b, axes))
+    b = pypocketfft.genuine_hartley(
+            pypocketfft.genuine_hartley(a.real.astype(np.float32), axes=axes,
+                                        nthreads=nthreads),
+            axes=axes, inorm=2, nthreads=nthreads)
+    err = update_err(err, "hmaxf", _l2error(a.real.astype(np.float32), b, axes))
+    if all(a.shape[i] > 1 for i in axes):
+        b = pypocketfft.dct(
+            pypocketfft.dct(a.real, axes=axes, nthreads=nthreads, type=1),
+            axes=axes, type=1, nthreads=nthreads, inorm=2)
+        err = update_err(err, "c1max", _l2error(a.real, b, axes))
+        b = pypocketfft.dct(
+            pypocketfft.dct(a_32.real, axes=axes, nthreads=nthreads, type=1),
+            axes=axes, type=1, nthreads=nthreads, inorm=2)
+        err = update_err(err, "c1maxf", _l2error(a_32.real, b, axes))
+    b = pypocketfft.dct(
+        pypocketfft.dct(a.real, axes=axes, nthreads=nthreads, type=2),
+        axes=axes, type=3, nthreads=nthreads, inorm=2)
+    err = update_err(err, "c23max", _l2error(a.real, b, axes))
+    b = pypocketfft.dct(
+        pypocketfft.dct(a_32.real, axes=axes, nthreads=nthreads, type=2),
+        axes=axes, type=3, nthreads=nthreads, inorm=2)
+    err = update_err(err, "c23maxf", _l2error(a_32.real, b, axes))
+    b = pypocketfft.dct(
+        pypocketfft.dct(a.real, axes=axes, nthreads=nthreads, type=4),
+        axes=axes, type=4, nthreads=nthreads, inorm=2)
+    err = update_err(err, "c4max", _l2error(a.real, b, axes))
+    b = pypocketfft.dct(
+        pypocketfft.dct(a_32.real, axes=axes, nthreads=nthreads, type=4),
+        axes=axes, type=4, nthreads=nthreads, inorm=2)
+    err = update_err(err, "c4maxf", _l2error(a_32.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a.real, axes=axes, nthreads=nthreads, type=1),
+        axes=axes, type=1, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s1max", _l2error(a.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a_32.real, axes=axes, nthreads=nthreads, type=1),
+        axes=axes, type=1, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s1maxf", _l2error(a_32.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a.real, axes=axes, nthreads=nthreads, type=2),
+        axes=axes, type=3, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s23max", _l2error(a.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a_32.real, axes=axes, nthreads=nthreads, type=2),
+        axes=axes, type=3, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s23maxf", _l2error(a_32.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a.real, axes=axes, nthreads=nthreads, type=4),
+        axes=axes, type=4, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s4max", _l2error(a.real, b, axes))
+    b = pypocketfft.dst(
+        pypocketfft.dst(a_32.real, axes=axes, nthreads=nthreads, type=4),
+        axes=axes, type=4, nthreads=nthreads, inorm=2)
+    err = update_err(err, "s4maxf", _l2error(a_32.real, b, axes))
+
+
+err = dict()
+while True:
+    test(err)