test.py 5.74 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 18
def _l2error(a,b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))

@pmp("len", len1D)
19 20
@pmp("inorm", [0,1,2])
def test1D(len, inorm):
Martin Reinecke's avatar
Martin Reinecke committed
21
    a=np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
Martin Reinecke's avatar
Martin Reinecke committed
22
    b=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
23
    c=a.astype(np.complex256)
24 25
    assert_(_l2error(a, pypocketfft.ifftn(pypocketfft.fftn(c,inorm=inorm), inorm=2-inorm))<1e-18)
    assert_(_l2error(a, pypocketfft.ifftn(pypocketfft.fftn(a,inorm=inorm), inorm=2-inorm))<1.5e-15)
26
    assert_(_l2error(a.real, pypocketfft.ifftn(pypocketfft.fftn(a.real,inorm=inorm), inorm=2-inorm))<1.5e-15)
27
    assert_(_l2error(a.real, pypocketfft.fftn(pypocketfft.ifftn(a.real,inorm=inorm), inorm=2-inorm))<1.5e-15)
28
    assert_(_l2error(a.real, pypocketfft.irfftn(pypocketfft.rfftn(a.real,inorm=inorm), inorm=2-inorm,lastsize=len))<1.5e-15)
29
    tmp=a.copy()
30 31 32
    assert_ (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True, inorm=inorm), inplace=True, inorm=2-inorm) is tmp)
    assert_(_l2error(tmp, a)<1.5e-15)
    assert_(_l2error(b, pypocketfft.ifftn(pypocketfft.fftn(b, inorm=inorm), inorm=2-inorm))<6e-7)
33
    assert_(_l2error(b.real, pypocketfft.ifftn(pypocketfft.fftn(b.real,inorm=inorm), inorm=2-inorm))<6e-7)
34
    assert_(_l2error(b.real, pypocketfft.fftn(pypocketfft.ifftn(b.real,inorm=inorm), inorm=2-inorm))<6e-7)
35
    assert_(_l2error(b.real, pypocketfft.irfftn(pypocketfft.rfftn(b.real, inorm=inorm), lastsize=len, inorm=2-inorm))<6e-7)
36
    tmp=b.copy()
37 38
    assert_ (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True, inorm=inorm), inplace=True, inorm=2-inorm) is tmp)
    assert_(_l2error(tmp, b)<6e-7)
Martin Reinecke's avatar
Martin Reinecke committed
39

Martin Reinecke's avatar
Martin Reinecke committed
40
@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
41 42
@pmp("nthreads", (0,1,2))
def test_fftn(shp, nthreads):
Martin Reinecke's avatar
Martin Reinecke committed
43
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
44
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a, nthreads=nthreads))<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
45
    a=a.astype(np.complex64)
46
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a, nthreads=nthreads))<5e-7)
Martin Reinecke's avatar
Martin Reinecke committed
47

Martin Reinecke's avatar
more  
Martin Reinecke committed
48 49 50 51
@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
52
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes))<1e-15)
Martin Reinecke's avatar
more  
Martin Reinecke committed
53
    a=a.astype(np.complex64)
54
    assert_(_l2error(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes))<5e-7)
Martin Reinecke's avatar
more  
Martin Reinecke committed
55 56 57 58

@pmp("shp", shapes)
def test_rfftn(shp):
    a=np.random.rand(*shp)-0.5
59
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a))<1e-15)
Martin Reinecke's avatar
more  
Martin Reinecke committed
60
    a=a.astype(np.float32)
61
    assert_(_l2error(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a))<5e-7)
62 63 64 65 66

@pmp("shp", shapes)
def test_rfft_scipy(shp):
    for i in range(len(shp)):
        a=np.random.rand(*shp)-0.5
67 68
        assert_(_l2error(pyfftw.interfaces.scipy_fftpack.rfft(a,axis=i), pypocketfft.rfft_scipy(a,axis=i))<1e-15)
        assert_(_l2error(pyfftw.interfaces.scipy_fftpack.irfft(a,axis=i), pypocketfft.irfft_scipy(a,axis=i,inorm=2))<1e-15)
Martin Reinecke's avatar
more  
Martin Reinecke committed
69 70 71 72 73

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

Martin Reinecke's avatar
Martin Reinecke committed
78 79 80
@pmp("shp", shapes)
def test_identity(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
81
    assert_(_l2error(pypocketfft.ifftn(pypocketfft.fftn(a),inorm=2), a)<1.5e-15)
82 83
    assert_(_l2error(pypocketfft.ifftn(pypocketfft.fftn(a.real),inorm=2), a.real)<1.5e-15)
    assert_(_l2error(pypocketfft.fftn(pypocketfft.ifftn(a.real),inorm=2), a.real)<1.5e-15)
84
    tmp=a.copy()
85 86
    assert_ (pypocketfft.ifftn(pypocketfft.fftn(tmp, inplace=True), inorm=2, inplace=True) is tmp)
    assert_(_l2error(tmp, a)<1.5e-15)
Martin Reinecke's avatar
Martin Reinecke committed
87
    a=a.astype(np.complex64)
88
    assert_(_l2error(pypocketfft.ifftn(pypocketfft.fftn(a),inorm=2), a)<6e-7)
Martin Reinecke's avatar
Martin Reinecke committed
89 90

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
91
def test_identity_r(shp):
Martin Reinecke's avatar
Martin Reinecke committed
92 93 94 95
    a=np.random.rand(*shp)-0.5
    b=a.astype(np.float32)
    for ax in range(a.ndim):
        n = a.shape[ax]
96 97
        assert_(_l2error(pypocketfft.irfftn(pypocketfft.rfftn(a,(ax,)),(ax,),lastsize=n,inorm=2), a)<1e-15)
        assert_(_l2error(pypocketfft.irfftn(pypocketfft.rfftn(b,(ax,)),(ax,),lastsize=n,inorm=2), b)<5e-7)
Martin Reinecke's avatar
more  
Martin Reinecke committed
98

Martin Reinecke's avatar
Martin Reinecke committed
99 100 101 102
@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))
103
    assert_(_l2error(pypocketfft.rfftn(pypocketfft.irfftn(a),inorm=2), a)<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
104

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
105 106
@pmp("shp", shapes2D+shapes3D)
def test_hartley2(shp):
Martin Reinecke's avatar
more  
Martin Reinecke committed
107
    a=np.random.rand(*shp)-0.5
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
108
    v1 = pypocketfft.hartley2(a)
Martin Reinecke's avatar
more  
Martin Reinecke committed
109 110
    v2 = pypocketfft.fftn(a.astype(np.complex128))
    v2 = v2.real+v2.imag
111
    assert_(_l2error(v1, v2)<1e-15)
Martin Reinecke's avatar
more  
Martin Reinecke committed
112

Martin Reinecke's avatar
fixes  
Martin Reinecke committed
113
@pmp("shp", shapes)
Martin Reinecke's avatar
more  
Martin Reinecke committed
114 115 116
def test_hartley_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley(pypocketfft.hartley(a))/a.size
117
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
fixes  
Martin Reinecke committed
118 119 120 121 122

@pmp("shp", shapes)
def test_hartley2_identity(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley2(pypocketfft.hartley2(a))/a.size
123
    assert_(_l2error(a, v1)<1e-15)
Martin Reinecke's avatar
Martin Reinecke committed
124
    v1 = a.copy()
125 126
    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
127 128 129 130 131

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