def __call__(self, power):
"""
Parameters
----------
power : xarray.DataArray
a DataArray holding the :math:`P(k,\mu)` values on a
coordinate grid with ``k`` and ``mu`` dimensions
Returns
-------
Pell : xarray.DataArray
a DataArray holding the :math:`P_\ell(k)` on a coordinate grid
with ``k`` and ``ell`` dimensions.
"""
self.power = power
# input coordinate grid
k = self.power['k']
mu = self.power['mu']
# return array
Nk = len(self.power)
shape = (len(k), len(self.ells))
Pell = xr.DataArray(np.empty(shape),
coords=[('k', k), ('ell', self.ells)])
# compute each Pell
for ell in self.ells:
kern = (2 * ell + 1.) * legendre(ell)(mu)
val = np.array(
[pygcl.SimpsIntegrate(mu, kern * d) for d in self.power])
Pell.loc[dict(ell=ell)] = val
return Pell
def __call__(self, power):
"""
Parameters
----------
power : xarray.DataArray
a DataArray holding the :math:`P(k,\mu)` values on a
coordinate grid with ``k`` and ``mu`` dimensions
Returns
-------
Pell : xarray.DataArray
a DataArray holding the :math:`P_\ell(k)` on a coordinate grid
with ``k`` and ``ell`` dimensions.
"""
self.power = power
# input coordinate grid
k = self.power['k']
mu = self.power['mu']
# the bin shape
Nk = self.grid.Nk
Nmu = len(self.mu_edges) - 1
binshape = (Nk + 2, Nmu + 2)
# compute the bin indices
k_idx = np.arange(self.grid.Nk, dtype=int)[:, None]
dig_k = np.repeat(k_idx, self.grid.Nmu, axis=1)[self.grid.notnull] + 1
dig_mu = np.digitize(self.grid.mu[self.grid.notnull], self.mu_edges)
indices = np.ravel_multi_index([dig_k, dig_mu], binshape)
# sum up power
toret = np.zeros(binshape)
minlength = np.prod(binshape)
P = self.power.values
toret.flat = np.bincount(indices,
weights=P[self.grid.notnull],
minlength=minlength)
toret = toret.reshape(binshape)[1:-1, 1:-1][..., ::2]
# and the normalization
N = np.zeros(binshape)
N.flat = np.bincount(indices, minlength=minlength)
N = N.reshape(binshape)[1:-1, 1:-1][..., ::2]
# normalize properly
toret /= N
# return array
shape = (len(k), len(self.mu_cen))
Pwedge = xr.DataArray(np.empty(shape),
coords=[('k', k), ('mu', self.mu_cen)])
for i in range(Pwedge.shape[1]):
Pwedge[:, i] = toret[:, i]
return Pwedge
def kmax(self, val):
"""
The maximum wavenumber ``k`` to include.
Shape can be either a float or array of length :attr:`N2`.
"""
if val is None: val = np.inf
toret = np.empty(self.N2)
toret[:] = val
return toret
def __call__(self, power, k_out=None, extrap=False, mcfit_kwargs={}, **kws):
"""
Evaluate the convolved multipoles.
Parameters
----------
power : xarray.DataArray
a DataArray holding the :math:`P(k,\mu)` values on a
coordinate grid with ``k`` and ``mu`` dimensions.
k_out : array_like, optional
if provided, evaluate the convolved multipoles at these
``k`` values using a spline
**kws :
additional keywords for testing purposes
Returns
-------
Pell : xarray.DataArray
a DataArray holding the convolved :math:`P_\ell(k)` on a
coordinate grid with ``k`` and ``ell`` dimensions.
"""
from pyRSD.extern import mcfit
# get testing keywords
dry_run = kws.get('dry_run', False)
no_convolution = kws.get('no_convolution', False)
# get the unconvovled theory multipoles
Pell0 = GriddedMultipoleTransfer.__call__(self, power)
# create additional logspaced k values for zero-padding up to k=100 h/Mpc
oldk = Pell0['k'].values
dk = np.diff(np.log10(oldk))[0]
newk = 10**(np.arange(np.log10(oldk.max()) + dk, 2 + 0.5*dk, dk))
newk = np.concatenate([oldk, newk])
# now copy over with zeros
Nk = len(newk); Nell = Pell0.shape[1]
Pell = xr.DataArray(np.zeros((Nk,Nell)), coords={'k':newk, 'ell':Pell0.ell}, dims=['k', 'ell'])
Pell.loc[dict(k=Pell0['k'])] = Pell0[:]
# do the convolution
if not no_convolution:
# FFT the input power multipoles
xi = np.empty((Nk, Nell), order='F') # column-continuous
for i, ell in enumerate(self.ells):
P2xi = mcfit.P2xi(newk, l=ell, **mcfit_kwargs)
rr, xi[:,i] = P2xi(Pell.sel(ell=ell).values, extrap=extrap)
# the linear combination of multipoles
if dry_run:
xi_conv = xi.copy()
else:
xi_conv = self.convolver(self.ells, rr, xi, order='F')
# FFTLog back
Pell_conv = np.empty((Nk, Nell), order='F')
for i, ell in enumerate(self.ells):
xi2P = mcfit.xi2P(rr, l=ell, **mcfit_kwargs)
kk, Pell_conv[:,i] = xi2P(xi_conv[:,i], extrap=extrap)
else:
Pell_conv = Pell
# interpolate to k_out
coords = coords={'ell':Pell0.ell}
if k_out is not None:
shape = (len(k_out), len(self.ells))
toret = np.ones(shape) * np.nan
for i, ell in enumerate(self.ells):
idx = np.isfinite(newk)
spl = spline(newk[idx], Pell_conv[idx,i])
toret[:,i] = spl(k_out)
coords['k'] = k_out
else:
toret = Pell_conv
coords['k'] = newk
return xr.DataArray(toret, coords=coords, dims=['k', 'ell'])