def test_FFT(FFT):
if FFT.rank == 0:
A = random((N, N, N)).astype(FFT.float)
if hasattr(FFT, 'params'):
if FFT.params['method'] == 'Nyquist':
C = rfftn(A, axes=(0, 1, 2))
C[:, :, -1] = 0 # Remove Nyquist frequency
A[:] = irfftn(C, axes=(0, 1, 2))
B2 = zeros((N, N, N / 2 + 1), dtype=FFT.complex)
B2[:] = rfftn(A, axes=(0, 1, 2))
else:
A = zeros((N, N, N), dtype=FFT.float)
B2 = zeros((N, N, N / 2 + 1), dtype=FFT.complex)
atol, rtol = (1e-10, 1e-8) if FFT.float is float64 else (5e-7, 1e-5)
FFT.comm.Bcast(A, root=0)
FFT.comm.Bcast(B2, root=0)
a = zeros(FFT.real_shape(), dtype=FFT.float)
c = zeros(FFT.complex_shape(), dtype=FFT.complex)
a[:] = A[FFT.real_local_slice()]
c = FFT.fftn(a, c)
assert all(abs((c - B2[FFT.complex_local_slice()]) / c.max()) < rtol)
#assert allclose(c, B2[FFT.complex_local_slice()], rtol, atol)
a = FFT.ifftn(c, a)
assert all(abs((a - A[FFT.real_local_slice()]) / a.max()) < rtol)
def test_FFT(FFT):
N = FFT.N
if FFT.rank == 0:
A = random(N).astype(FFT.float)
if FFT.communication == 'AlltoallN':
C = empty(FFT.global_complex_shape(), dtype=FFT.complex)
C = rfftn(A, C, axes=(0, 1, 2))
C[:, :, -1] = 0 # Remove Nyquist frequency
A = irfftn(C, A, axes=(0, 1, 2))
B2 = zeros(FFT.global_complex_shape(), dtype=FFT.complex)
B2 = rfftn(A, B2, axes=(0, 1, 2))
else:
A = zeros(N, dtype=FFT.float)
B2 = zeros(FFT.global_complex_shape(), dtype=FFT.complex)
atol, rtol = (1e-10, 1e-8) if FFT.float is float64 else (5e-7, 1e-4)
FFT.comm.Bcast(A, root=0)
FFT.comm.Bcast(B2, root=0)
a = zeros(FFT.real_shape(), dtype=FFT.float)
c = zeros(FFT.complex_shape(), dtype=FFT.complex)
a[:] = A[FFT.real_local_slice()]
c = FFT.fftn(a, c)
#print abs((c - B2[FFT.complex_local_slice()])/c.max()).max()
assert all(abs((c - B2[FFT.complex_local_slice()]) / c.max()) < rtol)
#assert allclose(c, B2[FFT.complex_local_slice()], rtol, atol)
a = FFT.ifftn(c, a)
#print abs((a - A[FFT.real_local_slice()])/a.max()).max()
assert all(abs((a - A[FFT.real_local_slice()]) / a.max()) < rtol)
def test_FFT(FFT):
if FFT.rank == 0:
A = random((N, N, N)).astype(FFT.float)
if hasattr(FFT, 'params'):
if FFT.params['method'] == 'Nyquist':
C = rfftn(A, axes=(0, 1, 2))
C[:, :, -1] = 0 # Remove Nyquist frequency
A[:] = irfftn(C, axes=(0, 1, 2))
else:
A = zeros((N, N, N), dtype=FFT.float)
FFT.comm.Bcast(A, root=0)
a = zeros(FFT.real_shape(), dtype=FFT.float)
c = zeros(FFT.complex_shape(), dtype=FFT.complex)
a[:] = A[FFT.real_local_slice()]
c = FFT.fftn(a, c)
B2 = rfftn(A, axes=(0, 1, 2))
assert allclose(c, B2[FFT.complex_local_slice()])
a = FFT.ifftn(c, a)
assert allclose(a, A[FFT.real_local_slice()], 5e-7, 5e-7)
def test_FFT_padded(FFT_padded):
FFT = FFT_padded
N = FFT.N
if FFT.rank == 0:
A = random(N).astype(FFT.float)
C = zeros((FFT.global_complex_shape()), dtype=FFT.complex)
C[:] = rfftn(A, axes=(0, 1, 2))
# Eliminate Nyquist, otherwise test will fail
C[-N / 2] = 0
C[:, -N / 2] = 0
A[:] = irfftn(C)
Cp = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 4 + 1),
dtype=FFT.complex)
Nf = N[2] / 2 + 1
ks = (fftfreq(N[1]) * N[1]).astype(int)
Cp[:N[0] / 2, ks, :Nf] = C[:N[0] / 2]
Cp[-N[0] / 2:, ks, :Nf] = C[N[0] / 2:]
# If Nyquist is retained then these are needed to symmetrize and pass test
###Cp[:, :, Nf-1] *= 0.5
#Cp[:, -N/2] *= 0.5
#Cp[-N/2] *= 0.5
#Cp[N/2] = Cp[-N/2]
#Cp[:, N/2] = Cp[:, -N/2]
Ap = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 2), dtype=FFT.float)
Ap[:] = irfftn(Cp * 1.5**3, axes=(0, 1, 2))
else:
C = zeros(FFT.global_complex_shape(), dtype=FFT.complex)
Ap = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 2), dtype=FFT.float)
A = zeros(N, dtype=FFT.float)
FFT.comm.Bcast(C, root=0)
FFT.comm.Bcast(Ap, root=0)
FFT.comm.Bcast(A, root=0)
ae = zeros(FFT.real_shape_padded(), dtype=FFT.float)
c = zeros(FFT.complex_shape(), dtype=FFT.complex)
c[:] = C[FFT.complex_local_slice()]
ae[:] = Ap[FFT.real_local_slice(padded=True)]
ap = zeros(FFT.real_shape_padded(), dtype=FFT.float)
cp = zeros(FFT.complex_shape(), dtype=FFT.complex)
ap = FFT.ifftn(c, ap, dealias="3/2-rule")
atol, rtol = (1e-10, 1e-8) if FFT.float is float64 else (5e-7, 1e-4)
assert allclose(ap, ae, rtol, atol)
def test_FFT_padded(FFT_padded):
FFT = FFT_padded
N = FFT.N
if FFT.rank == 0:
A = random(N).astype(FFT.float)
C = zeros((FFT.global_complex_shape()), dtype=FFT.complex)
C[:] = rfftn(A, axes=(0, 1, 2))
# Eliminate Nyquist, otherwise test will fail
C[-N[0] / 2] = 0
C[:, -N[1] / 2] = 0
if hasattr(FFT, 'params'):
if FFT.params['method'] == 'Nyquist':
C[:, :, -1] = 0 # Remove Nyquist frequency
A[:] = irfftn(C)
Cp = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 4 + 1),
dtype=FFT.complex)
Nf = N[2] / 2 + 1
ks = (fftfreq(N[1], 1. / N[1])).astype(int)
Cp[:N[0] / 2, ks, :Nf] = C[:N[0] / 2]
Cp[-N[0] / 2:, ks, :Nf] = C[N[0] / 2:]
# If Nyquist is retained then these are needed to symmetrize and pass test
Cp[-N[0] / 2] *= 0.5
Cp[:, -N[1] / 2] *= 0.5
Cp[N[0] / 2] = Cp[-N[0] / 2]
Cp[:, N[1] / 2] = Cp[:, -N[1] / 2]
Ap = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 2), dtype=FFT.float)
Ap[:] = irfftn(Cp, axes=(0, 1, 2)) * 1.5**3
else:
C = zeros(FFT.global_complex_shape(), dtype=FFT.complex)
Ap = zeros((3 * N[0] / 2, 3 * N[1] / 2, 3 * N[2] / 2), dtype=FFT.float)
A = zeros(N, dtype=FFT.float)
FFT.comm.Bcast(C, root=0)
FFT.comm.Bcast(Ap, root=0)
FFT.comm.Bcast(A, root=0)
ae = zeros(FFT.real_shape_padded(), dtype=FFT.float)
c = zeros(FFT.complex_shape(), dtype=FFT.complex)
c[:] = C[FFT.complex_local_slice()]
ae[:] = Ap[FFT.real_local_slice(padded=True)]
ap = zeros(FFT.real_shape_padded(), dtype=FFT.float)
cp = zeros(FFT.complex_shape(), dtype=FFT.complex)
ap = FFT.ifftn(c, ap, dealias="3/2-rule")
atol, rtol = (1e-10, 1e-8) if FFT.float is float64 else (5e-7, 1e-4)
#print np.linalg.norm(ap-ae)
assert allclose(ap, ae, rtol, atol)
cp = FFT.fftn(ap, cp, dealias="3/2-rule")
#from IPython import embed; embed()
#print np.linalg.norm(abs((cp-c)/cp.max()))
assert all(abs((cp - c) / cp.max()) < rtol)