def lpt1(dlin_k, pos, kvec=None, name="LTP1"):
""" Run first order LPT on linear density field, returns displacements of particles
reading out at q. The result has the same dtype as q.
Parameters:
-----------
dlin_k: TODO: @modichirag add documentation
Returns:
--------
displacement: tensor (batch_size, npart, 3)
Displacement field
"""
with tf.name_scope(name):
dlin_k = tf.convert_to_tensor(dlin_k, name="lineark")
pos = tf.convert_to_tensor(pos, name="pos")
shape = dlin_k.get_shape()
batch_size, nc = shape[0], shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
displacement = []
for d in range(3):
kweight = gradient_kernel(kvec, d) * lap
dispc = tf.multiply(dlin_k, kweight)
disp = c2r3d(dispc, norm=nc[0] * nc[1] * nc[2])
displacement.append(cic_readout(disp, pos))
displacement = tf.stack(displacement, axis=2)
return displacement
def lpt2_source(dlin_k, kvec=None, name="LPT2Source"):
""" Generate the second order LPT source term.
Parameters:
-----------
dlin_k: TODO: @modichirag add documentation
Returns:
--------
source: tensor (batch_size, nc, nc, nc)
Source term
"""
with tf.name_scope(name):
dlin_k = tf.convert_to_tensor(dlin_k, name="lineark")
shape = dlin_k.get_shape()
batch_size, nc = shape[0], shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
source = tf.zeros(tf.shape(dlin_k))
D1 = [1, 2, 0]
D2 = [2, 0, 1]
phi_ii = []
# diagnoal terms
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
for d in range(3):
grad = gradient_kernel(kvec, d)
kweight = lap * grad * grad
phic = tf.multiply(dlin_k, kweight)
phi_ii.append(c2r3d(phic, norm=nc[0] * nc[1] * nc[2]))
for d in range(3):
source = tf.add(source, tf.multiply(phi_ii[D1[d]], phi_ii[D2[d]]))
# free memory
phi_ii = []
# off-diag terms
for d in range(3):
gradi = gradient_kernel(kvec, D1[d])
gradj = gradient_kernel(kvec, D2[d])
kweight = lap * gradi * gradj
phic = tf.multiply(dlin_k, kweight)
phi = c2r3d(phic, norm=nc[0] * nc[1] * nc[2])
source = tf.subtract(source, tf.multiply(phi, phi))
source = tf.multiply(source, 3.0 / 7.)
return r2c3d(source, norm=nc[0] * nc[1] * nc[2])
def apply_pgd(x, delta_k, alpha, kl, ks, kvec=None, name="ApplyPGD"):
"""
Estimate the short range force on the particles given a state.
Parameters:
-----------
x: tensor
Input state tensor of shape (3, batch_size, npart, 3)
delta_k:
Density in the Fourier space
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
"""
with tf.name_scope(name):
x = tf.convert_to_tensor(x, name="pos")
delta_k = tf.convert_to_tensor(delta_k, name="delta_k")
shape = delta_k.get_shape()
nc = shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
ndim = 3
norm = nc[0] * nc[1] * nc[2]
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
PGD_range = tf.cast(PGD_kernel(kvec, kl, ks), tf.complex64)
kweight = lap * PGD_range
pot_k = tf.multiply(delta_k, kweight)
f = []
for d in range(ndim):
force_dc = tf.multiply(pot_k, gradient_kernel(kvec, d))
forced = c2r3d(force_dc, norm=norm)
force = cic_readout(forced, x)
f.append(force)
f = tf.stack(f, axis=2)
f = tf.multiply(f, alpha)
return f
def apply_longrange(x,
delta_k,
split=0,
factor=1,
kvec=None,
name="ApplyLongrange"):
""" like long range, but x is a list of positions
TODO: Better documentation, also better name?
"""
# use the four point kernel to suppresse artificial growth of noise like terms
with tf.name_scope(name):
x = tf.convert_to_tensor(x, name="pos")
delta_k = tf.convert_to_tensor(delta_k, name="delta_k")
shape = delta_k.get_shape()
nc = shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
ndim = 3
norm = nc[0] * nc[1] * nc[2]
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
fknlrange = longrange_kernel(kvec, split)
kweight = lap * fknlrange
pot_k = tf.multiply(delta_k, kweight)
f = []
for d in range(ndim):
force_dc = tf.multiply(pot_k, gradient_kernel(kvec, d))
forced = c2r3d(force_dc, norm=norm)
force = cic_readout(forced, x)
f.append(force)
f = tf.stack(f, axis=2)
f = tf.multiply(f, factor)
return f
def benchmark_model(mesh):
"""
Initializes a 3D volume with random noise, and execute a forward FFT
"""
# Setup parameters
bs = FLAGS.box_size
nc = FLAGS.cube_size
batch_size = FLAGS.batch_size
a0 = FLAGS.a0
a = 1.0
nsteps = FLAGS.pm_steps
# Compute a few things first, using simple tensorflow
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)
# Initialize the integration steps
stages = np.linspace(FLAGS.a0, 1.0, FLAGS.pm_steps, endpoint=True)
# Generate a batch of 3D initial conditions
initial_conditions = flowpm.linear_field(
nc, # size of the cube
bs, # Physical size of the cube
ipklin, # Initial power spectrum
batch_size=batch_size)
# 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)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = 8
n_block_y = 4
n_block_z = 1
halo_size = 4
# Parameters of the large scales decomposition
downsampling_factor = 2
lnc = nc // 2**downsampling_factor
fx_dim = mtf.Dimension("nx", nc)
fy_dim = mtf.Dimension("ny", nc)
fz_dim = mtf.Dimension("nz", nc)
tfx_dim = mtf.Dimension("tx", nc)
tfy_dim = mtf.Dimension("ty", nc)
tfz_dim = mtf.Dimension("tz", nc)
# Dimensions of the low resolution grid
tx_dim = mtf.Dimension("tx_lr", nc)
ty_dim = mtf.Dimension("ty_lr", nc)
tz_dim = mtf.Dimension("tz_lr", nc)
nx_dim = mtf.Dimension('nx_block', n_block_x)
ny_dim = mtf.Dimension('ny_block', n_block_y)
nz_dim = mtf.Dimension('nz_block', n_block_z)
sx_dim = mtf.Dimension('sx_block', nc // n_block_x)
sy_dim = mtf.Dimension('sy_block', nc // n_block_y)
sz_dim = mtf.Dimension('sz_block', nc // n_block_z)
batch_dim = mtf.Dimension("batch", batch_size)
pk_dim = mtf.Dimension("npk", len(plin))
pk = mtf.import_tf_tensor(mesh, plin.astype('float32'), shape=[pk_dim])
# Compute necessary Fourier kernels
kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False)
kx = mtf.import_tf_tensor(mesh,
kvec[0].squeeze().astype('float32'),
shape=[tfx_dim])
ky = mtf.import_tf_tensor(mesh,
kvec[1].squeeze().astype('float32'),
shape=[tfy_dim])
kz = mtf.import_tf_tensor(mesh,
kvec[2].squeeze().astype('float32'),
shape=[tfz_dim])
kv = [ky, kz, kx]
kvec_lr = flowpm.kernels.fftk([nc, nc, nc], symmetric=False)
kx_lr = mtf.import_tf_tensor(mesh,
kvec_lr[0].squeeze().astype('float32'),
shape=[tx_dim])
ky_lr = mtf.import_tf_tensor(mesh,
kvec_lr[1].squeeze().astype('float32'),
shape=[ty_dim])
kz_lr = mtf.import_tf_tensor(mesh,
kvec_lr[2].squeeze().astype('float32'),
shape=[tz_dim])
kv_lr = [ky_lr, kz_lr, kx_lr]
# kvec for high resolution blocks
shape = [batch_dim, fx_dim, fy_dim, fz_dim]
lr_shape = [batch_dim, fx_dim, fy_dim, fz_dim]
hr_shape = [batch_dim, nx_dim, ny_dim, nz_dim, sx_dim, sy_dim, sz_dim]
part_shape = [batch_dim, fx_dim, fy_dim, fz_dim]
initc = mtfpm.linear_field(mesh, shape, bs, nc, pk, kv)
state = mtfpm.lpt_init_single(
initc,
a0,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
#state = mtfpm.lpt_init(low, high, 0.1, kv_lr, kv_hr, halo_size, hr_shape, lr_shape,
# part_shape[1:], downsampling_factor=downsampling_factor, antialias=True,)
# Here we can run our nbody
final_state = state #mtfpm.nbody(state, stages, lr_shape, hr_shape, kv_lr, kv_hr, halo_size, downsampling_factor=downsampling_factor)
# paint the field
final_field = mtf.zeros(mesh, shape=hr_shape)
for block_size_dim in hr_shape[-3:]:
final_field = mtf.pad(final_field, [halo_size, halo_size],
block_size_dim.name)
final_field = mesh_utils.cic_paint(final_field, final_state[0], halo_size)
# Halo exchange
for blocks_dim, block_size_dim in zip(hr_shape[1:4], final_field.shape[-3:]):
final_field = mpm.halo_reduce(final_field, blocks_dim, block_size_dim,
halo_size)
# Remove borders
for block_size_dim in hr_shape[-3:]:
final_field = mtf.slice(final_field, halo_size, block_size_dim.size,
block_size_dim.name)
#final_field = mtf.reshape(final_field, [batch_dim, fx_dim, fy_dim, fz_dim])
# Hack usisng custom reshape because mesh is pretty dumb
final_field = mtf.slicewise(lambda x: x[:, 0, 0, 0], [final_field],
output_dtype=tf.float32,
output_shape=[batch_dim, fx_dim, fy_dim, fz_dim],
name='my_dumb_reshape',
splittable_dims=part_shape[:-1] + hr_shape[:4])
return mtf.reduce_sum(final_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)
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)
