test.py 6.36 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
10
11
12
shapes1D = ((10,), (127,))
shapes2D = ((128, 128), (128, 129), (1, 129), (129,1))
shapes3D = ((32,17,39),)
shapes = shapes1D+shapes2D+shapes3D
Martin Reinecke's avatar
Martin Reinecke committed
13
len1D=range(1,2048)
Martin Reinecke's avatar
Martin Reinecke committed
14

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

Martin Reinecke's avatar
Martin Reinecke committed
18
def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
19
    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
20
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
21

Martin Reinecke's avatar
Martin Reinecke committed
22
def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
Martin Reinecke's avatar
Martin Reinecke committed
23
    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
Martin Reinecke's avatar
Martin Reinecke committed
24
                           out=out, nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
25
26
27
28
29
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)

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)

Martin Reinecke's avatar
Martin Reinecke committed
34
def rfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
35
    return pypocketfft.r2r_fftpack(a, axes=(axis,), r2hc=True,
Martin Reinecke's avatar
Martin Reinecke committed
36
                                   forward=True, inorm=inorm, out=out,
Martin Reinecke's avatar
Martin Reinecke committed
37
                                   nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
38
def irfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
39
    return pypocketfft.r2r_fftpack(a, axes=(axis,), r2hc=False,
Martin Reinecke's avatar
Martin Reinecke committed
40
                                   forward=False, inorm=inorm, out=out,
Martin Reinecke's avatar
Martin Reinecke committed
41
                                   nthreads=nthreads)
Martin Reinecke's avatar
Martin Reinecke committed
42
@pmp("len", len1D)
43
44
@pmp("inorm", [0,1,2])
def test1D(len, inorm):
Martin Reinecke's avatar
Martin Reinecke committed
45
    a=np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
46
    b=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
47
    c=a.astype(np.complex256)
Martin Reinecke's avatar
Martin Reinecke committed
48
49
50
51
52
    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)
53
    tmp=a.copy()
Martin Reinecke's avatar
Martin Reinecke committed
54
    assert_ (ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm) is tmp)
55
    assert_(_l2error(tmp, a)<1.5e-15)
Martin Reinecke's avatar
Martin Reinecke committed
56
57
58
59
    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)
60
    tmp=b.copy()
Martin Reinecke's avatar
Martin Reinecke committed
61
    assert_ (ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm) is tmp)
62
    assert_(_l2error(tmp, b)<6e-7)
Martin Reinecke's avatar
Martin Reinecke committed
63

Martin Reinecke's avatar
Martin Reinecke committed
64
@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
65
66
@pmp("nthreads", (0,1,2))
def test_fftn(shp, nthreads):
Martin Reinecke's avatar
Martin Reinecke committed
67
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
68
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), fftn(a, nthreads=nthreads))<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
69
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
70
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), fftn(a, nthreads=nthreads))<5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
71

Martin Reinecke's avatar
more    
Martin Reinecke committed
72
73
74
75
@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
def test_fftn2D(shp, axes):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
76
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), fftn(a, axes=axes))<1e-15)
Martin Reinecke's avatar
more    
Martin Reinecke committed
77
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
78
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), fftn(a, axes=axes))<5e-7)
Martin Reinecke's avatar
more    
Martin Reinecke committed
79
80
81
82

@pmp("shp", shapes)
def test_rfftn(shp):
    a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
83
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a))<1e-15)
Martin Reinecke's avatar
more    
Martin Reinecke committed
84
    a=a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
85
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a))<5e-7)
86
87
88
89
90

@pmp("shp", shapes)
def test_rfft_scipy(shp):
    for i in range(len(shp)):
        a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
91
92
        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
93
94
95
96
97

@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
def test_rfftn2D(shp, axes):
    a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
Martin Reinecke committed
98
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), rfftn(a, axes=axes))<1e-15)
Martin Reinecke's avatar
more    
Martin Reinecke committed
99
    a=a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
100
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), rfftn(a, axes=axes))<5e-7)
Martin Reinecke's avatar
more    
Martin Reinecke committed
101

Martin Reinecke's avatar
Martin Reinecke committed
102
103
104
@pmp("shp", shapes)
def test_identity(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
105
106
107
    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)
108
    tmp=a.copy()
Martin Reinecke's avatar
Martin Reinecke committed
109
    assert_ (ifftn(fftn(tmp, out=tmp), inorm=2, out=tmp) is tmp)
110
    assert_(_l2error(tmp, a)<1.5e-15)
Martin Reinecke's avatar
Martin Reinecke committed
111
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
112
    assert_(_l2error(ifftn(fftn(a),inorm=2), a)<6e-7)
Martin Reinecke's avatar
Martin Reinecke committed
113
114

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
115
def test_identity_r(shp):
Martin Reinecke's avatar
Martin Reinecke committed
116
117
118
119
    a=np.random.rand(*shp)-0.5
    b=a.astype(np.float32)
    for ax in range(a.ndim):
        n = a.shape[ax]
Martin Reinecke's avatar
Martin Reinecke committed
120
121
        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
122

Martin Reinecke's avatar
Martin Reinecke committed
123
124
125
@pmp("shp", shapes)
def test_identity_r2(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
126
127
    a=rfftn(irfftn(a))
    assert_(_l2error(rfftn(irfftn(a),inorm=2), a)<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
128

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
129
@pmp("shp", shapes2D+shapes3D)
130
def test_genuine_hartley(shp):
Martin Reinecke's avatar
more    
Martin Reinecke committed
131
    a=np.random.rand(*shp)-0.5
132
    v1 = pypocketfft.genuine_hartley(a)
Martin Reinecke's avatar
Martin Reinecke committed
133
    v2 = fftn(a.astype(np.complex128))
Martin Reinecke's avatar
more    
Martin Reinecke committed
134
    v2 = v2.real+v2.imag
135
    assert_(_l2error(v1, v2)<1e-15)
Martin Reinecke's avatar
more    
Martin Reinecke committed
136

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
137
@pmp("shp", shapes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
138
139
def test_hartley_identity(shp):
    a=np.random.rand(*shp)-0.5
140
    v1 = pypocketfft.separable_hartley(pypocketfft.separable_hartley(a))/a.size
141
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
142
143

@pmp("shp", shapes)
144
def test_genuine_hartley_identity(shp):
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
145
    a=np.random.rand(*shp)-0.5
146
    v1 = pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(a))/a.size
147
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
148
    v1 = a.copy()
Martin Reinecke's avatar
Martin Reinecke committed
149
    assert_(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(v1, out=v1), inorm=2, out=v1) is v1)
150
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
151
152
153

@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
154
def test_genuine_hartley_2D(shp, axes):
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
155
    a=np.random.rand(*shp)-0.5
156
    assert_(_l2error(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(a, axes=axes), axes=axes, inorm=2),a)<1e-15)