diff --git a/nbodykit/algorithms/fftpower.py b/nbodykit/algorithms/fftpower.py index 4b316bb36..62d87670b 100644 --- a/nbodykit/algorithms/fftpower.py +++ b/nbodykit/algorithms/fftpower.py @@ -198,8 +198,8 @@ 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 mode not in ['1d', '2d'] and not isinstance(mode, AdvancedProjection): + raise ValueError("`mode` should be either '1d' or '2d', or an AdvancedProjection") if poles is None: poles = [] @@ -212,8 +212,9 @@ def __init__(self, first, mode, Nmesh=None, BoxSize=None, second=None, FFTBase.__init__(self, first, second, Nmesh, BoxSize) + self.mode = mode # save meta-data - self.attrs['mode'] = mode + self.attrs['mode'] = str(mode) self.attrs['los'] = los self.attrs['Nmu'] = Nmu self.attrs['poles'] = poles @@ -283,6 +284,11 @@ def run(self): # measure the 3D power (y3d is a ComplexField) y3d, attrs = self._compute_3d_power(self.first, self.second) + if isinstance(self.mode, AdvancedProjection): + # advanced projection + return self.mode.project(y3d, 'power'), None + + # binning in k out to the minimum nyquist frequency # (accounting for possibly anisotropic box) dk = self.attrs['dk'] @@ -297,6 +303,7 @@ def run(self): muedges = numpy.linspace(0, 1, self.attrs['Nmu']+1, endpoint=True) edges = [kedges, muedges] coords = [kcoords, None] + result, pole_result = project_to_basis(y3d, edges, poles=self.attrs['poles'], los=self.attrs['los']) @@ -502,6 +509,166 @@ def __setstate__(self, state): self.__dict__.update(state) self.power = BinnedStatistic(['k'], [self.edges], self.power) +class AdvancedProjection: + """ An advanced projection projects a field object to a certain set of + basis to create a BinnedStatistic. + + The bases of the basis are labeled by 'names'. The usual label + is given by the center of edges. But subclasses can override them + by modifying the coords attribute of the corresponding name. This + is for example done in MagPoleProjection to set the labels of + poles to the correct values. + + Subclass shall override the basis() function and / or the + get_coords() function. + + basis() generates a basis vector of the same type as the input variable + (y3d) for the given index in the label space. Usually + (the default implementation), the basis is a binning mask, which is set + to 1 when the point falls into the bin, and zero otherwise. + Override basis() to use more complicated basis, e.g. Legendre polynormials. + + get_coords() convert the coordinates on original field object to + the coords corresponding to the labels. + + """ + + def __init__(self, names, edges, **kwargs): + self.names = names + self.edges = dict(zip(names, edges)) + self.attrs = {} + self.attrs.update(kwargs) + + # use default centers + self.coords = [None for e in edges] + + def project(self, y3d, valuename): + data = numpy.zeros(shape=[len(self.edges[dim]) - 1 for dim in self.names], + # one column per dim + dtype=[(name, 'f8') for name in self.names] + + # one colum for the value and one column for the number of modes + [(valuename, 'c16'), ('modes', 'f8')] + ) + + unitary = y3d.copy() + unitary[:] = 1.0 + + # generate the full coords vectors for central values + full_coords = {} + for dim in self.names: + def filter(k, v): + return self.get_coords(k, dim) + full_coords[dim] = y3d.copy().apply(filter) + + for index in numpy.ndindex(data.shape): + basis = self.basis(index, y3d) + modes = basis.cdot(unitary).real + data[index]['modes'] = modes + if modes > 0: + data[index][valuename] = basis.cdot(y3d) / modes + # compute central coords + for dim in self.names: + data[index][dim] = basis.cdot(full_coords[dim]).real / modes + + return BinnedStatistic(self.names, [self.edges[dim] for dim in self.names], data, coords=self.coords, **self.attrs) + + def basis(self, index, y3d): + """ default basis are just binning. Override for fancier basis, e.g. P_leg(mu) """ + def filter(k, v): + mask = True + for i, dim in enumerate(self.edges): + min, max = self.get_range(index[i], dim) + coords = self.get_coords(k, dim) + mask = mask & ((coords >= min) & (coords < max)) + return 1.0 * mask + return y3d.copy().apply(filter) + + def get_range(self, ind, dim): + kmin = self.edges[dim][ind] + kmax = self.edges[dim][ind + 1] + return numpy.array([kmin, kmax]) + + def get_coords(self, k, dim): + if dim == 'k': + return k.normp(2) ** 0.5 + + if dim == 'pole': + return 0 + + if dim == 'mu': + # this uses special attr 'los', defined only by some of the Projections. + mu = 0 + for i in range(len(k)): + mu = mu + k[i] * self.attrs['los'][i] + + knorm = k.normp(2, zeromode=1.0) ** 0.5 + mu = mu / knorm + return mu + + +class MagProjection(AdvancedProjection): + def __init__(self, kedges): + """ Magnitude projection, edges=[0, 2 pi / L, ...] """ + AdvancedProjection.__init__(self, names=['k'], edges=[kedges]) + +class ZRhoProjection(AdvancedProjection): + def __init__(self, zedges, rhoedges, los): + """ Magnitude projection, edges=[0, 2 pi / L, ...] """ + AdvancedProjection.__init__(self, names=['k_z', 'k_r'], edges=[zedges, rhoedges], los=los) + + def get_coords(self, k, dim): + if dim == 'k_z': + return AdvancedProjection.get_coords(self, k, 'mu') * k.normp(2) ** 0.5 + if dim == 'k_r': + return (k.normp(2) - self.get_coords(k, 'k_z') ** 2) ** 0.5 + +class MagMuProjection(AdvancedProjection): + def __init__(self, kedges, muedges, los): + AdvancedProjection.__init__(self, + names=['k', 'mu'], + edges=[kedges, muedges], + los=los + ) + + def get_coords(self, k, dim): + if dim == 'mu': + return abs(AdvancedProjection.get_coords(self, k, dim)) + else: + return AdvancedProjection.get_coords(self, k, dim) + + +class MagPoleProjection(AdvancedProjection): + def __init__(self, kedges, poles, los): + # pad the poles at the end because we never touch the last mode. + poles = numpy.array(list(poles) + [-1]) + AdvancedProjection.__init__(self, + names=['k', 'pole'], + edges=[kedges, poles], + los=los + ) + self.coords[1] = poles[:-1] + + def basis(self, index, y3d): + from scipy.special import legendre + pole = self.edges['pole'][index[1]] + legpoly = legendre(pole) + min, max = self.get_range(index[0], 'k') + + def filter(k, v): + kcoords = self.get_coords(k, 'k') + mask = ((kcoords >= min) & (kcoords < max)) + mu = self.get_coords(k, 'mu') + L = legpoly(mu) + return L * mask + + return y3d.copy().apply(filter) + + def project(self, y3d, valuename): + r = AdvancedProjection.project(self, y3d, valuename) + # patch up pole value because we can't compute it with simple avaraging. + r['pole'][...] = self.coords[1][None, :] + return r + 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 @@ -534,11 +701,6 @@ def project_to_basis(y3d, edges, los=[0, 0, 1], poles=[]): 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 ------- @@ -567,6 +729,11 @@ def project_to_basis(y3d, edges, los=[0, 0, 1], poles=[]): """ comm = y3d.pm.comm x3d = y3d.x + + # 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 hermitian_symmetric = numpy.iscomplexobj(y3d) from scipy.special import legendre diff --git a/nbodykit/algorithms/tests/test_fftpower.py b/nbodykit/algorithms/tests/test_fftpower.py index 99d3cc210..ebf4285cf 100644 --- a/nbodykit/algorithms/tests/test_fftpower.py +++ b/nbodykit/algorithms/tests/test_fftpower.py @@ -43,6 +43,65 @@ def test_cic_aliasing(comm): red_chi2 = (residual**2).sum()/len(Pk) # should be about 0.5-0.6 assert red_chi2 < 1.0 +@MPITest([1]) +def test_fftpower_advanced(comm): + from nbodykit.algorithms.fftpower import MagProjection, MagMuProjection, MagPoleProjection, ZRhoProjection + + CurrentMPIComm.set(comm) + source = UniformCatalog(nbar=3e-3, BoxSize=512., seed=42) + + r = FFTPower(source, mode=MagProjection( + numpy.linspace(0, 16 * 2 * numpy.pi / 512., 17, endpoint=True) + ), Nmesh=32) + p = r.power + + r2 = FFTPower(source, mode='1d', Nmesh=32) + p2 = r2.power + + # FIXME: this is not expect to be the same because of round off errors. + # perhaps try some very narrow bins? + + #assert p.shape == p2.shape + #assert_allclose(p['modes'], p2['modes'], atol=16) + #assert_allclose(p['power'], p2['power'], rtol=1e-5) + + r = FFTPower(source, mode=MagMuProjection( + kedges = numpy.linspace(0, 16 * 2 * numpy.pi / 512., 17, endpoint=True), + muedges = numpy.linspace(0, 1, 6, endpoint=True), + los=[0, 0, 1], + ), Nmesh=32) + + p = r.power + + r2 = FFTPower(source, mode='2d', Nmesh=32) + p2 = r2.power + + assert p.shape == p2.shape + + # FIXME: this is not expect to be the same because of round off errors. + # perhaps try some very narrow bins? + + assert_allclose(p.coords['mu'], p2.coords['mu']) + assert_allclose(p.coords['k'], p2.coords['k']) + #assert_allclose(p['modes'], p2['modes'], atol=16) + #assert_allclose(p['power'], p2['power'], rtol=1e-5) + + r = FFTPower(source, mode=MagPoleProjection( + kedges = numpy.linspace(0, 16 * 2 * numpy.pi / 512., 17, endpoint=True), + poles = [0, 2, 4], + los=[0, 0, 1], + ), Nmesh=32) + + p = r.power + assert_array_equal(p.coords['pole'], [0, 2, 4]) + assert_allclose(p['pole'], [[0, 2, 4]] * len(p['pole'])) + + r = FFTPower(source, mode=ZRhoProjection( + numpy.linspace(0, 16 * 2 * numpy.pi / 512., 17, endpoint=True), + numpy.linspace(0, 16 * 2 * numpy.pi / 512., 17, endpoint=True), + los=[0, 0, 1], + ), Nmesh=32) + p = r.power @MPITest([1]) def test_fftpower_poles(comm):