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 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=1e-5)
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 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=5e-5)
dx1, dx2 = fastpm.lpt(rhok, q, pm) dx, p, f = fastpm.nbody(rhok, q, [0.1, 0.6, 1.0], Planck15, pm) dx0, p0, f0 = fastpm.nbody(rhok, q, [0.1], Planck15, pm) model.output(dx1=dx1, dx2=dx2, dx=dx, p=p, dx0=dx0, p0=p0, f0=f0, f=f) wn = pm.generate_whitenoise(555, unitary=True) x = wn[...] x = numpy.stack([x.real, x.imag], -1) from fastpm.core import Solver, leapfrog solver = Solver(pm, Planck15, B=1) linear = solver.linear(wn, powerspectrum) dx1_f = lpt1(linear, q) dx2_f = lpt1(lpt2source(linear), q) lpt = solver.lpt(linear, Q=q, a=0.1) print('comparing lpt order by order') dx1, dx2 = model.compute(['dx1', 'dx2'], init=dict(x=x)) print('model', dx1.std(axis=0), dx2.std(axis=0)) print('fastpm', dx1_f.std(axis=0), dx2_f.std(axis=0)) print('model', dx1[0], dx2[0]) print('fastpm', dx1_f[0], dx2_f[0]) print('comparing lpt dx and p ') dx0, p0 = model.compute(['dx0', 'p0'], init=dict(x=x)) print('model', dx0.std(axis=0), p0.std(axis=0))