def test_take_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.take(local, ind, comm, out=res)
r = heal(res, comm)
s = s[i]
assert_array_equal(r, s)
def test_take_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.take(local, ind, comm, out=res)
r = heal(res, comm)
s = s[i]
assert_array_equal(r, s)
def test_take(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.take(local, ind, comm)
r = heal(res, comm)
s = s[i]
assert res.size == ind.size
assert_array_equal(r, s)
def test_take(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.take(local, ind, comm)
r = heal(res, comm)
s = s[i]
assert res.size == ind.size
assert_array_equal(r, s)
def resample(self, out):
""" Resample the Field by filling 0 or truncating modes.
Convert from and between Real/Complex automatically.
Parameters
----------
out : Field
must be provided because it is a different PM. Can be RealField or ComplexField
"""
assert isinstance(out, Field)
if all(out.Nmesh == self.Nmesh):
# no resampling needed. Just do Fourier transforms.
if isinstance(self, RealField) and isinstance(out, ComplexField):
self.r2c(out)
if isinstance(self, RealField) and isinstance(out, RealField):
out.value[...] = self.value
if isinstance(self, ComplexField) and isinstance(out, RealField):
self.c2r(out)
if isinstance(self, ComplexField) and isinstance(out, ComplexField):
out.value[...] = self.value
return out
if isinstance(self, RealField):
self = self.r2c()
if isinstance(out, RealField):
complex = ComplexField(out.pm)
else:
complex = out
complex.value[...] = 0.0
tmp = numpy.empty_like(self.value)
self.ravel(out=tmp)
# indtable stores the index in pmsrc for the mode in pmdest
# since pmdest < pmsrc, all items are alright.
indtable = [reindex(self.Nmesh[d], out.Nmesh[d]) for d in range(self.value.ndim)]
ind = build_index(
[t[numpy.r_[s]]
for t, s in zip(indtable, complex.slices) ],
self.cshape)
# fill the points that has values in pmsrc
mask = ind >= 0
# their indices
argind = ind[mask]
# take the data
data = mpsort.take(tmp.flat, argind, self.pm.comm)
# fill in the value
complex[mask] = data
# ensure the down sample is real
for i, slab in zip(complex.slabs.i, complex.slabs):
mask = numpy.bitwise_and.reduce(
[(n - ii) % n == ii
for ii, n in zip(i, complex.Nmesh)])
slab.imag[mask] = 0
# remove the nyquist of the output
# FIXME: the nyquist is messy due to hermitian constraints
# let's do not touch them till we know they are important.
mask = numpy.bitwise_or.reduce(
[ ii == n // 2
for ii, n in zip(i, complex.Nmesh)])
slab[mask] = 0
# also remove the nyquist of the input
mask = numpy.bitwise_or.reduce(
[ ii == n // 2
for ii, n in zip(i, self.Nmesh)])
slab[mask] = 0
if isinstance(out, RealField):
complex.c2r(out)
return out
def resample(self, out):
""" Resample the Field by filling 0 or truncating modes.
Convert from and between Real/Complex automatically.
Parameters
----------
out : Field
must be provided because it is a different PM. Can be RealField or ComplexField
"""
assert isinstance(out, Field)
if all(out.Nmesh == self.Nmesh):
# no resampling needed. Just do Fourier transforms.
if isinstance(self, RealField) and isinstance(out, ComplexField):
self.r2c(out)
if isinstance(self, RealField) and isinstance(out, RealField):
out.value[...] = self.value
if isinstance(self, ComplexField) and isinstance(out, RealField):
self.c2r(out)
if isinstance(self, ComplexField) and isinstance(out, ComplexField):
out.value[...] = self.value
return out
if isinstance(self, RealField):
self = self.r2c()
if isinstance(out, RealField):
complex = ComplexField(out.pm)
else:
complex = out
complex.value[...] = 0.0
tmp = numpy.empty_like(self.value)
self.sort(out=tmp)
# indtable stores the index in pmsrc for the mode in pmdest
# since pmdest < pmsrc, all items are alright.
indtable = [reindex(self.Nmesh[d], out.Nmesh[d]) for d in range(self.value.ndim)]
ind = build_index(
[t[numpy.r_[s]]
for t, s in zip(indtable, complex.slices) ],
self.global_shape)
# fill the points that has values in pmsrc
mask = ind >= 0
# their indices
argind = ind[mask]
# take the data
data = mpsort.take(tmp.flat, argind, self.pm.comm)
# fill in the value
complex[mask] = data
if isinstance(out, RealField):
complex.c2r(out)
return out
