import os
import numpy
import logging
from nbodykit import CurrentMPIComm
from nbodykit.binned_statistic import BinnedStatistic
from nbodykit.meshtools import SlabIterator
from nbodykit.base.catalog import CatalogSourceBase
from nbodykit.base.mesh import MeshSource
from nbodykit.base.catalogmesh import CatalogMesh
[docs]class FFTBase(object):
"""
Base class provides functions for periodic FFT based Power spectrum code
Parameters
----------
first : CatalogSource
the first catalog source
second : CatalogSource, None
the second source, or None for auto-correlations
Nmesh : int, 3-vector
the number of cells per mesh size
BoxSize : 3-vector
the size of the box
kmin : float
the edge of the first ``k`` bin
dk : float
the size of the linearly spaced ``k`` bins
"""
def __init__(self, first, second, Nmesh, BoxSize):
from pmesh.pm import ParticleMesh
first = _cast_source(first, Nmesh=Nmesh, BoxSize=BoxSize)
if second is not None:
second = _cast_source(second, Nmesh=Nmesh, BoxSize=BoxSize)
else:
second = first
self.first = first
self.second = second
# grab comm from first source
self.comm = first.comm
# check for comm mismatch
assert second.comm is first.comm, "communicator mismatch between input sources"
# check box size
if not numpy.array_equal(first.attrs['BoxSize'], second.attrs['BoxSize']):
raise ValueError("'BoxSize' mismatch between sources in FFTPower")
# save meta-data
self.attrs = {}
self.attrs['Nmesh'] = first.attrs['Nmesh'].copy()
self.attrs['BoxSize'] = first.attrs['BoxSize'].copy()
self.attrs.update(zip(['Lx', 'Ly', 'Lz'], self.attrs['BoxSize']))
self.attrs.update({'volume':self.attrs['BoxSize'].prod()})
[docs] def save(self, output):
"""
Save the FFTPower result to disk.
The format is currently JSON.
"""
import json
from nbodykit.utils import JSONEncoder
# only the master rank writes
if self.comm.rank == 0:
self.logger.info('measurement done; saving result to %s' % output)
with open(output, 'w') as ff:
json.dump(self.__getstate__(), ff, cls=JSONEncoder)
[docs] @classmethod
@CurrentMPIComm.enable
def load(cls, output, comm=None):
"""
Load a saved FFTPower result.
The result has been saved to disk with :func:`save`.
"""
import json
from nbodykit.utils import JSONDecoder
if comm.rank == 0:
with open(output, 'r') as ff:
state = json.load(ff, cls=JSONDecoder)
else:
state = None
state = comm.bcast(state)
self = object.__new__(cls)
self.__setstate__(state)
self.comm = comm
return self
def _compute_3d_power(self):
"""
Compute and return the 3D power from two input sources
Returns
-------
p3d : array_like (complex)
the 3D complex array holding the power spectrum
"""
c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
# compute the auto power of single supplied field
if self.first is self.second:
c2 = c1
else:
c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
# calculate the 3d power spectrum, slab-by-slab to save memory
p3d = c1
for (s0, s1, s2) in zip(p3d.slabs, c1.slabs, c2.slabs):
s0[...] = s1 * s2.conj()
for i, s0 in zip(p3d.slabs.i, p3d.slabs):
# clear the zero mode.
mask = True
for i1 in i:
mask = mask & (i1 == 0)
s0[mask] = 0
# the complex field is dimensionless; power is L^3
# ref to http://icc.dur.ac.uk/~tt/Lectures/UA/L4/cosmology.pdf
p3d[...] *= self.attrs['BoxSize'].prod()
# get the number of objects (in a safe manner)
N1 = c1.attrs.get('N', 0)
N2 = c2.attrs.get('N', 0)
self.attrs.update({'N1':N1, 'N2':N2})
# add shotnoise (nonzero only for auto-spectra)
Pshot = 0
if self.first is self.second:
if 'shotnoise' not in c1.attrs:
if isinstance(self.first, CatalogMesh):
import warnings
warnings.warn(("no 'shotnoise' found for auto power spectrum "
"of discrete data in FFTPower"))
else:
Pshot = c1.attrs['shotnoise']
self.attrs['shotnoise'] = Pshot
return p3d
[docs]class FFTPower(FFTBase):
"""
Algorithm to compute the 1d or 2d power spectrum and/or multipoles
in a periodic box, using a Fast Fourier Transform (FFT).
This computes the power spectrum as the square of the Fourier modes of the
density field, which are computed via a FFT.
Results are computed when the object is inititalized. See the documenation
of :func:`~FFTPower.run` for the attributes storing the results.
.. note::
A full tutorial on the class is available in the documentation
:ref:`here <fftpower>`.
Parameters
----------
first : CatalogSource, MeshSource
the source for the first field; if a CatalogSource is provided, it
is automatically converted to MeshSource using the default painting
parameters (via :func:`~nbodykit.base.catalogmesh.CatalogMesh.to_mesh`)
mode : {'1d', '2d'}
compute either 1d or 2d power spectra
Nmesh : int, optional
the number of cells per side in the particle mesh used to paint the source
BoxSize : int, 3-vector, optional
the size of the box
second : CatalogSource, MeshSource, optional
the second source for cross-correlations
los : array_like , optional
the direction to use as the line-of-sight; must be a unit vector
Nmu : int, optional
the number of mu bins to use from :math:`\mu=[0,1]`;
if `mode = 1d`, then ``Nmu`` is set to 1
dk : float, optional
the linear spacing of ``k`` bins to use; if not provided, the
fundamental mode of the box is used
kmin : float, optional
the lower edge of the first ``k`` bin to use
poles : list of int, optional
a list of multipole numbers ``ell`` to compute :math:`P_\ell(k)`
from :math:`P(k,\mu)`
"""
logger = logging.getLogger('FFTPower')
def __init__(self, first, mode, Nmesh=None, BoxSize=None, second=None,
los=[0, 0, 1], Nmu=5, dk=None, kmin=0., poles=[]):
# mode is either '1d' or '2d'
if mode not in ['1d', '2d']:
raise ValueError("`mode` should be either '1d' or '2d'")
if poles is None:
poles = []
# check los
if numpy.isscalar(los) or len(los) != 3:
raise ValueError("line-of-sight ``los`` should be vector with length 3")
if not numpy.allclose(numpy.einsum('i,i', los, los), 1.0, rtol=1e-5):
raise ValueError("line-of-sight ``los`` must be a unit vector")
FFTBase.__init__(self, first, second, Nmesh, BoxSize)
# save meta-data
self.attrs['mode'] = mode
self.attrs['los'] = los
self.attrs['Nmu'] = Nmu
self.attrs['poles'] = poles
if dk is None:
dk = 2 * numpy.pi / self.attrs['BoxSize'].min()
self.attrs['dk'] = dk
self.attrs['kmin'] = kmin
self.run()
[docs] def run(self):
"""
Compute the power spectrum in a periodic box, using FFTs. This
function returns nothing, but attaches several attributes
to the class:
- :attr:`edges`
- :attr:`power`
- :attr:`poles`
Attributes
----------
edges : array_like
the edges of the wavenumber bins
power : :class:`~nbodykit.binned_statistic.BinnedStatistic`
a BinnedStatistic object that holds the measured :math:`P(k)` or
:math:`P(k,\mu)`. It stores the following variables:
- k :
the mean value for each ``k`` bin
- mu : ``mode=2d`` only
the mean value for each ``mu`` bin
- power :
complex array storing the real and imaginary components of the power
- modes :
the number of Fourier modes averaged together in each bin
poles : :class:`~nbodykit.binned_statistic.BinnedStatistic` or ``None``
a BinnedStatistic object to hold the multipole results
:math:`P_\ell(k)`; if no multipoles were requested by the user,
this is ``None``. It stores the following variables:
- k :
the mean value for each ``k`` bin
- power_L :
complex array storing the real and imaginary components for
the :math:`\ell=L` multipole
- modes :
the number of Fourier modes averaged together in each bin
attrs : dict
dictionary of meta-data; in addition to storing the input parameters,
it includes the following fields computed during the algorithm
execution:
- shotnoise : float
the power Poisson shot noise, equal to :math:`V/N`, where
:math:`V` is the volume of the box and `N` is the total
number of objects; if a cross-correlation is computed, this
will be equal to zero
- N1 : int
the total number of objects in the first source
- N2 : int
the total number of objects in the second source
"""
# only need one mu bin if 1d case is requested
if self.attrs['mode'] == "1d": self.attrs['Nmu'] = 1
# measure the 3D power (y3d is a ComplexField)
y3d = self._compute_3d_power()
# binning in k out to the minimum nyquist frequency
# (accounting for possibly anisotropic box)
dk = self.attrs['dk']
kmin = self.attrs['kmin']
kedges = numpy.arange(kmin, numpy.pi*y3d.Nmesh.min()/y3d.BoxSize.max() + dk/2, dk)
# project on to the desired basis
muedges = numpy.linspace(0, 1, self.attrs['Nmu']+1, endpoint=True)
edges = [kedges, muedges]
result, pole_result = project_to_basis(y3d, edges,
poles=self.attrs['poles'],
los=self.attrs['los'])
# format the power results into structured array
if self.attrs['mode'] == "1d":
cols = ['k', 'power', 'modes']
icols = [0, 2, 3]
edges = edges[0]
else:
cols = ['k', 'mu', 'power', 'modes']
icols = [0, 1, 2, 3]
# power results as a structured array
dtype = numpy.dtype([(name, result[icol].dtype.str) for icol,name in zip(icols,cols)])
power = numpy.squeeze(numpy.empty(result[0].shape, dtype=dtype))
for icol, col in zip(icols, cols):
power[col][:] = numpy.squeeze(result[icol])
# multipole results as a structured array
poles = None
if pole_result is not None:
k, poles, N = pole_result
cols = ['k'] + ['power_%d' %l for l in self.attrs['poles']] + ['modes']
result = [k] + [pole for pole in poles] + [N]
dtype = numpy.dtype([(name, result[icol].dtype.str) for icol,name in enumerate(cols)])
poles = numpy.empty(result[0].shape, dtype=dtype)
for icol, col in enumerate(cols):
poles[col][:] = result[icol]
# set all the necessary results
self.edges = edges
self.poles = poles
self.power = power
self._make_datasets()
def __getstate__(self):
state = dict(
edges=self.edges,
power=self.power.data,
poles=getattr(self.poles, 'data', None),
attrs=self.attrs)
return state
def __setstate__(self, state):
self.__dict__.update(state)
self._make_datasets()
def _make_datasets(self):
if self.attrs['mode'] == '1d':
self.power = BinnedStatistic(['k'], [self.edges], self.power, fields_to_sum=['modes'], **self.attrs)
else:
self.power = BinnedStatistic(['k', 'mu'], self.edges, self.power, fields_to_sum=['modes'], **self.attrs)
if self.poles is not None:
self.poles = BinnedStatistic(['k'], [self.power.edges['k']], self.poles, fields_to_sum=['modes'], **self.attrs)
def _compute_3d_power(self):
"""
Compute and return the 3D power from two input sources
Returns
-------
p3d : array_like (complex)
the 3D complex array holding the power spectrum
"""
c1 = self.first.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
# compute the auto power of single supplied field
if self.first is self.second:
c2 = c1
else:
c2 = self.second.paint(mode='complex', Nmesh=self.attrs['Nmesh'])
# calculate the 3d power spectrum, slab-by-slab to save memory
p3d = c1
for (s0, s1, s2) in zip(p3d.slabs, c1.slabs, c2.slabs):
s0[...] = s1 * s2.conj()
for i, s0 in zip(p3d.slabs.i, p3d.slabs):
# clear the zero mode.
mask = True
for i1 in i:
mask = mask & (i1 == 0)
s0[mask] = 0
# the complex field is dimensionless; power is L^3
# ref to http://icc.dur.ac.uk/~tt/Lectures/UA/L4/cosmology.pdf
p3d[...] *= self.attrs['BoxSize'].prod()
# get the number of objects (in a safe manner)
N1 = c1.attrs.get('N', 0)
N2 = c2.attrs.get('N', 0)
self.attrs.update({'N1':N1, 'N2':N2})
# add shotnoise (nonzero only for auto-spectra)
Pshot = 0
if self.first is self.second:
if 'shotnoise' not in c1.attrs:
if isinstance(self.first, CatalogMesh):
import warnings
warnings.warn(("no 'shotnoise' found for auto power spectrum "
"of discrete data in FFTPower"))
else:
Pshot = c1.attrs['shotnoise']
self.attrs['shotnoise'] = Pshot
return p3d
[docs]class ProjectedFFTPower(FFTBase):
"""
The power spectrum of a field in a periodic box, projected over certain axes.
This is not really always physically meaningful, but convenient for
making sense of Lyman-Alpha forest or lensing maps.
This is usually called the 1d power spectrum or 2d power spectrum.
Results are computed when the object is inititalized. See the documenation
of :func:`~ProjectedFFTPower.run` for the attributes storing the results.
Parameters
----------
first : CatalogSource, MeshSource
the source for the first field; if a CatalogSource is provided, it
is automatically converted to MeshSource using the default painting
parameters (via :func:`~nbodykit.base.catalogmesh.CatalogMesh.to_mesh`)
Nmesh : int, optional
the number of cells per side in the particle mesh used to paint the source
BoxSize : int, 3-vector, optional
the size of the box
second : CatalogSource, MeshSource, optional
the second source for cross-correlations
axes : tuple
axes to measure the power on. The axes not in the list will be averaged out.
For example:
- (0, 1) : project to x,y and measure power
- (0) : project to x and measure power.
dk : float, optional
the linear spacing of ``k`` bins to use; if not provided, the
fundamental mode of the box is used
kmin : float, optional
the lower edge of the first ``k`` bin to use
"""
logger = logging.getLogger('ProjectedFFTPower')
def __init__(self, first, Nmesh=None, BoxSize=None, second=None,
axes=(0, 1), dk=None, kmin=0.):
FFTBase.__init__(self, first, second, Nmesh, BoxSize)
# only deal with 1d and 2d projections.
assert len(axes) in (1, 2), "length of ``axes`` in ProjectedFFTPower should be 1 or 2"
if dk is None:
dk = 2 * numpy.pi / self.attrs['BoxSize'].min()
self.attrs['dk'] = dk
self.attrs['kmin'] = kmin
self.attrs['axes'] = axes
self.run()
[docs] def run(self):
"""
Run the algorithm. This attaches the following attributes to the class:
- :attr:`edges`
- :attr:`power`
Attributes
----------
edges : array_like
the edges of the wavenumber bins
power : :class:`~nbodykit.binned_statistic.BinnedStatistic`
a BinnedStatistic object that holds the projected power.
It stores the following variables:
- k :
the mean value for each ``k`` bin
- power :
complex array holding the real and imaginary components of the
projected power
- modes :
the number of Fourier modes averaged together in each bin
"""
c1 = self.first.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
r1 = c1.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
# average along projected axes;
# part of product is the rfftn vs r2c (for axes)
# the rest is for the mean (Nmesh - axes)
c1 = numpy.fft.rfftn(r1) / self.attrs['Nmesh'].prod()
# compute the auto power of single supplied field
if self.first is self.second:
c2 = c1
else:
c2 = self.second.paint(Nmesh=self.attrs['Nmesh'], mode='complex')
r2 = c2.preview(self.attrs['Nmesh'], axes=self.attrs['axes'])
c2 = numpy.fft.rfftn(r2) / self.attrs['Nmesh'].prod() # average along projected axes
pk = c1 * c2.conj()
# clear the zero mode
pk.flat[0] = 0
shape = numpy.array([self.attrs['Nmesh'][i] for i in self.attrs['axes']], dtype='int')
boxsize = numpy.array([self.attrs['BoxSize'][i] for i in self.attrs['axes']])
I = numpy.eye(len(shape), dtype='int') * -2 + 1
k = [numpy.fft.fftfreq(N, 1. / (N * 2 * numpy.pi / L))[:pkshape].reshape(kshape) for N, L, kshape, pkshape in zip(shape, boxsize, I, pk.shape)]
kmag = sum(ki ** 2 for ki in k) ** 0.5
W = numpy.empty(pk.shape, dtype='f4')
W[...] = 2.0
W[..., 0] = 1.0
W[..., -1] = 1.0
dk = self.attrs['dk']
kmin = self.attrs['kmin']
axes = list(self.attrs['axes'])
kedges = numpy.arange(kmin, numpy.pi * self.attrs['Nmesh'][axes].min() / self.attrs['BoxSize'][axes].max() + dk/2, dk)
xsum = numpy.zeros(len(kedges) + 1)
Psum = numpy.zeros(len(kedges) + 1, dtype='complex128')
Nsum = numpy.zeros(len(kedges) + 1)
dig = numpy.digitize(kmag.flat, kedges)
xsum.flat += numpy.bincount(dig, weights=(W * kmag).flat, minlength=xsum.size)
Psum.real.flat += numpy.bincount(dig, weights=(W * pk.real).flat, minlength=xsum.size)
Psum.imag.flat += numpy.bincount(dig, weights=(W * pk.imag).flat, minlength=xsum.size)
Nsum.flat += numpy.bincount(dig, weights=W.flat, minlength=xsum.size)
self.power = numpy.empty(len(kedges) - 1,
dtype=[('k', 'f8'), ('power', 'c16'), ('modes', 'f8')])
with numpy.errstate(invalid='ignore'):
self.power['k'] = (xsum / Nsum)[1:-1]
self.power['power'] = (Psum / Nsum)[1:-1] * boxsize.prod() # dimension is 'volume'
self.power['modes'] = Nsum[1:-1]
self.edges = kedges
self.power = BinnedStatistic(['k'], [self.edges], self.power)
def __getstate__(self):
state = dict(
edges=self.edges,
power=self.power.data,
attrs=self.attrs)
return state
def __setstate__(self, state):
self.__dict__.update(state)
self.power = BinnedStatistic(['k'], [self.edges], self.power)
[docs]def project_to_basis(y3d, edges, los=[0, 0, 1], poles=[]):
"""
Project a 3D statistic on to the specified basis. The basis will be one
of:
- 2D (`x`, `mu`) bins: `mu` is the cosine of the angle to the line-of-sight
- 2D (`x`, `ell`) bins: `ell` is the multipole number, which specifies
the Legendre polynomial when weighting different `mu` bins
.. note::
The 2D (`x`, `mu`) bins will be computed only if `poles` is specified.
See return types for further details.
Notes
-----
* the `mu` range extends from 0.0 to 1.0
* the `mu` bins are half-inclusive half-exclusive, except the last bin
is inclusive on both ends (to include `mu = 1.0`)
Parameters
----------
y3d : RealField or ComplexField
the 3D array holding the statistic to be projected to the specified basis
edges : list of arrays, (2,)
list of arrays specifying the edges of the desired `x` bins and `mu` bins
los : array_like,
the line-of-sight direction to use, which `mu` is defined with
respect to; default is [0, 0, 1] for z.
poles : list of int, optional
if provided, a list of integers specifying multipole numbers to
project the 2d `(x, mu)` bins on to
hermitian_symmetric : bool, optional
Whether the input array `y3d` is Hermitian-symmetric, i.e., the negative
frequency terms are just the complex conjugates of the corresponding
positive-frequency terms; if ``True``, the positive frequency terms
will be explicitly double-counted to account for this symmetry
Returns
-------
result : tuple
the 2D binned results; a tuple of ``(xmean_2d, mumean_2d, y2d, N_2d)``, where:
- xmean_2d : array_like, (Nx, Nmu)
the mean `x` value in each 2D bin
- mumean_2d : array_like, (Nx, Nmu)
the mean `mu` value in each 2D bin
- y2d : array_like, (Nx, Nmu)
the mean `y3d` value in each 2D bin
- N_2d : array_like, (Nx, Nmu)
the number of values averaged in each 2D bin
pole_result : tuple or `None`
the multipole results; if `poles` supplied it is a tuple of ``(xmean_1d, poles, N_1d)``,
where:
- xmean_1d : array_like, (Nx,)
the mean `x` value in each 1D multipole bin
- poles : array_like, (Nell, Nx)
the mean multipoles value in each 1D bin
- N_1d : array_like, (Nx,)
the number of values averaged in each 1D bin
"""
comm = y3d.pm.comm
x3d = y3d.x
hermitian_symmetric = numpy.iscomplexobj(y3d)
from scipy.special import legendre
# setup the bin edges and number of bins
xedges, muedges = edges
x2edges = xedges**2
Nx = len(xedges) - 1
Nmu = len(muedges) - 1
# always make sure first ell value is monopole, which
# is just (x, mu) projection since legendre of ell=0 is 1
do_poles = len(poles) > 0
_poles = [0]+sorted(poles) if 0 not in poles else sorted(poles)
legpoly = [legendre(l) for l in _poles]
ell_idx = [_poles.index(l) for l in poles]
Nell = len(_poles)
# valid ell values
if any(ell < 0 for ell in _poles):
raise ValueError("in `project_to_basis`, multipole numbers must be non-negative integers")
# initialize the binning arrays
musum = numpy.zeros((Nx+2, Nmu+2))
xsum = numpy.zeros((Nx+2, Nmu+2))
ysum = numpy.zeros((Nell, Nx+2, Nmu+2), dtype=y3d.dtype) # extra dimension for multipoles
Nsum = numpy.zeros((Nx+2, Nmu+2), dtype='i8')
# if input array is Hermitian symmetric, only half of the last
# axis is stored in `y3d`
symmetry_axis = -1 if hermitian_symmetric else None
# iterate over y-z planes of the coordinate mesh
for slab in SlabIterator(x3d, axis=0, symmetry_axis=symmetry_axis):
# the square of coordinate mesh norm
# (either Fourier space k or configuraton space x)
xslab = slab.norm2()
# if empty, do nothing
if len(xslab.flat) == 0: continue
# get the bin indices for x on the slab
dig_x = numpy.digitize(xslab.flat, x2edges)
# make xslab just x
xslab **= 0.5
# get the bin indices for mu on the slab
mu = slab.mu(los) # defined with respect to specified LOS
dig_mu = numpy.digitize(abs(mu).flat, muedges)
# make the multi-index
multi_index = numpy.ravel_multi_index([dig_x, dig_mu], (Nx+2,Nmu+2))
# sum up x in each bin (accounting for negative freqs)
xslab[:] *= slab.hermitian_weights
xsum.flat += numpy.bincount(multi_index, weights=xslab.flat, minlength=xsum.size)
# count number of modes in each bin (accounting for negative freqs)
Nslab = numpy.ones_like(xslab) * slab.hermitian_weights
Nsum.flat += numpy.bincount(multi_index, weights=Nslab.flat, minlength=Nsum.size)
# compute multipoles by weighting by Legendre(ell, mu)
for iell, ell in enumerate(_poles):
# weight the input 3D array by the appropriate Legendre polynomial
weighted_y3d = legpoly[iell](mu) * y3d[slab.index]
# add conjugate for this kx, ky, kz, corresponding to
# the (-kx, -ky, -kz) --> need to make mu negative for conjugate
# Below is identical to the sum of
# Leg(ell)(+mu) * y3d[:, nonsingular] (kx, ky, kz)
# Leg(ell)(-mu) * y3d[:, nonsingular].conj() (-kx, -ky, -kz)
# or
# weighted_y3d[:, nonsingular] += (-1)**ell * weighted_y3d[:, nonsingular].conj()
# but numerically more accurate.
if hermitian_symmetric:
if ell % 2: # odd, real part cancels
weighted_y3d.real[slab.nonsingular] = 0.
weighted_y3d.imag[slab.nonsingular] *= 2.
else: # even, imag part cancels
weighted_y3d.real[slab.nonsingular] *= 2.
weighted_y3d.imag[slab.nonsingular] = 0.
# sum up the weighted y in each bin
weighted_y3d *= (2.*ell + 1.)
ysum[iell,...].real.flat += numpy.bincount(multi_index, weights=weighted_y3d.real.flat, minlength=Nsum.size)
if numpy.iscomplexobj(ysum):
ysum[iell,...].imag.flat += numpy.bincount(multi_index, weights=weighted_y3d.imag.flat, minlength=Nsum.size)
# sum up the absolute mag of mu in each bin (accounting for negative freqs)
mu[:] *= slab.hermitian_weights
musum.flat += numpy.bincount(multi_index, weights=abs(mu).flat, minlength=musum.size)
# sum binning arrays across all ranks
xsum = comm.allreduce(xsum)
musum = comm.allreduce(musum)
ysum = comm.allreduce(ysum)
Nsum = comm.allreduce(Nsum)
# add the last 'internal' mu bin (mu == 1) to the last visible mu bin
# this makes the last visible mu bin inclusive on both ends.
ysum[..., -2] += ysum[..., -1]
musum[:, -2] += musum[:, -1]
xsum[:, -2] += xsum[:, -1]
Nsum[:, -2] += Nsum[:, -1]
# reshape and slice to remove out of bounds points
sl = slice(1, -1)
with numpy.errstate(invalid='ignore'):
# 2D binned results
y2d = (ysum[0,...] / Nsum)[sl,sl] # ell=0 is first index
xmean_2d = (xsum / Nsum)[sl,sl]
mumean_2d = (musum / Nsum)[sl, sl]
N_2d = Nsum[sl,sl]
# 1D multipole results (summing over mu (last) axis)
if do_poles:
N_1d = Nsum[sl,sl].sum(axis=-1)
xmean_1d = xsum[sl,sl].sum(axis=-1) / N_1d
poles = ysum[:, sl,sl].sum(axis=-1) / N_1d
poles = poles[ell_idx,...]
# return y(x,mu) + (possibly empty) multipoles
result = (xmean_2d, mumean_2d, y2d, N_2d)
pole_result = (xmean_1d, poles, N_1d) if do_poles else None
return result, pole_result
def _cast_source(source, BoxSize, Nmesh):
"""
Cast an object to a MeshSource. BoxSize and Nmesh is used
only on CatalogSource
"""
from pmesh.pm import Field
from nbodykit.source.mesh import FieldMesh
if isinstance(source, Field):
# if input is a Field object, wrap it as a MeshSource.
source = FieldMesh(source)
elif isinstance(source, CatalogSourceBase):
# if input is CatalogSource, use defaults to make it into a mesh
if not isinstance(source, MeshSource):
source = source.to_mesh(BoxSize=BoxSize, Nmesh=Nmesh, dtype='f8', compensated=True)
# now we are having a MeshSource
# paint the density field to the mesh
if not isinstance(source, MeshSource):
raise TypeError("Unknown type of source in FFTPower: %s" % str(type(source)))
if BoxSize is not None and any(source.attrs['BoxSize'] != BoxSize):
raise ValueError("Mismatched Boxsize between __init__ and source.attrs")
if Nmesh is not None and any(source.attrs['Nmesh'] != Nmesh):
raise ValueError(("Mismatched Nmesh between __init__ and source.attrs; "
"if trying to re-sample with a different mesh, specify "
"`Nmesh` as keyword of to_mesh()"))
return source