In [1]:
import numpy as np
import subprocess
%matplotlib nbagg
import matplotlib.pyplot as plt
import pyfftw
from mayavi import mlab
In [2]:
def generate_data_3D(
n,
dtype = np.complex128,
p = 1.5):
"""
generate something that has the proper shape
"""
assert(n % 2 == 0)
a = np.zeros((n, n, n/2+1), dtype = dtype)
a[:] = np.random.randn(*a.shape) + 1j*np.random.randn(*a.shape)
k, j, i = np.mgrid[-n/2+1:n/2+1, -n/2+1:n/2+1, 0:n/2+1]
k = (k**2 + j**2 + i**2)**.5
k = np.roll(k, n//2+1, axis = 0)
k = np.roll(k, n//2+1, axis = 1)
a /= k**p
a[0, :, :] = 0
a[:, 0, :] = 0
a[:, :, 0] = 0
ii = np.where(k == 0)
a[ii] = 0
ii = np.where(k > n/3)
a[ii] = 0
return a
n = 13*4
N = 128
Kdata0 = generate_data_3D(n, p = 2).astype(np.complex64)
Kdata1 = generate_data_3D(n, p = 2).astype(np.complex64)
Kdata1 = Kdata0.copy()
Kdata0.T.copy().tofile("Kdata0")
Kdata1.T.copy().tofile("Kdata1")
In [3]:
def padd_with_zeros(
a,
n,
odtype = None):
if (type(odtype) == type(None)):
odtype = a.dtype
assert(a.shape[0] <= n)
b = np.zeros((n, n, n/2 + 1), dtype = odtype)
m = a.shape[0]
b[ :m/2, :m/2, :m/2+1] = a[ :m/2, :m/2, :m/2+1]
b[ :m/2, n-m/2: , :m/2+1] = a[ :m/2, m-m/2: , :m/2+1]
b[n-m/2: , :m/2, :m/2+1] = a[m-m/2: , :m/2, :m/2+1]
b[n-m/2: , n-m/2: , :m/2+1] = a[m-m/2: , m-m/2: , :m/2+1]
return b
def transform_py(bla):
b = padd_with_zeros(bla, N)
c = np.zeros((N, N, N), np.float32)
t = pyfftw.FFTW(
b, c,
axes = (0, 1, 2),
direction = 'FFTW_BACKWARD',
flags = ('FFTW_ESTIMATE',),
threads = 2)
t.execute()
return c
import chichi_misc as cm
def array_to_8cubes(
a,
odtype = None):
assert(len(a.shape) >= 3)
assert((a.shape[0] % 8 == 0) and
(a.shape[1] % 8 == 0) and
(a.shape[2] % 8 == 0))
if type(odtype) == type(None):
odtype = a.dtype
c = np.zeros(
((((a.shape[0] // 8)*(a.shape[1] // 8)*(a.shape[2] // 8)),) +
(8, 8, 8) +
a.shape[3:]),
dtype = odtype)
bk, bj, bi = np.mgrid[0:a.shape[0]//8,
0:a.shape[1]//8,
0:a.shape[2]//8]
bindices = np.array([bi, bj, bk])
cindices = cm.grid3D_to_zindex(bindices)
plist = []
zlist = []
for k in range(a.shape[0]//8):
for j in range(a.shape[1]//8):
for i in range(a.shape[2]//8):
z = cm.grid3D_to_zindex(np.array([k, j, i]))
c[z] = a[8*k:8*(k+1), 8*j:8*(j+1), 8*i:8*(i+1)]
plist.append([i, j, k])
zlist.append(z)
plist = np.array(plist)
zlist = np.array(zlist)
i = np.argsort(zlist)
return c, zlist[i], plist[i]
d0 = transform_py(Kdata0)
d1 = transform_py(Kdata1)
Rdata_py_tmp = np.array([d0, d1]).transpose((1, 2, 3, 0))
Rdata_py, zlist, plist = array_to_8cubes(Rdata_py_tmp)
# i0 = np.random.randint(16)
# i1 = np.random.randint(16)
# i2 = np.random.randint(16)
# z = cm.grid3D_to_zindex(np.array([i0, i1, i2]))
# mlab.contour3d(Rdata_py[z, :, :, :, 0], color = (1, 0, 0))
# mlab.contour3d(Rdata_py_tmp[i0*8:i0*8 + 8, i1*8:i1*8+8, i2*8:i2*8+8, 0],
# color = (0, 0, 1))
# mlab.show()
In [23]:
def get_cpp_data(branch = None):
if not (type(branch) == type(None)):
subprocess.call(['git', 'checkout', branch])
subprocess.call(['rm', 'Rdata_z0000000', 'Rdata_z0000800'])
if subprocess.call(['make', 'full.elf']) == 0:
subprocess.call(['time',
'mpirun.mpich',
'-np',
'8',
'./full.elf',
'{0}'.format(n/2+1),
'{0}'.format(n),
'{0}'.format(n),
'{0}'.format(N),
'{0}'.format(N),
'{0}'.format(N)])
Rdata0 = np.fromfile(
"Rdata_z0000000",
dtype = np.float32).reshape(-1, 8, 8, 8, 2)
Rdata1 = np.fromfile(
"Rdata_z0000800",
dtype = np.float32).reshape(-1, 8, 8, 8, 2)
return np.concatenate([Rdata0, Rdata1])
else:
print ('compilation error')
return None
#Rdata = get_cpp_data(branch = 'develop')
# develop says 30 secs, inplace fft is 28 secs
#Rdata = get_cpp_data(branch = 'feature-inplace_fft')
Rdata = get_cpp_data()
In [15]:
distance = np.max(np.abs(Rdata_py - Rdata))
print distance
if distance > 1e-5:
i0 = np.random.randint(16)
i1 = np.random.randint(16)
i2 = np.random.randint(16)
z = cm.grid3D_to_zindex(np.array([i0, i1, i2]))
#z = 0
print cm.zindex_to_grid3D(z)
s = np.max(np.abs(Rdata_py[None, z, :, :, :, 0] - Rdata[..., 0]),
axis = (1, 2, 3))
z1 = np.argmin(s)
print z, z1, s[z1]
#print Rdata[z1]
In [16]:
# i0 = np.random.randint(16)
# i1 = np.random.randint(16)
# i2 = np.random.randint(16)
i0 = 6
i1 = 11
i2 = 15
z = cm.grid3D_to_zindex(np.array([i0, i1, i2]))
mlab.contour3d(Rdata_py[z, :, :, :, 0],
color = (0, 1, 0))
# mlab.contour3d(Rdata_py_tmp[i0*8:i0*8 + 8,
# i1*8:i1*8 + 8,
# i2*8:i2*8 + 8, 0],
# color = (0, 0, 1))
mlab.contour3d(Rdata[z, :, :, :, 0], color = (1, 0, 0))
mlab.show()
In [ ]: