test.py 6.39 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
130
@pmp("shp", shapes2D+shapes3D)
def test_hartley2(shp):
Martin Reinecke's avatar
more    
Martin Reinecke committed
131
    a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
132
    v1 = pypocketfft.hartley2(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
140
def test_hartley_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley(pypocketfft.hartley(a))/a.size
141
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
142
143
144
145
146

@pmp("shp", shapes)
def test_hartley2_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley2(pypocketfft.hartley2(a))/a.size
147
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
148
    v1 = a.copy()
149
150
    assert_ (pypocketfft.hartley2(pypocketfft.hartley2(v1, inplace=True), inorm=2, inplace=True) is v1)
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
151
152
153
154
155

@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
def test_hartley2_2D(shp, axes):
    a=np.random.rand(*shp)-0.5
156
    assert_(_l2error(pypocketfft.hartley2(pypocketfft.hartley2(a, axes=axes), axes=axes, inorm=2),a)<1e-15)