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)