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))