def whitenoise_to_linear(nc,
boxsize,
whitec,
pk,
kvec=None,
batch_size=1,
seed=None,
dtype=tf.float32,
name="LinearField"):
"""Generates a linear field with a given linear power spectrum and whitenoise realization
"""
with tf.name_scope(name):
# Transform nc to a list of necessary
if isinstance(nc, int):
nc = [nc, nc, nc]
if isinstance(boxsize, int) or isinstance(boxsize, float):
boxsize = [boxsize, boxsize, boxsize]
if kvec is None:
kvec = fftk(nc, symmetric=False)
kmesh = sum((kk / boxsize[i] * nc[i])**2 for i, kk in enumerate(kvec))**0.5
pkmesh = pk(kmesh)
lineark = tf.multiply(whitec, (pkmesh /
(boxsize[0] * boxsize[1] * boxsize[2]))**0.5)
linear = c2r3d(lineark, norm=nc[0] * nc[1] * nc[2], name=name, dtype=dtype)
return linear
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 linear_field(nc,
boxsize,
pk,
kvec=None,
batch_size=1,
seed=None,
dtype=tf.float32,
name="LinearField"):
"""Generates a linear field with a given linear power spectrum.
Parameters:
-----------
nc: int, or list of ints
Number of cells in the field. If a list is provided, number of cells per
dimension.
boxsize: float, or list of floats
Physical size of the cube, in Mpc/h.
pk: interpolator
Power spectrum to use for the field
kvec: array
k_vector corresponding to the cube, optional
batch_size: int
Size of batches
seed: int
Seed to initialize the gaussian random field
dtype: tf.dtype
Type of the sampled field, e.g. tf.float32 or tf.float64
Returns
------
linfield: tensor (batch_size, nc, nc, nc)
Realization of the linear field with requested power spectrum
"""
with tf.name_scope(name):
# Transform nc to a list of necessary
if isinstance(nc, int):
nc = [nc, nc, nc]
if isinstance(boxsize, int) or isinstance(boxsize, float):
boxsize = [boxsize, boxsize, boxsize]
if kvec is None:
kvec = fftk(nc, symmetric=False)
kmesh = sum(
(kk / boxsize[i] * nc[i])**2 for i, kk in enumerate(kvec))**0.5
pkmesh = pk(kmesh)
whitec = white_noise(nc,
batch_size=batch_size,
seed=seed,
type='complex')
lineark = tf.multiply(
whitec, (pkmesh / (boxsize[0] * boxsize[1] * boxsize[2]))**0.5)
linear = c2r3d(lineark,
norm=nc[0] * nc[1] * nc[2],
name=name,
dtype=dtype)
return linear
def lpt2_source(dlin_k, kvec=None, name="LPT2Source"):
""" Generate the second order LPT source term.
Parameters:
-----------
dlin_k: TODO: @modichirag add documentation
Returns:
--------
source: tensor (batch_size, nc, nc, nc)
Source term
"""
with tf.name_scope(name):
dlin_k = tf.convert_to_tensor(dlin_k, name="lineark")
shape = dlin_k.get_shape()
batch_size, nc = shape[0], shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
source = tf.zeros(tf.shape(dlin_k))
D1 = [1, 2, 0]
D2 = [2, 0, 1]
phi_ii = []
# diagnoal terms
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
for d in range(3):
grad = gradient_kernel(kvec, d)
kweight = lap * grad * grad
phic = tf.multiply(dlin_k, kweight)
phi_ii.append(c2r3d(phic, norm=nc[0] * nc[1] * nc[2]))
for d in range(3):
source = tf.add(source, tf.multiply(phi_ii[D1[d]], phi_ii[D2[d]]))
# free memory
phi_ii = []
# off-diag terms
for d in range(3):
gradi = gradient_kernel(kvec, D1[d])
gradj = gradient_kernel(kvec, D2[d])
kweight = lap * gradi * gradj
phic = tf.multiply(dlin_k, kweight)
phi = c2r3d(phic, norm=nc[0] * nc[1] * nc[2])
source = tf.subtract(source, tf.multiply(phi, phi))
source = tf.multiply(source, 3.0 / 7.)
return r2c3d(source, norm=nc[0] * nc[1] * nc[2])
def apply_pgd(x, delta_k, alpha, kl, ks, kvec=None, name="ApplyPGD"):
"""
Estimate the short range force on the particles given a state.
Parameters:
-----------
x: tensor
Input state tensor of shape (3, batch_size, npart, 3)
delta_k:
Density in the Fourier space
alpha: float
Free parameter. Factor of proportionality between the displacement and the Particle-mesh force.
kl: float
Long range scale parameter
ks: float
Short range scale parameter
"""
with tf.name_scope(name):
x = tf.convert_to_tensor(x, name="pos")
delta_k = tf.convert_to_tensor(delta_k, name="delta_k")
shape = delta_k.get_shape()
nc = shape[1:]
if kvec is None:
kvec = fftk(nc, symmetric=False)
ndim = 3
norm = nc[0] * nc[1] * nc[2]
lap = tf.cast(laplace_kernel(kvec), tf.complex64)
PGD_range = tf.cast(PGD_kernel(kvec, kl, ks), tf.complex64)
kweight = lap * PGD_range
pot_k = tf.multiply(delta_k, kweight)
f = []
for d in range(ndim):
force_dc = tf.multiply(pot_k, gradient_kernel(kvec, d))
forced = c2r3d(force_dc, norm=norm)
force = cic_readout(forced, x)
f.append(force)
f = tf.stack(f, axis=2)
f = tf.multiply(f, alpha)
return f
def compensate_cic(field, name="CompensateCiC"):
"""
Compensate for CiC painting
Args:
field: input 3D cic-painted field
Returns:
compensated_field
"""
with tf.name_scope(name):
shape = field.get_shape()
batch_size, nc = shape[0], shape[1:]
kvec = fpk.fftk(nc, symmetric=False)
delta_k = r2c3d(field, norm=nc[0] * nc[1] * nc[2])
delta_k = tf.cast(fpk.cic_compensation(kvec), tf.complex64) * delta_k
return c2r3d(delta_k, norm=nc[0] * nc[1] * nc[2])
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