def run_test_c2c_impl(self, shape, axes, inverse=False, fftshift=False):
shape = list(shape)
shape[-1] *= 2 # For complex
known_data = np.random.normal(size=shape).astype(np.float32).view(np.complex64)
idata = bf.ndarray(known_data, space='cuda')
odata = bf.empty_like(idata)
fft = Fft()
fft.init(idata, odata, axes=axes, apply_fftshift=fftshift)
fft.execute(idata, odata, inverse)
if inverse:
if fftshift:
known_data = np.fft.ifftshift(known_data, axes=axes)
# Note: Numpy applies normalization while CUFFT does not
norm = reduce(lambda a, b: a * b, [known_data.shape[d]
for d in axes])
known_result = gold_ifftn(known_data, axes=axes) * norm
else:
known_result = gold_fftn(known_data, axes=axes)
if fftshift:
known_result = np.fft.fftshift(known_result, axes=axes)
x = (np.abs(odata.copy('system') - known_result) / known_result > RTOL).astype(np.int32)
a = odata.copy('system')
b = known_result
compare(odata.copy('system'), known_result)
python类ifftn()的实例源码
def correlate(x, y):
from numpy import fft
sx = numpy.array(x.shape)
sy = numpy.array(y.shape)
if (sx >= sy).sum():
slices = [slice(None, sx[i] - sy[i] + 1) for i in range(len(sx))]
X = fft.fftn(x)
Y = fft.fftn(zerofill(y, sx))
else:
sf = sx + sy - 1
slices = [slice(None, sf[i]) for i in range(len(sf))]
X = fft.fftn(x, sf)
Y = fft.fftn(zerofill(y, sf), sf)
return fft.ifftn(X.conjugate() * Y)[slices].real
def kfilter(ar, kf, lx=2*pi, ly=2*pi, lz=2*pi):
"""
Function to filter a 3D array by zeroing out larger k values
"""
while len(ar.shape) < 3:
ar = ar.reshape(ar.shape + (1,))
# COMPUTE THE ARRAY SIZE
kx, ky, kz, km = create_kgrid(*ar.shape, lx=lx, ly=ly, lz=lz)
# FOURIER TRANSFORM THE ARRAY
far = nf.fftshift(nf.fftn(ar))
# SET VALUES ABOVE kf AS 0+0i
far = (np.sign(km - kf) - 1.)/(-2.)*far
# BACK TRANSFORM TO REAL SPACE
arf = np.real(nf.ifftn(nf.ifftshift(far)))
return arf
#==================================================
def kdiv(arx,ary,arz,kf=None,lx=2*pi,ly=2*pi,lz=2*pi):
"""
Function to compute divergence in Fourier space
"""
# COMPUTE THE ARRAY SIZE
kx, ky, kz, km = create_kgrid(*np.shape(arx), lx=lx, ly=ly, lz=lz)
# FOURIER TRANSFORM THE ARRAY
far = [nf.fftshift(nf.fftn(a)) for a in (arx, ary, arz)]
# COMPUTE div=i*(kx*ax+ky*ay+kz*az)
mg = np.meshgrid(kx,ky,kz)
divarf = 1.j*reduce(operator.add, [a*b for a,b in zip(mg, far)])
# SET VALUES ABOVE kf AS 0+0i if kf non zero
if kf is not None:
divarf = (np.sign(km - kf) - 1.)/(-2.)*divarf
divar = np.real(nf.ifftn(nf.ifftshift(divarf)))
return divar
def kurl(arx,ary,arz, kf, lx=2*np.pi, ly=2*np.pi, lz=2*np.pi):
kx, ky, kz, km = create_kgrid(*arx.shape, lx=lx, ly=ly, lz=lz)
# FOURIER TRANSFORM THE ARRAY
farx = nf.fftshift(nf.fftn(arx))
fary = nf.fftshift(nf.fftn(ary))
farz = nf.fftshift(nf.fftn(arz))
# SET VALUES ABOVE kf AS 0+0i
farx = (np.sign(km - kf) - 1.)/(-2.)*farx
fary = (np.sign(km - kf) - 1.)/(-2.)*fary
farz = (np.sign(km - kf) - 1.)/(-2.)*farz
# COMPUTE VORTICITY
axf = eye*(ky*farz-kz*fary)
ayf = eye*(kz*farx-kx*farz)
azf = eye*(kx*fary-ky*farx)
# BACK TRANSFORM TO REAL SPACE
wx = np.real(nf.ifftn(nf.ifftshift(axf)))
wy = np.real(nf.ifftn(nf.ifftshift(ayf)))
wz = np.real(nf.ifftn(nf.ifftshift(azf)))
return wx,wy,wz
def convolve(x, f):
from numpy import fft, all
sx = numpy.array(x.shape)
sf = numpy.array(f.shape)
if not all(sx >= sf): return convolve(f, x)
y = fft.ifftn(fft.fftn(x) * fft.fftn(f, sx)).real
slices = [slice(sf[i] - 1, sx[i]) for i in range(len(sf))]
return y[slices]
def test_UnscaledFFT_3d(backend, M, N, K, B ):
b = backend()
# forward
x = b.rand_array( (M*N*K, B) )
y = b.rand_array( (M*N*K, B) )
x_h = x.to_host().reshape( (M,N,K,B), order='F' )
A = b.UnscaledFFT( (M,N,K), dtype=x.dtype )
A.eval(y, x)
y_exp = np.fft.fftn( x_h, axes=(0,1,2) )
y_act = y.to_host().reshape( (M,N,K,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
# adjoint
x = b.rand_array( (M*N*K, B) )
y = b.rand_array( (M*N*K, B) )
x_h = x.to_host().reshape( (M,N,K,B), order='F' )
A.H.eval(y, x)
y_exp = np.fft.ifftn( x_h, axes=(0,1,2) ) * (M*N*K)
y_act = y.to_host().reshape( (M,N,K,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
def test_UnscaledFFT_2d(backend, M, N, B ):
b = backend()
# forward
x = b.rand_array( (M*N, B) )
y = b.rand_array( (M*N, B) )
x_h = x.to_host().reshape( (M,N,B), order='F' )
A = b.UnscaledFFT( (M,N), dtype=x.dtype )
A.eval(y, x)
y_exp = np.fft.fftn( x_h, axes=(0,1) )
y_act = y.to_host().reshape( (M,N,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
# adjoint
x = b.rand_array( (M*N, B) )
y = b.rand_array( (M*N, B) )
x_h = x.to_host().reshape( (M,N,B), order='F' )
A.H.eval(y, x)
y_exp = np.fft.ifftn( x_h, axes=(0,1) ) * (M*N)
y_act = y.to_host().reshape( (M,N,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
def test_UnscaledFFT_1d(backend, M, B ):
b = backend()
# forward
x = b.rand_array( (M, B) )
y = b.rand_array( (M, B) )
x_h = x.to_host().reshape( (M,B), order='F' )
A = b.UnscaledFFT( (M,), dtype=x.dtype )
A.eval(y, x)
y_exp = np.fft.fftn( x_h, axes=(0,) )
y_act = y.to_host().reshape( (M,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
# adjoint
x = b.rand_array( (M, B) )
y = b.rand_array( (M, B) )
x_h = x.to_host().reshape( (M,B), order='F' )
A.H.eval(y, x)
y_exp = np.fft.ifftn( x_h, axes=(0,) ) * M
y_act = y.to_host().reshape( (M,B), order='F' )
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
def test_CenteredFFT(backend, M, N, K, B ):
from numpy.fft import fftshift, ifftshift, fftn, ifftn
b = backend()
A = b.FFTc( (M,N,K), dtype=np.dtype('complex64') )
# forward
ax = (0,1,2)
x = b.rand_array( (M*N*K,B) )
y = b.rand_array( (M*N*K,B) )
x_h = x.to_host().reshape( (M,N,K,B), order='F' )
A.eval(y, x)
y_act = y.to_host().reshape( (M,N,K,B), order='F' )
y_exp = fftshift( fftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
# adjoint
x = b.rand_array( (M*N*K,B) )
y = b.rand_array( (M*N*K,B) )
x_h = x.to_host().reshape( (M,N,K,B), order='F' )
A.H.eval(y, x)
y_act = y.to_host().reshape( (M,N,K,B), order='F' )
y_exp = fftshift( ifftn( ifftshift(x_h, axes=ax), axes=ax, norm='ortho'), axes=ax)
npt.assert_allclose(y_act, y_exp, rtol=1e-2)
def vecpot(arx,ary,arz, kf, lx=2*np.pi, ly=2*np.pi, lz=2*np.pi):
"""
Function to compute vector potential of a 3D array
"""
nx,ny,nz=arx.shape
# COMPUTE THE ARRAY SIZE
kx, ky, kz, km = create_kgrid(*arx.shape, lx=lx, ly=ly, lz=lz)
#kx, ky, kz, k2 = kx[:,nna,nna],ky[nna,:,nna],kz[nna,nna,:], km**2
k2=km**2
k2[nx/2,ny/2,nz/2]=1.
# FOURIER TRANSFORM THE ARRAY
farx = nf.fftshift(nf.fftn(arx))
fary = nf.fftshift(nf.fftn(ary))
farz = nf.fftshift(nf.fftn(arz))
# SET VALUES ABOVE kf AS 0+0i
farx = (np.sign(km - kf) - 1.)/(-2.)*farx
fary = (np.sign(km - kf) - 1.)/(-2.)*fary
farz = (np.sign(km - kf) - 1.)/(-2.)*farz
# FIND THE CORRESPONDING VECTOR POTENTIAL A = -ik x B /k^2
axf = -eye*(ky*farz-kz*fary)/k2
ayf = -eye*(kz*farx-kx*farz)/k2
azf = -eye*(kx*fary-ky*farx)/k2
# BACK TRANSFORM TO REAL SPACE
ax = np.real(nf.ifftn(nf.ifftshift(axf)))
ay = np.real(nf.ifftn(nf.ifftshift(ayf)))
az = np.real(nf.ifftn(nf.ifftshift(azf)))
return ax,ay,az