def standardrecon(base, pos, bias, R):
#base = base.astype(np.float32)
#pos = pos.astype(base.dtype)
smwts = tf.exp(tf.multiply(-kmesh**2, R**2))
basek = utils.r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
basesm = utils.c2r3d(basek, norm=nc**3)
grid = bs/nc*np.indices((nc, nc, nc)).reshape(3, -1).T.astype(np.float32)
grid = tf.constant(np.expand_dims(grid, 0))
grid = grid *nc/bs
pos = pos *nc/bs
mesh = basesm #tf.constant(basesm.astype(np.float32))
meshk = utils.r2c3d(mesh, norm=nc**3)
DX = tfpm.lpt1(meshk, pos, kvec=kvec)
DX = tf.multiply(DX, -1/bias)
pos = tf.add(pos, DX)
displaced = tf.zeros_like(mesh)
displaced = utils.cic_paint(displaced, pos, name='displaced')
DXrandom = tfpm.lpt1(meshk, grid, kvec)
DXrandom = tf.multiply(DXrandom, -1/bias)
posrandom = tf.add(grid, DXrandom)
random = tf.zeros_like(mesh)
random = utils.cic_paint(random, posrandom, name='random')
return displaced, random
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)