from nbodykit.source.mesh.catalog import CatalogMesh
from nbodykit import _global_options
import numpy
import logging
import warnings
# for converting from particle to mesh
from pmesh.pm import RealField
from nbodykit.utils import attrs_to_dict
[docs]class MultipleSpeciesCatalogMesh(CatalogMesh):
"""
A subclass of :class:`~nbodykit.base.catalogmesh.CatalogMesh`
designed to paint the density field from a sum of multiple types
of particles.
The :func:`compute` function paints the density field summed over
all particle species.
Parameters
----------
source : CatalogSource
the input catalog that we wish to interpolate to a mesh
BoxSize :
the size of the box
Nmesh : int, 3-vector
the number of cells per mesh side
dtype : str
the data type of the values stored on mesh
weight : str
column in ``source`` that specifies the weight value for each
particle in the ``source`` to use when gridding
value : str
column in ``source`` that specifies the field value for each particle;
the mesh stores a weighted average of this column
selection : str
column in ``source`` that selects the subset of particles to grid
position : str, optional
column in ``source`` specifying the position coordinates; default
is ``Position``
"""
logger = logging.getLogger('MultipleSpeciesCatalogMesh')
def __init__(self, source, Nmesh, BoxSize, dtype,
selection, position, weight, value,
interlaced, compensated, resampler):
from nbodykit.source.catalog import MultipleSpeciesCatalog
if not isinstance(source, MultipleSpeciesCatalog):
raise TypeError(("the input source for MultipleSpeciesCatalogMesh "
"must be a MultipleSpeciesCatalog"))
CatalogMesh.__init__(self,
source=source,
Nmesh=Nmesh, BoxSize=BoxSize, dtype=dtype,
Position=None,
Selection=None,
Weight=None,
Value=None,
resampler=resampler,
compensated=compensated,
interlaced=interlaced,
comm=source.comm)
self.species = source.species
self.position = position
self.selection = selection
self.weight = weight
self.value = value
def __iter__(self):
return iter(self.species)
[docs] def __getitem__(self, key):
"""
If indexed by a species name, return a CatalogMesh object holding
only the data columns for that species with the same parameters as
the current object.
If not a species name, this has the same behavior as
:func:`CatalogSource.__getitem__`.
"""
from nbodykit.source.mesh import CatalogMesh
if key in self.source.species:
cat = self.source[key]
# view as a catalog mesh
mesh = CatalogMesh(cat,
BoxSize=self.attrs['BoxSize'],
Nmesh=self.attrs['Nmesh'],
dtype=self.dtype,
Weight=cat[self.weight],
Value=cat[self.value],
Selection=cat[self.selection],
Position=cat[self.position],
interlaced=self.interlaced,
compensated=self.compensated,
resampler=self.resampler,
)
# attach attributes from self
return mesh.__finalize__(self)
else:
raise KeyError("%s is not a species defined in the source" % key)
[docs] def to_real_field(self, normalize=True):
r"""
Paint the density field holding the sum of all particle species,
returning a :class:`~pmesh.pm.RealField` object.
Meta-data computed for each particle is stored in the :attr:`attrs`
attribute of the returned RealField, with keys that are prefixed by
the species name. In particular, the total shot noise for the
mesh is defined as:
.. math::
P_\mathrm{shot} = \sum_i (W_i/W_\mathrm{tot})^2 P_{\mathrm{shot},i},
where the sum is over all species in the catalog, ``W_i`` is the
sum of the ``Weight`` column for the :math:`i^\mathrm{th}` species,
and :math:`W_\mathrm{tot}` is the sum of :math:`W_i` across all species.
Parameters
----------
normalize : bool, optional
if ``True``, normalize the density field as :math:`1+\delta`,
dividing by the total mean number of objects per cell, as given
by the ``num_per_cell`` meta-data value in :attr:`attrs`
Returns
-------
RealField :
the RealField holding the painted density field, with a
:attr:`attrs` dictionary attribute holding the meta-data
"""
# track the sum of the mean number of objects per cell across species
attrs = {'num_per_cell':0., 'N':0}
# initialize an empty real field
real = self.pm.create(type='real', value=0)
# loop over each species
for name in self.source.species:
if self.pm.comm.rank == 0:
self.logger.info("painting the '%s' species" %name)
# get a CatalogMesh for this species
species_mesh = self[name]
# paint (in-place) the un-normalized density field for this species
species_mesh.to_real_field(out=real, normalize=False)
# add to the mean number of objects per cell and total number
attrs['num_per_cell'] += real.attrs['num_per_cell']
attrs['N'] += real.attrs['N']
# store the meta-data for this species, with a prefix
attrs.update(attrs_to_dict(real, name+'.'))
# # normalize the field by nbar -> this is now 1+delta
if normalize:
real[:] /= attrs['num_per_cell']
# update the meta-data
real.attrs.clear()
real.attrs.update(attrs)
# compute total shot noise by summing of shot noise each species
real.attrs['shotnoise'] = 0
total_weight = sum(real.attrs['%s.W' %name] for name in self.source.species)
for name in self.source.species:
this_Pshot = real.attrs['%s.shotnoise' %name]
this_weight = real.attrs['%s.W' %name]
real.attrs['shotnoise'] += (this_weight/total_weight)**2 * this_Pshot
return real