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 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 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)