test.py 6.95 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
def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
21
    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
22
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
23

Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
24

Martin Reinecke's avatar
Martin Reinecke committed
25
def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
26
    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
27
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
28

Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
29

Martin Reinecke's avatar
Martin Reinecke committed
30
31
32
33
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
34

Martin Reinecke's avatar
Martin Reinecke committed
35
def irfftn(a, axes=None, lastsize=0, inorm=0, nthreads=1):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
36
    return pypocketfft.c2r(a, axes=axes, lastsize=lastsize, forward=False,
Martin Reinecke's avatar
Martin Reinecke committed
37
38
                           inorm=inorm, nthreads=nthreads)

Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
39

Martin Reinecke's avatar
Martin Reinecke committed
40
def rfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
41
    return pypocketfft.r2r_fftpack(a, axes=(axis,), real2hermitian=True,
Martin Reinecke's avatar
Martin Reinecke committed
42
                                   forward=True, inorm=inorm, out=out,
Martin Reinecke's avatar
Martin Reinecke committed
43
                                   nthreads=nthreads)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
44
45


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


Martin Reinecke's avatar
Martin Reinecke committed
52
@pmp("len", len1D)
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
53
@pmp("inorm", [0, 1, 2])
54
def test1D(len, inorm):
Martin Reinecke's avatar
PEP8    
Martin Reinecke committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    a = np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
    b = a.astype(np.complex64)
    c = a.astype(np.complex256)
    assert_(_l2error(a, ifftn(fftn(c, inorm=inorm), inorm=2-inorm)) < 1e-18)
    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
82

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

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

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

112
113
114
115

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

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

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

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

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

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

Martin Reinecke's avatar
more    
Martin Reinecke committed
174

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

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
181
182

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

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