test.py 6.48 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
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def fftn(a, axes=None, inorm=0, inplace=False, nthreads=1):
    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
                           inplace=inplace, nthreads=nthreads)

def ifftn(a, axes=None, inorm=0, inplace=False, nthreads=1):
    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
                           inplace=inplace, nthreads=nthreads)

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)

def rfft_scipy(a, axis, inorm=0, inplace=False, nthreads=1):
    return pypocketfft.r2r_fftpack(a, axes=(axis,), input_halfcomplex=False,
                                   forward=True, inorm=inorm, inplace=inplace,
                                   nthreads=nthreads)
def irfft_scipy(a, axis, inorm=0, inplace=False, nthreads=1):
    return pypocketfft.r2r_fftpack(a, axes=(axis,), input_halfcomplex=True,
                                   forward=False, inorm=inorm, inplace=inplace,
                                   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, inplace=True, inorm=inorm), inplace=True, 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, inplace=True, inorm=inorm), inplace=True, 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, inplace=True), inorm=2, inplace=True) 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
cleanup    
Martin Reinecke committed
149
    assert_(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(v1, inplace=True), inorm=2, inplace=True) 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)