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_nody():
""" Checking end to end nbody
"""
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)
stages = np.linspace(0.1, 1.0, 10, endpoint=True)
# 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)
finalstate = solver.nbody(statelpt, leapfrog(stages))
final_cube = pm.paint(finalstate.X)
# Same thing with flowpm
tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
state = tfpm.lpt_init(tlinear, a0, order=1)
state = tfpm.nbody(state, stages, nc)
tfread = pmutils.cic_paint(tf.zeros_like(tlinear), state[0]).numpy()
assert_allclose(final_cube, tfread[0], atol=1.2)
def pgd_loss(pgdparams, state, target_pk, return_pk=False):
"""
Defines the loss function for the PGD parameters
"""
shape = state.get_shape()
batch_size = shape[1]
# Step I: Apply PGD to the state vector
pdgized_state = tfpm.pgd(state,
pgdparams,
nc=[FLAGS.nc] * 3,
pm_nc_factor=FLAGS.B)
# Step II: Painting and compensating for cic window
field = cic_paint(tf.zeros([batch_size] + [FLAGS.nc] * 3),
pdgized_state[0])
field = compensate_cic(field)
# Step III: Compute power spectrum
k, pk = flowpm.power_spectrum(field,
boxsize=np.array([FLAGS.box_size] * 3),
kmin=np.pi / FLAGS.box_size,
dk=2 * np.pi / FLAGS.box_size)
# Averaging pk over realisations
pk = tf.reduce_mean(pk, axis=0)
# Step IV: compute loss
loss = tf.reduce_sum((1. - pk / target_pk)**2)
if return_pk:
return loss, pk
else:
return loss
def force(cosmo,
state,
nc,
pm_nc_factor=1,
kvec=None,
dtype=tf.float32,
name="Force",
**kwargs):
"""
Estimate force on the particles given a state.
Parameters:
-----------
state: tensor
Input state tensor of shape (3, batch_size, npart, 3)
boxsize: float
Size of the simulation volume (Mpc/h) TODO: check units
cosmology: astropy.cosmology
Cosmology object
pm_nc_factor: int
TODO: @modichirag please add doc
"""
with tf.name_scope(name):
state = tf.convert_to_tensor(state, name="state")
shape = state.get_shape()
batch_size = shape[1]
ncf = [n * pm_nc_factor for n in nc]
rho = tf.zeros([batch_size] + ncf)
wts = tf.ones((batch_size, nc[0] * nc[1] * nc[2]))
nbar = nc[0] * nc[1] * nc[2] / (ncf[0] * ncf[1] * ncf[2])
rho = cic_paint(rho, tf.multiply(state[0], pm_nc_factor), wts)
rho = tf.multiply(rho, 1. /
nbar) # I am not sure why this is not needed here
delta_k = r2c3d(rho, norm=ncf[0] * ncf[1] * ncf[2])
fac = tf.cast(1.5 * cosmo.Omega_m, dtype=dtype)
update = apply_longrange(tf.multiply(state[0], pm_nc_factor),
delta_k,
split=0,
factor=fac)
update = tf.expand_dims(update, axis=0)
indices = tf.constant([[2]])
shape = state.shape
update = tf.scatter_nd(indices, update, shape)
mask = tf.stack((tf.ones_like(state[0]), tf.ones_like(
state[0]), tf.zeros_like(state[0])),
axis=0)
state = tf.multiply(state, mask)
state = tf.add(state, update)
return state
def pgd(state, pgdparams, nc, pm_nc_factor=1, name="PGD", **kwargs):
"""
Estimate the short range force on the particles given a state.
Parameters:
-----------
state: tensor
Input state tensor of shape (3, batch_size, npart, 3)
nc: int, or list of ints
Number of cells
alpha: float
Free parameter. Factor of proportionality between the displacement and the Particle-mesh force.
kl: float
Long range scale parameter
ks: float
Short range scale parameter
pm_nc_factor: int
Resolution factor. Ratio of the size per side of the mesh and the number of particles per side
"""
with tf.name_scope(name):
state = tf.convert_to_tensor(state, name="state")
shape = state.get_shape()
batch_size = shape[1]
ncf = [n * pm_nc_factor for n in nc]
rho = tf.zeros([batch_size] + ncf)
wts = tf.ones((batch_size, nc[0] * nc[1] * nc[2]))
nbar = nc[0] * nc[1] * nc[2] / (ncf[0] * ncf[1] * ncf[2])
rho = cic_paint(rho, tf.multiply(state[0], pm_nc_factor), wts)
rho = tf.multiply(rho, 1. /
nbar) # I am not sure why this is not needed here
delta_k = r2c3d(rho, norm=ncf[0] * ncf[1] * ncf[2])
alpha, kl, ks = tf.split(pgdparams, 3)
update = apply_pgd(
tf.multiply(state[0], pm_nc_factor),
delta_k,
alpha,
kl,
ks,
)
update = tf.expand_dims(update, axis=0) / pm_nc_factor
indices = tf.constant([[0]])
shape = state.shape
update = tf.scatter_nd(indices, update, shape)
state = tf.add(state, update)
return state
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)
mesh = cic_paint(tf.zeros((1, nc, nc, nc), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)),
weight=wts.reshape(1, nparticle))
tfmesh = mesh.numpy()
assert_allclose(pmmesh, tfmesh[0], atol=1e-06)
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)
