def test_c2r_vjp(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') real = pm.generate_whitenoise(1234, type='real', mean=1.0) comp = real.r2c() def objective(comp): real = comp.c2r() obj = (real.value ** 2).sum() return comm.allreduce(obj) grad_real = RealField(pm) grad_real[...] = real[...] * 2 grad_comp = ComplexField(pm) grad_comp = grad_real.c2r_vjp(grad_real) grad_comp.decompress_vjp(grad_comp) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])): dx1, c1 = perturb(comp, ind1, dx) ng1 = (objective(c1) - objective(comp)) / dx ag1 = grad_comp.cgetitem(ind1) * dx1 / dx comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5)
def test_cnorm_log(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0) norm2 = comp1.cnorm(norm = lambda x: numpy.log(x.real ** 2 + x.imag ** 2)) norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2)).sum() assert_allclose(norm2, norm3)
def test_whitenoise_mean(comm): # the whitenoise shall preserve the large scale. pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8') complex1 = pm0.generate_whitenoise(seed=8, unitary=True, mean=1.0) assert_allclose(complex1.c2r().cmean(), 1.0)
def test_paint(comm): from pmesh.pm import ParticleMesh # initialize a random white noise field in real space pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm) real = pm.generate_whitenoise(type='real', seed=3333) # a RealField # FFT to a ComplexField complex = real.r2c() # gather to all ranks --> shape is now (8,8,8) on all ranks data = numpy.concatenate(comm.allgather(numpy.array(real.ravel())))\ .reshape(real.cshape) # mesh from data array realmesh = ArrayMesh(data, BoxSize=128., comm=comm) # paint() must yield the RealField assert_array_equal(real, realmesh.to_field(mode='real')) # gather complex to all ranks --> shape is (8,8,5) on all ranks cdata = numpy.concatenate(comm.allgather(numpy.array(complex.ravel())))\ .reshape(complex.cshape) # mesh from complex array cmesh = ArrayMesh(cdata, BoxSize=128., comm=comm) # paint() yields the ComplexField assert_allclose(complex, cmesh.to_field(mode='complex'), atol=1e-8)
def test_cnorm_grad(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') def objective(comp1): return comp1.cnorm() grad_comp1 = comp1 * 2 grad_comp1.decompress_vjp(grad_comp1) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c1 = perturb(comp1, ind1, dx) ng1 = (objective(c1) - objective(comp1)) / dx ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
def test_decompose(comm): pm = ParticleMesh(BoxSize=4.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') numpy.random.seed(1234) if comm.rank == 0: Npar = 1000 else: Npar = 0 pos = 4.0 * (numpy.random.uniform(size=(Npar, 3))) pos = pm.generate_uniform_particle_grid(shift=0.5) all_pos = numpy.concatenate(comm.allgather(pos), axis=0) for resampler in ['cic', 'tsc', 'db12']: def test(resampler): print(resampler) truth = numpy.zeros(pm.Nmesh, dtype='f8') affine = window.Affine(ndim=3, period=4) window.FindResampler(resampler).paint(truth, all_pos, transform=affine) truth = comm.bcast(truth) layout = pm.decompose(pos, smoothing=resampler) npos = layout.exchange(pos) real = pm.paint(npos, resampler=resampler) full = numpy.zeros(pm.Nmesh, dtype='f8') full[real.slices] = real full = comm.allreduce(full) #print(full.sum(), pm.Nmesh.prod()) #print(truth.sum(), pm.Nmesh.prod()) #print(comm.rank, npos, real.slices) assert_almost_equal(full, truth) # can't yield! test(resampler)
def test_c2c(comm): # this test requires pfft-python 0.1.16. pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128') numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() real2 = complex.c2r() assert numpy.iscomplexobj(real) assert numpy.iscomplexobj(real2) assert numpy.iscomplexobj(complex) assert_array_equal(complex.cshape, pm.Nmesh) assert_array_equal(real2.cshape, pm.Nmesh) assert_array_equal(real.cshape, pm.Nmesh) real.readout(npos) assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_inplace_fft(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') real = RealField(pm) numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() complex2 = real.r2c(out=Ellipsis) assert real._base in complex2._base assert_almost_equal(numpy.asarray(complex), numpy.asarray(complex2), decimal=7) real = complex2.c2r() real2 = complex2.c2r(out=Ellipsis) assert real2._base in complex2._base assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_operators(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4') numpy.random.seed(1234) real = pm.create(type='real', value=0) complex = pm.create(type='complex', value=0) real = real + 1 real = 1 + real real = real + real.value real = real * real.value real = real * real assert isinstance(real, RealField) complex = 1 + complex assert isinstance(complex, ComplexField) complex = complex + 1 assert isinstance(complex, ComplexField) complex = complex + complex.value assert isinstance(complex, ComplexField) complex = numpy.conj(complex) * complex assert isinstance(complex, ComplexField) assert (real == real).dtype == numpy.dtype('?') assert not isinstance(real == real, RealField) assert not isinstance(complex == complex, ComplexField) assert not isinstance(numpy.sum(real), RealField) complex = numpy.conj(complex)
def test_preview(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='real') preview = comp1.preview(axes=(0, 1, 2)) preview = comp1.preview(Nmesh=4, axes=(0, 1, 2)) for ind1 in numpy.ndindex(*(list(comp1.cshape))): assert_allclose(preview[ind1], comp1.cgetitem(ind1)) preview1 = comp1.preview(Nmesh=4, axes=(0, 1)) previewsum1 = preview.sum(axis=2) assert_allclose(preview1, previewsum1) preview2 = comp1.preview(Nmesh=4, axes=(1, 2)) previewsum2 = preview.sum(axis=0) assert_allclose(preview2, previewsum2) preview3 = comp1.preview(Nmesh=4, axes=(0, 2)) previewsum3 = preview.sum(axis=1) assert_allclose(preview3, previewsum3) preview4 = comp1.preview(Nmesh=4, axes=(2, 0)) previewsum4 = preview.sum(axis=1).T assert_allclose(preview4, previewsum4) preview5 = comp1.preview(Nmesh=4, axes=(0,)) previewsum5 = preview.sum(axis=(1, 2)) assert_allclose(preview5, previewsum5) preview6 = comp1.preview(Nmesh=8, axes=(0,))
def test_negnyquist(comm): # the nyquist mode wave number in the hermitian complex field must be negative. # nbodykit depends on this behavior. # see https://github.com/bccp/nbodykit/pull/459 pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') c = pm.create(type='complex') assert (c.x[-1][0][-1] < 0).all() assert (c.x[-1][0][:-1] >= 0).all()
def test_indices_c2c(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='c16') comp = pm.create(type='complex') real = pm.create(type='real') assert_almost_equal(comp.x[0], [[0], [0.785], [-1.571], [-0.785]], decimal=3) assert_almost_equal(comp.x[1], [[0, 0.785, -1.571, -0.785]], decimal=3) assert_almost_equal(real.x[0], [[0], [2], [-4], [-2]], decimal=3) assert_almost_equal(real.x[1], [[0, 2, -4, -2]], decimal=3)
def test_c2c_r2c_edges(comm): pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='c16') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='f8') real1 = pm1.create(type='real') real2 = pm2.create(type='real') assert_allclose(real1.x[0], real2.x[0]) assert_allclose(real1.x[1], real2.x[1]) assert_allclose(real1.x[2], real2.x[2])
def test_cdot_cnorm(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') norm1 = comp1.cdot(comp1) norm2 = comp1.cnorm() norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2).sum() assert_allclose(norm2, norm3) assert_allclose(norm2, norm1)
def test_grid(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') grid, id = pm.generate_uniform_particle_grid(shift=0.5, return_id=True) assert len(id) == len(grid) allid = numpy.concatenate(comm.allgather(id), axis=0) # must be all unique assert len(numpy.unique(allid)) == len(allid) assert numpy.max(allid) == len(allid) - 1 assert numpy.min(allid) == 0
def test_create_typenames(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4') numpy.random.seed(1234) real = pm.create(type=RealField, value=0) from pmesh.pm import _typestr_to_type real = pm.create(type=RealField, value=0) real = pm.create(type=_typestr_to_type('real'), value=0) real.cast(type=_typestr_to_type('real')) real.cast(type=RealField) real.cast(type=RealField)
def test_lic(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(Nmesh=[8, 8], comm=comm, BoxSize=8.0) vx = pm.create(type='real') vy = pm.create(type='real') vx = vx.apply(lambda r, v: r[0]) vy = vy.apply(lambda r, v: 1.0) # FIXME: This only tests if the code 'runs', but no correctness # is tested. r = lic([vx, vy], kernel=lambda s: 1.0, length=2.0, ds=0.1)
def test_cdot_grad(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex', mean=1) comp2 = pm.generate_whitenoise(1235, type='complex', mean=1) def objective(comp1, comp2): return comp1.cdot(comp2).real y = objective(comp1, comp2) grad_comp2 = comp1.cdot_vjp(1.0) grad_comp1 = comp2.cdot_vjp(1.0) grad_comp1.decompress_vjp(grad_comp1) grad_comp2.decompress_vjp(grad_comp2) print(grad_comp1) print("comp1") ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c1 = perturb(comp1, ind1, dx) ng1 = (objective(c1, comp2) - objective(comp1, comp2)) / dx ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5) print("comp2") ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c2 = perturb(comp2, ind1, dx) ng1 = (objective(comp1, c2) - objective(comp1, comp2)) / dx ag1 = grad_comp2.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5)
def test_fof_parallel_no_merge(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm) Q = pm.generate_uniform_particle_grid() cat = ArrayCatalog({'Position' : Q}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm) fof = FOF(cat, linking_length=0.9, nmin=0) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) # one particle per group assert max(labels) == cat.csize - 1
def test_2d_2d(comm): import pfft pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8') assert pm._use_padded == False real = pm.generate_whitenoise(seed=123, type='real') complex = pm.generate_whitenoise(seed=123, type='complex') assert_array_equal(real, complex.c2r()) real2 = pm2.generate_whitenoise(seed=123, type='real') assert real2.shape[:2] == real.shape
def test_cdot_types(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') comp2 = pm.generate_whitenoise(1239, type='untransposedcomplex') with pytest.raises(TypeError): norm1 = comp1.cdot(comp2) norm2 = comp2.cdot(comp1) # this should work though fragile, because only on the Nmesh and 1 rank # the shape of values is the same. norm1 = comp1.cdot(comp2.value) norm2 = comp2.cdot(comp1.value)
def test_fft(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4') numpy.random.seed(1234) real = pm.create(type='real', value=0) raw = real._base.view_raw() real[...] = 2 real[::2, ::2] = -2 real3 = real.copy() complex = real.r2c() assert_almost_equal(numpy.asarray(real), numpy.asarray(real3), decimal=7) real2 = complex.c2r() assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_real_resample(comm): from functools import reduce pmh = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') pml = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') reall = pml.create(type='real') reall.apply(lambda i, v: (i[0] % 2) * (i[1] %2 ), kind='index', out=Ellipsis) for resampler in ['nearest', 'cic', 'tsc', 'cubic']: realh = pmh.upsample(reall, resampler=resampler, keep_mean=False) reall2 = pml.downsample(realh, resampler=resampler) # print(resampler, comm.rank, comm.size, reall, realh) assert_almost_equal(reall.csum(), realh.csum()) assert_almost_equal(reall.csum(), reall2.csum())
def test_cdot(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') comp2 = pm.generate_whitenoise(1239, type='complex') norm1 = comp1.cdot(comp2) norm2 = comp2.cdot(comp1) norm_r = comp1.c2r().cdot(comp2.c2r()) / pm.Nmesh.prod() assert_allclose(norm2.real, norm_r) assert_allclose(norm1.real, norm_r) assert_allclose(norm1.real, norm2.real) assert_allclose(norm1.imag, -norm2.imag)
def test_paint(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm) real = pm.generate_whitenoise(type='real', seed=3333) complex = real.r2c() realmesh = FieldMesh(real) complexmesh = FieldMesh(complex) assert_array_equal(real, realmesh.to_field(mode='real')) assert_array_equal(complex, complexmesh.to_field(mode='complex')) assert_allclose(complex, realmesh.to_field(mode='complex')) assert_allclose(real, complexmesh.to_field(mode='real'))
def __init__(self, comm, Nmesh, BoxSize, dtype): # ensure self.comm is set, though usually already set by the child. self.comm = comm self.dtype = dtype if Nmesh is None or BoxSize is None: raise ValueError("both Nmesh and BoxSize must not be None to initialize ParticleMesh") Nmesh = numpy.array(Nmesh) if Nmesh.ndim == 0: ndim = 3 else: ndim = len(Nmesh) _Nmesh = numpy.empty(ndim, dtype='i8') _Nmesh[:] = Nmesh self.pm = ParticleMesh(BoxSize=BoxSize, Nmesh=_Nmesh, dtype=self.dtype, comm=self.comm) self.attrs['BoxSize'] = self.pm.BoxSize.copy() self.attrs['Nmesh'] = self.pm.Nmesh.copy() # modify the underlying method # actions may have been overriden! self._actions = [] self.base = None
def test_whitenoise(comm): # the whitenoise shall preserve the large scale. pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8') pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8') complex1_down = ComplexField(pm0) complex2_down = ComplexField(pm0) complex1 = pm1.generate_whitenoise(seed=8, unitary=True) complex2 = pm2.generate_whitenoise(seed=8, unitary=True) complex1.resample(complex1_down) complex2.resample(complex2_down) mask1 = complex1_down.value != complex2_down.value assert_array_equal(complex1_down.value, complex2_down.value)
def main(): from argparse import ArgumentParser ap = ArgumentParser() ap.add_argument("--ndim", type=int, choices=[2, 3], default=2, help="Number of dimensions, 2 or 3") ap.add_argument("--nmesh", type=int, default=256, help="Size of FFT mesh") ns = ap.parse_args() pm = ParticleMesh(BoxSize=32.0, Nmesh=[ns.nmesh] * ns.ndim, comm=MPI.COMM_WORLD) u = pm.create(mode='real') def transfer(i, v): r = [(ii - 0.5 * ni) * (Li / ni) for ii, ni, Li in zip(i, v.Nmesh, v.BoxSize)] r2 = sum(ri ** 2 for ri in r) return 4.0 * numpy.arctan(numpy.exp(3 - r2)) u = u.apply(transfer, kind='index') du = pm.create(mode='real') du[...] = 0 steps = numpy.linspace(0, 16, 321, endpoint=True) tmonitor = [0, 4, 8, 11.5, 15] def monitor(t, dt, u_k, dv_k): norm = u_k.cnorm() if pm.comm.rank == 0: print("---- timestep %5.3f, step size %5.4f" % (t, dt)) print("norm of u_k is %g." % norm) for tm in tmonitor.copy(): if abs(t - tm) > dt * 0.5: continue preview = u_k.c2r().preview(Nmesh=min([512, ns.nmesh]), axes=(0, 1)) if pm.comm.rank == 0: print("writing a snapshot") from matplotlib.figure import Figure from matplotlib.backends.backend_agg import FigureCanvasAgg fig = Figure(figsize=(8, 8)) ax = fig.add_subplot(111) ax.imshow(preview.T, origin='lower', extent=(0, pm.BoxSize[0], 0, pm.BoxSize[1])) canvas = FigureCanvasAgg(fig) fig.savefig('klein-gordon-result-%05.3f.png' % t, dpi=128) tmonitor.remove(tm) kgsolver(steps, u, du, lambda u : numpy.sin(u), monitor=monitor)
def __init__(self, field, points, Nmesh, smoothing=None): from pmesh.pm import ParticleMesh self.field = field self.points = points self.Nmesh = Nmesh self.smoothing = smoothing self.pm = ParticleMesh(BoxSize=self.field.BoxSize, Nmesh=[self.Nmesh]*3, dtype='f4', comm=self.comm)
def test_fof_parallel_merge(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm) Q = pm.generate_uniform_particle_grid(shift=0) Q1 = Q.copy() Q1[:] += 0.01 Q2 = Q.copy() Q2[:] -= 0.01 Q3 = Q.copy() Q3[:] += 0.02 cat = ArrayCatalog({'Position' : numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm) fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) assert max(labels) == pm.Nmesh.prod() - 1 assert all(numpy.bincount(labels) == 4)
def test_complex_iter(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') complex = ComplexField(pm) for x, slab in zip(complex.slabs.x, complex.slabs): assert_array_equal(slab.shape, sum(x[d]**2 for d in range(len(pm.Nmesh))).shape) for a, b in zip(slab.x, x): assert_almost_equal(a, b)
def test_fof_parallel_merge(comm): CurrentMPIComm.set(comm) from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[32, 32, 32], Nmesh=[32, 32, 32], comm=comm) Q = pm.generate_uniform_particle_grid(shift=0) Q1 = Q.copy() Q1[:] += 0.01 Q2 = Q.copy() Q2[:] -= 0.01 Q3 = Q.copy() Q3[:] += 0.02 cat = ArrayCatalog({'Position' : numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) assert max(labels) == pm.Nmesh.prod() - 1 assert all(numpy.bincount(labels) == 4)
def __init__(self, pm, cosmology, B=1): """ """ if not isinstance(cosmology, Cosmology): raise TypeError("only nbodykit.cosmology object is supported") fpm = ParticleMesh(Nmesh=pm.Nmesh * B, BoxSize=pm.BoxSize, dtype=pm.dtype, comm=pm.comm, resampler=pm.resampler) self.pm = pm self.fpm = fpm self.cosmology = cosmology
def test_resample(comm): if comm.size == 4: comm.barrier() from pmesh.pm import ParticleMesh # testing the gradient of down sampling, which we use the most. pm = ParticleMesh(BoxSize=128.0, Nmesh=(8, 8), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') dlink, ampl = _addampl(dlink) code = vm.code() code.Decompress(C='dlin_k') code.C2R(C='dlin_k', R='mesh') code.Resample(Neff=4) code.Chi2(variable='mesh') _test_model(code, dlink, ampl)
def test_whitenoise_untransposed(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4') f1 = pm.generate_whitenoise(seed=3333, type='untransposedcomplex') f2 = pm.generate_whitenoise(seed=3333, type='transposedcomplex') f1r = numpy.concatenate(comm.allgather(numpy.array(f1.ravel()))) f2r = numpy.concatenate(comm.allgather(numpy.array(f2.ravel()))) assert_array_equal(f1r, f2r) # this should have asserted r2c transforms as well. r1 = f1.c2r() r2 = f2.c2r() r1r = numpy.concatenate(comm.allgather(numpy.array(r1.ravel()))) r2r = numpy.concatenate(comm.allgather(numpy.array(r2.ravel()))) assert_array_equal(r1r, r2r)
def test_field_compressed(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='c16') comp = pm.create(type='complex') real = pm.create(type='real') assert comp.compressed == False assert real.compressed == False pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') comp = pm.create(type='complex') real = pm.create(type='real') assert comp.compressed == True assert real.compressed == False
def test_lpt2(): """ Checking lpt2_source, this also checks the laplace and gradient kernels """ pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order source = fpmops.lpt2source(lineark).c2r() # Same thing from tensorflow tfsource = tfpm.lpt2_source( pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0))) tfread = pmutils.c2r3d(tfsource).numpy() assert_allclose(source, tfread[0], atol=1e-5)
def _makesource(self, BoxSize, Nmesh): from nbodykit import mockmaker from pmesh.pm import ParticleMesh # the particle mesh for gridding purposes _Nmesh = numpy.empty(3, dtype='i8') _Nmesh[:] = Nmesh pm = ParticleMesh(BoxSize=BoxSize, Nmesh=_Nmesh, dtype='f4', comm=self.comm) # growth rate to do RSD in the Zel'dovich approx f = self.cosmo.scale_independent_growth_rate(self.attrs['redshift']) if self.comm.rank == 0: self.logger.info("Growth Rate is %g" % f) # compute the linear overdensity and displacement fields delta, disp = mockmaker.gaussian_real_fields(pm, self.Plin, self.attrs['seed'], unitary_amplitude=self.attrs['unitary_amplitude'], inverted_phase=self.attrs['inverted_phase'], compute_displacement=True, logger=self.logger) if self.comm.rank == 0: self.logger.info("gaussian field is generated") # poisson sample to points # this returns position and velocity offsets kws = {'bias':self.attrs['bias'], 'seed':self.attrs['seed'], 'logger' : self.logger} pos, disp = mockmaker.poisson_sample_to_points(delta, disp, pm, self.attrs['nbar'], **kws) if self.comm.rank == 0: self.logger.info("poisson sampling is generated") # move particles from initial position based on the Zeldovich displacement pos[:] = (pos + disp) % BoxSize # velocity from displacement (assuming Mpc/h) # this is f * H(z) * a / h = f 100 E(z) a --> converts from Mpc/h to km/s z = self.attrs['redshift'] velocity_norm = f * 100 * self.cosmo.efunc(z) / (1+z) vel = velocity_norm * disp # return data dtype = numpy.dtype([ ('Position', ('f4', 3)), ('Velocity', ('f4', 3)), ('VelocityOffset', ('f4', 3)) ]) source = numpy.empty(len(pos), dtype) source['Position'][:] = pos[:] # in Mpc/h source['Velocity'][:] = vel[:] # in km/s source['VelocityOffset'][:] = f*disp[:] # in Mpc/h return source, pm
def __init__(self, datasource, Nmesh, seed=12345, ratio=0.01, smoothing=None, format='hdf5'): from pmesh.pm import ParticleMesh self.datasource = datasource self.Nmesh = Nmesh self.seed = seed self.ratio = ratio self.smoothing = smoothing self.format = format self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype='f4', comm=self.comm)
def test_cic_readout(): bs = 50 nc = 16 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') nparticle = 100 pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype( np.float32) base = 100 * np.random.random(nc**3).reshape(nc, nc, nc).astype(np.float32) pmmesh = pm.create(mode='real', value=base) pmread = pmmesh.readout(pos) with tf.Session() as sess: mesh = cic_readout( tf.constant(base.reshape((1, nc, nc, nc)), dtype=tf.float32), (pos * nc / bs).reshape((1, nparticle, 3))) sess.run(tf.global_variables_initializer()) tfread = sess.run(mesh) assert_allclose(pmread, tfread[0], rtol=1e-06)
def test_real_apply(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') real = RealField(pm) def filter(x, v): return x[0] * 10 + x[1] real.apply(filter, out=Ellipsis) for i, x, slab in zip(real.slabs.i, real.slabs.x, real.slabs): assert_array_equal(slab, x[0] * 10 + x[1])
def test_complex_apply(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') complex = ComplexField(pm) def filter(k, v): return k[0] + k[1] * 1j complex.apply(filter, out=Ellipsis) for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs): assert_array_equal(slab, x[0] + x[1] * 1j)
def test_cic_paint(): bs = 50 nc = 16 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') nparticle = 100 pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype( np.float32) wts = np.random.random(nparticle).astype(np.float32) # Painting with pmesg pmmesh = pm.paint(pos, mass=wts) with tf.Session() as sess: mesh = cic_paint(tf.zeros((1, nc, nc, nc), dtype=tf.float32), (pos * nc / bs).reshape((1, nparticle, 3)), weight=wts.reshape(1, nparticle)) sess.run(tf.global_variables_initializer()) tfmesh = sess.run(mesh) assert_allclose(pmmesh, tfmesh[0], atol=1e-06)
def test_coords(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') grid_x = pm.create_coords('real') assert len(grid_x) == 3 assert grid_x[0].dtype == pm.dtype grid_i = pm.create_coords('real', return_indices=True) assert len(grid_i) == 3 grid_x = pm.create_coords('complex') grid_i = pm.create_coords('complex', return_indices=True) assert len(grid_x) == 3 assert grid_x[0].dtype == pm.dtype assert len(grid_i) == 3 pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f4') grid_x = pm.create_coords('transposedcomplex') grid_i = pm.create_coords('transposedcomplex', return_indices=True) assert len(grid_x) == 3 assert grid_x[0].dtype == pm.dtype assert len(grid_i) == 3
def __init__(self, pm, B=1, shift=0, dtype='f8'): self.fpm = ParticleMesh(Nmesh=pm.Nmesh * B, BoxSize=pm.BoxSize, dtype=dtype, comm=pm.comm) q = operators.create_grid(pm, shift=shift, dtype=dtype) N = pm.comm.allreduce(len(q)) self.mean_number_count = 1.0 * N / (1.0 * self.fpm.Nmesh.prod()) ParticleMeshVM.__init__(self, pm, q) MPINumeric.__init__(self, pm.comm)
def test_lpt1(): """ Checking lpt1, this also checks the laplace and gradient kernels """ pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order lpt = fpmops.lpt1(lineark, grid) # Same thing from tensorflow with tf.Session() as sess: state = tfpm.lpt1( pmutils.r2c3d(tf.expand_dims(tf.constant(lineark.c2r()), axis=0)), grid.reshape((1, -1, 3)) * nc / bs) tfread = sess.run(state) assert_allclose(lpt, tfread[0] * bs / nc, atol=1e-5)
def test_fft(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape( -1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() real2 = complex.c2r() real.readout(npos) assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_lpt_prior(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') code = vm.code() code.Decompress(C='dlin_k') code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0) code.Decompose() code.Paint() code.Chi2(variable='mesh') code.Prior(powerspectrum=lambda k : 1.0) code.Multiply(a='prior', b='prior', f=0.1) code.Add(a='prior', b='chi2', c='chi2') dlink, ampl = _addampl(dlink) _test_model(code, dlink, ampl)
def test_lpt_init(): """ Checking lpt init """ a0 = 0.1 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) solver = Solver(pm, Planck15, B=1) # Generate initial state with fastpm whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) statelpt = solver.lpt(lineark, grid, a0, order=1) # Same thing with flowpm tlinear = tf.expand_dims(np.array(lineark.c2r()), 0) tfread = tfpm.lpt_init(tlinear, a0, order=1).numpy() assert_allclose(statelpt.X, tfread[0, 0] * bs / nc, rtol=1e-2)
def test_lpt1_64(): """ Checking lpt1, this also checks the laplace and gradient kernels This variant of the test checks that it works for cubes of size 64 """ nc = 64 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order lpt = fpmops.lpt1(lineark, grid) # Same thing from tensorflow tfread = tfpm.lpt1( pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)), grid.reshape((1, -1, 3)) * nc / bs).numpy() assert_allclose(lpt, tfread[0] * bs / nc, atol=5e-5)
def test_wrong_ndim(comm): numpy.random.seed(42) pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') rfield = RealField(pm) data = numpy.random.random(size=rfield.shape) rfield[...] = data[:] # SlabIterator only works for 2D or 3D coordinate meshes with pytest.raises(NotImplementedError): for slab in SlabIterator([rfield.x[0]], axis=0, symmetry_axis=None): pass
def main(ns): comm = MPI.COMM_WORLD result = simulate(comm, ns) pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm) report = analyze(pm, result) if comm.rank == 0: write_report(ns.output, report)
def test_fdownsample(comm): """ fourier space resample, deprecated """ pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') numpy.random.seed(3333) truth = numpy.fft.rfftn(numpy.random.normal(size=(8, 8))) complex1 = ComplexField(pm1) for ind in numpy.ndindex(*complex1.cshape): complex1.csetitem(ind, truth[ind]) assert_almost_equal(complex1[...], complex1.c2r().r2c()) complex2 = ComplexField(pm2) for ind in numpy.ndindex(*complex2.cshape): newind = tuple([i if i <= 2 else 8 - (4 - i) for i in ind]) if any(i == 2 for i in ind): complex2.csetitem(ind, 0) else: complex2.csetitem(ind, truth[newind]) tmpr = RealField(pm2) tmp = ComplexField(pm2) complex1.resample(tmp) assert_almost_equal(complex2[...], tmp[...], decimal=5) complex1.c2r().resample(tmp) assert_almost_equal(complex2[...], tmp[...], decimal=5) complex1.resample(tmpr) assert_almost_equal(tmpr.r2c(), tmp[...]) complex1.c2r().resample(tmpr) assert_almost_equal(tmpr.r2c(), tmp[...])
def test_untransposed_complex_apply(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8') complex = UntransposedComplexField(pm) def filter(k, v): knormp = k.normp() assert_allclose(knormp, sum(ki**2 for ki in k)) return k[0] + k[1] * 1j + k[2] complex = complex.apply(filter) for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs): assert_array_equal(slab, x[0] + x[1] * 1j + x[2])
def test_real_iter(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') real = RealField(pm) for i, x, slab in zip(real.slabs.i, real.slabs.x, real.slabs): assert_same_base(slab, real.value) assert_almost_equal(slab.shape, sum(x[d]**2 for d in range(len(pm.Nmesh))).shape) for a, b in zip(slab.x, x): assert_array_equal(a, b) for a, b in zip(slab.i, i): assert_array_equal(a, b)
def fiddlebiasgal(aa, suff, nc=nc, mcfv=[1.], saveb=False, bname='h1bias', ofolder=None): '''Fiddle bias for galaxies''' if ofolder is None: ofolder = project + '/%s/fastpm_%0.4f/'%(sim, aa) pm = ParticleMesh(BoxSize = bs, Nmesh = [nc, nc, nc]) print('Read in catalogs') cencat = BigFileCatalog(project + sim + '/fastpm_%0.4f/cencat'%aa) satcat = BigFileCatalog(project + sim + '/fastpm_%0.4f/satcat'%aa+suff) cpos, spos = cencat['Position'], satcat['Position'] cmass, smass = cencat['Mass'], satcat['Mass'] pos = np.concatenate((cpos, spos), axis=0) dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint() pkm = FFTPower(dm/dm.cmean(), mode='1d').power k, pkm = pkm['k'], pkm['power'] b1, b1sq = np.zeros((k.size, len(mcfv))), np.zeros((k.size, len(mcfv))) for imc, mcf in enumerate(mcfv): print(mcf) ch1mass = HI_masscutfiddle(cmass, aa, mcutf=mcf) sh1mass = HI_masscutfiddle(smass, aa, mcutf=mcf) h1mass = np.concatenate((ch1mass, sh1mass), axis=0) # h1mesh = pm.paint(pos, mass=h1mass) pkh1 = FFTPower(h1mesh/h1mesh.cmean(), mode='1d').power['power'] pkh1m = FFTPower(h1mesh/h1mesh.cmean(), second=dm/dm.cmean(), mode='1d').power['power'] #Bias b1[:, imc] = pkh1m/pkm b1sq[:, imc] = pkh1/pkm np.savetxt(ofolder+bname+'auto'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1sq**0.5), axis=1), header='mcut factors = %s\nk, pkh1xm/pkm, pkh1/pkm^0.5'%mcfv) np.savetxt(ofolder+bname+'cross'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1), axis=1), header='mcut factors = %s\nk, pkh1xm/pkm, pkh1mx/pkm'%mcfv) return k, b1, b1sq
def finalize(self): self['aout'] = numpy.array(self['aout']) self.pm = ParticleMesh(BoxSize=self['boxsize'], Nmesh=[self['nc']] * self['ndim'], resampler=self['resampler'], dtype='f4') mask = numpy.array([a not in self['stages'] for a in self['aout']], dtype='?') missing_stages = self['aout'][mask] if len(missing_stages): raise ValueError( 'Some stages are requested for output but missing: %s' % str(missing_stages))
def set_up(): params = get_params() pm = ParticleMesh(Nmesh=params['Nmesh'], BoxSize=params['BoxSize'], comm=MPI.COMM_WORLD, resampler='cic') # generate initial conditions cosmo = Planck15.clone(P_k_max=30) x = pm.generate_uniform_particle_grid(shift=0.1) BoxSize2D = [deg / 180. * np.pi for deg in params['BoxSize2D']] logger = None sim = lightcone.WLSimulation(stages=numpy.linspace(0.1, 1.0, params['N_steps'], endpoint=True), cosmology=cosmo, pm=pm, boxsize2D=BoxSize2D, params=params, logger=None) kmaps = [sim.mappm.create('real', value=0.) for ii in range(1)] return pm, cosmo, x, kmaps, sim.DriftFactor, sim.mappm, sim