test.py 1.33 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import pypocketfft
import pyfftw
import numpy as np
import pytest
from numpy.testing import assert_allclose

pmp = pytest.mark.parametrize

shapes = ((10,), (127,), (128, 128), (128, 129), (1, 129), (129,1), (32,17,39))

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

@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))
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a))/a.size, a, atol=2e-15*vmax, rtol=0)
    a=a.astype(np.complex64)
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a))/a.size, a, atol=6e-7*vmax, rtol=0)

@pmp("shp", shapes)
def test_identity_r(shp):
    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]
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(a,ax),ax)/n, a, atol=2e-15*vmax, rtol=0)
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(b,ax),ax)/n, b, atol=6e-7*vmax, rtol=0)