def recon_dm(linear, data):
print('new graph')
final_field = pm(linear)
residual = final_field - data #.astype(np.float32)
chisq = tf.multiply(residual, residual)
chisq = tf.reduce_mean(chisq)
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_mean(tf.multiply(priormesh, 1/priorwt))
loss = chisq + prior
return loss, chisq, prior
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 test_r2c2r():
bs = 50
nc = 16
batch_size = 3
base = 100 * np.random.randn(batch_size, nc, nc, nc).astype(np.float64)
with tf.Session() as sess:
cfield = r2c3d(tf.constant(base, dtype=tf.float64),
dtype=tf.complex128)
rfield = c2r3d(cfield, dtype=tf.float64)
sess.run(tf.global_variables_initializer())
rec = sess.run(rfield)
assert_allclose(base, rec, rtol=1e-09)
def recon_prototype(linear):
"""
"""
linear = tf.reshape(linear, minimum.shape)
#loss = tf.reduce_sum(tf.square(linear - minimum))
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
Rsm = tf.placeholder(tf.float32, name='smoothing')
if FLAGS.anneal:
print("\nAdd annealing section to graph\n")
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
#
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
chisq = tf.multiply(chisq, 1 / nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
prior = tf.multiply(prior, 1 / nc**3, name='prior')
#
loss = chisq + prior
return loss
def recon_prototype(linearflat):
"""
"""
linear = tf.reshape(linearflat, data.shape)
#
#loss = tf.reduce_sum(tf.square(linear - minimum))
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
#
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
chisq = tf.multiply(chisq, 1 / nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
prior = tf.multiply(prior, 1 / nc**3, name='prior')
#
loss = chisq + prior
grad = tf.gradients(loss, linearflat)
print(grad)
return loss, grad[0]
def recon(self, linear, data):
"""
"""
args = self.args
print('new recon graph')
base = self.pm(linear)
galmean = tfp.distributions.Poisson(rate=args.plambda * (1 + base))
logprob = -tf.reduce_mean(galmean.log_prob(data))
#logprob = tf.multiply(logprob, 1/nc**3, name='logprob')
#Prior
lineark = r2c3d(linear, norm=args.nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_mean(tf.multiply(priormesh, 1 / args.priorwt))
#prior = tf.multiply(prior, 1/nc**3, name='prior')
loss = logprob + prior
return loss, logprob, prior
def test_lpt2():
""" Checking lpt2_source, this also checks the laplace and gradient kernels
"""
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
source = fpmops.lpt2source(lineark).c2r()
# Same thing from tensorflow
tfsource = tfpm.lpt2_source(
pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)))
tfread = pmutils.c2r3d(tfsource).numpy()
assert_allclose(source, tfread[0], atol=1e-5)
def recon_dm(linear, data):
"""
"""
print('new graph')
final_field = pm(linear)
residual = final_field - data
chisq = tf.multiply(residual, residual)
chisq = tf.reduce_mean(chisq)
# chisq = tf.multiply(chisq, 1/nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_mean(tf.multiply(priormesh, 1 / priorwt))
# prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = chisq + prior
return loss, chisq, prior
def pm_poisson():
print("PM graph")
linear = flowpm.linear_field(nc, bs, ipklin, batch_size=batch_size)
if args.nbody:
print('Nobdy sim')
state = lpt_init(linear, a0=a0, order=args.lpt_order)
final_state = nbody(state, stages, nc)
else:
print('ZA/2LPT sim')
final_state = lpt_init(linear, a0=af, order=args.lpt_order)
tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0])
base = tfinal_field
if Rsm != 0:
basek = r2c3d(tfinal_field, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
galmean = tfp.distributions.Poisson(rate = plambda * (1 + base))
result = galmean.sample()
return linear, tfinal_field, result, base
def test_lpt1_64():
""" Checking lpt1, this also checks the laplace and gradient kernels
This variant of the test checks that it works for cubes of size 64
"""
nc = 64
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
lpt = fpmops.lpt1(lineark, grid)
# Same thing from tensorflow
tfread = tfpm.lpt1(
pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)),
grid.reshape((1, -1, 3)) * nc / bs).numpy()
assert_allclose(lpt, tfread[0] * bs / nc, atol=5e-5)
def test_lpt1():
""" Checking lpt1, this also checks the laplace and gradient kernels
"""
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
lpt = fpmops.lpt1(lineark, grid)
# Same thing from tensorflow
with tf.Session() as sess:
state = tfpm.lpt1(
pmutils.r2c3d(tf.expand_dims(tf.constant(lineark.c2r()), axis=0)),
grid.reshape((1, -1, 3)) * nc / bs)
tfread = sess.run(state)
assert_allclose(lpt, tfread[0] * bs / nc, atol=1e-5)
def tfpowerspec(x, boxsize):
nc = x.shape[-1]
kvec = flowpm.kernels.fftk([nc, nc, nc], symmetric=False)
kmesh = sum((kk / boxsize * nc)**2 for kk in kvec)**0.5
kmesh = np.expand_dims(kmesh, 0).astype(float32)
kbinmap = np.digitize(kmesh, kedges, right=False).astype(int32)
kbinmap[kbinmap == kbinmap.max()] = kbinmap.max() - 1
kbinmap -= 1
kbinmap = tf.constant(kbinmap)
kbincount = tfp.stats.count_integers(kbinmap)
tfkmesh = tf.constant(kmesh)
kvals = tfp.stats.count_integers(kbinmap, weights=tfkmesh) / tf.cast(
kbincount, tf.float32)
kvals = tf.repeat(tf.reshape(kvals, (1, nc)), x.shape[0], 0)
#
pmesh = tf.square(tf.abs(r2c3d(x, norm=nc**3))) * boxsize**3
tfpower = tfp.stats.count_integers(tf.repeat(kbinmap, pmesh.shape[0], 0),
weights=pmesh,
axis=[1, 2, 3])
tfpower = tf.reshape(tfpower, (nc, x.shape[0]))
tfpower = tf.transpose(tfpower) / tf.cast(kbincount, tf.float32)
return kvals, tfpower
def recon_prototype(data,
anneal=True,
nc=FLAGS.nc,
bs=FLAGS.box_size,
batch_size=FLAGS.batch_size,
a0=FLAGS.a0,
a=FLAGS.af,
nsteps=FLAGS.nsteps,
dtype=tf.float32):
"""
Prototype of function computing LPT deplacement.
Returns output tensorflow and mesh tensorflow tensors
"""
if dtype == tf.float32:
npdtype = "float32"
cdtype = tf.complex64
elif dtype == tf.float64:
npdtype = "float64"
cdtype = tf.complex128
print(dtype, npdtype)
#graph = mtf.Graph()
#mesh = mtf.Mesh(graph, "my_mesh")
linear = tf.get_variable('linmesh',
shape=(1, nc, nc, nc),
dtype=tf.float32,
initializer=tf.random_normal_initializer(),
trainable=True)
state = lpt_init(linear, a0=0.1, order=1)
final_state = nbody(state, stages, FLAGS.nc)
final_field = cic_paint(tf.zeros_like(linear), final_state[0])
residual = final_field - data.astype(np.float32)
base = residual
##Anneal
Rsm = tf.placeholder(tf.float32, name='smoothing')
if anneal:
#def anneal
Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
basek = r2c3d(base, norm=nc**3)
basek = tf.multiply(basek, tf.cast(smwts, tf.complex64))
base = c2r3d(basek, norm=nc**3)
chisq = tf.multiply(base, base)
chisq = tf.reduce_sum(chisq)
#chisq = tf.multiply(chisq, 1/nc**3, name='chisq')
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt))
#prior = tf.multiply(prior, 1/nc**3, name='prior')
#
loss = chisq + prior
## #optimizer = tf.optimize.AdamWeightDecayOptimizer(0.01)
## opt = tf.train.GradientDescentOptimizer(learning_rate=0.01)
##
## # Compute the gradients for a list of variables.
## grads_and_vars = opt.compute_gradients(loss, [linear])
## print("\ngradients : ", grads_and_vars)
## update_ops = opt.apply_gradients(grads_and_vars)
##
## #optimizer = tf.keras.optimizers.Adam(0.01)
## #var_grads = tf.gradients([loss], [linear])
##
##
## #update_ops = optimizer.apply_gradients(var_grads, linear)
## #update_ops = optimizer.apply_gradients(zip(var_grads, [linear]))
## #update_ops = None
## #lr = tf.placeholder(tf.float32, shape=())
## #update_op = mtf.assign(fieldvar, fieldvar - var_grads[0]*lr)
##
return linear, final_field, loss, chisq, prior
def recon_model(data, sigma=0.01**0.5, maxiter=100, anneal=False, dataovd=False, gtol=1e-5):
#bs, nc = config['boxsize'], config['nc']
kvec = flowpm.kernels.fftk([nc, nc, nc], symmetric=False)
kmesh = sum(kk**2 for kk in kvec)**0.5
priorwt = ipklin(kmesh) * bs ** -3
g = tf.Graph()
with g.as_default():
initlin = tf.placeholder(tf.float32, data.shape, name='initlin')
linear = tf.get_variable('linmesh', shape=(nc, nc, nc),
initializer=tf.random_normal_initializer(), trainable=True)
initlin_op = linear.assign(initlin, name='initlin_op')
#PM
icstate = tfpm.lptinit(linear, FLAGS.a0, name='icstate')
fnstate = tfpm.nbody(icstate, stages, nc, name='fnstate')
final = tf.zeros_like(linear)
final = cic_paint(final, fnstate[0], name='final')
if dataovd:
print('\Converting final density to overdensity because data is that\n')
fmean = tf.reduce_mean(final)
final = tf.multiply(final, 1/ fmean)
final = final - 1
#
#Prior
lineark = r2c3d(linear, norm=nc**3)
priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32))
prior = tf.reduce_sum(tf.multiply(priormesh, 1/priorwt))
prior = tf.multiply(prior, 1/nc**3, name='prior')
likelihood = tf.subtract(final, data)
likelihood = tf.multiply(likelihood, 1/sigma)
#galmean = tfp.distributions.Poisson(rate = plambda * (1 + finalfield))
#logprob = galmean.log_prob(data)
##Anneal
Rsm = tf.placeholder(tf.float32, name='smoothing')
if anneal :
print('\nAdding annealing part to graph\n')
Rsm = tf.multiply(Rsm, bs/nc)
Rsmsq = tf.multiply(Rsm, Rsm)
smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq))
likelihood = tf.squeeze(likelihood)
likelihoodk = r2c3d(likelihood, norm=nc**3)
likelihoodk = tf.multiply(likelihoodk, tf.cast(smwts, tf.complex64))
residual = c2r3d(likelihoodk, norm=nc**3)
else:
residual = tf.identity(likelihood)
chisq = tf.multiply(residual, residual)
chisq = tf.reduce_sum(chisq)
chisq = tf.multiply(chisq, 1/nc**3, name='chisq')
loss = tf.add(chisq, prior, name='loss')
#optimizer = ScipyOptimizerInterface(loss, var_list=[linear], method='L-BFGS-B',
# options={'maxiter': maxiter, 'gtol':gtol})
optimizer = tf.optimize.AdamWeightDecayOptimizer(0.01)
var_grads = tf.gradients(
[loss], [linear])
update_ops = optimizer.apply_grads(var_grads, linear)
tf.add_to_collection('inits', [initlin_op, initlin])
#tf.add_to_collection('opt', optimizer)
tf.add_to_collection('opt', update_ops)
tf.add_to_collection('diagnostics', [prior, chisq, loss])
tf.add_to_collection('reconpm', [linear, final, fnstate])
tf.add_to_collection('data', data)
return g
