def test_sort(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 6], comm=comm, dtype='f8')
real = RealField(pm)
truth = numpy.arange(8 * 6)
real[...] = truth.reshape(8, 6)[real.slices]
unsorted = real.copy()
real.sort(out=Ellipsis)
conjecture = numpy.concatenate(comm.allgather(real.value.ravel()))
assert_array_equal(conjecture, truth)
real.unravel(real)
assert_array_equal(real, unsorted)
complex = ComplexField(pm)
truth = numpy.arange(8 * 4)
complex[...] = truth.reshape(8, 4)[complex.slices]
complex.ravel(out=Ellipsis)
conjecture = numpy.concatenate(comm.allgather(complex.value.ravel()))
assert_array_equal(conjecture, truth)
def test_sort(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 6], comm=comm, dtype='f8')
real = RealField(pm)
truth = numpy.arange(8 * 6)
real[...] = truth.reshape(8, 6)[real.slices]
unsorted = real.copy()
real.sort(out=Ellipsis)
conjecture = numpy.concatenate(comm.allgather(real.value.ravel()))
assert_array_equal(conjecture, truth)
real.unravel(real)
assert_array_equal(real, unsorted)
complex = ComplexField(pm)
truth = numpy.arange(8 * 4)
complex[...] = truth.reshape(8, 4)[complex.slices]
complex.ravel(out=Ellipsis)
conjecture = numpy.concatenate(comm.allgather(complex.value.ravel()))
assert_array_equal(conjecture, truth)
def main():
ns = ap.parse_args()
comm = MPI.COMM_WORLD
ff = bigfile.BigFileMPI(comm, ns.fastpm)
with ff['.'] as bb:
BoxSize = bb.attrs['BoxSize'][0]
Redshift = 1 / bb.attrs['ScalingFactor'][0] - 1
Nmesh = int(BoxSize / ns.resolution * 2)
# round it to 8.
Nmesh -= Nmesh % 8
if comm.rank == 0:
logger.info("source = %s", ns.fastpm)
logger.info("output = %s", ns.output)
logger.info("BoxSize = %g", BoxSize)
logger.info("Redshift = %g", Redshift)
logger.info("Nmesh = %g", Nmesh)
pm = ParticleMesh([Nmesh, Nmesh, Nmesh], BoxSize, comm=comm)
real = RealField(pm)
real[...] = 0
with ff['Position'] as ds:
logger.info(ds.size)
for i in range(0, ds.size, ns.chunksize):
sl = slice(i, i + ns.chunksize)
pos = ds[sl]
layout = pm.decompose(pos)
lpos = layout.exchange(pos)
real.paint(lpos, hold=True)
mean = real.cmean()
if comm.rank == 0:
logger.info("mean particle per cell = %s", mean)
real[...] /= mean
real[...] -= 1
complex = real.r2c()
for k, i, slab in zip(complex.slabs.x, complex.slabs.i, complex.slabs):
k2 = sum(kd**2 for kd in k)
# tophat
f = tophat(ns.filtersize, k2**0.5)
slab[...] *= f
# zreion
slab[...] *= Bk(k2**0.5)
slab[...] *= (1 + Redshift)
real = complex.c2r()
real[...] += Redshift
mean = real.cmean()
if comm.rank == 0:
logger.info("zreion.mean = %s", mean)
buffer = numpy.empty(real.size, real.dtype)
real.sort(out=buffer)
if comm.rank == 0:
logger.info("sorted for output")
with bigfile.BigFileMPI(comm, ns.output, create=True) as ff:
with ff.create_from_array(ns.dataset, buffer) as bb:
bb.attrs['BoxSize'] = BoxSize
bb.attrs['Redshift'] = Redshift
bb.attrs['TopHatFilterSize'] = ns.filtersize
bb.attrs['Nmesh'] = Nmesh
#
# hack: compatible with current MPGadget. This is not really needed
# we'll remove the bins later, since BoxSize and Nmesh are known.
with ff.create("XYZ_bins", dtype='f8', size=Nmesh) as bb:
if comm.rank == 0:
bins = numpy.linspace(0, BoxSize * 1000., Nmesh, dtype='f8')
bb.write(0, bins)
if comm.rank == 0:
logger.info("done. written at %s", ns.output)
