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