test.py 7.28 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
59
tol = {np.float32: 6e-7, np.float64: 1.5e-15, np.longfloat: 1e-18}
ctype = {np.float32: np.complex64, np.float64: np.complex128, np.longfloat: np.longcomplex}
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
60

Martin Reinecke's avatar
Martin Reinecke committed
61
@pmp("len", len1D)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
62
@pmp("inorm", [0, 1, 2])
63
64
@pmp("dtype", [np.float32, np.float64, np.longfloat])
def test1D(len, inorm, dtype):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
65
    a = np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
66
67
68
    a = a.astype(ctype[dtype])
    eps = tol[dtype]
    assert_(_l2error(a, ifftn(fftn(a, inorm=inorm), inorm=2-inorm)) < eps)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
69
    assert_(_l2error(a.real, ifftn(fftn(a.real, inorm=inorm), inorm=2-inorm))
70
            < eps)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
71
    assert_(_l2error(a.real, fftn(ifftn(a.real, inorm=inorm), inorm=2-inorm))
72
            < eps)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
73
    assert_(_l2error(a.real, irfftn(rfftn(a.real, inorm=inorm),
74
                                    inorm=2-inorm, lastsize=len)) < eps)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
75
76
77
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
78
    assert_(_l2error(tmp, a) < eps)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
79

Martin Reinecke's avatar
Martin Reinecke committed
80

Martin Reinecke's avatar
Martin Reinecke committed
81
@pmp("shp", shapes)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
82
@pmp("nthreads", (0, 1, 2))
Martin Reinecke's avatar
Martin Reinecke committed
83
def test_fftn(shp, nthreads):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
84
85
86
87
88
89
90
    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
91

Martin Reinecke's avatar
more    
Martin Reinecke committed
92
@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
93
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
Martin Reinecke's avatar
more    
Martin Reinecke committed
94
def test_fftn2D(shp, axes):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
95
96
97
98
99
100
101
    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
102
103
104

@pmp("shp", shapes)
def test_rfftn(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
105
106
107
108
109
    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)

110
111
112
113

@pmp("shp", shapes)
def test_rfft_scipy(shp):
    for i in range(len(shp)):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
114
115
116
117
118
119
        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
120
121

@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
122
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
Martin Reinecke's avatar
more    
Martin Reinecke committed
123
def test_rfftn2D(shp, axes):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
124
125
126
127
128
129
130
    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
131

Martin Reinecke's avatar
Martin Reinecke committed
132
133
@pmp("shp", shapes)
def test_identity(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
134
135
136
137
138
139
140
141
142
143
    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
144
145

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
146
def test_identity_r(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
147
148
    a = np.random.rand(*shp)-0.5
    b = a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
149
150
    for ax in range(a.ndim):
        n = a.shape[ax]
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
151
152
153
154
155
        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
156

Martin Reinecke's avatar
Martin Reinecke committed
157
158
@pmp("shp", shapes)
def test_identity_r2(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
159
160
161
162
    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
163

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
164
@pmp("shp", shapes2D+shapes3D)
165
def test_genuine_hartley(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
166
    a = np.random.rand(*shp)-0.5
167
    v1 = pypocketfft.genuine_hartley(a)
Martin Reinecke's avatar
Martin Reinecke committed
168
    v2 = fftn(a.astype(np.complex128))
Martin Reinecke's avatar
more    
Martin Reinecke committed
169
    v2 = v2.real+v2.imag
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
170
171
    assert_(_l2error(v1, v2) < 1e-15)

Martin Reinecke's avatar
more    
Martin Reinecke committed
172

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
173
@pmp("shp", shapes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
174
def test_hartley_identity(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
175
    a = np.random.rand(*shp)-0.5
176
    v1 = pypocketfft.separable_hartley(pypocketfft.separable_hartley(a))/a.size
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
177
178
    assert_(_l2error(a, v1) < 1e-15)

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
179
180

@pmp("shp", shapes)
181
def test_genuine_hartley_identity(shp):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
182
    a = np.random.rand(*shp)-0.5
183
    v1 = pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(a))/a.size
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
184
    assert_(_l2error(a, v1) < 1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
185
    v1 = a.copy()
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
186
187
188
189
    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
190
191

@pmp("shp", shapes2D)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
192
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
193
def test_genuine_hartley_2D(shp, axes):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
194
195
196
    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
197
198
199


@pmp("len", len1D)
200
201
202
203
204
205
@pmp("inorm", [0, 1])  # inorm==2 not needed, tested via inverse
@pmp("type", [1, 2, 3, 4])
@pmp("dtype", [np.float32, np.float64, np.longfloat])
def testdcst1D(len, inorm, type, dtype):
    a = (np.random.rand(len)-0.5).astype(dtype)
    eps = tol[dtype]
206
207
    itp = (0, 1, 3, 2, 4)
    itype = itp[type]
208
209
210
    if type != 1 or len > 1:  # there are no length-1 type 1 DCTs
        _assert_close(a, pypocketfft.dct(pypocketfft.dct(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), eps)
    _assert_close(a, pypocketfft.dst(pypocketfft.dst(a, inorm=inorm, type=type), inorm=2-inorm, type=itype), eps)