def alltoallv_array(self, send, total_size, offsets, sizes):
if len(send.shape) > 1:
recv = []
for i in range(send.shape[0]):
recv.append(
self.alltoallv_array(send[i, :].copy(), total_size,
offsets, sizes))
recv = np.array(recv)
return recv
offset = offsets[self.comm.rank]
tmp_send = send.view(self.__tocast)
recv = np.empty(total_size, dtype=send.dtype)
if isinstance(send, YTArray):
# We assume send.units is consistent with the units
# on the receiving end.
if isinstance(send, ImageArray):
recv = ImageArray(recv, units=send.units)
else:
recv = YTArray(recv, send.units)
recv[offset:offset + send.size] = send[:]
dtr = send.dtype.itemsize / tmp_send.dtype.itemsize # > 1
roff = [off * dtr for off in offsets]
rsize = [siz * dtr for siz in sizes]
tmp_recv = recv.view(self.__tocast)
self.comm.Allgatherv((tmp_send, tmp_send.size, MPI.CHAR),
(tmp_recv, (rsize, roff), MPI.CHAR))
return recv
def recv_array(self, source, tag=0):
metadata = self.comm.recv(source=source, tag=tag)
dt, ne = metadata[:2]
if ne is None and dt is None:
return self.comm.recv(source=source, tag=tag)
arr = np.empty(ne, dtype=dt)
if len(metadata) == 5:
registry = UnitRegistry(lut=metadata[3], add_default_symbols=False)
if metadata[-1] == "ImageArray":
arr = ImageArray(arr, units=metadata[2], registry=registry)
else:
arr = YTArray(arr, metadata[2], registry=registry)
tmp = arr.view(self.__tocast)
self.comm.Recv([tmp, MPI.CHAR], source=source, tag=tag)
return arr
