def val_and_grad(Omega_c, whitec):
params = tf.stack([Omega_c])
with tf.GradientTape() as tape:
tape.watch(params)
cosmology = flowpm.cosmology.Planck15(Omega_c=params[0])
k = tf.constant(np.logspace(-4, 1, 256), dtype=tf.float32)
pk = tfpower.linear_matter_power(cosmology, k)
pk_fun = lambda x: tf.cast(
tf.reshape(
interpolate.interp_tf(
tf.reshape(tf.cast(x, tf.float32), [-1]), k, pk), x.shape), tf.
complex64)
initial_conditions = whitenoise_to_linear([nc, nc, nc],
[box_size, box_size, box_size],
whitec,
pk_fun,
seed=100,
batch_size=1)
state = flowpm.lpt_init(cosmology, initial_conditions, 0.1)
final_state = flowpm.nbody(cosmology, state, stages, [nc, nc, nc])
final_field = flowpm.cic_paint(
tf.zeros_like(initial_conditions), final_state[0])
final_field = tf.reshape(final_field, [nc, nc, nc])
loss = tf.reduce_mean(final_field**2)
return loss, tape.gradient(loss, params)
def main(_):
cosmology = flowpm.cosmology.Planck15()
# Compute the k vectora that will be needed in the PGD fit
k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3),
boxsize=np.array([FLAGS.box_size] * 3),
kmin=np.pi / FLAGS.box_size,
dk=2 * np.pi / FLAGS.box_size)
# Create some initial conditions
klin = tf.constant(np.logspace(-4, 1, 512), dtype=tf.float32)
pk = linear_matter_power(cosmology, klin)
pk_fun = lambda x: tf.cast(
tf.reshape(
interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]),
klin, pk), x.shape), tf.complex64)
initial_conditions = flowpm.linear_field(
[FLAGS.nc, FLAGS.nc, FLAGS.nc],
[FLAGS.box_size, FLAGS.box_size, FLAGS.box_size],
pk_fun,
batch_size=FLAGS.batch_size)
initial_state = flowpm.lpt_init(cosmology, initial_conditions,
FLAGS.a_init)
stages = np.linspace(FLAGS.a_init, 1., FLAGS.nsteps, endpoint=True)
print('Starting simulation')
state, scale_factors, pgdparams = fit_nbody(cosmology,
initial_state,
stages,
[FLAGS.nc, FLAGS.nc, FLAGS.nc],
pm_nc_factor=FLAGS.B)
print('Simulation done')
pickle.dump(
{
'B': FLAGS.B,
'nsteps': FLAGS.nsteps,
'params': pgdparams,
'scale_factors': scale_factors,
'cosmology': cosmology.to_dict(),
'boxsize': FLAGS.box_size,
'nc': FLAGS.nc
}, open(FLAGS.filename, "wb"))
def a_of_chi(cosmo, chi):
r"""Computes the scale factor for corresponding (array) of radial comoving
distance by reverse linear interpolation.
Parameters:
-----------
cosmo: Cosmology
Cosmological parameters
chi: array_like or tf.TensorArray
radial comoving distance to query.
Returns:
--------
a : tf.TensorArray
Scale factors corresponding to requested distances
"""
# Check if distances have already been computed, force computation otherwise
if "background.radial_comoving_distance" not in cosmo._workspace.keys():
rad_comoving_distance(cosmo, 1.0)
cache = cosmo._workspace["tfbackground.radial_comoving_distance"]
chi = tf.cast(chi, dtype=tf.float32)
return interp_tf(chi, cache["chi"], cache["a"])
def main(_):
cosmology = flowpm.cosmology.Planck15()
# Create a simple Planck15 cosmology without neutrinos, and makes sure sigma8
# is matched
nbdykit_cosmo = Cosmology.from_astropy(Planck15.clone(m_nu=0 * u.eV))
nbdykit_cosmo = nbdykit_cosmo.match(sigma8=cosmology.sigma8.numpy())
# Compute the k vectora that will be needed in the PGD fit
k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3),
boxsize=np.array([FLAGS.box_size] * 3),
kmin=np.pi / FLAGS.box_size,
dk=2 * np.pi / FLAGS.box_size)
# Create some initial conditions
klin = tf.constant(np.logspace(-4, 1, 512), dtype=tf.float32)
pk = linear_matter_power(cosmology, klin)
pk_fun = lambda x: tf.cast(
tf.reshape(
interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]),
klin, pk), x.shape), tf.complex64)
initial_conditions = flowpm.linear_field(
[FLAGS.nc, FLAGS.nc, FLAGS.nc],
[FLAGS.box_size, FLAGS.box_size, FLAGS.box_size],
pk_fun,
batch_size=FLAGS.batch_size)
initial_state = flowpm.lpt_init(cosmology, initial_conditions,
FLAGS.a_init)
stages = np.linspace(FLAGS.a_init, 1., FLAGS.nsteps, endpoint=True)
print('Starting simulation')
# Run the Nbody
states = flowpm.nbody(cosmology,
initial_state,
stages, [FLAGS.nc, FLAGS.nc, FLAGS.nc],
pm_nc_factor=FLAGS.B,
return_intermediate_states=True)
print('Simulation done')
# Initialize PGD params
alpha = tf.Variable([FLAGS.alpha0], dtype=tf.float32)
scales = tf.Variable([FLAGS.kl0, FLAGS.ks0], dtype=tf.float32)
optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate)
params = []
scale_factors = []
# We begin by fitting the last time step
for j, (a, state) in enumerate(states[::-1]):
# Let's compute the target power spectrum at that scale factor
target_pk = HalofitPower(nbdykit_cosmo,
1. / a - 1.)(k).astype('float32')
for i in range(FLAGS.niter if j == 0 else FLAGS.niter_refine):
optimizer.minimize(
partial(pgd_loss, alpha, scales, state, target_pk), [alpha] if
(FLAGS.fix_scales and j > 0) else [alpha, scales])
if i % 10 == 0:
loss, pk = pgd_loss(alpha,
scales,
state,
target_pk,
return_pk=True)
if i == 0:
pk0 = pk
print("step %d, loss:" % i, loss)
params.append(np.concatenate([alpha.numpy(), scales.numpy()]))
scale_factors.append(a)
print("Fitted params (alpha, kl, ks)", params[-1])
plt.loglog(k, target_pk, "k")
plt.loglog(k, pk0, ':', label='starting')
plt.loglog(k, pk, '--', label='after n steps')
plt.grid(which='both')
plt.savefig('PGD_fit_%0.2f.png' % a)
plt.close()
pickle.dump(
{
'B': FLAGS.B,
'nsteps': FLAGS.nsteps,
'params': params,
'scale_factors': scale_factors,
'cosmology': cosmology.to_dict(),
'boxsize': FLAGS.box_size,
'nc': FLAGS.nc
}, open(FLAGS.filename, "wb"))