def get_phh(self,ks,logdens1,logdens2,redshift):
"""get_phh
Compute the halo-halo power spectrum :math:`P_{hh}(k;n_1,n_2)` between 2 mass threshold halo samples specified by the corresponding cumulative number densities.
Args:
ks (numpy array): Wavenumbers in :math:`[h\mathrm{Mpc}^{-1}]`
logdens1 (float): Logarithm of the cumulative halo number density of the first halo sample taken from the most massive, :math:`\log_{10}[n_1/(h^{-1}\mathrm{Mpc})^3]`
logdens2 (float): Logarithm of the cumulative halo number density of the second halo sample taken from the most massive, :math:`\log_{10}[n_2/(h^{-1}\mathrm{Mpc})^3]`
redshift (float): redshift at which the power spectrum is evaluated
Returns:
numpy array: halo power spectrum in :math:`[(h^{-1}\mathrm{Mpc})^{3}]`
"""
xs = np.logspace(-3,3,4000)
xi_tree = self._get_xiauto_tree(xs,logdens1,logdens2,redshift)
rswitch = min(60.,0.5 * self.cosmo.get_BAO_approx())
if logdens1 >= -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(xs,logdens1,logdens2,redshift)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + xi_tree * (1-np.exp(-(xs/rswitch)**4))
elif logdens1 >= -5.75 and logdens2 < -5.75:
xi_dir = self._get_xiauto_direct(xs,logdens1,-5.75,redshift) * self.g1.bias_ratio(redshift,logdens2)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + xi_tree * (1-np.exp(-(xs/rswitch)**4))
elif logdens1 < -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(xs,-5.75,logdens2,redshift) * self.g1.bias_ratio(redshift,logdens1)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + xi_tree * (1-np.exp(-(xs/rswitch)**4))
else:
xi_dir = self._get_xiauto_direct(xs,-5.75,-5.75,redshift) * self.g1.bias_ratio(redshift,logdens1)*self.g1.bias_ratio(redshift,logdens2)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + xi_tree * (1-np.exp(-(xs/rswitch)**4))
return pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs,xi_tot))(ks)
def __init__(self,
nc,
bs,
bias,
errormesh,
a0=0.1,
af=1.0,
nsteps=5,
nbody=True,
lpt_order=2,
anneal=True,
prior=True):
self.nc = nc
self.bs = bs
self.bias = bias
self.errormesh = errormesh
self.a0, self.af, self.nsteps = a0, af, nsteps
self.stages = np.linspace(a0, af, nsteps, endpoint=True)
self.nbody = nbody
self.lpt_order = lpt_order
self.anneal = True
self.klin = np.loadtxt('../../data/Planck15_a1p00.txt').T[0].astype(
np.float32)
self.plin = np.loadtxt('../../data/Planck15_a1p00.txt').T[1].astype(
np.float32)
self.ipklin = iuspline(self.klin, self.plin)
# Compute necessary Fourier kernels
self.kvec = tools.fftk((nc, nc, nc), boxsize=bs, symmetric=False)
self.kmesh = (sum(k**2 for k in self.kvec)**0.5).astype(np.float32)
self.priorwt = self.ipklin(self.kmesh)
self.R0 = tf.constant(0.)
self.prior = prior
def _get_xiauto_spl(self, logdens1, logdens2, redshift):
xs = np.logspace(-1, 3., 2000)
xi_tree = self._get_xiauto_tree(xs, logdens1, logdens2, redshift)
if logdens1 >= -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(xs, logdens1, logdens2, redshift)
rswitch = min(60., 0.5 * self.cosmo.get_BAO_approx())
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + \
xi_tree * (1-np.exp(-(xs/rswitch)**4))
elif logdens1 >= -5.75 and logdens2 < -5.75:
xi_dir = self._get_xiauto_direct(
xs, logdens1, -5.75, redshift) * self.g1.bias_ratio(redshift, logdens2)
rswitch = min(60., 0.5 * self.cosmo.get_BAO_approx())
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + \
xi_tree * (1-np.exp(-(xs/rswitch)**4))
elif logdens1 < -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(
xs, -5.75, logdens2, redshift) * self.g1.bias_ratio(redshift, logdens1)
rswitch = min(60., 0.5 * self.cosmo.get_BAO_approx())
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + \
xi_tree * (1-np.exp(-(xs/rswitch)**4))
else:
xi_dir = self._get_xiauto_direct(xs, -5.75, -5.75, redshift) * self.g1.bias_ratio(
redshift, logdens1)*self.g1.bias_ratio(redshift, logdens2)
rswitch = min(60., 0.5 * self.cosmo.get_BAO_approx())
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4) + \
xi_tree * (1-np.exp(-(xs/rswitch)**4))
return iuspline(xs, xi_tot)
def _get_xiauto_tree(self, xs, logdens1, logdens2, redshift):
ks = np.logspace(-3, 3, 2000)
g1 = self.g1.get(ks, redshift, logdens1)
g2 = self.g1.get(ks, redshift, logdens2)
pm_lin = self.get_pklin(ks)
ph_tree = g1 * g2 * pm_lin
return pyfftlog_interface.pk2xi_pyfftlog(iuspline(ks, ph_tree))(xs)
def test_linear_field_shape():
""" Testing just the shape of the sampled linear field
"""
klin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
tfread = tfpm.linear_field(nc, bs, ipklin, batch_size=5).numpy()
assert tfread.shape == (5, 16, 16, 16)
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 get_xilin(self, xs):
"""get_xilin
Compute the linear matter correlation function at z=0.
Args:
xs (numpy array): Separations in :math:`[h^{-1}\mathrm{Mpc}]`
Returns:
numpy array: Correlation function at separations given in the argument xs.
"""
ks = np.logspace(-3, 3, 300)
return pyfftlog_interface.pk2xi_pyfftlog(iuspline(ks, self.pkL.get(ks)))(xs)
def test_linear_field_shape():
bs = 50
nc = 16
klin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
with tf.Session() as sess:
field = linear_field(nc, bs, ipklin, batch_size=5)
sess.run(tf.global_variables_initializer())
tfread = sess.run(field)
assert tfread.shape == (5, 16, 16, 16)
def get_pknl(self, k, z):
"""get_pknl
Compute the nonlinear matter power spectrum. Note that this is still in a development phase, and the accuracy has not yet been fully evaluated.
Args:
k (numpy array): Wavenumbers in :math:`[h\mathrm{Mpc}^{-1}]`
z (float): redshift
Returns:
numpy array: Nonlinear matter power spectrum at wavenumbers given in the argument k.
"""
xs = np.logspace(-3, 3, 2000)
xinl = self.get_xinl(xs, z)
return pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs, xinl))(k)
def _get_phh_direct_cut(self,ks,logdens1,logdens2,redshift):
xs = np.logspace(-3,3,4000)
rswitch = min(60.,0.5 * self.cosmo.get_BAO_approx())
if logdens1 >= -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(xs,logdens1,logdens2,redshift)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4)
elif logdens1 >= -5.75 and logdens2 < -5.75:
xi_dir = self._get_xiauto_direct(xs,logdens1,-5.75,redshift) * self.g1.bias_ratio(redshift,logdens2)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4)
elif logdens1 < -5.75 and logdens2 >= -5.75:
xi_dir = self._get_xiauto_direct(xs,-5.75,logdens2,redshift) * self.g1.bias_ratio(redshift,logdens1)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4)
else:
xi_dir = self._get_xiauto_direct(xs,-5.75,-5.75,redshift) * self.g1.bias_ratio(redshift,logdens1)*self.g1.bias_ratio(redshift,logdens2)
xi_tot = xi_dir * np.exp(-(xs/rswitch)**4)
return pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs,xi_tot))(ks)
def get_Sigma(self, R2d, logdens, redshift):
"""get_Sigma
Compute the surface mass density, :math:`\Sigma(R;n_h)`, for a mass threshold halo sample specified by the corresponding cumulative number density.
Args:
R2d (numpy array): 2 dimensional projected separation in :math:`h^{-1}\mathrm{Mpc}`
logdens (float): Logarithm of the cumulative halo number density taken from the most massive, :math:`\log_{10}[n_h/(h^{-1}\mathrm{Mpc})^3]`
redshift (float): redshift at which the lens halos are located
Returns:
numpy array: surface mass density in :math:`[h M_\odot \mathrm{pc}^{-2}]`
"""
xs = np.logspace(-3, 3, 2000)
xi_tot = self._get_xicross(xs, logdens, redshift)
pk_spl = pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs, xi_tot))
return self.cosmo.get_Omega0() * self.cosmo.rho_cr / 1e12 * pyfftlog_interface.pk2xiproj_J0_pyfftlog(pk_spl)(R2d)
def finalize(self):
self['aout'] = numpy.array(self['aout'])
self['klin'] = np.loadtxt(self['pkfile']).T[0]
self['plin'] = np.loadtxt(self['pkfile']).T[1]
self['ipklin'] = iuspline(self['klin'], self['plin'])
#
bs, nc, B = self['boxsize'], self['nc'], self['pm_nc_factor']
ncf = int(nc*B)
self['f_config'] = {'boxsize':bs,
'nc':int(ncf),
'kvec':fftk(shape=(ncf, ncf, ncf), boxsize=bs, symmetric=False, dtype=self['dtype'])}
#
#self.pm = ParticleMesh(BoxSize=self['boxsize'], Nmesh= [self['nc']] * self['ndim'], resampler=self['resampler'], dtype='f4')
mask = numpy.array([ a not in self['stages'] for a in self['aout']], dtype='?')
missing_stages = self['aout'][mask]
if len(missing_stages):
raise ValueError('Some stages are requested for output but missing: %s' % str(missing_stages))
def get_wauto(self, R2d, logdens1, logdens2, redshift):
"""get_wauto
Compute the projected halo-halo correlation function :math:`w_{hh}(R;n_1,n_2)` for 2 mass threshold halo samples specified by the corresponding cumulative number densities.
Args:
R2d (numpy array): 2 dimensional projected separation in :math:`[h^{-1}\mathrm{Mpc}]`
logdens1 (float): Logarithm of the cumulative halo number density of the first halo sample taken from the most massive, :math:`\log_{10}[n_1/(h^{-1}\mathrm{Mpc})^3]`
logdens2 (float): Logarithm of the cumulative halo number density of the second halo sample taken from the most massive, :math:`\log_{10}[n_2/(h^{-1}\mathrm{Mpc})^3]`
redshift (float): redshift at which the power spectrum is evaluated
Returns:
numpy array: projected halo correlation function in :math:`[h^{-1}\mathrm{Mpc}]`
"""
xs = np.logspace(-3, 3, 1000)
xi_auto = self.get_xiauto(xs, logdens1, logdens2, redshift)
pk_spl = pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs, xi_auto))
return pyfftlog_interface.pk2xiproj_J0_pyfftlog(pk_spl, logkmin=-3.0, logkmax=3.0)(R2d)
def get_phm(self,ks,logdens,redshift):
"""get_phm
Compute the halo-matter cross power spectrum :math:`P_{hm}(k;n_h)` for a mass threshold halo sample specified by the corresponding cumulative number density.
Args:
ks (numpy array): Wavenumbers in :math:`[h\mathrm{Mpc}^{-1}]`
logdens (float): Logarithm of the cumulative halo number density of the halo sample taken from the most massive, :math:`\log_{10}[n_h/(h^{-1}\mathrm{Mpc})^3]`
redshift (float): redshift at which the power spectrum is evaluated
Returns:
numpy array: Halo-matter cross power spectrum in :math:`[(h^{-1}\mathrm{Mpc})^{3}]`
"""
xs = np.logspace(-4,3,4000)
xi_dir = self._get_xicross_direct(xs,logdens,redshift)
xi_tree = self._get_xicross_tree(xs,logdens,redshift)
rswitch = min(60.,0.5 * self.cosmo.get_BAO_approx())
xi = xi_dir * np.exp(-(xs/rswitch)**4) + xi_tree * (1-np.exp(-(xs/rswitch)**4))
return pyfftlog_interface.xi2pk_pyfftlog(iuspline(xs,xi),logrmin = -4.0, logrmax = 3.0)(ks)
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 lpt_prototype(mesh,
nc=FLAGS.nc,
bs=FLAGS.box_size,
batch_size=FLAGS.batch_size,
a0=FLAGS.a0,
a=FLAGS.af,
nsteps=FLAGS.nsteps):
"""
Prototype of function computing LPT deplacement.
Returns output tensorflow and mesh tensorflow tensors
"""
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)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = FLAGS.nx
n_block_y = FLAGS.ny
n_block_z = 1
halo_size = FLAGS.hsize
if halo_size >= 0.5 * min(nc // n_block_x, nc // n_block_y,
nc // n_block_z):
new_size = int(0.5 *
min(nc // n_block_x, nc // n_block_y, nc // n_block_z))
print('WARNING: REDUCING HALO SIZE from %d to %d' %
(halo_size, new_size))
halo_size = new_size
# Parameters of the large scales decomposition
downsampling_factor = 0
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)
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)
k_dims = [tx_dim, ty_dim, tz_dim]
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 for low resolution grid
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]
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]
# Begin simulation
initc = mtfpm.linear_field(mesh, shape, bs, nc, pk, kv)
# # Reshaping array into high resolution mesh
# field = mtf.slicewise(lambda x:tf.expand_dims(tf.expand_dims(tf.expand_dims(x, axis=1),axis=1),axis=1),
# [initc],
# output_dtype=tf.float32,
# output_shape=hr_shape,
# name='my_reshape',
# splittable_dims=lr_shape[:-1]+hr_shape[1:4]+part_shape[1:3])
#
state = mtfpm.lpt_init_single(
initc,
a0,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
# Here we can run our nbody
final_state = state #mtfpm.nbody(state, stages, lr_shape, hr_shape, k_dims, 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 initc, final_field
import tensorflow as tf
import flowpm
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as iuspline
sys.path.append('./utils')
import tools
bs, nc = 400, 128
nsteps = 10
ainit = 0.1
stages = np.linspace(ainit, 1.0, nsteps, endpoint=True)
pk = np.loadtxt('../data/ics_matterpow_0.dat')
ipklin = iuspline(pk[:, 0], pk[:, 1])
print('loaded')
##
##for seed in range(10, 10000, 10):
##
## print(seed)
## path = '../data/make_data_code/L%d-N%d-B1-T%d/S%d/'%(bs, nc, nsteps, seed)
##
## if os.path.isdir(path ):
## if not os.path.isfile(path + '/fpm-d'):
## #ick = tools.readbigfile(path + '/linear/LinearDensityK/')
## #ic = np.fft.irfftn(ick)*nc**3
## ic = tools.readbigfile(path + '/mesh/s/')
## initial_conditions = tf.cast(tf.expand_dims(tf.constant(ic), 0), tf.float32) - 1.
##
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 lpt_prototype(mesh,
initial_conditions,
derivs,
nc=FLAGS.nc,
bs=FLAGS.box_size,
batch_size=FLAGS.batch_size,
a0=FLAGS.a0,
a=FLAGS.af,
nsteps=FLAGS.nsteps):
"""
Prototype of function computing LPT deplacement.
Returns output tensorflow and mesh tensorflow tensors
"""
stages = np.linspace(a0, a, nsteps, endpoint=True)
lap, grad_x, grad_y, grad_z = derivs
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)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = FLAGS.nx
n_block_y = FLAGS.ny
n_block_z = 1
halo_size = FLAGS.hsize
# Parameters of the large scales decomposition
downsampling_factor = 0
lnc = nc // 2**downsampling_factor
#
ffx_dim = mtf.Dimension("fnx", nc)
ffy_dim = mtf.Dimension("fny", nc)
ffz_dim = mtf.Dimension("fnz", nc)
fx_dim = mtf.Dimension("nx", nc)
fy_dim = mtf.Dimension("ny", nc)
fz_dim = mtf.Dimension("nz", nc)
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)
k_dims = [tx_dim, ty_dim, tz_dim]
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])
# kvec for low resolution grid
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]
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]
field = mtf.import_tf_tensor(mesh, initial_conditions, shape=part_shape)
state = mtfpm.lpt_init_single(
field,
a,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
print('TOTO', state)
# Here we can run our nbody
final_state = state
# final_state = mtfpm.nbody_single(state, stages, lr_shape, hr_shape,
# kv_lr, halo_size)
# 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])
final_field = mtf.reshape(final_field,
[batch_dim, ffx_dim, ffy_dim, ffz_dim])
return final_field
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 recon_model(mesh,
data,
bparams,
ipkerror,
R0,
x0,
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
"""
b1, b2, bs2 = bparams
kerror, perror = ipkerror[0].astype(np.float32), ipkerror[1].astype(
np.float32)
if dtype == tf.float32:
npdtype = "float32"
cdtype = tf.complex64
elif dtype == tf.float64:
npdtype = "float64"
cdtype = tf.complex128
print("Dtype : ", dtype, npdtype)
# Compute a few things first, using simple tensorflow
kny = 1 * np.pi * nc / bs
R1, R2 = 3., 3 * 1.2
stages = np.linspace(a0, a, nsteps, endpoint=True)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = FLAGS.nx
n_block_y = FLAGS.ny
n_block_z = 1
halo_size = FLAGS.hsize
if halo_size >= 0.5 * min(nc // n_block_x, nc // n_block_y,
nc // n_block_z):
new_size = int(0.5 *
min(nc // n_block_x, nc // n_block_y, nc // n_block_z))
print('WARNING: REDUCING HALO SIZE from %d to %d' %
(halo_size, new_size))
halo_size = new_size
# Parameters of the large scales decomposition
scalar = mtf.Dimension("scalar", 1)
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)
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)
#k_dims = [tx_dim, ty_dim, tz_dim]
batch_dim = mtf.Dimension("batch", batch_size)
klin = np.loadtxt('..//data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('..//data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
pk_dim = mtf.Dimension("npk", len(plin))
pk = mtf.import_tf_tensor(mesh, plin.astype(npdtype), shape=[pk_dim])
pke_dim = mtf.Dimension("epk", len(perror))
pkerror = mtf.import_tf_tensor(mesh,
perror.astype(npdtype),
shape=[pke_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 for low resolution grid
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]
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]
splittables = lr_shape[:-1] + hr_shape[1:4] + part_shape[1:3]
#
# Begin simulation
if x0 is None:
fieldvar = mtf.get_variable(mesh,
'linear',
part_shape,
initializer=tf.random_normal_initializer(
mean=0.0, stddev=1, seed=None))
else:
fieldvar = mtf.get_variable(mesh,
'linear',
part_shape,
initializer=tf.constant_initializer(x0))
state = mtfpm.lpt_init_single(
fieldvar,
a0,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
final_state = mtfpm.nbody_single(state, stages, lr_shape, hr_shape, kv_lr,
halo_size)
# paint the field
final_field = mtf.zeros(mesh, shape=part_shape)
final_field = mcomp.cic_paint_fr(final_field, final_state, part_shape,
hr_shape, halo_size, splittables, mesh)
##
#Get the fields for bias
hr_field = mcomp.fr_to_hr(final_field, hr_shape, halo_size, splittables,
mesh)
mstate = mpm.mtf_indices(hr_field.mesh,
shape=part_shape[1:],
dtype=tf.float32)
X = mtf.einsum([mtf.ones(hr_field.mesh, [batch_dim]), mstate],
output_shape=[batch_dim] + mstate.shape[:])
k_dims_pr = [d.shape[0] for d in kv]
k_dims_pr = [k_dims_pr[2], k_dims_pr[0], k_dims_pr[1]]
tfnc, tfbs = cswisef.float_to_mtf(nc * 1., mesh,
scalar), cswisef.float_to_mtf(
bs, mesh, scalar)
#
initc = fieldvar
d0 = initc - mtf.reduce_mean(initc)
#
d2 = initc * initc
d2 = d2 - mtf.reduce_mean(d2)
#
cfield = mesh_utils.r2c3d(d0, k_dims_pr, dtype=cdtype)
shearfield = mtf.zeros(mesh, shape=part_shape)
shearfield = shear(shearfield, cfield, kv, tfnc, tfbs)
s2 = shearfield - mtf.reduce_mean(shearfield)
dread = mcomp.cic_readout_fr(d0, [X],
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh)
d2read = mcomp.cic_readout_fr(d2, [X],
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh)
s2read = mcomp.cic_readout_fr(s2, [X],
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh)
ed, ed2, es2 = mtf.zeros(mesh, shape=part_shape), mtf.zeros(
mesh, shape=part_shape), mtf.zeros(mesh, shape=part_shape)
ed = mcomp.cic_paint_fr(ed,
final_state,
output_shape=part_shape,
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh,
weights=dread)
ed2 = mcomp.cic_paint_fr(ed2,
final_state,
output_shape=part_shape,
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh,
weights=d2read)
es2 = mcomp.cic_paint_fr(es2,
final_state,
output_shape=part_shape,
hr_shape=hr_shape,
halo_size=halo_size,
splittables=splittables,
mesh=mesh,
weights=s2read)
model = ed * b1 + ed2 * b2 + es2 * bs2
mtfdata = mtf.import_tf_tensor(mesh,
tf.convert_to_tensor(data),
shape=shape)
diff = model - mtfdata
# Get prior
k_dims_pr = [d.shape[0] for d in kv]
k_dims_pr = [k_dims_pr[2], k_dims_pr[0], k_dims_pr[1]]
cfield = mesh_utils.r2c3d(fieldvar, k_dims_pr, dtype=cdtype)
def _cwise_prior(kfield, pk, kx, ky, kz):
kx = tf.reshape(kx, [-1, 1, 1])
ky = tf.reshape(ky, [1, -1, 1])
kz = tf.reshape(kz, [1, 1, -1])
kk = tf.sqrt((kx / bs * nc)**2 + (ky / bs * nc)**2 + (kz / bs * nc)**2)
kshape = kk.shape
kk = tf.reshape(kk, [-1])
pkmesh = tfp.math.interp_regular_1d_grid(
x=kk,
x_ref_min=1e-05,
x_ref_max=1000.0,
y_ref=pk,
grid_regularizing_transform=tf.log)
priormesh = tf.reshape(pkmesh, kshape)
return tf.abs(kfield) / priormesh**0.5
cpfield = mtf.cwise(_cwise_prior, [cfield, pk] + kv,
output_dtype=tf.float32)
prior = mtf.reduce_sum(mtf.square(cpfield)) * bs**3 #* nc**3
# Total loss
cdiff = mesh_utils.r2c3d(diff, k_dims_pr, dtype=cdtype)
def _cwise_diff(kfield, pk, kx, ky, kz):
kx = tf.reshape(kx, [-1, 1, 1])
ky = tf.reshape(ky, [1, -1, 1])
kz = tf.reshape(kz, [1, 1, -1])
kk = tf.sqrt((kx / bs * nc)**2 + (ky / bs * nc)**2 + (kz / bs * nc)**2)
kshape = kk.shape
kk = tf.reshape(kk, [-1])
pkmesh = tfp.math.interp_regular_1d_grid(x=kk,
x_ref_min=kerror.min(),
x_ref_max=kerror.max(),
y_ref=pk)
priormesh = tf.reshape(pkmesh, kshape)
priormesh = tf.cast(priormesh**0.5, kfield.dtype)
return kfield / priormesh
cdiff = mtf.cwise(_cwise_diff, [cdiff, pkerror] + kv, output_dtype=cdtype)
def _cwise_smooth(kfield, kx, ky, kz):
kx = tf.reshape(kx, [-1, 1, 1])
ky = tf.reshape(ky, [1, -1, 1])
kz = tf.reshape(kz, [1, 1, -1])
kk = (kx / bs * nc)**2 + (ky / bs * nc)**2 + (kz / bs * nc)**2
wts = tf.cast(tf.exp(-kk * (R0 * bs / nc)**2), kfield.dtype)
return kfield * wts
cdiff = mtf.cwise(_cwise_smooth, [cdiff] + kv, output_dtype=cdtype)
diff = mesh_utils.c2r3d(cdiff, diff.shape[-3:], dtype=dtype)
chisq = mtf.reduce_sum(mtf.square(diff))
loss = chisq + prior
fields = [fieldvar, final_field, model]
metrics = [chisq, prior, loss]
return fields, metrics, kv
def recon_model(mesh,
data,
R0,
x0,
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)
# Compute a few things first, using simple tensorflow
stages = np.linspace(a0, a, nsteps, endpoint=True)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = FLAGS.nx
n_block_y = FLAGS.ny
n_block_z = 1
halo_size = FLAGS.hsize
if halo_size >= 0.5 * min(nc // n_block_x, nc // n_block_y,
nc // n_block_z):
new_size = int(0.5 *
min(nc // n_block_x, nc // n_block_y, nc // n_block_z))
print('WARNING: REDUCING HALO SIZE from %d to %d' %
(halo_size, new_size))
halo_size = new_size
# Parameters of the large scales decomposition
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)
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)
k_dims = [tx_dim, ty_dim, tz_dim]
batch_dim = mtf.Dimension("batch", batch_size)
klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
pk_dim = mtf.Dimension("npk", len(plin))
pk = mtf.import_tf_tensor(mesh, plin.astype(npdtype), 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 for low resolution grid
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]
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]
# Begin simulation
if x0 is None:
fieldvar = mtf.get_variable(mesh,
'linear',
part_shape,
initializer=tf.random_normal_initializer(
mean=0.0, stddev=1, seed=None))
else:
fieldvar = mtf.get_variable(mesh,
'linear',
part_shape,
initializer=tf.constant_initializer(x0))
print("\nfieldvar : \n", fieldvar)
# Here we can run our nbody
if FLAGS.nbody:
state = mtfpm.lpt_init_single(
fieldvar,
a0,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
# Here we can run our nbody
final_state = mtfpm.nbody_single(state, stages, lr_shape, hr_shape,
kv_lr, halo_size)
else:
final_state = mtfpm.lpt_init_single(
fieldvar,
stages[-1],
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
# 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.slicewise(
lambda x: x[:, 0, 0, 0], [final_field],
output_dtype=dtype,
output_shape=[batch_dim, fx_dim, fy_dim, fz_dim],
name='my_dumb_reshape',
splittable_dims=part_shape[:-1] + hr_shape[:4])
mtfdata = mtf.import_tf_tensor(mesh,
tf.convert_to_tensor(data),
shape=shape)
# Get prior
k_dims_pr = [d.shape[0] for d in kv]
k_dims_pr = [k_dims_pr[2], k_dims_pr[0], k_dims_pr[1]]
cfield = mesh_utils.r2c3d(fieldvar, k_dims_pr, dtype=cdtype)
def _cwise_prior(kfield, pk, kx, ky, kz):
kx = tf.reshape(kx, [-1, 1, 1])
ky = tf.reshape(ky, [1, -1, 1])
kz = tf.reshape(kz, [1, 1, -1])
kk = tf.sqrt((kx / bs * nc)**2 + (ky / bs * nc)**2 + (kz / bs * nc)**2)
kshape = kk.shape
kk = tf.reshape(kk, [-1])
pkmesh = tfp.math.interp_regular_1d_grid(
x=kk,
x_ref_min=1e-05,
x_ref_max=1000.0,
y_ref=pk,
grid_regularizing_transform=tf.log)
priormesh = tf.reshape(pkmesh, kshape)
return tf.abs(kfield) / priormesh**0.5
cpfield = mtf.cwise(_cwise_prior, [cfield, pk] + kv,
output_dtype=tf.float32)
prior = mtf.reduce_sum(mtf.square(cpfield)) * bs**3 #*nc**3
# Total loss
diff = (final_field - mtfdata)
R0 = tf.constant(R0)
print("R0 in the recon_model : ", R0)
def _cwise_smooth(kfield, kx, ky, kz):
kx = tf.reshape(kx, [-1, 1, 1])
ky = tf.reshape(ky, [1, -1, 1])
kz = tf.reshape(kz, [1, 1, -1])
kk = (kx / bs * nc)**2 + (ky / bs * nc)**2 + (kz / bs * nc)**2
wts = tf.cast(tf.exp(-kk * (R0 * bs / nc)**2), kfield.dtype)
return kfield * wts
# Element-wise function that applies a Fourier kernel
plambda = FLAGS.plambda
def _cwise_logprob(finalfield, data):
galmean = tfp.distributions.Poisson(rate=plambda * (1 + finalfield))
logprob = galmean.log_prob(data)
return -1 * logprob
cfield = mesh_utils.r2c3d(final_field, k_dims_pr, dtype=cdtype)
cfield = mtf.cwise(_cwise_smooth, [cfield] + kv, output_dtype=cdtype)
final_fieldsm = mesh_utils.c2r3d(cfield, diff.shape[-3:], dtype=dtype)
chisq = mtf.cwise(_cwise_logprob, [final_fieldsm, mtfdata],
output_dtype=tf.float32) #
chisq = mtf.reduce_sum(chisq)
## #
loss = chisq + prior
def _cwise_sample(finalfield, data):
galmean = tfp.distributions.Poisson(rate=plambda * (1 + finalfield))
sample = galmean.sample()
return sample
sample = mtf.cwise(_cwise_sample, [final_fieldsm, mtfdata],
output_dtype=tf.float32) #
fields = [fieldvar, sample]
metrics = [chisq, prior, loss]
return fields, metrics, kv
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
plambda = FLAGS.plambda
ic, fin = np.load('../data/poisson_N%03d/ic.npy' % nc), np.load(
'../data/poisson_N%03d/psample_%0.2f.npy' % (nc, plambda))
print('Data loaded')
########################################################
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 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)
def nbody_prototype(mesh,
infield=False,
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
"""
# Compute a few things first, using simple tensorflow
stages = np.linspace(a0, a, nsteps, endpoint=True)
# Define the named dimensions
# Parameters of the small scales decomposition
n_block_x = FLAGS.nx
n_block_y = FLAGS.ny
n_block_z = 1
halo_size = FLAGS.hsize
# Parameters of the large scales decomposition
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)
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)
k_dims = [tx_dim, ty_dim, tz_dim]
batch_dim = mtf.Dimension("batch", batch_size)
klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0]
plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1]
ipklin = iuspline(klin, plin)
pk_dim = mtf.Dimension("npk", len(plin))
pk = mtf.import_tf_tensor(mesh, plin, 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 for low resolution grid
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]
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]
# Begin simulation
## Compute initial initial conditions distributed
input_field = tf.placeholder(dtype, [batch_size, nc, nc, nc])
if infield:
initc = mtf.import_tf_tensor(mesh, input_field, shape=part_shape)
else:
initc = mtfpm.linear_field(mesh, shape, bs, nc, pk, kv)
# Here we can run our nbody
if FLAGS.nbody:
state = mtfpm.lpt_init_single(
initc,
a0,
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
# Here we can run our nbody
final_state = mtfpm.nbody_single(state, stages, lr_shape, hr_shape,
kv_lr, halo_size)
else:
final_state = mtfpm.lpt_init_single(
initc,
stages[-1],
kv_lr,
halo_size,
lr_shape,
hr_shape,
part_shape[1:],
antialias=True,
)
# 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.slicewise(
lambda x: x[:, 0, 0, 0], [final_field],
output_dtype=dtype,
output_shape=[batch_dim, fx_dim, fy_dim, fz_dim],
name='my_dumb_reshape',
splittable_dims=part_shape[:-1] + hr_shape[:4])
return initc, final_field, input_field
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 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 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)
np.random.seed(100)
#tf.random.set_random_seed(100)
cscratch = "../figs_recon/"
nc, bs = 64, 200
a0, a, nsteps = 0.1, 1.0, 5
stages = np.linspace(a0, a, nsteps, endpoint=True)
anneal = True
niter = 200
optimizer = 'adam'
lr = 0.01
RRs = [2, 1, 0.5, 0]
klin = np.loadtxt('..//data/Planck15_a1p00.txt').T[0].astype(np.float32)
plin = np.loadtxt('..//data/Planck15_a1p00.txt').T[1].astype(np.float32)
ipklin = iuspline(klin, plin)
# Compute necessary Fourier kernels
kvec = tools.fftk((nc, nc, nc), boxsize=bs, symmetric=False)
kmesh = (sum(k**2 for k in kvec)**0.5).astype(np.float32)
priorwt = ipklin(kmesh)
fpath = "./tmp/dm-tf2-%s-%d/" % (optimizer, nc)
for ff in [fpath]:
#for ff in [fpath, fpath + '/figs']:
try:
os.makedirs(ff)
except Exception as e:
print(e)
dtype = tf.float32
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)
