コード例 #1
0
    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
コード例 #2
0
ファイル: wedges.py プロジェクト: ybh0822/pyRSD
    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
コード例 #3
0
    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
コード例 #4
0
    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'])