def pmpos(linear):
if args.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=a0, order=args.lpt_order)
final_state = nbody(state, stages, nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=af, order=args.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
return tfinal_field, final_state[0]
def forward(self, linear):
if self.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=self.a0, order=self.lpt_order)
final_state = nbody(state, self.stages, self.nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=self.af, order=self.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
return tfinal_field
def pm_data_test(dummy):
print("PM graph")
linear = flowpm.linear_field(nc, bs, ipklin, batch_size=world_size)
if args.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=a0, order=args.lpt_order)
final_state = nbody(state, stages, nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=af, order=args.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
return linear, tfinal_field
def pm(self, linear):
args = self.args
nc, bs = args.nc, args.bs
if args.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=args.a0, order=args.lpt_order)
final_state = nbody(state, args.stages, nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=args.af, order=args.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
return tfinal_field
def sample(self, linear):
if self.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=self.a0, order=self.lpt_order)
final_state = nbody(state, self.stages, self.nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=self.af, order=self.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
galmean = tfp.distributions.Poisson(rate=self.plambda *
(1 + tfinal_field))
sample = galmean.sample()
return sample
def recon_prototype(linearflat):
"""
"""
linear = tf.reshape(linearflat, data.shape)
#
#loss = tf.reduce_sum(tf.square(linear - minimum))
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
#final_field = pmgraph(linear)
base = final_field
if FLAGS.anneal:
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
galmean = tfp.distributions.Poisson(rate=plambda * (1 + base))
logprob = -tf.reduce_sum(galmean.log_prob(data))
#logprob = tf.multiply(logprob, 1/nc**3, name='logprob')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
#prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = logprob + prior
grad = tf.gradients(loss, linearflat)
print(grad)
return loss, grad[0]
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 pm(lin, params):
""" FastPM forward simulation
"""
state = flowpm.lpt_init(lin, a0=params['a0'])
final_state = flowpm.nbody(state, params['stages'], params['nc'])
final_field = flowpm.cic_paint(tf.zeros_like(lin), final_state[0])
return final_field
def recon_prototype(linear):
"""
"""
linear = tf.reshape(linear, minimum.shape)
#loss = tf.reduce_sum(tf.square(linear - minimum))
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
## Rsmsq = tf.multiply(Rsm*bs/nc, Rsm*bs/nc)
## smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
## basek = r2c3d(base, norm=nc**3)
## basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
## base = c2r3d(basek, norm=nc**3)
#### #
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
chisq = tf.multiply(chisq, 1/nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1/priorwt))
prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = chisq + prior
return loss
def forward(self, linear):
b1, b2 = self.bias[0], self.bias[1]
if self.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=self.a0, order=self.lpt_order)
final_state = nbody(state, self.stages, self.nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=self.af, order=self.lpt_order)
#final_field = cic_paint(tf.zeros_like(linear), final_state[0])
fpos = final_state[0]
w0 = tf.reshape(linear, (linear.shape[0], -1))
w0 = w0 - tf.expand_dims(tf.reduce_mean(w0, 1), -1)
w2 = w0 * w0
w2 = w2 - tf.expand_dims(tf.reduce_mean(w2, 1), -1)
weight = b1 * w0 + b2 * w2
bmodel = cic_paint(tf.zeros_like(linear), fpos, weight=weight)
return bmodel
def recon_prototype(x0=None):
"""
"""
# linear = tf.get_variable('linmesh', shape=(1, nc, nc, nc), dtype=tf.float32,
# initializer=tf.random_normal_initializer(), trainable=True)
if x0 is None:
linear = tf.get_variable('linmesh', shape=(1, nc, nc, nc), dtype=tf.float32,
initializer=tf.random_normal_initializer(), trainable=True)
else:
linear = tf.get_variable('linmesh', shape=(1, nc, nc, nc), dtype=tf.float32,
initializer=tf.constant_initializer(x0), trainable=True)
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
base = final_field
if FLAGS.anneal:
print('\nAdd annealing graph\n')
Rsmsq = tf.multiply(Rsm*bs/nc, Rsm*bs/nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
galmean = tfp.distributions.Poisson(rate = plambda * (1 + base))
sample = galmean.sample()
logprob = -tf.reduce_sum(galmean.log_prob(data))
#logprob = tf.multiply(logprob, 1/nc**3, name='logprob')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1/priorwt))
#prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = logprob + prior
#opt = tf.train.GradientDescentOptimizer(learning_rate=0.1)
opt = tf.train.AdamOptimizer(learning_rate=FLAGS.lr)
#step = tf.Variable(0, trainable=False)
#schedule = tf.optimizers.schedules.PiecewiseConstantDecay(
# [10000, 15000], [1e-0, 1e-1, 1e-2])
## lr and wd can be a function or a tensor
#lr = 1e-1 * schedule(step)
#wd = lambda: 1e-4 * schedule(step)
#opt = tfa.optimizers.AdamW(learning_rate=FLAGS.lr, weight_decay=1e-1)
# Compute the gradients for a list of variables.
grads_and_vars = opt.compute_gradients(loss, [linear])
print("\ngradients : ", grads_and_vars)
update_ops = opt.apply_gradients(grads_and_vars)
return linear, sample, update_ops, loss, logprob, prior
def pm_poisson():
print("PM graph")
linear = flowpm.linear_field(nc, bs, ipklin, batch_size=batch_size)
if args.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=a0, order=args.lpt_order)
final_state = nbody(state, stages, nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=af, order=args.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
base = tfinal_field
if Rsm != 0:
basek = r2c3d(tfinal_field, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
galmean = tfp.distributions.Poisson(rate = plambda * (1 + base))
result = galmean.sample()
return linear, tfinal_field, result, base
def test_save_state():
"""
Tests the BigFile saving function
"""
klin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
a0 = 0.1
nc = [16, 16, 16]
boxsize = [100., 100., 100.]
cosmo = flowpm.cosmology.Planck15()
initial_conditions = flowpm.linear_field(
nc, # size of the cube
boxsize, # Physical size of the cube
ipklin, # Initial powerspectrum
batch_size=2)
# Sample particles
state = flowpm.lpt_init(cosmo, initial_conditions, a0)
with tempfile.TemporaryDirectory() as tmpdirname:
filename = tmpdirname + '/testsave'
save_state(cosmo, state, a0, nc, boxsize, filename)
# Now try to reload the information using BigFile
bf = bigfile.BigFile(filename)
# Testing recovery of header
header = bf['Header']
assert_allclose(np.array(header.attrs['NC']), np.array(nc))
assert_allclose(np.array(header.attrs['BoxSize']), np.array(boxsize))
assert_allclose(np.array(header.attrs['OmegaCDM']),
np.array(cosmo.Omega_c))
assert_allclose(np.array(header.attrs['OmegaB']),
np.array(cosmo.Omega_b))
assert_allclose(np.array(header.attrs['OmegaK']),
np.array(cosmo.Omega_k))
assert_allclose(np.array(header.attrs['h']), np.array(cosmo.h))
assert_allclose(np.array(header.attrs['Sigma8']),
np.array(cosmo.sigma8))
assert_allclose(np.array(header.attrs['w0']), np.array(cosmo.w0))
assert_allclose(np.array(header.attrs['wa']), np.array(cosmo.wa))
assert_allclose(np.array(header.attrs['Time']), np.array(a0))
# Testing recovery of data
pos = bf['1/Position']
assert_allclose(pos[:], state[0, 1].numpy() / nc[0] * boxsize[0])
vel = bf['1/Velocity']
assert_allclose(vel[:], state[1, 1].numpy() / nc[0] * boxsize[0])
# Closing file
bf.close()
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 test_convergence_Born(return_results=False):
""" This function tests that given a set of density planes,
both lenstools and flowpm recover the same convergence maps in
angular coordinates.
"""
klin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
cosmo = flowpm.cosmology.Planck15()
a0 = 0.9
# Create a state vector
initial_conditions = flowpm.linear_field([nc, nc, 10 * nc],
[bs, bs, 10 * bs],
ipklin,
batch_size=2)
state = flowpm.lpt_init(cosmo, initial_conditions, a0)
r = tf.linspace(0., 10 * bs, 11)
r_center = 0.5 * (r[1:] + r[:-1])
a_center = flowpm.background.a_of_chi(cosmo, r_center)
constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
# To make it convenient to access simulation properties in lenstools
# let's quicly export and reload the sim
# TODO: remove the need for this!
flowpm.io.save_state(cosmo,
state,
a0, [nc, nc, 10 * nc], [bs, bs, 10 * bs],
'snapshot_born_testing',
attrs={'comoving_distance': r_center[0]})
# Reload the snapshot with lenstools
snapshot = FlowPMSnapshot.open('snapshot_born_testing')
# Get some density planes and create lenstool tracer
lensplanes = []
tracer = lt.simulations.RayTracer(lens_type=lt.simulations.DensityPlane)
for i in range(len(r_center)):
plane = flowpm.raytracing.density_plane(
state, [nc, nc, 10 * nc],
r_center[i] / bs * nc,
width=nc,
plane_resolution=plane_resolution)
r, a, p = r_center[i], a_center[i], plane[0]
lensplanes.append((r, a, plane))
density_normalization = bs * r / a
# We upsample the lensplanes before giving them to lenstools because
# lentools is using a weird kind of interpolation when converting from
# comoving coordinates to angular coords. with a larger
p = tf.image.resize(
tf.reshape(p, [1, plane_resolution, plane_resolution, 1]),
[2048, 2048])
p = (p[0, :, :, 0] * constant_factor * density_normalization).numpy()
p = p - np.mean(p)
lt_plane = lt.simulations.DensityPlane(
p,
angle=snapshot.header["box_size"],
redshift=1 / a - 1,
cosmology=snapshot.cosmology)
tracer.addLens(lt_plane)
# Adding dummy lensplane at the end
tracer.addLens(
lt.simulations.DensityPlane(np.zeros((2048, 2048)),
angle=snapshot.header["box_size"],
redshift=0.99,
cosmology=snapshot.cosmology))
tracer.addLens(
lt.simulations.DensityPlane(np.zeros((2048, 2048)),
angle=snapshot.header["box_size"],
redshift=2,
cosmology=snapshot.cosmology))
tracer.reorderLenses()
# Create an array of coordinates at which to retrieve the convernge maps
xgrid, ygrid = np.meshgrid(
np.linspace(0, field, npix, endpoint=False), # range of X coordinates
np.linspace(0, field, npix, endpoint=False)) # range of Y coordinates
coords = np.stack([xgrid, ygrid], axis=0) * u.deg
c = coords.reshape([2, -1]).T
# Compute convergence map with lenstool
lt_map = tracer.convergenceBorn(coords, z=1.0)
# Compute convergemce map with flowpm
fpm_map = flowpm.raytracing.convergenceBorn(cosmo,
lensplanes,
bs / nc,
bs,
c.to(u.rad),
z_source=tf.ones([1]),
field_npix=npix)
# Comparing the final maps
assert_allclose(lt_map,
fpm_map[0].numpy().reshape([npix, npix, -1])[:, :, -1],
atol=5e-4)
if return_results:
return lt_map, fpm_map
def test_density_plane(return_results=False):
""" Tests cutting density planes from snapshots against lenstools
"""
klin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
cosmo = flowpm.cosmology.Planck15()
a0 = 0.9
r0 = flowpm.background.rad_comoving_distance(cosmo, a0)
# Create a state vector
initial_conditions = flowpm.linear_field(nc, bs, ipklin, batch_size=2)
state = flowpm.lpt_init(cosmo, initial_conditions, a0)
# Export the snapshot
flowpm.io.save_state(cosmo,
state,
a0, [nc, nc, nc], [bs, bs, bs],
'snapshot_density_testing',
attrs={'comoving_distance': r0})
# Reload the snapshot with lenstools
snapshot = FlowPMSnapshot.open('snapshot_density_testing')
# Cut a lensplane in the middle of the volume
lt_plane, resolution, NumPart = snapshot.cutPlaneGaussianGrid(
normal=2,
plane_resolution=plane_resolution,
center=(bs / 2) * snapshot.Mpc_over_h,
thickness=(bs / 4) * snapshot.Mpc_over_h,
left_corner=np.zeros(3) * snapshot.Mpc_over_h,
smooth=None,
kind='density')
# Cut the same lensplane with flowpm
fpm_plane = flowpm.raytracing.density_plane(
state,
nc,
center=nc / 2,
width=nc / 4,
plane_resolution=plane_resolution)
# Apply additional normalization terms to match lenstools definitions
constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
density_normalization = bs / 4 * r0 / a0
fpm_plane = fpm_plane * density_normalization * constant_factor
# Checking first the mean value, which accounts for any normalization
# issues
assert_allclose(np.mean(fpm_plane[0]), np.mean(lt_plane), rtol=1e-5)
# To check pixelwise difference, we need to do some smoothing as lenstools and
# flowpm use different painting kernels
smooth_lt_plane = np.fft.ifft2(fourier_gaussian(np.fft.fft2(lt_plane),
3)).real
smooth_fpm_plane = np.fft.ifft2(
fourier_gaussian(np.fft.fft2(fpm_plane[0]), 3)).real
assert_allclose(smooth_fpm_plane, smooth_lt_plane, rtol=2e-2)
if return_results:
return fpm_plane, lt_plane, smooth_fpm_plane, smooth_lt_plane
def pmgraph(lin):
state = lpt_init(lin, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(lin), final_state[0])
return final_field
def main(_):
infield = True
mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
startw = time.time()
print(mesh_shape)
#layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
#mesh_shape = [("row", FLAGS.nx), ("col", FLAGS.ny)]
layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"),
("ny", "col"), ("ty", "row"), ("tz", "col"),
("ty_lr", "row"), ("tz_lr", "col"), ("nx_block", "row"),
("ny_block", "col")]
# Resolve the cluster from SLURM environment
cluster = tf.distribute.cluster_resolver.SlurmClusterResolver(
{"mesh": mesh_shape.size // FLAGS.gpus_per_task},
port_base=8822,
gpus_per_node=FLAGS.gpus_per_node,
gpus_per_task=FLAGS.gpus_per_task,
tasks_per_node=FLAGS.tasks_per_node)
cluster_spec = cluster.cluster_spec()
print(cluster_spec)
# Create a server for all mesh members
server = tf.distribute.Server(cluster_spec, "mesh", cluster.task_id)
print(server)
if cluster.task_id > 0:
server.join()
# Otherwise we are the main task, let's define the devices
devices = [
"/job:mesh/task:%d/device:GPU:%d" % (i, j)
for i in range(cluster_spec.num_tasks("mesh"))
for j in range(FLAGS.gpus_per_task)
]
print("List of devices", devices)
mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
mesh_shape, layout_rules, devices)
##Begin here
klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
#If initc, run normal flowpm to generate data
tf.reset_default_graph()
if infield:
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
start = time.time()
with tf.Session(server.target) as sess:
ic, fin = sess.run([tfic, tfinal_field])
print("\nTime taken for the vanilla flowpm thingy :\n ",
time.time() - start)
else:
ic = None
tf.reset_default_graph()
print('ic constructed')
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
initial_conditions, final_field, input_field = nbody_prototype(
mesh, infield, nc=FLAGS.nc, batch_size=FLAGS.batch_size)
# Lower mesh computation
start = time.time()
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
restore_hook = mtf.MtfRestoreHook(lowering)
end = time.time()
print('\n Time for lowering : %f \n' % (end - start))
tf_initc = lowering.export_to_tf_tensor(initial_conditions)
tf_final = lowering.export_to_tf_tensor(final_field)
nc = FLAGS.nc
with tf.Session(server.target) as sess:
start = time.time()
if infield:
ic_check, fin_check = sess.run([tf_initc, tf_final],
feed_dict={input_field: ic})
else:
ic_check, fin_check = sess.run([tf_initc, tf_final])
ic, fin = ic_check, fin_check
print('\n Time for the mesh run : %f \n' % (time.time() - start))
plt.figure(figsize=(15, 3))
plt.subplot(141)
plt.imshow(ic_check[0].sum(axis=2))
plt.title('Initial Conditions')
plt.subplot(142)
plt.imshow(fin[0].sum(axis=2))
plt.title('TensorFlow (single GPU)')
plt.colorbar()
plt.subplot(143)
plt.imshow(fin_check[0].sum(axis=2))
plt.title('Mesh TensorFlow')
plt.colorbar()
plt.subplot(144)
plt.imshow((fin_check[0] - fin[0]).sum(axis=2))
plt.title('Residuals')
plt.colorbar()
plt.savefig("comparison_mesh.png")
exit(0)
##
exit(0)
#tf.reset_default_graph()
# Run normal flowpm to generate data
try:
ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy')
print('Data loaded')
except Exception as e:
print('Exception occured', e)
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100,
dtype=dtype)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
ic, fin = sess.run([tfic, tfinal_field])
np.save(fpath + 'ic', ic)
np.save(fpath + 'final', fin)
################################################################
#tf.reset_default_graph()
print('ic constructed')
noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(np.float32)
data_noised = fin + noise
def main(_):
dtype = tf.float32
startw = time.time()
tf.random.set_random_seed(100)
np.random.seed(100)
# Compute a few things first, using simple tensorflow
a0 = FLAGS.a0
a = FLAGS.af
nsteps = FLAGS.nsteps
bs, nc = FLAGS.box_size, FLAGS.nc
klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
stages = np.linspace(a0, a, nsteps, endpoint=True)
tf.reset_default_graph()
# Run normal flowpm to generate data
try:
ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy')
print('Data loaded')
except Exception as e:
print('Exception occured', e)
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100,
dtype=dtype)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
ic, fin = sess.run([tfic, tfinal_field])
np.save(fpath + 'ic', ic)
np.save(fpath + 'final', fin)
k, pic = tools.power(ic[0] + 1, boxsize=bs)
k, pfin = tools.power(fin[0], boxsize=bs)
plt.plot(k, pic)
plt.plot(k, pfin)
plt.loglog()
plt.grid(which='both')
plt.savefig('pklin.png')
plt.close()
print(pic)
print(pfin)
#sys.exit(-1)
################################################################
tf.reset_default_graph()
print('ic constructed')
noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(
np.float32) * 1
data_noised = fin + noise
data = data_noised
startpos = noise.copy().flatten().astype(np.float32)
x0 = tf.placeholder(dtype=tf.float32,
shape=data.flatten().shape,
name='initlin')
Rsm = tf.placeholder(tf.float32, name='smoothing')
def recon_prototype(linearflat):
"""
"""
linear = tf.reshape(linearflat, data.shape)
#
#loss = tf.reduce_sum(tf.square(linear - minimum))
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
#
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
chisq = tf.multiply(chisq, 1 / nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
prior = tf.multiply(prior, 1 / nc**3, name='prior')
#
loss = chisq + prior
grad = tf.gradients(loss, linearflat)
print(grad)
return loss, grad[0]
@tf.function
def min_lbfgs():
return tfp.optimizer.lbfgs_minimize(
#make_val_and_grad_fn(recon_prototype),
recon_prototype,
initial_position=x0,
tolerance=1e-10,
max_iterations=100)
with tf.Session() as sess:
start = time.time()
results = sess.run(min_lbfgs(), {Rsm: 2, x0: startpos})
print("\n")
print(results)
print("\n")
minimum = results.position
print(minimum)
print("\nTime taken : ", time.time() - start)
start = time.time()
results = sess.run(min_lbfgs(), {Rsm: 1, x0: minimum})
print("\n")
print(results)
minimum = results.position
print("\n")
print(minimum)
print("\nTime taken : ", time.time() - start)
start = time.time()
results = sess.run(min_lbfgs(), {Rsm: 0, x0: minimum})
print("\n")
print(results)
minimum = results.position
print("\n")
print(minimum)
print("\nTime taken : ", time.time() - start)
tf.reset_default_graph()
print("\n")
print('\nminimized\n')
tfic = linear_field(
FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100,
dtype=dtype) * 0 + minimum.reshape(data_noised.shape)
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
minic, minfin = sess.run([tfic, tfinal_field])
dg.saveimfig(0, [minic, minfin], [ic, fin], fpath + '')
dg.save2ptfig(0, [minic, minfin], [ic, fin], fpath + '', bs)
np.save(fpath + 'recon0ic', minic)
np.save(fpath + 'recon-final', minfin)
##
exit(0)
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"))
def main(_):
dtype = tf.float32
startw = time.time()
tf.random.set_random_seed(100)
np.random.seed(100)
# Compute a few things first, using simple tensorflow
a0 = FLAGS.a0
a = FLAGS.af
nsteps = FLAGS.nsteps
bs, nc = FLAGS.box_size, FLAGS.nc
klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
stages = np.linspace(a0, a, nsteps, endpoint=True)
tf.reset_default_graph()
# Run normal flowpm to generate data
try:
ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy')
print('Data loaded')
except Exception as e:
print('Exception occured', e)
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100,
dtype=dtype)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
ic, fin = sess.run([tfic, tfinal_field])
np.save(fpath + 'ic', ic)
np.save(fpath + 'final', fin)
tf.reset_default_graph()
print('ic constructed')
linear, final_field, update_ops, loss, chisq, prior, Rsm = recon_prototype(
fin)
#initial_conditions = recon_prototype(mesh, fin, nc=FLAGS.nc, batch_size=FLAGS.batch_size, dtype=dtype)
# Lower mesh computation
with tf.Session() as sess:
#ic_check, fin_check = sess.run([tf_initc, tf_final])
#sess.run(tf_linear_op, feed_dict={input_field:ic})
#ic_check, fin_check = sess.run([linear, final_field])
#dg.saveimfig('-check', [ic_check, fin_check], [ic, fin], fpath)
#dg.save2ptfig('-check', [ic_check, fin_check], [ic, fin], fpath, bs)
#sess.run(tf_linear_op, feed_dict={input_field:np.random.normal(size=ic.size).reshape(ic.shape)})
sess.run(tf.global_variables_initializer())
ic0, fin0 = sess.run([linear, final_field])
dg.saveimfig('-init', [ic0, fin0], [ic, fin], fpath)
start = time.time()
titer = 20
niter = 201
iiter = 0
start0 = time.time()
RRs = [4, 2, 1, 0.5, 0]
lrs = np.array([0.1, 0.1, 0.1, 0.1, 0.1]) * 2
#lrs = [0.1, 0.05, 0.01, 0.005, 0.001]
for iR, zlR in enumerate(zip(RRs, lrs)):
RR, lR = zlR
for ff in [fpath + '/figs-R%02d' % (10 * RR)]:
try:
os.makedirs(ff)
except Exception as e:
print(e)
for i in range(niter):
iiter += 1
sess.run(update_ops, {Rsm: RR})
print(sess.run([loss, chisq, prior], {Rsm: RR}))
if (i % titer == 0):
end = time.time()
print('Iter : ', i)
print('Time taken for %d iterations: ' % titer,
end - start)
start = end
##
#ic1, fin1, cc, pp = sess.run([tf_initc, tf_final, tf_chisq, tf_prior], {R0:RR})
#ic1, fin1, cc, pp = sess.run([tf_initc, tf_final, tf_chisq, tf_prior], {R0:RR})
ic1, fin1 = sess.run([linear, final_field])
#print('Chisq and prior are : ', cc, pp)
dg.saveimfig(i, [ic1, fin1], [ic, fin],
fpath + '/figs-R%02d' % (10 * RR))
dg.save2ptfig(i, [ic1, fin1], [ic, fin],
fpath + '/figs-R%02d' % (10 * RR), bs)
dg.saveimfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs')
dg.save2ptfig(i * (iR + 1), [ic1, fin1], [ic, fin],
fpath + '/figs', bs)
ic1, fin1 = sess.run([linear, final_field])
print('Total time taken for %d iterations is : ' % iiter,
time.time() - start0)
dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath)
dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath, bs)
np.save(fpath + 'ic_recon', ic1)
np.save(fpath + 'final_recon', fin1)
print('Total wallclock time is : ', time.time() - start0)
##
exit(0)
def main(_):
infield = True
dtype = tf.float32
mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
nc, bs = FLAGS.nc, FLAGS.box_size
a0, a, nsteps = FLAGS.a0, FLAGS.af, FLAGS.nsteps
stages = np.linspace(a0, a, nsteps, endpoint=True)
numd = 1e-3
##Begin here
klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
#pypath = '/global/cscratch1/sd/chmodi/cosmo4d/output/version2/L0400_N0128_05step-fof/lhd_S0100/n10/opt_s999_iM12-sm3v25off/meshes/'
final = tools.readbigfile('../data//L0400_N0128_S0100_05step/mesh/d/')
ic = tools.readbigfile('../data/L0400_N0128_S0100_05step/mesh/s/')
fpos = tools.readbigfile(
'../data/L0400_N0128_S0100_05step/dynamic/1/Position/')
hpos = tools.readbigfile(
'../data/L0400_N0512_S0100_40step/FOF/PeakPosition//')[1:int(bs**3 *
numd)]
hmass = tools.readbigfile(
'../data/L0400_N0512_S0100_40step/FOF/Mass//')[1:int(bs**3 *
numd)].flatten()
meshpos = tools.paintcic(hpos, bs, nc)
meshmass = tools.paintcic(hpos, bs, nc, hmass.flatten() * 1e10)
data = meshmass
data /= data.mean()
data -= 1
kv = tools.fftk([nc, nc, nc], bs, symmetric=True, dtype=np.float32)
datasm = tools.fingauss(data, kv, 3, np.pi * nc / bs)
ic, data = np.expand_dims(ic, 0), np.expand_dims(data,
0).astype(np.float32)
datasm = np.expand_dims(datasm, 0).astype(np.float32)
print("Min in data : %0.4e" % datasm.min())
np.save(fpath + 'ic', ic)
np.save(fpath + 'data', data)
####################################################
#
tf.reset_default_graph()
tfic = tf.constant(ic.astype(np.float32))
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
state = sess.run(final_state)
fpos = state[0, 0] * bs / nc
bparams, bmodel = getbias(bs, nc, data[0] + 1, ic[0], fpos)
#bmodel += 1 #np.expand_dims(bmodel, 0) + 1
errormesh = data - np.expand_dims(bmodel, 0)
kerror, perror = tools.power(errormesh[0] + 1, boxsize=bs)
kerror, perror = kerror[1:], perror[1:]
print("Error power spectra", kerror, perror)
print("\nkerror", kerror.min(), kerror.max(), "\n")
print("\nperror", perror.min(), perror.max(), "\n")
suff = "-error"
dg.saveimfig(suff, [ic, errormesh], [ic, data], fpath + '/figs/')
dg.save2ptfig(suff, [ic, errormesh], [ic, data], fpath + '/figs/', bs)
ipkerror = iuspline(kerror, perror)
####################################################
#stdinit = srecon.standardinit(bs, nc, meshpos, hpos, final, R=8)
recon_estimator = tf.estimator.Estimator(model_fn=model_fn,
model_dir=fpath)
def predict_input_fn(data=data,
M0=0.,
w=3.,
R0=0.,
off=None,
istd=None,
x0=None):
features = {}
features['datasm'] = data
features['R0'] = R0
features['x0'] = x0
features['bparams'] = bparams
features['ipkerror'] = [kerror, perror] #ipkerror
return features, None
eval_results = recon_estimator.predict(
input_fn=lambda: predict_input_fn(x0=ic), yield_single_examples=False)
for i, pred in enumerate(eval_results):
if i > 0: break
suff = '-model'
dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/')
dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/', bs)
np.save(fpath + '/reconmeshes/ic_true' + suff, pred['ic'])
np.save(fpath + '/reconmeshes/fin_true' + suff, pred['final'])
np.save(fpath + '/reconmeshes/model_true' + suff, pred['model'])
#
randominit = np.random.normal(size=data.size).reshape(data.shape)
#eval_results = recon_estimator.predict(input_fn=lambda : predict_input_fn(x0 = np.expand_dims(stdinit, 0)), yield_single_examples=False)
eval_results = recon_estimator.predict(
input_fn=lambda: predict_input_fn(x0=randominit),
yield_single_examples=False)
for i, pred in enumerate(eval_results):
if i > 0: break
suff = '-init'
dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/')
dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/', bs)
np.save(fpath + '/reconmeshes/ic_init' + suff, pred['ic'])
np.save(fpath + '/reconmeshes/fin_init' + suff, pred['final'])
np.save(fpath + '/reconmeshes/model_init' + suff, pred['model'])
#
# Train and evaluate model.
RRs = [4., 2., 1., 0.5, 0.]
niter = 100
iiter = 0
for R0 in RRs:
print('\nFor iteration %d\n' % iiter)
print('With R0=%0.2f \n' % (R0))
def train_input_fn():
features = {}
features['datasm'] = data
features['R0'] = R0
features['bparams'] = bparams
features['ipkerror'] = [kerror, perror] #ipkerror
#features['x0'] = np.expand_dims(stdinit, 0)
features['x0'] = randominit
features['lr'] = 0.01
return features, None
recon_estimator.train(input_fn=train_input_fn, max_steps=iiter + niter)
eval_results = recon_estimator.predict(input_fn=predict_input_fn,
yield_single_examples=False)
for i, pred in enumerate(eval_results):
if i > 0: break
iiter += niter #
suff = '-%d-R%d' % (iiter, R0)
dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/')
dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data],
fpath + '/figs/', bs)
np.save(fpath + '/reconmeshes/ic' + suff, pred['ic'])
np.save(fpath + '/reconmeshes/fin' + suff, pred['final'])
np.save(fpath + '/reconmeshes/model' + suff, pred['model'])
sys.exit(0)
##
exit(0)
def main(_):
dtype = tf.float32
mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
print(mesh_shape)
#layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
#mesh_shape = [("row", FLAGS.nx), ("col", FLAGS.ny)]
layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"),
("ny", "col"), ("ty", "row"), ("tz", "col"),
("ty_lr", "row"), ("tz_lr", "col"), ("nx_block", "row"),
("ny_block", "col")]
# Resolve the cluster from SLURM environment
cluster = tf.distribute.cluster_resolver.SlurmClusterResolver(
{"mesh": mesh_shape.size // FLAGS.gpus_per_task},
port_base=8822,
gpus_per_node=FLAGS.gpus_per_node,
gpus_per_task=FLAGS.gpus_per_task,
tasks_per_node=FLAGS.tasks_per_node)
cluster_spec = cluster.cluster_spec()
print(cluster_spec)
# Create a server for all mesh members
server = tf.distribute.Server(cluster_spec, "mesh", cluster.task_id)
print(server)
if cluster.task_id > 0:
server.join()
# Otherwise we are the main task, let's define the devices
devices = [
"/job:mesh/task:%d/device:GPU:%d" % (i, j)
for i in range(cluster_spec.num_tasks("mesh"))
for j in range(FLAGS.gpus_per_task)
]
print("List of devices", devices)
mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
mesh_shape, layout_rules, devices)
##Begin here
klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
tf.reset_default_graph()
# Run normal flowpm to generate data
try:
ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy')
print('Data loaded')
except Exception as e:
print('Exception occured', e)
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100,
dtype=dtype)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session(server.target) as sess:
ic, fin = sess.run([tfic, tfinal_field])
np.save(fpath + 'ic', ic)
np.save(fpath + 'final', fin)
tf.reset_default_graph()
print('ic constructed')
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
initial_conditions, final_field, loss, var_grads, update_op, linear_op, input_field, lr, R0 = recon_prototype(
mesh, fin, nc=FLAGS.nc, batch_size=FLAGS.batch_size, dtype=dtype)
# Lower mesh computation
start = time.time()
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
restore_hook = mtf.MtfRestoreHook(lowering)
end = time.time()
print('\n Time for lowering : %f \n' % (end - start))
tf_initc = lowering.export_to_tf_tensor(initial_conditions)
tf_final = lowering.export_to_tf_tensor(final_field)
tf_grads = lowering.export_to_tf_tensor(var_grads[0])
tf_linear_op = lowering.lowered_operation(linear_op)
tf_update_ops = lowering.lowered_operation(update_op)
n_block_x, n_block_y, n_block_z = FLAGS.nx, FLAGS.ny, 1
nc = FLAGS.nc
ic_hrshape = ic.reshape([
FLAGS.batch_size, n_block_x, nc // n_block_x, n_block_y,
nc // n_block_y, n_block_z, nc // n_block_z
])
ic_hrshape = np.transpose(ic_hrshape, [0, 1, 3, 5, 2, 4, 6])
with tf.Session(server.target) as sess:
#ic_check, fin_check = sess.run([tf_initc, tf_final])
sess.run(tf_linear_op, feed_dict={input_field: ic_hrshape})
ic_check, fin_check = sess.run([tf_initc, tf_final])
dg.saveimfig('-check', [ic_check, fin_check], [ic, fin], fpath)
dg.save2ptfig('-check', [ic_check, fin_check], [ic, fin], fpath, bs)
sess.run(tf_linear_op,
feed_dict={
input_field:
np.random.normal(size=ic.size).reshape(ic_hrshape.shape)
})
ic0, fin0 = sess.run([tf_initc, tf_final])
dg.saveimfig('-init', [ic0, fin0], [ic, fin], fpath)
start = time.time()
niter = 5
iiter = 0
start0 = time.time()
RRs = [4, 2, 1, 0.5, 0]
lrs = np.array([0.2, 0.15, 0.1, 0.1, 0.1])
#lrs = [0.1, 0.05, 0.01, 0.005, 0.001]
for iR, zlR in enumerate(zip(RRs, lrs)):
RR, lR = zlR
#for ff in [fpath + '/figs-R%02d'%(10*RR)]:
for ff in [fpath + '/figsiter']:
try:
os.makedirs(ff)
except Exception as e:
print(e)
for i in range(301):
if (i % niter == 0):
end = time.time()
print('Iter : ', i)
print('Time taken for %d iterations: ' % niter,
end - start)
start = end
##
ic1, fin1 = sess.run([tf_initc, tf_final])
#dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath+'/figs-R%02d'%(10*RR))
#dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath+'/figs-R%02d'%(10*RR), bs)
dg.saveimfig2x2(iiter, [ic1, fin1], [ic, fin],
fpath + '/figsiter')
#
sess.run(tf_update_ops, {lr: lR, R0: RR})
iiter += 1
dg.saveimfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs')
dg.save2ptfig(i * (iR + 1), [ic1, fin1], [ic, fin],
fpath + '/figs', bs)
ic1, fin1 = sess.run([tf_initc, tf_final])
print('Total time taken for %d iterations is : ' % iiter,
time.time() - start0)
dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath)
dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath, bs)
np.save(fpath + 'ic_recon', ic1)
np.save(fpath + 'final_recon', fin1)
print('Total wallclock time is : ', time.time() - start0)
##
exit(0)
def recon_prototype(data,
anneal=True,
nc=FLAGS.nc,
bs=FLAGS.box_size,
batch_size=FLAGS.batch_size,
a0=FLAGS.a0,
a=FLAGS.af,
nsteps=FLAGS.nsteps,
dtype=tf.float32):
"""
Prototype of function computing LPT deplacement.
Returns output tensorflow and mesh tensorflow tensors
"""
if dtype == tf.float32:
npdtype = "float32"
cdtype = tf.complex64
elif dtype == tf.float64:
npdtype = "float64"
cdtype = tf.complex128
print(dtype, npdtype)
#graph = mtf.Graph()
#mesh = mtf.Mesh(graph, "my_mesh")
linear = tf.get_variable('linmesh',
shape=(1, nc, nc, nc),
dtype=tf.float32,
initializer=tf.random_normal_initializer(),
trainable=True)
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
##Anneal
Rsm = tf.placeholder(tf.float32, name='smoothing')
if anneal:
#def anneal
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
#chisq = tf.multiply(chisq, 1/nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
#prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = chisq + prior
## #optimizer = tf.optimize.AdamWeightDecayOptimizer(0.01)
## opt = tf.train.GradientDescentOptimizer(learning_rate=0.01)
##
## # Compute the gradients for a list of variables.
## grads_and_vars = opt.compute_gradients(loss, [linear])
## print("\ngradients : ", grads_and_vars)
## update_ops = opt.apply_gradients(grads_and_vars)
##
## #optimizer = tf.keras.optimizers.Adam(0.01)
## #var_grads = tf.gradients([loss], [linear])
##
##
## #update_ops = optimizer.apply_gradients(var_grads, linear)
## #update_ops = optimizer.apply_gradients(zip(var_grads, [linear]))
## #update_ops = None
## #lr = tf.placeholder(tf.float32, shape=())
## #update_op = mtf.assign(fieldvar, fieldvar - var_grads[0]*lr)
##
return linear, final_field, loss, chisq, prior
for ss in range(100, 1000, 100):
path = '../data/z00/L%04d_N%04d_S%04d_%dstep/'%(bs, nc, ss, nsteps)
ic = np.expand_dims(tools.readbigfile(path + '/mesh/s/').astype(np.float32), axis=0)
print(ic.shape)
initial_conditions = tf.cast(tf.constant(ic), tf.float32)
print(initial_conditions)
# Sample particles
state = flowpm.lpt_init(initial_conditions, a0=ainit)
# Evolve particles down to z=0
final_state = flowpm.nbody(state, stages, nc)
# Retrieve final density field
final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions), final_state[0])
with tf.Session() as sess:
#ic, sim = sess.run([initial_conditions, final_field])
sim = sess.run(final_field)
print(sim.shape)
np.save(path + '/mesh/d', np.squeeze(sim))
def main(_):
mesh_shape = [("row", 2), ("col", 2)]
layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"),
("ny", "col"), ("ty_lr", "row"), ("tz_lr", "col"),
("nx_block", "row"), ("ny_block", "col")]
mesh_hosts = ["localhost:%d" % (8222 + j) for j in range(4)]
# Create a cluster from the mesh hosts.
cluster = tf.train.ClusterSpec({
"mesh": mesh_hosts,
"master": ["localhost:8488"]
})
# Create a server for local mesh members
server = tf.train.Server(cluster, job_name="master", task_index=0)
mesh_devices = [
'/job:mesh/task:%d' % i for i in range(cluster.num_tasks("mesh"))
]
print("List of devices", mesh_devices)
mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
mesh_shape, layout_rules, mesh_devices)
# Build the model
# Create computational graphs and some initializations
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "nbody_mesh")
# Compute a few things first, using simple tensorflow
a0 = FLAGS.a0
a = FLAGS.af
nsteps = FLAGS.nsteps
bs, nc = FLAGS.box_size, FLAGS.nc
klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
stages = np.linspace(a0, a, nsteps, endpoint=True)
#pt = PerturbationGrowth(cosmology, a=[a], a_normalize=1.0)
# Generate a batch of 3D initial conditions
initial_conditions = flowpm.linear_field(
FLAGS.nc, # size of the cube
FLAGS.box_size, # Physical size of the cube
ipklin, # Initial power spectrum
batch_size=FLAGS.batch_size)
state = lpt_init(initial_conditions, a0=a0, order=1)
#final_state = state
final_state = nbody(state, stages, nc)
tfinal_field = cic_paint(tf.zeros_like(initial_conditions), final_state[0])
# Compute necessary Fourier kernels
kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False)
from flowpm.kernels import laplace_kernel, gradient_kernel
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
grad_x = gradient_kernel(kvec, 0)
grad_y = gradient_kernel(kvec, 1)
grad_z = gradient_kernel(kvec, 2)
derivs = [lap, grad_x, grad_y, grad_z]
mesh_final_field = lpt_prototype(mesh,
initial_conditions,
derivs,
bs=FLAGS.box_size,
nc=FLAGS.nc,
batch_size=FLAGS.batch_size)
# Lower mesh computation
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
# Retrieve output of computation
result = lowering.export_to_tf_tensor(mesh_final_field)
with tf.Session(server.target,
config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)) as sess:
a, b, c = sess.run([initial_conditions, tfinal_field, result])
np.save('init', a)
np.save('reference_final', b)
np.save('mesh_pyramid', c)
plt.figure(figsize=(15, 3))
plt.subplot(141)
plt.imshow(a[0].sum(axis=2))
plt.title('Initial Conditions')
plt.subplot(142)
plt.imshow(b[0].sum(axis=2))
plt.title('TensorFlow (single GPU)')
plt.colorbar()
plt.subplot(143)
plt.imshow(c[0].sum(axis=2))
plt.title('Mesh TensorFlow Single')
plt.colorbar()
plt.subplot(144)
plt.imshow((b[0] - c[0]).sum(axis=2))
plt.title('Residuals')
plt.colorbar()
plt.savefig("comparison-single.png")
exit(0)
def main(_):
dtype = tf.float32
mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
startw = time.time()
print(mesh_shape)
##
##
##Begin here
klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('..//data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
tf.reset_default_graph()
# Run normal flowpm to generate data
try:
ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy')
print('Data loaded')
except Exception as e:
print('Exception occured', e)
tfic = linear_field(FLAGS.nc,
FLAGS.box_size,
ipklin,
batch_size=1,
seed=100,
dtype=dtype)
if FLAGS.nbody:
state = lpt_init(tfic, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
else:
final_state = lpt_init(tfic, a0=stages[-1], order=1)
tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0])
with tf.Session() as sess:
ic, fin = sess.run([tfic, tfinal_field])
np.save(fpath + 'ic', ic)
np.save(fpath + 'final', fin)
print(ic.shape, fin.shape)
########################################################
print(ic.shape, fin.shape)
recon_estimator = tf.estimator.Estimator(model_fn=model_fn,
model_dir=fpath)
def eval_input_fn():
features = {}
features['data'] = fin
features['R0'] = 0
features['x0'] = None
features['lr'] = 0
return features, None
# Train and evaluate model.
RRs = [4., 2., 1., 0.5, 0.]
niter = 200
iiter = 0
for R0 in RRs:
print('\nFor iteration %d and R=%0.1f\n' % (iiter, R0))
def train_input_fn():
features = {}
features['data'] = fin
features['R0'] = R0
features['x0'] = np.random.normal(size=fin.size).reshape(fin.shape)
features['lr'] = 0.01
return features, None
for _ in range(1):
recon_estimator.train(input_fn=train_input_fn,
max_steps=iiter + niter)
eval_results = recon_estimator.predict(input_fn=eval_input_fn,
yield_single_examples=False)
for i, pred in enumerate(eval_results):
if i > 0: break
iiter += niter #
dg.saveimfig(iiter, [pred['ic'], pred['data']], [ic, fin],
fpath + '/figs/')
dg.save2ptfig(iiter, [pred['ic'], pred['data']], [ic, fin],
fpath + '/figs/', bs)
sys.exit(0)
def pm(linear):
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, nc)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
return tfinal_field
def main(_):
mesh_shape = [("row", FLAGS.nx), ("col", FLAGS.ny)]
layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"),
("ny", "col"), ("ty_lr", "row"), ("tz_lr", "col"),
("nx_block", "row"), ("ny_block", "col")]
mesh_impl = HvdSimdMeshImpl(mtf.convert_to_shape(mesh_shape),
mtf.convert_to_layout_rules(layout_rules))
# Build the model
# Create computational graphs and some initializations
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "nbody_mesh")
# Compute a few things first, using simple tensorflow
a0 = FLAGS.a0
a = FLAGS.af
nsteps = FLAGS.nsteps
bs, nc = FLAGS.box_size, FLAGS.nc
klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
stages = np.linspace(a0, a, nsteps, endpoint=True)
#pt = PerturbationGrowth(cosmology, a=[a], a_normalize=1.0)
# Generate a batch of 3D initial conditions
initial_conditions = flowpm.linear_field(
FLAGS.nc, # size of the cube
FLAGS.box_size, # Physical size of the cube
ipklin, # Initial power spectrum
batch_size=FLAGS.batch_size)
cosmo = flowpm.cosmology.Planck15()
state = lpt_init(cosmo, initial_conditions, a, order=1)
final_state = state
#final_state = nbody(cosmo, state, stages, nc)
tfinal_field = cic_paint(tf.zeros_like(initial_conditions), final_state[0])
# Compute necessary Fourier kernels
kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False)
from flowpm.kernels import laplace_kernel, gradient_kernel
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
grad_x = gradient_kernel(kvec, 0)
grad_y = gradient_kernel(kvec, 1)
grad_z = gradient_kernel(kvec, 2)
derivs = [lap, grad_x, grad_y, grad_z]
mesh_final_field = lpt_prototype(mesh,
initial_conditions,
derivs,
bs=FLAGS.box_size,
nc=FLAGS.nc,
batch_size=FLAGS.batch_size)
# Lower mesh computation
lowering = mtf.Lowering(graph, {mesh: mesh_impl})
# Retrieve output of computation
result = lowering.export_to_tf_tensor(mesh_final_field)
with tf.Session() as sess:
a, b, c = sess.run([initial_conditions, tfinal_field, result])
if comm.rank == 0:
np.save('init', a)
np.save('reference_final', b)
np.save('mesh_pyramid', c)
plt.figure(figsize=(15, 3))
plt.subplot(141)
plt.imshow(a[0].sum(axis=2))
plt.title('Initial Conditions')
plt.subplot(142)
plt.imshow(b[0].sum(axis=2))
plt.title('TensorFlow (single GPU)')
plt.colorbar()
plt.subplot(143)
plt.imshow(c[0].sum(axis=2))
plt.title('Mesh TensorFlow Single')
plt.colorbar()
plt.subplot(144)
plt.imshow((b[0] - c[0]).sum(axis=2))
plt.title('Residuals')
plt.colorbar()
plt.savefig("comparison-single.png")
exit(0)
