test.py 4.91 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
import pypocketfft
import pyfftw
import numpy as np
import pytest

pmp = pytest.mark.parametrize

Martin Reinecke's avatar
more    
Martin Reinecke committed
8
9
10
11
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
12
len1D=range(1,2048)
Martin Reinecke's avatar
Martin Reinecke committed
13

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

@pmp("len", len1D)
def test1D(len):
    a=np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
20
    b=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
21
    assert(_l2error(a, pypocketfft.ifftn(pypocketfft.fftn(a),fct=1./len))<1.5e-15)
Martin Reinecke's avatar
Martin Reinecke committed
22
    assert(_l2error(a.real, pypocketfft.irfftn(pypocketfft.rfftn(a.real),fct=1./len,lastsize=len))<1.5e-15)
23
24
25
    tmp=a.copy()
    assert (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True), fct=1./len, inplace=True) is tmp)
    assert(_l2error(tmp, a)<1.5e-15)
Martin Reinecke's avatar
Martin Reinecke committed
26
27
    assert(_l2error(b, pypocketfft.ifftn(pypocketfft.fftn(b),fct=1./len))<6e-7)
    assert(_l2error(b.real, pypocketfft.irfftn(pypocketfft.rfftn(b.real),fct=1./len,lastsize=len))<6e-7)
28
29
    tmp=b.copy()
    assert (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True), fct=1./len, inplace=True) is tmp)
Martin Reinecke's avatar
Martin Reinecke committed
30
    assert(_l2error(tmp, b)<6e-7)
Martin Reinecke's avatar
Martin Reinecke committed
31

Martin Reinecke's avatar
Martin Reinecke committed
32
33
34
35
@pmp("shp", shapes)
def test_fftn(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    vmax = np.max(np.abs(pyfftw.interfaces.numpy_fft.fftn(a)))
Martin Reinecke's avatar
Martin Reinecke committed
36
    assert(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a))<1e-15*vmax)
Martin Reinecke's avatar
Martin Reinecke committed
37
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
38
    assert(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a))<5e-7*vmax)
Martin Reinecke's avatar
Martin Reinecke committed
39

Martin Reinecke's avatar
more    
Martin Reinecke committed
40
41
42
43
44
@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
    vmax = np.max(np.abs(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes)))
Martin Reinecke's avatar
Martin Reinecke committed
45
    assert(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes))<1e-15*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
46
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
47
    assert(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes))<5e-7*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
48
49
50
51
52

@pmp("shp", shapes)
def test_rfftn(shp):
    a=np.random.rand(*shp)-0.5
    vmax = np.max(np.abs(pyfftw.interfaces.numpy_fft.fftn(a)))
Martin Reinecke's avatar
Martin Reinecke committed
53
    assert(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a))<1e-15*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
54
    a=a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
55
    assert(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a))<5e-7*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
56
57
58
59
60
61

@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
def test_rfftn2D(shp, axes):
    a=np.random.rand(*shp)-0.5
    vmax = np.max(np.abs(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes)))
Martin Reinecke's avatar
Martin Reinecke committed
62
    assert(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), pypocketfft.rfftn(a, axes=axes))<1e-15*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
63
    a=a.astype(np.float32)
Martin Reinecke's avatar
Martin Reinecke committed
64
    assert(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), pypocketfft.rfftn(a, axes=axes))<5e-7*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
65

Martin Reinecke's avatar
Martin Reinecke committed
66
67
68
69
@pmp("shp", shapes)
def test_identity(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    vmax = np.max(np.abs(a))
Martin Reinecke's avatar
Martin Reinecke committed
70
    assert(_l2error(pypocketfft.ifftn(pypocketfft.fftn(a),fct=1./a.size), a)<1.5e-15*vmax)
71
72
73
    tmp=a.copy()
    assert (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True), fct=1./a.size, inplace=True) is tmp)
    assert(_l2error(tmp, a)<1.5e-15*vmax)
Martin Reinecke's avatar
Martin Reinecke committed
74
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
75
    assert(_l2error(pypocketfft.ifftn(pypocketfft.fftn(a),fct=1./a.size), a)<6e-7*vmax)
Martin Reinecke's avatar
Martin Reinecke committed
76
77

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
78
def test_identity_r(shp):
Martin Reinecke's avatar
Martin Reinecke committed
79
80
81
82
83
    a=np.random.rand(*shp)-0.5
    b=a.astype(np.float32)
    vmax = np.max(np.abs(a))
    for ax in range(a.ndim):
        n = a.shape[ax]
Martin Reinecke's avatar
Martin Reinecke committed
84
85
        assert(_l2error(pypocketfft.irfftn(pypocketfft.rfftn(a,(ax,)),(ax,),lastsize=n,fct=1./n), a)<1e-15*vmax)
        assert(_l2error(pypocketfft.irfftn(pypocketfft.rfftn(b,(ax,)),(ax,),lastsize=n,fct=1./n), b)<5e-7*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
86

Martin Reinecke's avatar
Martin Reinecke committed
87
88
89
90
91
92
@pmp("shp", shapes)
def test_identity_r2(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    a=pypocketfft.rfftn(pypocketfft.irfftn(a))
    fct = 1./pypocketfft.irfftn(a).size
    vmax = np.max(np.abs(a))
Martin Reinecke's avatar
Martin Reinecke committed
93
    assert(_l2error(pypocketfft.rfftn(pypocketfft.irfftn(a),fct=fct), a)<1e-15*vmax)
Martin Reinecke's avatar
Martin Reinecke committed
94

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
95
96
@pmp("shp", shapes2D+shapes3D)
def test_hartley2(shp):
Martin Reinecke's avatar
more    
Martin Reinecke committed
97
    a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
98
    v1 = pypocketfft.hartley2(a)
Martin Reinecke's avatar
more    
Martin Reinecke committed
99
100
101
    v2 = pypocketfft.fftn(a.astype(np.complex128))
    vmax = np.max(np.abs(v1))
    v2 = v2.real+v2.imag
Martin Reinecke's avatar
Martin Reinecke committed
102
    assert(_l2error(v1, v2)<1e-15*vmax)
Martin Reinecke's avatar
more    
Martin Reinecke committed
103

Martin Reinecke's avatar
fixes    
Martin Reinecke committed
104
@pmp("shp", shapes)
Martin Reinecke's avatar
more    
Martin Reinecke committed
105
106
107
108
def test_hartley_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley(pypocketfft.hartley(a))/a.size
    vmax = np.max(np.abs(a))
Martin Reinecke's avatar
Martin Reinecke committed
109
    assert(_l2error(a, v1)<1e-15*vmax)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
110
111
112
113
114
115

@pmp("shp", shapes)
def test_hartley2_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley2(pypocketfft.hartley2(a))/a.size
    vmax = np.max(np.abs(a))
Martin Reinecke's avatar
Martin Reinecke committed
116
    assert(_l2error(a, v1)<1e-15*vmax)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
117
118
119
120
121
122
123

@pmp("shp", shapes2D)
@pmp("axes", ((0,),(1,),(0,1),(1,0)))
def test_hartley2_2D(shp, axes):
    a=np.random.rand(*shp)-0.5
    fct = 1./np.prod(np.take(shp,axes))
    vmax = np.max(np.abs(a))
Martin Reinecke's avatar
Martin Reinecke committed
124
    assert(_l2error(pypocketfft.hartley2(pypocketfft.hartley2(a, axes=axes), axes=axes, fct=fct),a)<1e-15*vmax)