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'])