def create_random_dataset(DS):
pprint('Populating random data on disk', rank=0)
if mpi_root():
f = tempfile.NamedTemporaryFile(dir=TMP)
data = np.random.rand(DS, DS, DS).astype(np.float32)
fname = f.name + '.h5'
with h5.File(fname, 'w') as g:
g.create_dataset('data', data=data)
else:
fname = None
fname = mpi_comm().bcast(fname, root=0)
pprint(fname)
h5in = h5.File(fname, 'r')
h5data = h5in['data']
return h5in, h5data
def convolve1(ds, sigma):
mpi_barrier()
pprint('Convolving in 3D', rank=0)
with MpiTimer('Separable 3D convolution') as T:
ds3d_ = ds.clone('gaussian_3d')
for z in range(mpi_rank(), ds.chunk_grid[0], mpi_size()):
zS = z * ds.chunk_size[0]
zE = min(zS + ds.chunk_size[0], ds.shape[0])
for y in range(ds.chunk_grid[1]):
yS = y * ds.chunk_size[1]
yE = min(yS + ds.chunk_size[1], ds.shape[1])
for x in range(ds.chunk_grid[2]):
xS = x * ds.chunk_size[2]
xE = min(xS + ds.chunk_size[2], ds.shape[2])
out = gaussian_filter(ds[zS:zE, yS:yE, xS:xE], sigma=sigma)
ds3d_[zS:zE, yS:yE, xS:xE] = out
mpi_barrier()
if mpi_root():
try:
imsave(
join(
OUT_PATH,
'conv3d_{}-{}-{}.png'.format(mpi_size(), ds.shape[0],
ds.chunk_size[0])),
ds3d_[ds.shape[0] // 2])
except NameError:
pass
mpi_barrier()
ds3d_.delete()
return T.time
pprint)
log.getLogger().setLevel(log.INFO)
dn.use(engine='mpi', backend='hdf5')
###############################################################################
DS = 256
CS = 50
engine, backend = dn.status()
pprint('Loading data', rank=0)
if mpi_root():
f = tempfile.NamedTemporaryFile()
data = np.random.rand(DS, DS, DS).astype(np.float32)
with h5.File(f.name, 'w') as g:
g.create_dataset('data', data=data)
fname = f.name
else:
fname = None
fname = mpi_comm().bcast(fname, root=0)
h5in = h5.File(fname, 'r')
h5data = h5in['data']
###############################################################################
with MpiTimer('DosNa %s (%s)' % (engine.name, backend.name)):
def convolve2(ds, sigma):
mpi_barrier()
with MpiTimer('Three separable 1D convolutions') as T:
pprint('Convolving axis Z', rank=0)
ds1_ = ds.clone('gaussian_z')
for z in range(mpi_rank(), ds.chunk_grid[0], mpi_size()):
zS = z * ds.chunk_size[0]
zE = min(zS + ds.chunk_size[0], ds.shape[0])
out = gaussian_filter1d(ds[zS:zE], sigma=sigma, axis=0)
ds1_[zS:zE] = out
mpi_barrier()
# Gaussian second axis
pprint('Convolving axis Y', rank=0)
ds2_ = ds.clone('gaussian_y')
for y in range(mpi_rank(), ds.chunk_grid[1], mpi_size()):
yS = y * ds.chunk_size[1]
yE = min(yS + ds.chunk_size[1], ds.shape[1])
out = gaussian_filter1d(ds1_[:, yS:yE], sigma=sigma, axis=1)
ds2_[:, yS:yE] = out
mpi_barrier()
# Gaussian second axis
pprint('Convolving axis X', rank=0)
ds3_ = ds.clone('gaussian_x')
for x in range(mpi_rank(), ds.chunk_grid[2], mpi_size()):
xS = x * ds.chunk_size[2]
xE = min(xS + ds.chunk_size[2], ds.shape[2])
out = gaussian_filter1d(ds2_[..., xS:xE], sigma=sigma, axis=2)
ds3_[..., xS:xE] = out
mpi_barrier()
if mpi_root():
try:
imsave(
join(
OUT_PATH,
'conv3x1d_{}-{}-{}.png'.format(mpi_size(), ds.shape[0],
ds.chunk_size[0])),
ds3_[ds.shape[0] // 2])
except NameError:
pass
mpi_barrier()
ds1_.delete()
ds2_.delete()
ds3_.delete()
return T.time