test.py 10 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
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
    if type != 1 or len > 1:
218
219
220
221
222
223
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-18)
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-15)
        _assert_close(b, pypocketfft.dct(pypocketfft.dct(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-7)
    _assert_close(a, pypocketfft.dst(pypocketfft.dst(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-18)
    _assert_close(a, pypocketfft.dst(pypocketfft.dst(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-15)
    _assert_close(b, pypocketfft.dst(pypocketfft.dst(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-7)
224

Martin Reinecke's avatar
Martin Reinecke committed
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
@pmp("len", len1D)
@pmp("type", [1, 2, 3])
def testdcst1Dortho(len, type):
    a = np.random.rand(len)-0.5
    b = a.astype(np.float32)
    c = a.astype(np.float128)
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
    if type != 1 or len > 1:
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(c, ortho=True, type=type), ortho=True, type=itype), 2e-18)
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(a, ortho=True, type=type), ortho=True, type=itype), 1.5e-15)
        _assert_close(b, pypocketfft.dct(pypocketfft.dct(b, ortho=True, type=type), ortho=True, type=itype), 6e-7)
    if type != 1:
        _assert_close(a, pypocketfft.dst(pypocketfft.dst(c, ortho=True, type=type), ortho=True, type=itype), 2e-18)
        _assert_close(a, pypocketfft.dst(pypocketfft.dst(a, ortho=True, type=type), ortho=True, type=itype), 1.5e-15)
        _assert_close(b, pypocketfft.dst(pypocketfft.dst(b, ortho=True, type=type), ortho=True, type=itype), 6e-7)

242
243
244
245
246
247
248
249
250
251
252
253

# TEMPORARY: separate test for DCT/DST IV, since they are less accurate
@pmp("len", len1D)
@pmp("inorm", [0, 1, 2])
@pmp("type", [4])
def testdcst1D4(len, inorm, type):
    a = np.random.rand(len)-0.5
    b = a.astype(np.float32)
    c = a.astype(np.float128)
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
    if type != 1 or len > 1:
254
255
256
257
258
259
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-16)
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-13)
        _assert_close(b, pypocketfft.dct(pypocketfft.dct(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-5)
    _assert_close(a, pypocketfft.dst(pypocketfft.dst(c, inorm=inorm, type=type), inorm=2-inorm, type=itype), 2e-16)
    _assert_close(a, pypocketfft.dst(pypocketfft.dst(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), 1.5e-13)
    _assert_close(b, pypocketfft.dst(pypocketfft.dst(b, inorm=inorm, type=type), inorm=2-inorm, type=itype), 6e-5)