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)