test.py 8.11 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1 2 3 4
import pypocketfft
import pyfftw
import numpy as np
import pytest
5
from numpy.testing import assert_
Martin Reinecke's avatar
Martin Reinecke committed
6 7 8

pmp = pytest.mark.parametrize

Martin Reinecke's avatar
more  
Martin Reinecke committed
9
shapes1D = ((10,), (127,))
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
10 11
shapes2D = ((128, 128), (128, 129), (1, 129), (129, 1))
shapes3D = ((32, 17, 39),)
Martin Reinecke's avatar
more  
Martin Reinecke committed
12
shapes = shapes1D+shapes2D+shapes3D
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
13
len1D = range(1, 2048)
Martin Reinecke's avatar
Martin Reinecke committed
14

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
15 16

def _l2error(a, b):
Martin Reinecke's avatar
Martin Reinecke committed
17 18
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
19

Martin Reinecke's avatar
Martin Reinecke committed
20 21 22 23 24 25 26
def _assert_close(a, b, epsilon):
    err = _l2error(a, b)
    if (err >= epsilon):
        print("Error: {} > {}".format(err, epsilon))
    assert_(err<epsilon)


Martin Reinecke's avatar
Martin Reinecke committed
27
def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
28
    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
29
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
30

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
31

Martin Reinecke's avatar
Martin Reinecke committed
32
def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
33
    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
34
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
35

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
36

Martin Reinecke's avatar
Martin Reinecke committed
37 38 39 40
def rfftn(a, axes=None, inorm=0, nthreads=1):
    return pypocketfft.r2c(a, axes=axes, forward=True, inorm=inorm,
                           nthreads=nthreads)

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
41

Martin Reinecke's avatar
Martin Reinecke committed
42
def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
43
    return pypocketfft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
Martin Reinecke's avatar
Martin Reinecke committed
44 45
                           inorm=inorm, nthreads=nthreads)

Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
46

Martin Reinecke's avatar
Martin Reinecke committed
47
def rfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
48
    return pypocketfft.r2r_fftpack(a, axes=(axis,), real2hermitian=True,
Martin Reinecke's avatar
Martin Reinecke committed
49
                                   forward=True, inorm=inorm, out=out,
Martin Reinecke's avatar
Martin Reinecke committed
50
                                   nthreads=nthreads)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
51 52


Martin Reinecke's avatar
Martin Reinecke committed
53
def irfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
54
    return pypocketfft.r2r_fftpack(a, axes=(axis,), real2hermitian=False,
Martin Reinecke's avatar
Martin Reinecke committed
55
                                   forward=False, inorm=inorm, out=out,
Martin Reinecke's avatar
Martin Reinecke committed
56
                                   nthreads=nthreads)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
57 58


Martin Reinecke's avatar
Martin Reinecke committed
59
@pmp("len", len1D)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
60
@pmp("inorm", [0, 1, 2])
61
def test1D(len, inorm):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
62 63 64
    a = np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
    b = a.astype(np.complex64)
    c = a.astype(np.complex256)
Martin Reinecke's avatar
Martin Reinecke committed
65
    _assert_close(a, ifftn(fftn(c, inorm=inorm), inorm=2-inorm), 1e-18)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
    assert_(_l2error(a, ifftn(fftn(a, inorm=inorm), inorm=2-inorm)) < 1.5e-15)
    assert_(_l2error(a.real, ifftn(fftn(a.real, inorm=inorm), inorm=2-inorm))
            < 1.5e-15)
    assert_(_l2error(a.real, fftn(ifftn(a.real, inorm=inorm), inorm=2-inorm))
            < 1.5e-15)
    assert_(_l2error(a.real, irfftn(rfftn(a.real, inorm=inorm),
                                    inorm=2-inorm, lastsize=len)) < 1.5e-15)
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
    assert_(_l2error(tmp, a) < 1.5e-15)
    assert_(_l2error(b, ifftn(fftn(b, inorm=inorm), inorm=2-inorm)) < 6e-7)
    assert_(_l2error(b.real, ifftn(fftn(b.real, inorm=inorm), inorm=2-inorm))
            < 6e-7)
    assert_(_l2error(b.real, fftn(ifftn(b.real, inorm=inorm), inorm=2-inorm))
            < 6e-7)
    assert_(_l2error(b.real, irfftn(rfftn(b.real, inorm=inorm), lastsize=len,
                                    inorm=2-inorm)) < 6e-7)
    tmp = b.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
    assert_(_l2error(tmp, b) < 6e-7)

Martin Reinecke's avatar
Martin Reinecke committed
89

Martin Reinecke's avatar
Martin Reinecke committed
90
@pmp("shp", shapes)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
91
@pmp("nthreads", (0, 1, 2))
Martin Reinecke's avatar
Martin Reinecke committed
92
def test_fftn(shp, nthreads):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
93 94 95 96 97 98 99
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a),
                     fftn(a, nthreads=nthreads)) < 1e-15)
    a = a.astype(np.complex64)
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a),
                     fftn(a, nthreads=nthreads)) < 5e-7)

Martin Reinecke's avatar
Martin Reinecke committed
100

Martin Reinecke's avatar
more  
Martin Reinecke committed
101
@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
102
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
Martin Reinecke's avatar
more  
Martin Reinecke committed
103
def test_fftn2D(shp, axes):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
104 105 106 107 108 109 110
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes),
                     fftn(a, axes=axes)) < 1e-15)
    a = a.astype(np.complex64)
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes),
                     fftn(a, axes=axes)) < 5e-7)

Martin Reinecke's avatar
more  
Martin Reinecke committed
111 112 113

@pmp("shp", shapes)
def test_rfftn(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
114 115 116 117 118
    a = np.random.rand(*shp)-0.5
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a)) < 1e-15)
    a = a.astype(np.float32)
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a)) < 5e-7)

