test.py

import pypocketfft
import pyfftw
import numpy as np
import pytest
from numpy.testing import assert_, assert_array_less

pmp = pytest.mark.parametrize

shapes1D = ((10,), (127,))
shapes2D = ((128, 128), (128, 129), (1, 129), (129, 1))
shapes3D = ((32, 17, 39),)
shapes = shapes1D+shapes2D+shapes3D
len1D = range(1, 2048)

ftol=6e-7
dtol=1.5e-15
ldtol=1e-18
tol = {np.float32: 6e-7, np.complex64: 6e-7,
       np.float64: 1.5e-15, np.complex128: 1.5e-15,
       np.longdouble: 1e-18, np.longcomplex: 1e-18}


def _l2error(a, b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))


def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
    return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
                           out=out, nthreads=nthreads)


def ifftn(a, axes=None, inorm=0, out=None, nthreads=1):
    return pypocketfft.c2c(a, axes=axes, forward=False, inorm=inorm,
                           out=out, 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, out=None, nthreads=1):
    return pypocketfft.r2r_fftpack(a, axes=(axis,), real2hermitian=True,
                                   forward=True, inorm=inorm, out=out,
                                   nthreads=nthreads)


def irfft_scipy(a, axis, inorm=0, out=None, nthreads=1):
    return pypocketfft.r2r_fftpack(a, axes=(axis,), real2hermitian=False,
                                   forward=False, inorm=inorm, out=out,
                                   nthreads=nthreads)


def chk(v1, v2):
    eps = tol[v1.dtype.type]
    err =_l2error(v1,v2)
    if err > eps:
        assert_array_less(err, eps)

@pmp("len", len1D)
@pmp("inorm", [0, 1, 2])
def test1D(len, inorm):
    a = np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
    b = a.astype(np.complex64)
    c = a.astype(np.longcomplex)
    chk(c, ifftn(fftn(c, inorm=inorm), inorm=2-inorm))
    chk(a, ifftn(fftn(a, inorm=inorm), inorm=2-inorm))
    chk(a.real, ifftn(fftn(a.real, inorm=inorm), inorm=2-inorm))
    chk(a.real, fftn(ifftn(a.real, inorm=inorm), inorm=2-inorm))
    chk(a.real, irfftn(rfftn(a.real, inorm=inorm), inorm=2-inorm,
                       lastsize=len))
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
    chk(tmp, a)
    chk(b, ifftn(fftn(b, inorm=inorm), inorm=2-inorm))
    chk(b.real, ifftn(fftn(b.real, inorm=inorm), inorm=2-inorm))
    chk(b.real, fftn(ifftn(b.real, inorm=inorm), inorm=2-inorm))
    chk(b.real, irfftn(rfftn(b.real, inorm=inorm), lastsize=len,
                       inorm=2-inorm))
    tmp = b.copy()
    assert_(ifftn(fftn(tmp, out=tmp, inorm=inorm), out=tmp, inorm=2-inorm)
            is tmp)
    chk(tmp, b)


@pmp("shp", shapes)
@pmp("nthreads", (0, 1, 2))
def test_fftn(shp, nthreads):
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    chk(pyfftw.interfaces.numpy_fft.fftn(a), fftn(a, nthreads=nthreads))
    a = a.astype(np.complex64)
    chk(pyfftw.interfaces.numpy_fft.fftn(a), fftn(a, nthreads=nthreads))


@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
    chk(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), fftn(a, axes=axes))
    a = a.astype(np.complex64)
    chk(pyfftw.interfaces.numpy_fft.fftn(a, axes=axes), fftn(a, axes=axes))


@pmp("shp", shapes)
def test_rfftn(shp):
    a = np.random.rand(*shp)-0.5
    chk(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a))
    a = a.astype(np.float32)
    chk(pyfftw.interfaces.numpy_fft.rfftn(a), rfftn(a))


@pmp("shp", shapes)
def test_rfft_scipy(shp):
    for i in range(len(shp)):
        a = np.random.rand(*shp)-0.5
        chk(pyfftw.interfaces.scipy_fftpack.rfft(a, axis=i),
            rfft_scipy(a, axis=i))
        chk(pyfftw.interfaces.scipy_fftpack.irfft(a, axis=i),
            irfft_scipy(a, axis=i, inorm=2))


@pmp("shp", shapes2D)
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
def test_rfftn2D(shp, axes):
    a = np.random.rand(*shp)-0.5
    chk(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), rfftn(a, axes=axes))
    a = a.astype(np.float32)
    chk(pyfftw.interfaces.numpy_fft.rfftn(a, axes=axes), rfftn(a, axes=axes))


@pmp("shp", shapes)
def test_identity(shp):
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    chk(ifftn(fftn(a), inorm=2), a)
    chk(ifftn(fftn(a.real), inorm=2), a.real)
    chk(fftn(ifftn(a.real), inorm=2), a.real)
    tmp = a.copy()
    assert_(ifftn(fftn(tmp, out=tmp), inorm=2, out=tmp) is tmp)
    chk(tmp, a)
    a = a.astype(np.complex64)
    chk(ifftn(fftn(a), inorm=2), a)


@pmp("shp", shapes)
def test_identity_r(shp):
    a = np.random.rand(*shp)-0.5
    b = a.astype(np.float32)
    for ax in range(a.ndim):
        n = a.shape[ax]
        chk(irfftn(rfftn(a, (ax,)), (ax,), lastsize=n, inorm=2), a)
        chk(irfftn(rfftn(b, (ax,)), (ax,), lastsize=n, inorm=2), b)


@pmp("shp", shapes)
def test_identity_r2(shp):
    a = np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    a = rfftn(irfftn(a))
    chk(rfftn(irfftn(a), inorm=2), a)


@pmp("shp", shapes2D+shapes3D)
def test_genuine_hartley(shp):
    a = np.random.rand(*shp)-0.5
    v1 = pypocketfft.genuine_hartley(a)
    v2 = fftn(a.astype(np.complex128))
    v2 = v2.real+v2.imag
    chk(v1, v2)


@pmp("shp", shapes)
def test_hartley_identity(shp):
    a = np.random.rand(*shp)-0.5
    v1 = pypocketfft.separable_hartley(pypocketfft.separable_hartley(a))/a.size
    chk(a, v1)


@pmp("shp", shapes)
def test_genuine_hartley_identity(shp):
    a = np.random.rand(*shp)-0.5
    v1 = pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(a))/a.size
    chk(a, v1)
    v1 = a.copy()
    assert_(pypocketfft.genuine_hartley(
        pypocketfft.genuine_hartley(v1, out=v1), inorm=2, out=v1) is v1)
    chk(a, v1)


@pmp("shp", shapes2D)
@pmp("axes", ((0,), (1,), (0, 1), (1, 0)))
def test_genuine_hartley_2D(shp, axes):
    a = np.random.rand(*shp)-0.5
    chk(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(
        a, axes=axes), axes=axes, inorm=2), a)