Source code for nbodykit.mockmaker

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 lognormal_transform(density, bias=1.): r""" Apply a (biased) lognormal transformation of the density field by computing: .. math:: F(\delta) = \frac{1}{N} e^{b*\delta} where :math:`\delta` is the initial overdensity field and the normalization :math:`N` is chosen such that :math:`\langle F(\delta) \rangle = 1` Parameters ---------- density : array_like the input density field to apply the transformation to bias : float, optional optionally apply a linear bias to the density field; default is unbiased (1.0) Returns ------- toret : RealField the real field holding the transformed density field """ toret = density.copy() toret[:] = numpy.exp(bias * density.value) toret[:] /= toret.cmean(dtype='f8') return toret
[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