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
[docs]def gaussian_complex_fields(pm, linear_power, seed,
unitary_amplitude=False, inverted_phase=False,
compute_displacement=False):
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")
# 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, mode='complex', unitary=unitary_amplitude)
if inverted_phase: delta_k[...] *= -1
# initialize the displacement fields for (x,y,z)
if compute_displacement:
disp_k = [ComplexField(pm) 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)
# 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)
zero_idx = k2 == 0.
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'):
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
# 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):
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)
# FFT the density to real-space
delta = delta_k.c2r()
# FFT the velocity back to real space
if compute_displacement:
disp = [disp_k[i].c2r() for i in range(delta.ndim)]
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, comm=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
"""
if comm is None:
comm = MPI.COMM_WORLD
# create a random state with the input seed
rng = numpy.random.RandomState(seed)
# 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)
# mean number of objects per cell
H = delta.BoxSize / delta.Nmesh
overallmean = H.prod() * nbar
# number of objects in each cell (per rank)
cellmean = delta.value*overallmean
cellmean = GatherArray(cellmean.flatten(), comm, root=0)
# rank 0 computes the poisson sampling
if comm.rank == 0:
N = rng.poisson(cellmean)
else:
N = None
# scatter N back evenly across the ranks
counts = comm.allgather(delta.value.size)
N = ScatterArray(N, comm, root=0, counts=counts).reshape(delta.shape)
Nlocal = N.sum() # local number of particles
Ntot = comm.allreduce(Nlocal) # the collective number of particles
nonzero_cells = N.nonzero() # indices of nonzero cells
# initialize the mesh of particle positions and displacement
# this has the shape: (number of dimensions, number of nonzero cells)
pos_mesh = numpy.empty(numpy.shape(nonzero_cells), dtype=delta.dtype)
disp_mesh = numpy.empty_like(pos_mesh)
# generate the coordinates for each nonzero cell
for i in range(delta.ndim):
# particle positions initially on the coordinate grid
pos_mesh[i] = numpy.squeeze(delta.pm.x[i])[nonzero_cells[i]]
# displacements for each particle
disp_mesh[i] = displacement[i][nonzero_cells]
# rank 0 computes the in-cell uniform offsets
if comm.rank == 0:
in_cell_shift = numpy.empty((Ntot, delta.ndim), dtype=delta.dtype)
for i in range(delta.ndim):
in_cell_shift[:,i] = rng.uniform(0, H[i], size=Ntot)
else:
in_cell_shift = None
# scatter the in-cell uniform offsets back to the ranks
counts = comm.allgather(Nlocal)
in_cell_shift = ScatterArray(in_cell_shift, comm, root=0, counts=counts)
# initialize the output array of particle positions and displacement
# this has shape: (local number of particles, number of dimensions)
pos = numpy.zeros((Nlocal, delta.ndim), dtype=delta.dtype)
disp = numpy.zeros_like(pos)
# coordinates of each object (placed randomly in each cell)
for i in range(delta.ndim):
pos[:,i] = numpy.repeat(pos_mesh[i], N[nonzero_cells]) + in_cell_shift[:,i]
pos[:,i] %= delta.BoxSize[i]
# displacements of each object
for i in range(delta.ndim):
disp[:,i] = numpy.repeat(disp_mesh[i], N[nonzero_cells])
return pos, disp