import numpy
import numbers
from mpi4py import MPI
from pmesh.pm import RealField, ComplexField
from nbodykit.meshtools import SlabIterator
from nbodykit.utils import GatherArray, ScatterArray
from nbodykit.mpirng import MPIRandomState
import mpsort
[docs]def gaussian_complex_fields(pm, linear_power, seed,
unitary_amplitude=False, inverted_phase=False,
compute_displacement=False, logger=None):
r"""
Make a Gaussian realization of a overdensity field, :math:`\delta(x)`.
If specified, also compute the corresponding 1st order Lagrangian
displacement field (Zel'dovich approximation) :math:`\psi(x)`,
which is related to the linear velocity field via:
.. math::
v(x) = f a H \psi(x)
Notes
-----
This computes the overdensity field using the following steps:
#. Generate complex variates with unity variance
#. Scale the Fourier field by :math:`(P(k) / V)^{1/2}`
After step 2, the complex field has unity variance. This
is equivalent to generating real-space normal variates
with mean and unity variance, calling r2c() and dividing by :math:`N^3`
since the variance of the complex FFT (with no additional normalization)
is :math:`N^3 \times \sigma^2_\mathrm{real}`.
Furthermore, the power spectrum is defined as V * variance.
So a normalization factor of 1 / V shows up in step 2,
cancels this factor such that the power spectrum is P(k).
The linear displacement field is computed as:
.. math::
\psi_i(k) = i \frac{k_i}{k^2} \delta(k)
.. note::
To recover the linear velocity in proper units, i.e., km/s,
from the linear displacement, an additional factor of
:math:`f \times a \times H(a)` is required
Parameters
----------
pm : pmesh.pm.ParticleMesh
the mesh object
linear_power : callable
a function taking wavenumber as its only argument, which returns
the linear power spectrum
seed : int
the random seed used to generate the random field
compute_displacement : bool, optional
if ``True``, also return the linear Zel'dovich displacement field;
default is ``False``
unitary_amplitude : bool, optional
if ``True``, the seed gaussian has unitary_amplitude.
inverted_phase: bool, optional
if ``True``, the phase of the seed gaussian is inverted
Returns
-------
delta_k : ComplexField
the real-space Gaussian overdensity field
disp_k : ComplexField or ``None``
if requested, the Gaussian displacement field
"""
if not isinstance(seed, numbers.Integral):
raise ValueError("the seed used to generate the linear field must be an integer")
if logger and pm.comm.rank == 0:
logger.info("Generating whitenoise")
# use pmesh to generate random complex white noise field (done in parallel)
# variance of complex field is unity
# multiply by P(k)**0.5 to get desired variance
delta_k = pm.generate_whitenoise(seed, type='untransposedcomplex', unitary=unitary_amplitude)
if logger and pm.comm.rank == 0:
logger.info("Write noise generated")
if inverted_phase: delta_k[...] *= -1
# initialize the displacement fields for (x,y,z)
if compute_displacement:
disp_k = [pm.create(type='untransposedcomplex') for i in range(delta_k.ndim)]
for i in range(delta_k.ndim): disp_k[i][:] = 1j
else:
disp_k = None
# volume factor needed for normalization
norm = 1.0 / pm.BoxSize.prod()
# iterate in slabs over fields
slabs = [delta_k.slabs.x, delta_k.slabs]
if compute_displacement:
slabs += [d.slabs for d in disp_k]
# loop over the mesh, slab by slab
for islabs in zip(*slabs):
kslab, delta_slab = islabs[:2] # the k arrays and delta slab
# the square of the norm of k on the mesh
k2 = sum(kk**2 for kk in kslab)
zero_idx = k2 == 0.
k2[zero_idx] = 1. # avoid dividing by zero
# the linear power (function of k)
power = linear_power((k2**0.5).flatten())
# multiply complex field by sqrt of power
delta_slab[...].flat *= (power*norm)**0.5
# set k == 0 to zero (zero config-space mean)
delta_slab[zero_idx] = 0.
# compute the displacement
if compute_displacement:
# ignore division where k==0 and set to 0
with numpy.errstate(invalid='ignore', divide='ignore'):
for i in range(delta_k.ndim):
disp_slab = islabs[2+i]
disp_slab[...] *= kslab[i] / k2 * delta_slab[...]
disp_slab[zero_idx] = 0. # no bulk displacement
if logger and pm.comm.rank == 0:
logger.info("Displacement computed in fourier space")
# return Fourier-space density and displacement (which could be None)
return delta_k, disp_k
[docs]def gaussian_real_fields(pm, linear_power, seed,
unitary_amplitude=False,
inverted_phase=False, compute_displacement=False, logger=None):
r"""
Make a Gaussian realization of a overdensity field in
real-space :math:`\delta(x)`.
If specified, also compute the corresponding linear Zel'dovich
displacement field :math:`\psi(x)`, which is related to the
linear velocity field via:
Notes
-----
See the docstring for :func:`gaussian_complex_fields` for the
steps involved in generating the fields.
Parameters
----------
pm : pmesh.pm.ParticleMesh
the mesh object
linear_power : callable
a function taking wavenumber as its only argument, which returns
the linear power spectrum
seed : int
the random seed used to generate the random field
compute_displacement : bool, optional
if ``True``, also return the linear Zel'dovich displacement field;
default is ``False``
unitary_amplitude : bool, optional
if ``True``, the seed gaussian has unitary_amplitude.
inverted_phase: bool, optional
if ``True``, the phase of the seed gaussian is inverted
Returns
-------
delta : RealField
the real-space Gaussian overdensity field
disp : RealField or ``None``
if requested, the Gaussian displacement field
"""
# make fourier fields
delta_k, disp_k = gaussian_complex_fields(pm, linear_power, seed,
inverted_phase=inverted_phase,
unitary_amplitude=unitary_amplitude,
compute_displacement=compute_displacement,
logger=logger)
# FFT the density to real-space
delta = delta_k.c2r()
std = (delta ** 2).cmean() ** 0.5
if logger and pm.comm.rank == 0:
logger.info("Overdensity computed in configuration space: std = %s" % str(std))
# FFT the velocity back to real space
if compute_displacement:
disp = [disp_k[i].c2r() for i in range(delta.ndim)]
std = [(disp[i] ** 2).cmean() ** 0.5 for i in range(delta.ndim)]
if logger and pm.comm.rank == 0:
logger.info("Displacement computed in configuration space: std = %s" % str(std))
else:
disp = None
# return density and displacement (which could be None)
return delta, disp
[docs]def poisson_sample_to_points(delta, displacement, pm, nbar, bias=1., seed=None, logger=None):
"""
Poisson sample the linear delta and displacement fields to points.
The steps in this function:
#. Apply a biased, lognormal transformation to the input ``delta`` field
#. Poisson sample the overdensity field to discrete points
#. Disribute the positions of particles uniformly within the mesh cells,
and assign the displacement field at each cell to the particles
Parameters
----------
delta : RealField
the linear overdensity field to sample
displacement : list of RealField (3,)
the linear displacement fields which is used to move the particles
nbar : float
the desired number density of the output catalog of objects
bias : float, optional
apply a linear bias to the overdensity field (default is 1.)
seed : int, optional
the random seed used to Poisson sample the field to points
Returns
-------
pos : array_like, (N, 3)
the Cartesian positions of each of the generated particles
displ : array_like, (N, 3)
the displacement field sampled for each of the generated particles in the
same units as the ``pos`` array
"""
comm = delta.pm.comm
# seed1 used for poisson sampling
# seed2 used for uniform shift within a cell.
seed1, seed2 = numpy.random.RandomState(seed).randint(0, 0xfffffff, size=2)
# apply the lognormal transformation to the initial conditions density
# this creates a positive-definite delta (necessary for Poisson sampling)
lagrangian_bias = bias - 1.
delta = lognormal_transform(delta, bias=lagrangian_bias)
if logger and pm.comm.rank == 0:
logger.info("Lognormal transformation done")
# mean number of objects per cell
H = delta.BoxSize / delta.Nmesh
overallmean = H.prod() * nbar
# number of objects in each cell (per rank, as a RealField)
cellmean = delta * overallmean
# create a random state with the input seed
rng = MPIRandomState(seed=seed1, comm=comm, size=delta.size)
# generate poissons. Note that we use ravel/unravel to
# maintain MPI invariane.
Nravel = rng.poisson(lam=cellmean.ravel())
N = delta.pm.create(type='real')
N.unravel(Nravel)
Ntot = N.csum()
if logger and pm.comm.rank == 0:
logger.info("Poisson sampling done, total number of objects is %d" % Ntot)
pos_mesh = delta.pm.generate_uniform_particle_grid(shift=0.0)
disp_mesh = numpy.empty_like(pos_mesh)
# no need to do decompose because pos_mesh is strictly within the
# local volume of the RealField.
N_per_cell = N.readout(pos_mesh, resampler='nnb')
for i in range(N.ndim):
disp_mesh[:, i] = displacement[i].readout(pos_mesh, resampler='nnb')
# fight round off errors, if any
N_per_cell = numpy.int64(N_per_cell + 0.5)
pos = pos_mesh.repeat(N_per_cell, axis=0)
disp = disp_mesh.repeat(N_per_cell, axis=0)
del pos_mesh
del disp_mesh
if logger and pm.comm.rank == 0:
logger.info("catalog produced. Assigning in cell shift.")
# generate linear ordering of the positions.
# this should have been a method in pmesh, e.g. argument
# to genereate_uniform_particle_grid(return_id=True);
# FIXME: after pmesh update, remove this
orderby = numpy.int64(pos[:, 0] / H[0] + 0.5)
for i in range(1, delta.ndim):
orderby[...] *= delta.Nmesh[i]
orderby[...] += numpy.int64(pos[:, i] / H[i] + 0.5)
# sort by ID to maintain MPI invariance.
pos = mpsort.sort(pos, orderby=orderby, comm=comm)
disp = mpsort.sort(disp, orderby=orderby, comm=comm)
if logger and pm.comm.rank == 0:
logger.info("sorting done")
rng_shift = MPIRandomState(seed=seed2, comm=comm, size=len(pos))
in_cell_shift = rng_shift.uniform(0, H[i], itemshape=(delta.ndim,))
pos[...] += in_cell_shift
pos[...] %= delta.BoxSize
if logger and pm.comm.rank == 0:
logger.info("catalog shifted.")
return pos, disp