def test_permute_out(comm): s = numpy.arange(10) local = split(s, comm) i = numpy.arange(9, -1, -1) ind = split(i, comm, adjustsize(local.size, comm)) res = numpy.empty(local.size, local.dtype) mpsort.permute(local, ind, comm, out=res) r = heal(res, comm) s = s[i] assert_array_equal(r, s)
def unsort(self, flatiter): """ Unsort c-ordered field values to the field. Parameters ---------- flatiter : numpy.flatiter Notes ----- self is updated. `array` does not have to be C_CONTIGUOUS flat iterator of array is used. """ if not isinstance(flatiter, numpy.flatiter): flatiter = flatiter.flat assert isinstance(flatiter, numpy.flatiter) assert len(flatiter) == self.size ind = numpy.ravel_multi_index(numpy.mgrid[self.slices], self.global_shape) mpsort.permute(flatiter, argindex=ind.flat, comm=self.pm.comm, out=self.flat)
def test_permute(comm): s = numpy.arange(10) local = split(s, comm) i = numpy.arange(9, -1, -1) ind = split(i, comm, adjustsize(local.size, comm)) res = mpsort.permute(local, ind, comm) r = heal(res, comm) s = s[i] assert res.size == ind.size assert_array_equal(r, s)
def unravel(self, flatiter): """ Unsort c-ordered field values to the field. Parameters ---------- flatiter : numpy.flatiter Notes ----- self is updated. `array` does not have to be C_CONTIGUOUS flat iterator of array is used. """ if not isinstance(flatiter, numpy.flatiter): flatiter = flatiter.flat assert isinstance(flatiter, numpy.flatiter) assert self.pm.comm.allreduce(len(flatiter)) == self.csize if self.pm.comm.size > 1: ind = numpy.ravel_multi_index(numpy.mgrid[self.slices], self.cshape) mpsort.permute(flatiter, argindex=ind.flat, comm=self.pm.comm, out=self.flat) else: # optimize for a single rank -- directly copy the result self.flat[...] = flatiter