def operate(self, out):
arg = self.args[0]
axis0 = self.axis0
axis1 = self.axis1
arg.require_coeff_space()
out.layout = arg.layout
k0 = self.basis0.elements[self.slices[self.axis0]]
k1 = self.basis0.elements[self.slices[self.axis1]]
y = self.basis0.grid(scale=self.domain.dealias[self.axis0])
phi = np.einsum('ik,jk->ijk', np.sin(k0[:, None] * y),
np.sin(k1[:, None] * y))
di = np.einsum('ijk,jkl', arg.data, phi)
parity = arg.meta[axis0]['parity'] * arg.meta[axis1]['parity']
out.meta[axis0]['parity'] = parity
out.meta[axis1]['parity'] = parity
out.require_grid_space(axis=1)
out.data[:, :, :] = 0.
out.data[axslice(axis0, 0, 1)] = np.expand_dims(di, axis=1)
out.data[axslice(axis1, 0, 1)] = np.expand_dims(di, axis=2)
# out.data[axslice(axis0,1,None)] = 0
# out.data[axslice(axis0,None, None)] = out.data[axslice(axis0,0,1)]
# out.require_coeff_space(axis=axis1)
# out.data *= self.filter_mask
out.require_layout(arg.layout)
def tree_merge_distributed_sets(set_paths,
executor,
blocksize=2,
cleanup=False,
startlevel=0,
maxlevel=None):
set_paths = [pathlib.Path(sp) for sp in set_paths]
logger.info("Merging sets {}".format(set_paths))
if maxlevel is None:
maxlevel = np.inf
# Get process mesh
set_stem = set_paths[0].stem
proc_path = set_paths[0].joinpath(f"{set_stem}_p0.h5")
with h5py.File(str(proc_path), mode='r') as proc_file:
proc_dset = proc_file['tasks']['T']
global_shape = np.array(proc_dset.attrs['global_shape'])
local_shape = np.array(proc_dset.attrs['count'])
mesh = np.ceil(global_shape / local_shape).astype(int)
# Loop backwards over process mesh
procs = np.arange(np.prod(mesh)).reshape(mesh)
level = 0
for D in reversed(range(len(mesh))):
M = procs.shape[D]
# Recursively merge blocks
while M > 1:
if level >= maxlevel:
break
futures = []
chunks = int(np.ceil(M / blocksize))
proc_blocks = np.array_split(procs, chunks, axis=D)
proc_blocks = [
procs.reshape((-1, procs.shape[D])) for procs in proc_blocks
]
proc_blocks = [procs.tolist() for procs in proc_blocks]
proc_blocks = list(itertools.chain(*zip(*proc_blocks)))
if level >= startlevel:
# Loop over blocks within sets
for set_path in set_paths:
for proc_block in proc_blocks:
f = executor.submit(merge_level_procs, set_path, level,
proc_block, D)
futures.append(f)
wait(futures)
procs = procs[axslice(D, None, None, blocksize)]
M = procs.shape[D]
level += 1
# Copy final output
if np.prod(procs.shape) == 1:
for set_path in set_paths:
set_stem = set_path.stem
proc_path = set_path.joinpath(f"{set_stem}_p0_l{level}.h5")
joint_path = set_path.parent.joinpath(f"{set_stem}.h5")
shutil.copy(str(proc_path), str(joint_path))
def operate(self, out):
arg = self.args[0]
axis0 = self.axis0
axis1 = self.axis1
# Enforce conditions for shearing space
arg.require_grid_space(axis=axis1)
arg.require_coeff_space(axis=axis0)
arg.require_local(axis=axis0)
# Apply Fourier shear to flatten the diagonal
# s.t out(y0, y1) = arg(y0+y1-a, y1)
out.layout = arg.layout
np.multiply(arg.data, self.shear, out=out.data)
# Interpolate on flattened diagonal
# s.t. out(y0, y1) = arg(y1, y1)
self.basis0.Interpolate(out, 'left', out=out).evaluate()
out.meta['y0']['constant'] = False
# Broadcast and filter coefficients
# s.t. out(y0, y1) = arg(y0+y1-a, y0+y1-a)
out.data[axslice(axis0, None, None)] = out.data[axslice(axis0, 0, 1)]
out.require_coeff_space(axis=axis1)
out.data *= self.filter_mask
# Move back to starting layout
out.require_layout(arg.layout)
def __init__(self, set_path, mesh, level=0, blocksize=2):
self.set_path = set_path = pathlib.Path(set_path)
set_stem = set_path.stem
# Get process mesh
procs = np.arange(np.prod(mesh)).reshape(mesh)
D = len(mesh) - 1
for i in range(level):
if procs.shape[D] == 1:
D = D - 1
procs = procs[axslice(D, None, None, blocksize)]
self.procs = procs
# Load proc files
self.proc_files = []
for proc in self.procs.ravel():
if level == 0:
proc_path = set_path.joinpath(f"{set_stem}_p{proc}.h5")
else:
proc_path = set_path.joinpath(f"{set_stem}_p{proc}_l{level}.h5")
proc_file = h5py.File(str(proc_path), 'r')
self.proc_files.append(proc_file)