119 120 121 122

@pmp("shp", shapes)
def test_rfft_scipy(shp):
    for i in range(len(shp)):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
123 124 125 126 127 128
        a = np.random.rand(*shp)-0.5
        assert_(_l2error(pyfftw.interfaces.scipy_fftpack.rfft(a, axis=i),
                         rfft_scipy(a, axis=i)) < 1e-15)
        assert_(_l2error(pyfftw.interfaces.scipy_fftpack.irfft(a, axis=i),
                         irfft_scipy(a, axis=i, inorm=2)) < 1e-15)

Martin Reinecke's avatar
more  
Martin Reinecke committed
129 130

@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
131
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
Martin Reinecke's avatar
more  
Martin Reinecke committed
132
def test_rfftn2D(shp, axes):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
133 134 135 136 137 138 139
    a = np.random.rand(*shp)-0.5
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes),
                     rfftn(a, axes=axes)) < 1e-15)
    a = a.astype(np.float32)
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes),
                     rfftn(a, axes=axes)) < 5e-7)

Martin Reinecke's avatar
more  
Martin Reinecke committed
140

Martin Reinecke's avatar
Martin Reinecke committed
141 142
@pmp("shp", shapes)
def test_identity(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
143 144 145 146 147 148 149 150 151 152
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    assert_(_l2error(ifftn(fftn(a), inorm=2), a) < 1.5e-15)
    assert_(_l2error(ifftn(fftn(a.real), inorm=2), a.real) < 1.5e-15)
    assert_(_l2error(fftn(ifftn(a.real), inorm=2), a.real) < 1.5e-15)
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp), inorm=2, out=tmp) is tmp)
    assert_(_l2error(tmp, a) < 1.5e-15)
    a = a.astype(np.complex64)
    assert_(_l2error(ifftn(fftn(a), inorm=2), a) < 6e-7)

Martin Reinecke's avatar
Martin Reinecke committed
153 154

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
155
def test_identity_r(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
156 157
    a = np.random.rand(*shp)-0.5
    b = a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
158 159
    for ax in range(a.ndim):
        n = a.shape[ax]
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
160 161 162 163 164
        assert_(_l2error(irfftn(rfftn(a, (ax,)), (ax,), lastsize=n, inorm=2),
                         a) < 1e-15)
        assert_(_l2error(irfftn(rfftn(b, (ax,)), (ax,), lastsize=n, inorm=2),
                         b) < 5e-7)

Martin Reinecke's avatar
more  
Martin Reinecke committed
165

Martin Reinecke's avatar
Martin Reinecke committed
166 167
@pmp("shp", shapes)
def test_identity_r2(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
168 169 170 171
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    a = rfftn(irfftn(a))
    assert_(_l2error(rfftn(irfftn(a), inorm=2), a) < 1e-15)

Martin Reinecke's avatar
Martin Reinecke committed
172

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
173
@pmp("shp", shapes2D+shapes3D)
174
def test_genuine_hartley(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
175
    a = np.random.rand(*shp)-0.5
176
    v1 = pypocketfft.genuine_hartley(a)
Martin Reinecke's avatar
Martin Reinecke committed
177
    v2 = fftn(a.astype(np.complex128))
Martin Reinecke's avatar
more  
Martin Reinecke committed
178
    v2 = v2.real+v2.imag
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
179 180
    assert_(_l2error(v1, v2) < 1e-15)

Martin Reinecke's avatar
more  
Martin Reinecke committed
181

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
182
@pmp("shp", shapes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
183
def test_hartley_identity(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
184
    a = np.random.rand(*shp)-0.5
185
    v1 = pypocketfft.separable_hartley(pypocketfft.separable_hartley(a))/a.size
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
186 187
    assert_(_l2error(a, v1) < 1e-15)

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
188 189

@pmp("shp", shapes)
190
def test_genuine_hartley_identity(shp):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
191
    a = np.random.rand(*shp)-0.5
192
    v1 = pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(a))/a.size
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
193
    assert_(_l2error(a, v1) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
194
    v1 = a.copy()
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
195 196 197 198
    assert_(pypocketfft.genuine_hartley(
        pypocketfft.genuine_hartley(v1, out=v1), inorm=2, out=v1) is v1)
    assert_(_l2error(a, v1) < 1e-15)

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
199 200

@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
201
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
202
def test_genuine_hartley_2D(shp, axes):
Martin Reinecke's avatar
PEP8  
Martin Reinecke committed
203 204 205
    a = np.random.rand(*shp)-0.5
    assert_(_l2error(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(
        a, axes=axes), axes=axes, inorm=2), a) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
206 207 208 209


@pmp("len", len1D)
@pmp("inorm", [0, 1, 2])
210 211
@pmp("type", [1, 2, 3])
def testdcst1D(len, inorm, type):
Martin Reinecke's avatar
Martin Reinecke committed
212 213 214
    a = np.random.rand(len)-0.5
    b = a.astype(np.float32)
    c = a.astype(np.float128)
215 216 217 218 219 220 221 222 223
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
    if type != 1 or len > 1:
        _assert_close(a, pypocketfft.r2r_dct(pypocketfft.r2r_dct(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-18)
        _assert_close(a, pypocketfft.r2r_dct(pypocketfft.r2r_dct(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-15)
        _assert_close(b, pypocketfft.r2r_dct(pypocketfft.r2r_dct(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-7)
    _assert_close(a, pypocketfft.r2r_dst(pypocketfft.r2r_dst(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-18)
    _assert_close(a, pypocketfft.r2r_dst(pypocketfft.r2r_dst(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-15)
    _assert_close(b, pypocketfft.r2r_dst(pypocketfft.r2r_dst(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-7)