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

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
14
15
16
17
18

@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)))
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a), atol=2e-15*vmax, rtol=0)
Martin Reinecke's avatar
Martin Reinecke committed
19
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a,norm="ortho"), pypocketfft.fftn(a,norm="ortho"), atol=2e-15*vmax, rtol=0)
Martin Reinecke's avatar
Martin Reinecke committed
20
21
22
    a=a.astype(np.complex64)
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a), atol=6e-7*vmax, rtol=0)

Martin Reinecke's avatar
more    
Martin Reinecke committed
23
24
25
26
27
28
29
30
31
32
33
34
35
36
@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)))
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes), atol=2e-15*vmax, rtol=0)
    a=a.astype(np.complex64)
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), pypocketfft.fftn(a, axes=axes), atol=6e-7*vmax, rtol=0)

@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)))
    assert_allclose(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a), atol=2e-15*vmax, rtol=0)
Martin Reinecke's avatar
Martin Reinecke committed
37
    assert_allclose(pyfftw.interfaces.numpy_fft.rfftn(a,norm="ortho"), pypocketfft.rfftn(a,norm="ortho"), atol=2e-15*vmax, rtol=0)
Martin Reinecke's avatar
more    
Martin Reinecke committed
38
39
40
41
42
43
44
45
46
47
48
49
    a=a.astype(np.float32)
    assert_allclose(pyfftw.interfaces.numpy_fft.rfftn(a), pypocketfft.rfftn(a), atol=6e-7*vmax, rtol=0)

@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)))
    assert_allclose(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), pypocketfft.rfftn(a, axes=axes), atol=2e-15*vmax, rtol=0)
    a=a.astype(np.float32)
    assert_allclose(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), pypocketfft.rfftn(a, axes=axes), atol=6e-7*vmax, rtol=0)

Martin Reinecke's avatar
Martin Reinecke committed
50
51
52
53
@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
54
55
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a)), a, atol=2e-15*vmax, rtol=0)
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a,norm="ortho"), norm="ortho"), a, atol=2e-15*vmax, rtol=0)
Martin Reinecke's avatar
Martin Reinecke committed
56
    a=a.astype(np.complex64)
Martin Reinecke's avatar
Martin Reinecke committed
57
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a)), a, atol=6e-7*vmax, rtol=0)
Martin Reinecke's avatar
Martin Reinecke committed
58
59

@pmp("shp", shapes)
Martin Reinecke's avatar
Martin Reinecke committed
60
def test_identity_r(shp):
Martin Reinecke's avatar
Martin Reinecke committed
61
62
63
64
65
    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
66
67
68
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(a,(ax,)),(ax,),n), a, atol=2e-15*vmax, rtol=0)
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(a,(ax,),norm="ortho"),(ax,),n,norm="ortho"), a, atol=2e-15*vmax, rtol=0)
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(b,(ax,)),(ax,),n), b, atol=6e-7*vmax, rtol=0)
Martin Reinecke's avatar
more    
Martin Reinecke committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

@pmp("shp", shapes3D)
def xtest_hartley(shp):
    a=np.random.rand(*shp)-0.5
    v1 = pypocketfft.hartley(a)
    v2 = pypocketfft.fftn(a.astype(np.complex128))
    vmax = np.max(np.abs(v1))
    v2 = v2.real+v2.imag
    assert_allclose(v1, v2, atol=2e-15*vmax, rtol=0)
@pmp("shp", shapes3D)

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))
    assert_allclose(a, v1, atol=2e-15*vmax, rtol=0)