Beispiel #1
0
    def to_complex_field(self):
        """
        Return the ComplexField stored on disk.

        .. note::
            The mesh stored on disk must be stored with ``mode=complex``

        Returns
        -------
        real : pmesh.pm.ComplexField
            an array-like object holding the mesh loaded from disk in Fourier
            space
        """
        if not self.isfourier:
            return NotImplemented
        pmread = self.pm

        if self.comm.rank == 0:
            self.logger.info("reading complex field from %s" % self.path)

        with FileMPI(comm=self.comm, filename=self.path)[self.dataset] as ds:
            complex2 = ComplexField(pmread)
            assert self.comm.allreduce(complex2.size) == ds.size
            start = numpy.sum(self.comm.allgather(complex2.size)[:self.comm.rank], dtype='intp')
            end = start + complex2.size
            complex2.unravel(ds[start:end])

        return complex2
Beispiel #2
0
    def read(self, real):
        import bigfile
        if self.comm.rank == 0:
            self.logger.info("Reading from Nmesh = %d to Nmesh = %d" %(self.Nmesh, real.Nmesh[0]))

        if any(real.Nmesh != self.Nmesh):
            pmread = ParticleMesh(BoxSize=real.BoxSize, Nmesh=(self.Nmesh, self.Nmesh, self.Nmesh),
                    dtype='f4', comm=self.comm)
        else:
            pmread = real.pm

        f = bigfile.BigFileMPI(self.comm, self.path)

        with f[self.dataset] as ds:
            if self.isfourier:
                if self.comm.rank == 0:
                    self.logger.info("reading complex field")
                complex2 = ComplexField(pmread)
                assert self.comm.allreduce(complex2.size) == ds.size
                start = sum(self.comm.allgather(complex2.size)[:self.comm.rank])
                end = start + complex2.size
                complex2.unsort(ds[start:end])
                complex2.resample(real)
            else:
                if self.comm.rank == 0:
                    self.logger.info("reading real field")
                real2 = RealField(pmread)
                start = sum(self.comm.allgather(real2.size)[:self.comm.rank])
                end = start + real2.size
                real2.unsort(ds[start:end])
                real2.resample(real)
Beispiel #3
0
    def to_complex_field(self):
        """
        Return the ComplexField stored on disk.

        .. note::
            The mesh stored on disk must be stored with ``mode=complex``

        Returns
        -------
        real : pmesh.pm.ComplexField
            an array-like object holding the mesh loaded from disk in Fourier
            space
        """
        if not self.isfourier:
            return NotImplemented
        pmread = self.pm

        if self.comm.rank == 0:
            self.logger.info("reading complex field from %s" % self.path)

        with BigFileMPI(comm=self.comm,
                        filename=self.path)[self.dataset] as ds:
            complex2 = ComplexField(pmread)
            assert self.comm.allreduce(complex2.size) == ds.size
            start = sum(self.comm.allgather(complex2.size)[:self.comm.rank])
            end = start + complex2.size
            complex2.unsort(ds[start:end])

        return complex2
Beispiel #4
0
def test_c2r_vjp(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    real = pm.generate_whitenoise(1234, type='real', mean=1.0)
    comp = real.r2c()

    def objective(comp):
        real = comp.c2r()
        obj = (real.value ** 2).sum()
        return comm.allreduce(obj)

    grad_real = RealField(pm)
    grad_real[...] = real[...] * 2
    grad_comp = ComplexField(pm)
    grad_comp = grad_real.c2r_vjp(grad_real)
    grad_comp.decompress_vjp(grad_comp)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])):
        dx1, c1 = perturb(comp, ind1, dx)
        ng1 = (objective(c1) - objective(comp)) / dx
        ag1 = grad_comp.cgetitem(ind1) * dx1 / dx
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)
Beispiel #5
0
    def run(self):
        """
        Run the algorithm, which computes and returns the grid in C_CONTIGUOUS order partitioned by ranks.
        """
        from nbodykit import measurestats

        if self.comm.rank == 0:
            self.logger.info("importing done")
            self.logger.info("Resolution Nmesh : %d" % self.paintNmesh)
            self.logger.info("paintbrush : %s" % self.painter.paintbrush)

        # setup the particle mesh object, taking BoxSize from the painters
        pmpaint = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.paintNmesh] * 3, dtype="f4", comm=self.comm)
        pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype="f4", comm=self.comm)

        real, stats = self.painter.paint(pmpaint, self.datasource)

        if self.writeFourier:
            result = ComplexField(pm)
        else:
            result = RealField(pm)
        real.resample(result)

        # reuses the memory
        result.sort(out=result)
        result = result.ravel()

        # return all the necessary results
        return result, stats
Beispiel #6
0
def test_c2r_vjp(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    real = pm.generate_whitenoise(1234, mode='real')
    comp = real.r2c()

    def objective(comp):
        real = comp.c2r()
        obj = (real.value ** 2).sum()
        return comm.allreduce(obj)

    grad_real = RealField(pm)
    grad_real[...] = real[...] * 2
    grad_comp = ComplexField(pm)
    grad_comp = grad_real.c2r_vjp(grad_real)
    grad_comp.decompress_vjp(grad_comp)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])):
        dx1, c1 = perturb(comp, ind1, dx)
        ng1 = (objective(c1) - objective(comp)) / dx
        ag1 = grad_comp.cgetitem(ind1) * dx1 / dx
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)
Beispiel #7
0
def test_complex_apply(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    complex = ComplexField(pm)

    def filter(k, v):
        return k[0] + k[1] * 1j

    complex.apply(filter, out=Ellipsis)

    for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs):
        assert_array_equal(slab, x[0] + x[1] * 1j)
Beispiel #8
0
def test_complex_apply(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    complex = ComplexField(pm)
    def filter(k, v):
        knormp = k.normp()
        assert_allclose(knormp, sum(ki ** 2 for ki in k))
        return k[0] + k[1] * 1j

    complex.apply(filter, out=Ellipsis)

    for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs):
        assert_array_equal(slab, x[0] + x[1] * 1j)
Beispiel #9
0
def test_cmean(comm):
    # this tests cmean (collective mean) along with resampling preseves it.

    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    complex1 = ComplexField(pm1)
    complex2 = ComplexField(pm2)
    real2 = RealField(pm2)
    real1 = RealField(pm1)
    for i, kk, slab in zip(complex1.slabs.i, complex1.slabs.x, complex1.slabs):
        slab[...] = sum([k**2 for k in kk]) **0.5

    complex1.c2r(real1)
    real1.resample(real2)
    assert_almost_equal(real1.cmean(), real2.cmean())
Beispiel #10
0
def test_hermitian_weights(comm):

    numpy.random.seed(42)

    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
    cfield = ComplexField(pm)
    data = numpy.random.random(size=cfield.shape)
    data = data[:] + 1j * data[:]

    cfield[...] = data[:]
    x = cfield.x

    # iterate over symmetry axis
    for i, slab in enumerate(SlabIterator(x, axis=2, symmetry_axis=2)):

        # nonsingular weights give indices of positive frequencies
        nonsig = slab.nonsingular
        weights = slab.hermitian_weights

        # weights == 2 when iterating frequency is positive
        if numpy.float(slab.coords(2)) > 0.:
            assert weights > 1
            assert numpy.all(nonsig == True)
        else:
            assert weights == 1.0
            assert numpy.all(nonsig == False)
Beispiel #11
0
def test_cmean(comm):
    # this tests cmean (collective mean) along with resampling preseves it.

    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    complex1 = ComplexField(pm1)
    complex2 = ComplexField(pm2)
    real2 = RealField(pm2)
    real1 = RealField(pm1)
    for i, kk, slab in zip(complex1.slabs.i, complex1.slabs.x, complex1.slabs):
        slab[...] = sum([k**2 for k in kk])**0.5

    complex1.c2r(real1)
    real1.resample(real2)
    assert_almost_equal(real1.cmean(), real2.cmean())
Beispiel #12
0
def test_whitenoise(comm):
    # the whitenoise shall preserve the large scale.
    pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8')
    complex1_down = ComplexField(pm0)
    complex2_down = ComplexField(pm0)

    complex1 = pm1.generate_whitenoise(seed=8, unitary=True)
    complex2 = pm2.generate_whitenoise(seed=8, unitary=True)

    complex1.resample(complex1_down)
    complex2.resample(complex2_down)

    mask1 = complex1_down.value != complex2_down.value
    assert_array_equal(complex1_down.value, complex2_down.value)
Beispiel #13
0
def test_complex_iter(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    complex = ComplexField(pm)

    for x, slab in zip(complex.slabs.x, complex.slabs):
        assert_array_equal(slab.shape,
                           sum(x[d]**2 for d in range(len(pm.Nmesh))).shape)
        for a, b in zip(slab.x, x):
            assert_almost_equal(a, b)
Beispiel #14
0
def test_sort(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 6], comm=comm, dtype='f8')
    real = RealField(pm)
    truth = numpy.arange(8 * 6)
    real[...] = truth.reshape(8, 6)[real.slices]
    unsorted = real.copy()
    real.sort(out=Ellipsis)
    conjecture = numpy.concatenate(comm.allgather(real.value.ravel()))
    assert_array_equal(conjecture, truth)

    real.unravel(real)
    assert_array_equal(real, unsorted)

    complex = ComplexField(pm)
    truth = numpy.arange(8 * 4)
    complex[...] = truth.reshape(8, 4)[complex.slices]
    complex.ravel(out=Ellipsis)
    conjecture = numpy.concatenate(comm.allgather(complex.value.ravel()))
    assert_array_equal(conjecture, truth)
Beispiel #15
0
def test_sort(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 6], comm=comm, dtype='f8')
    real = RealField(pm)
    truth = numpy.arange(8 * 6)
    real[...] = truth.reshape(8, 6)[real.slices]
    unsorted = real.copy()
    real.sort(out=Ellipsis)
    conjecture = numpy.concatenate(comm.allgather(real.value.ravel()))
    assert_array_equal(conjecture, truth)

    real.unravel(real)
    assert_array_equal(real, unsorted)

    complex = ComplexField(pm)
    truth = numpy.arange(8 * 4)
    complex[...] = truth.reshape(8, 4)[complex.slices]
    complex.ravel(out=Ellipsis)
    conjecture = numpy.concatenate(comm.allgather(complex.value.ravel()))
    assert_array_equal(conjecture, truth)
Beispiel #16
0
    def run(self):
        """
        Run the algorithm, which computes and returns the grid in C_CONTIGUOUS order partitioned by ranks.
        """
        from nbodykit import measurestats

        if self.comm.rank == 0:
            self.logger.info('importing done')
            self.logger.info('Resolution Nmesh : %d' % self.paintNmesh)
            self.logger.info('paintbrush : %s' % self.painter.paintbrush)

        # setup the particle mesh object, taking BoxSize from the painters
        pmpaint = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.paintNmesh] * 3, dtype='f4', comm=self.comm)
        pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype='f4', comm=self.comm)

        real, stats = self.painter.paint(pmpaint, self.datasource)

        if self.writeFourier:
            result = ComplexField(pm)
        else:
            result = RealField(pm)
        real.resample(result)

        # reuses the memory
        result.sort(out=result)
        result = result.ravel()

        # return all the necessary results
        return result, stats
Beispiel #17
0
    def read(self, real):
        import bigfile
        if self.comm.rank == 0:
            self.logger.info("Reading from Nmesh = %d to Nmesh = %d" %
                             (self.Nmesh, real.Nmesh[0]))

        if any(real.Nmesh != self.Nmesh):
            pmread = ParticleMesh(BoxSize=real.BoxSize,
                                  Nmesh=(self.Nmesh, self.Nmesh, self.Nmesh),
                                  dtype='f4',
                                  comm=self.comm)
        else:
            pmread = real.pm

        f = bigfile.BigFileMPI(self.comm, self.path)

        with f[self.dataset] as ds:
            if self.isfourier:
                if self.comm.rank == 0:
                    self.logger.info("reading complex field")
                complex2 = ComplexField(pmread)
                assert self.comm.allreduce(complex2.size) == ds.size
                start = sum(
                    self.comm.allgather(complex2.size)[:self.comm.rank])
                end = start + complex2.size
                complex2.unsort(ds[start:end])
                complex2.resample(real)
            else:
                if self.comm.rank == 0:
                    self.logger.info("reading real field")
                real2 = RealField(pmread)
                start = sum(self.comm.allgather(real2.size)[:self.comm.rank])
                end = start + real2.size
                real2.unsort(ds[start:end])
                real2.resample(real)
Beispiel #18
0
    def _compute_multipoles(self, kedges):
        """
        Compute the window-convoled power spectrum multipoles, for a data set
        with non-trivial survey geometry.

        This estimator builds upon the work presented in Bianchi et al. 2015
        and Scoccimarro et al. 2015, but differs in the implementation. This
        class uses the spherical harmonic addition theorem such that
        only :math:`2\ell+1` FFTs are required per multipole, rather than the
        :math:`(\ell+1)(\ell+2)/2` FFTs in the implementation presented by
        Bianchi et al. and Scoccimarro et al.

        References
        ----------
        * Bianchi, Davide et al., `Measuring line-of-sight-dependent Fourier-space clustering using FFTs`,
          MNRAS, 2015
        * Scoccimarro, Roman, `Fast estimators for redshift-space clustering`, Phys. Review D, 2015
        """
        # clear compensation from the actions
        for source in [self.first, self.second]:
            source.actions[:] = []
            source.compensated = False
            assert len(source.actions) == 0

        # compute the compensations
        compensation = {}
        for name, mesh in zip(['first', 'second'], [self.first, self.second]):
            compensation[name] = get_compensation(mesh)
            if self.comm.rank == 0:
                if compensation[name] is not None:
                    args = (compensation[name]['func'].__name__, name)
                    self.logger.info(
                        "using compensation function %s for source '%s'" %
                        args)
                else:
                    self.logger.warning(
                        "no compensation applied for source '%s'" % name)

        rank = self.comm.rank
        pm = self.first.pm

        # setup the 1D-binning
        muedges = numpy.linspace(-1, 1, 2, endpoint=True)
        edges = [kedges, muedges]

        # make a structured array to hold the results
        cols = ['k'] + ['power_%d' % l
                        for l in sorted(self.attrs['poles'])] + ['modes']
        dtype = ['f8'] + ['c8'] * len(self.attrs['poles']) + ['i8']
        dtype = numpy.dtype(list(zip(cols, dtype)))
        result = numpy.empty(len(kedges) - 1, dtype=dtype)

        # offset the box coordinate mesh ([-BoxSize/2, BoxSize]) back to
        # the original (x,y,z) coords
        offset = self.attrs['BoxCenter'] + 0.5 * pm.BoxSize / pm.Nmesh

        # always need to compute ell=0
        poles = sorted(self.attrs['poles'])
        if 0 not in poles:
            poles = [0] + poles
        assert poles[0] == 0

        # spherical harmonic kernels (for ell > 0)
        Ylms = [[get_real_Ylm(l, m) for m in range(-l, l + 1)]
                for l in poles[1:]]

        # paint the 1st FKP density field to the mesh (paints: data - alpha*randoms, essentially)
        rfield1 = self.first.compute(Nmesh=self.attrs['Nmesh'])
        meta1 = rfield1.attrs.copy()
        if rank == 0:
            self.logger.info("%s painting of 'first' done" %
                             self.first.resampler)

        # store alpha: ratio of data to randoms
        self.attrs['alpha'] = meta1['alpha']

        # FFT 1st density field and apply the resampler transfer kernel
        cfield = rfield1.r2c()
        if compensation['first'] is not None:
            cfield.apply(out=Ellipsis, **compensation['first'])
        if rank == 0: self.logger.info('ell = 0 done; 1 r2c completed')

        # monopole A0 is just the FFT of the FKP density field
        # NOTE: this holds FFT of density field #1
        volume = pm.BoxSize.prod()
        A0_1 = ComplexField(pm)
        A0_1[:] = cfield[:] * volume  # normalize with a factor of volume

        # paint second mesh too?
        if self.first is not self.second:

            # paint the second field
            rfield2 = self.second.compute(Nmesh=self.attrs['Nmesh'])
            meta2 = rfield2.attrs.copy()
            if rank == 0:
                self.logger.info("%s painting of 'second' done" %
                                 self.second.resampler)

            # need monopole of second field
            if 0 in self.attrs['poles']:

                # FFT density field and apply the resampler transfer kernel
                A0_2 = rfield2.r2c()
                A0_2[:] *= volume
                if compensation['second'] is not None:
                    A0_2.apply(out=Ellipsis, **compensation['second'])
        else:
            rfield2 = rfield1
            meta2 = meta1

            # monopole of second field is first field
            if 0 in self.attrs['poles']:
                A0_2 = A0_1

        # ensure alpha from first mesh is equal to alpha from second mesh
        # NOTE: this is mostly just a sanity check, and should always be true if
        # we made it this far already
        if not numpy.allclose(
                rfield1.attrs['alpha'], rfield2.attrs['alpha'], rtol=1e-3):
            msg = (
                "ConvolvedFFTPower cross-correlations currently require the same"
                " FKPCatalog (data/randoms), such that only the weight column can vary;"
                " different ``alpha`` values found for first/second meshes")
            raise ValueError(msg)

        # save the painted density field #2 for later
        density2 = rfield2.copy()

        # initialize the memory holding the Aell terms for
        # higher multipoles (this holds sum of m for fixed ell)
        # NOTE: this will hold FFTs of density field #2
        Aell = ComplexField(pm)

        # the real-space grid
        xgrid = [
            xx.astype('f8') + offset[ii]
            for ii, xx in enumerate(density2.slabs.optx)
        ]
        xnorm = numpy.sqrt(sum(xx**2 for xx in xgrid))
        xgrid = [x / xnorm for x in xgrid]

        # the Fourier-space grid
        kgrid = [kk.astype('f8') for kk in cfield.slabs.optx]
        knorm = numpy.sqrt(sum(kk**2 for kk in kgrid))
        knorm[knorm == 0.] = numpy.inf
        kgrid = [k / knorm for k in kgrid]

        # proper normalization: same as equation 49 of Scoccimarro et al. 2015
        for name in ['data', 'randoms']:
            self.attrs[name + '.norm'] = self.normalization(
                name, self.attrs['alpha'])

        if self.attrs['randoms.norm'] > 0:
            norm = 1.0 / self.attrs['randoms.norm']

            # check normalization
            Adata = self.attrs['data.norm']
            Aran = self.attrs['randoms.norm']
            if not numpy.allclose(Adata, Aran, rtol=0.05):
                msg = "normalization in ConvolvedFFTPower different by more than 5%; "
                msg += ",algorithm requires they must be similar\n"
                msg += "\trandoms.norm = %.6f, data.norm = %.6f\n" % (Aran,
                                                                      Adata)
                msg += "\tpossible discrepancies could be related to normalization "
                msg += "of n(z) column ('%s')\n" % self.first.nbar
                msg += "\tor the consistency of the FKP weight column for 'data' "
                msg += "and 'randoms';\n"
                msg += "\tn(z) columns for 'data' and 'randoms' should be "
                msg += "normalized to represent n(z) of the data catalog"
                raise ValueError(msg)

            if rank == 0:
                self.logger.info(
                    "normalized power spectrum with `randoms.norm = %.6f`" %
                    Aran)
        else:
            # an empty random catalog is provides, so we will ignore the normalization.
            norm = 1.0
            if rank == 0:
                self.logger.info(
                    "normalization of power spectrum is neglected, as no random is provided."
                )

        # loop over the higher order multipoles (ell > 0)
        start = time.time()
        for iell, ell in enumerate(poles[1:]):

            # clear 2D workspace
            Aell[:] = 0.

            # iterate from m=-l to m=l and apply Ylm
            substart = time.time()
            for Ylm in Ylms[iell]:

                # reset the real-space mesh to the original density #2
                rfield2[:] = density2[:]

                # apply the config-space Ylm
                for islab, slab in enumerate(rfield2.slabs):
                    slab[:] *= Ylm(xgrid[0][islab], xgrid[1][islab],
                                   xgrid[2][islab])

                # real to complex of field #2
                rfield2.r2c(out=cfield)

                # apply the Fourier-space Ylm
                for islab, slab in enumerate(cfield.slabs):
                    slab[:] *= Ylm(kgrid[0][islab], kgrid[1][islab],
                                   kgrid[2][islab])

                # add to the total sum
                Aell[:] += cfield[:]

                # and this contribution to the total sum
                substop = time.time()
                if rank == 0:
                    self.logger.debug("done term for Y(l=%d, m=%d) in %s" %
                                      (Ylm.l, Ylm.m, timer(substart, substop)))

            # apply the compensation transfer function
            if compensation['second'] is not None:
                Aell.apply(out=Ellipsis, **compensation['second'])

            # factor of 4*pi from spherical harmonic addition theorem + volume factor
            Aell[:] *= 4 * numpy.pi * volume

            # log the total number of FFTs computed for each ell
            if rank == 0:
                args = (ell, len(Ylms[iell]))
                self.logger.info('ell = %d done; %s r2c completed' % args)

            # calculate the power spectrum multipoles, slab-by-slab to save memory
            # NOTE: this computes (A0 of field #1) * (Aell of field #2).conj()
            for islab in range(A0_1.shape[0]):
                Aell[islab, ...] = norm * A0_1[islab] * Aell[islab].conj()

            # project on to 1d k-basis (averaging over mu=[0,1])
            proj_result, _ = project_to_basis(Aell, edges)
            result['power_%d' % ell][:] = numpy.squeeze(proj_result[2])

        # summarize how long it took
        stop = time.time()
        if rank == 0:
            self.logger.info(
                "higher order multipoles computed in elapsed time %s" %
                timer(start, stop))

        # also compute ell=0
        if 0 in self.attrs['poles']:

            # the 3D monopole
            for islab in range(A0_1.shape[0]):
                A0_1[islab, ...] = norm * A0_1[islab] * A0_2[islab].conj()

            # the 1D monopole
            proj_result, _ = project_to_basis(A0_1, edges)
            result['power_0'][:] = numpy.squeeze(proj_result[2])

        # save the number of modes and k
        result['k'][:] = numpy.squeeze(proj_result[0])
        result['modes'][:] = numpy.squeeze(proj_result[-1])

        # compute shot noise
        self.attrs['shotnoise'] = self.shotnoise(self.attrs['alpha'])

        # copy over any painting meta data
        if self.first is self.second:
            copy_meta(self.attrs, meta1)
        else:
            copy_meta(self.attrs, meta1, prefix='first')
            copy_meta(self.attrs, meta2, prefix='second')

        return result
Beispiel #19
0
def test_cgetitem(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
    for i in numpy.ndindex((4, 4)):
        complex = RealField(pm)
        complex[...] = 0
        v2 = complex.csetitem(i, 100.)
        v1 = complex.cgetitem(i)
        assert v2 == 100.
        assert_array_equal(v1, v2)

    for i in numpy.ndindex((4, 3)):
        complex = ComplexField(pm)
        complex[...] = 0
        v2 = complex.csetitem(i, 100. + 10j)
        complex.c2r(out=Ellipsis).r2c(out=Ellipsis)
        v1 = complex.cgetitem(i)
        if i == (0, 0):
            assert v2 == 100.
            assert comm.allreduce(complex.value.sum()) == 100.
        elif i == (0, 2):
            assert v2 == 100.
            assert comm.allreduce(complex.value.sum()) == 100.
        elif i == (2, 0):
            assert v2 == 100.
            assert comm.allreduce(complex.value.sum()) == 100.
        elif i == (1, 0):
            assert v2 == 100 + 10j
            assert comm.allreduce(complex.value.sum()) == 200.
        elif i == (3, 0):
            assert v2 == 100 + 10j
            assert comm.allreduce(complex.value.sum()) == 200.
        elif i == (3, 2):
            assert v2 == 100 + 10j
            assert comm.allreduce(complex.value.sum()) == 200.
        elif i == (1, 2):
            assert v2 == 100 + 10j
            assert comm.allreduce(complex.value.sum()) == 200.
        elif i == (2, 2):
            assert v2 == 100.
            assert comm.allreduce(complex.value.sum()) == 100.
        else:
            assert v2 == 100. + 10j
            assert_array_equal(comm.allreduce(complex.value.sum()), 100. + 10j)
        assert_array_equal(v1, v2)

    for i in numpy.ndindex((4, 3, 2)):
        complex = ComplexField(pm)
        complex[...] = 0
        v2 = complex.csetitem(i, 100.)
        complex.c2r(out=Ellipsis).r2c(out=Ellipsis)
        v1 = complex.cgetitem(i)
        if i == (0, 0, 0):
            assert v2 == 100.
        if i == (0, 0, 1):
            assert v2 == 0.
        elif i == (0, 2, 0):
            assert v2 == 100.
        elif i == (0, 2, 1):
            assert v2 == 0.
        elif i == (2, 0, 0):
            assert v2 == 100.
        elif i == (2, 0, 1):
            assert v2 == 0.
        elif i == (2, 2, 0):
            assert v2 == 100.
        elif i == (2, 2, 1):
            assert v2 == 0.
        else:
            assert v2 == 100.
        assert_array_equal(v1, v2)
Beispiel #20
0
def test_ctol(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
    complex = ComplexField(pm)
    value, local = complex._ctol((3, 3))
    assert local is None
Beispiel #21
0
def test_fupsample(comm):
    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    numpy.random.seed(3333)
    truth = numpy.fft.rfftn(numpy.random.normal(size=(8, 8)))

    complex1 = ComplexField(pm1)
    for ind in numpy.ndindex(*complex1.cshape):
        complex1.csetitem(ind, truth[ind])
        if any(i == 4 for i in ind):
            complex1.csetitem(ind, 0)
        else:
            complex1.csetitem(ind, truth[ind])

        if any(i >= 2 and i < 7 for i in ind):
            complex1.csetitem(ind, 0)

    assert_almost_equal(complex1[...], complex1.c2r().r2c())
    complex2 = ComplexField(pm2)
    for ind in numpy.ndindex(*complex2.cshape):
        newind = tuple([i if i <= 2 else 8 - (4 - i) for i in ind])
        if any(i == 2 for i in ind):
            complex2.csetitem(ind, 0)
        else:
            complex2.csetitem(ind, truth[newind])

    tmpr = RealField(pm1)
    tmp = ComplexField(pm1)

    complex2.resample(tmp)

    assert_almost_equal(complex1[...], tmp[...], decimal=5)

    complex2.c2r().resample(tmp)

    assert_almost_equal(complex1[...], tmp[...], decimal=5)

    complex2.resample(tmpr)

    assert_almost_equal(tmpr.r2c(), tmp[...])

    complex2.c2r().resample(tmpr)

    assert_almost_equal(tmpr.r2c(), tmp[...])
Beispiel #22
0
def test_fupsample(comm):
    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    numpy.random.seed(3333)
    truth = numpy.fft.rfftn(numpy.random.normal(size=(8, 8)))

    complex1 = ComplexField(pm1)
    for ind in numpy.ndindex(*complex1.cshape):
        complex1.csetitem(ind, truth[ind])
        if any(i == 4 for i in ind):
            complex1.csetitem(ind, 0)
        else:
            complex1.csetitem(ind, truth[ind])

        if any(i >= 2 and i < 7 for i in ind):
            complex1.csetitem(ind, 0)

    assert_almost_equal(complex1[...], complex1.c2r().r2c())
    complex2 = ComplexField(pm2)
    for ind in numpy.ndindex(*complex2.cshape):
        newind = tuple([i if i <= 2 else 8 - (4 - i) for i in ind])
        if any(i == 2 for i in ind):
            complex2.csetitem(ind, 0)
        else:
            complex2.csetitem(ind, truth[newind])

    tmpr = RealField(pm1)
    tmp = ComplexField(pm1)

    complex2.resample(tmp)

    assert_almost_equal(complex1[...], tmp[...], decimal=5)

    complex2.c2r().resample(tmp)

    assert_almost_equal(complex1[...], tmp[...], decimal=5)

    complex2.resample(tmpr)

    assert_almost_equal(tmpr.r2c(), tmp[...])

    complex2.c2r().resample(tmpr)

    assert_almost_equal(tmpr.r2c(), tmp[...])
Beispiel #23
0
def test_ctol(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
    complex = ComplexField(pm)
    value, local = complex._ctol((3, 3))
    assert local is None
Beispiel #24
0
def test_shape_complex(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
    comp = ComplexField(pm)
    assert (tuple(comp.cshape) == (8, 5))
Beispiel #25
0
def gaussian_complex_fields(pm,
                            linear_power,
                            seed,
                            unitary_amplitude=False,
                            inverted_phase=False,
                            compute_displacement=False):
    r"""
    Make a Gaussian realization of a overdensity field, :math:`\delta(x)`.

    If specified, also compute the corresponding 1st order Lagrangian
    displacement field (Zel'dovich approximation) :math:`\psi(x)`,
    which is related to the linear velocity field via:

    .. math::

        v(x) = f a H \psi(x)

    Notes
    -----
    This computes the overdensity field using the following steps:

    #. Generate complex variates with unity variance
    #. Scale the Fourier field by :math:`(P(k) / V)^{1/2}`

    After step 2, the complex field has unity variance. This
    is equivalent to generating real-space normal variates
    with mean and unity variance, calling r2c() and dividing by :math:`N^3`
    since the variance of the complex FFT (with no additional normalization)
    is :math:`N^3 \times \sigma^2_\mathrm{real}`.

    Furthermore, the power spectrum is defined as V * variance.
    So a normalization factor of 1 / V shows up in step 2,
    cancels this factor such that the power spectrum is P(k).

    The linear displacement field is computed as:

    .. math::

        \psi_i(k) = i \frac{k_i}{k^2} \delta(k)

    .. note::

        To recover the linear velocity in proper units, i.e., km/s,
        from the linear displacement, an additional factor of
        :math:`f \times a \times H(a)` is required

    Parameters
    ----------
    pm : pmesh.pm.ParticleMesh
        the mesh object
    linear_power : callable
        a function taking wavenumber as its only argument, which returns
        the linear power spectrum
    seed : int
        the random seed used to generate the random field
    compute_displacement : bool, optional
        if ``True``, also return the linear Zel'dovich displacement field;
        default is ``False``
    unitary_amplitude : bool, optional
        if ``True``, the seed gaussian has unitary_amplitude.
    inverted_phase: bool, optional
        if ``True``, the phase of the seed gaussian is inverted

    Returns
    -------
    delta_k : ComplexField
        the real-space Gaussian overdensity field
    disp_k : ComplexField or ``None``
        if requested, the Gaussian displacement field
    """
    if not isinstance(seed, numbers.Integral):
        raise ValueError(
            "the seed used to generate the linear field must be an integer")

    # use pmesh to generate random complex white noise field (done in parallel)
    # variance of complex field is unity
    # multiply by P(k)**0.5 to get desired variance
    delta_k = pm.generate_whitenoise(seed,
                                     mode='complex',
                                     unitary=unitary_amplitude)

    if inverted_phase: delta_k[...] *= -1

    # initialize the displacement fields for (x,y,z)
    if compute_displacement:
        disp_k = [ComplexField(pm) for i in range(delta_k.ndim)]
        for i in range(delta_k.ndim):
            disp_k[i][:] = 1j
    else:
        disp_k = None

    # volume factor needed for normalization
    norm = 1.0 / pm.BoxSize.prod()

    # iterate in slabs over fields
    slabs = [delta_k.slabs.x, delta_k.slabs]
    if compute_displacement:
        slabs += [d.slabs for d in disp_k]

    # loop over the mesh, slab by slab
    for islabs in zip(*slabs):
        kslab, delta_slab = islabs[:2]  # the k arrays and delta slab

        # the square of the norm of k on the mesh
        k2 = sum(kk**2 for kk in kslab)
        zero_idx = k2 == 0.

        k2[zero_idx] = 1.  # avoid dividing by zero

        # the linear power (function of k)
        power = linear_power((k2**0.5).flatten())

        # multiply complex field by sqrt of power
        delta_slab[...].flat *= (power * norm)**0.5

        # set k == 0 to zero (zero config-space mean)
        delta_slab[zero_idx] = 0.

        # compute the displacement
        if compute_displacement:

            # ignore division where k==0 and set to 0
            with numpy.errstate(invalid='ignore'):
                for i in range(delta_k.ndim):
                    disp_slab = islabs[2 + i]
                    disp_slab[...] *= kslab[i] / k2 * delta_slab[...]
                    disp_slab[zero_idx] = 0.  # no bulk displacement

    # return Fourier-space density and displacement (which could be None)
    return delta_k, disp_k
Beispiel #26
0
	def _compute_multipoles(self):
		"""
		Compute the window-convoled power spectrum multipoles, for a data set
		with non-trivial survey geometry.

		This estimator builds upon the work presented in Bianchi et al. 2015
		and Scoccimarro et al. 2015, but differs in the implementation. This
		class uses the spherical harmonic addition theorem such that
		only :math:`2\ell+1` FFTs are required per multipole, rather than the
		:math:`(\ell+1)(\ell+2)/2` FFTs in the implementation presented by
		Bianchi et al. and Scoccimarro et al.

		References
		----------
		* Bianchi, Davide et al., `Measuring line-of-sight-dependent Fourier-space clustering using FFTs`,
		  MNRAS, 2015
		* Scoccimarro, Roman, `Fast estimators for redshift-space clustering`, Phys. Review D, 2015
		"""
		# clear compensation from the actions
		for source in [self.first, self.second]:
			source.actions[:] = []; source.compensated = False
			assert len(source.actions) == 0

		# compute the compensations
		compensation = {}
		for name, mesh in zip(['first', 'second'], [self.first, self.second]):
			compensation[name] = get_compensation(mesh)
			if self.comm.rank == 0:
				if compensation[name] is not None:
					args = (compensation[name]['func'].__name__, name)
					self.logger.info("using compensation function %s for source '%s'" % args)
				else:
					self.logger.warning("no compensation applied for source '%s'" % name)

		rank = self.comm.rank
		pm   = self.first.pm

		# setup the 1D-binning
		muedges = numpy.linspace(0, 1, 2, endpoint=True)
		edges = [self.edges, muedges]

		# make a structured array to hold the results
		cols   = ['k'] + ['power_%d' %l for l in sorted(self.attrs['poles'])] + ['A_%d' %l for l in sorted(self.attrs['poles'])] + ['modes']
		dtype  = ['f8'] + ['c8']*len(self.attrs['poles'])*2 + ['i8']
		dtype  = numpy.dtype(list(zip(cols, dtype)))
		result = numpy.empty(len(self.edges)-1, dtype=dtype)

		# offset the box coordinate mesh ([-BoxSize/2, BoxSize]) back to
		# the original (x,y,z) coords
		offset = self.attrs['BoxCenter'] + 0.5*pm.BoxSize / pm.Nmesh

		# always need to compute ell=0
		poles = sorted(self.attrs['poles'])
		if 0 not in poles:
			poles = [0] + poles
		assert poles[0] == 0

		# spherical harmonic kernels (for ell > 0)
		Ylms = [[get_real_Ylm(l,m) for m in range(-l, l+1)] for l in poles[1:]]

		# paint the 1st FKP density field to the mesh (paints: data - alpha*randoms, essentially)
		rfield1 = self.first.paint(Nmesh=self.attrs['Nmesh'])

		vol_per_cell = (pm.BoxSize/pm.Nmesh).prod()/rfield1.attrs['num_per_cell'] #to compensate the default normalization by num_per_cell
		rfield1[:] /= vol_per_cell
		meta1 = rfield1.attrs.copy()
		if rank == 0: self.logger.info("%s painting of 'first' done" %self.first.window)

		# FFT 1st density field and apply the paintbrush window transfer kernel
		cfield = rfield1.r2c()
		if compensation['first'] is not None:
			cfield.apply(out=Ellipsis, **compensation['first'])
		if rank == 0: self.logger.info('ell = 0 done; 1 r2c completed')

		# monopole A0 is just the FFT of the FKP density field
		# NOTE: this holds FFT of density field #1
		volume = pm.BoxSize.prod()
		A0_1 = ComplexField(pm)
		A0_1[:] = cfield[:] * volume # normalize with a factor of volume

		# paint second mesh too?
		if self.first is not self.second:

			# paint the second field
			rfield2 = self.second.paint(Nmesh=self.attrs['Nmesh'],normalize=False)
			rfield2[:] /= vol_per_cell
			meta2 = rfield2.attrs.copy()
			if rank == 0: self.logger.info("%s painting of 'second' done" %self.second.window)

			# need monopole of second field
			if 0 in self.attrs['poles']:

				# FFT density field and apply the paintbrush window transfer kernel
				A0_2 = rfield2.r2c()
				A0_2[:] *= volume
				if compensation['second'] is not None: A0_2.apply(out=Ellipsis, **compensation['second'])
				
		else:
			rfield2 = rfield1
			meta2 = meta1

			# monopole of second field is first field
			if 0 in self.attrs['poles']:
				A0_2 = A0_1

		# save the painted density field #2 for later
		density2 = rfield2.copy()

		# initialize the memory holding the Aell terms for
		# higher multipoles (this holds sum of m for fixed ell)
		# NOTE: this will hold FFTs of density field #2
		Aell = ComplexField(pm)

		# the real-space grid
		xgrid = [xx.astype('f8') + offset[ii] for ii, xx in enumerate(density2.slabs.optx)]
		xnorm = numpy.sqrt(sum(xx**2 for xx in xgrid))
		xgrid = [x/xnorm for x in xgrid]

		# the Fourier-space grid
		kgrid = [kk.astype('f8') for kk in cfield.slabs.optx]
		knorm = numpy.sqrt(sum(kk**2 for kk in kgrid)); knorm[knorm==0.] = numpy.inf
		kgrid = [k/knorm for k in kgrid]

		# loop over the higher order multipoles (ell > 0)
		start = time.time()
		for iell, ell in enumerate(poles[1:]):

			# clear 2D workspace
			Aell[:] = 0.

			# iterate from m=-l to m=l and apply Ylm
			substart = time.time()
			for Ylm in Ylms[iell]:

				# reset the real-space mesh to the original density #2
				rfield2[:] = density2[:]

				# apply the config-space Ylm
				for islab, slab in enumerate(rfield2.slabs):
					slab[:] *= Ylm(xgrid[0][islab], xgrid[1][islab], xgrid[2][islab])

				# real to complex of field #2
				rfield2.r2c(out=cfield)

				# apply the Fourier-space Ylm
				for islab, slab in enumerate(cfield.slabs):
					slab[:] *= Ylm(kgrid[0][islab], kgrid[1][islab], kgrid[2][islab])

				# add to the total sum
				Aell[:] += cfield[:]

				# and this contribution to the total sum
				substop = time.time()
				if rank == 0:
					self.logger.debug("done term for Y(l=%d, m=%d) in %s" %(Ylm.l, Ylm.m, timer(substart, substop)))

			# apply the compensation transfer function
			if compensation['second'] is not None:
				Aell.apply(out=Ellipsis, **compensation['second'])

			# factor of 4*pi from spherical harmonic addition theorem + volume factor
			Aell[:] *= 4*numpy.pi*volume

			# log the total number of FFTs computed for each ell
			if rank == 0:
				args = (ell, len(Ylms[iell]))
				self.logger.info('ell = %d done; %s r2c completed' %args)
				
			# project on to 1d k-basis (averaging over mu=[0,1])
			proj_result, _ = project_to_basis(Aell, edges)
			result['A_%d' %ell][:] = numpy.squeeze(proj_result[2])

			# calculate the power spectrum multipoles, slab-by-slab to save memory
			# NOTE: this computes (A0 of field #1) * (Aell of field #2).conj()
			for islab in range(A0_1.shape[0]):
				Aell[islab,...] = A0_1[islab] * Aell[islab].conj()

			# project on to 1d k-basis (averaging over mu=[0,1])
			proj_result, _ = project_to_basis(Aell, edges)
			result['power_%d' %ell][:] = numpy.squeeze(proj_result[2])

		# summarize how long it took
		stop = time.time()
		if rank == 0:
			self.logger.info("higher order multipoles computed in elapsed time %s" %timer(start, stop))

		# also compute ell=0
		if 0 in self.attrs['poles']:
		
			# the 1D monopole
			proj_result, _ = project_to_basis(A0_2, edges)
			result['A_0'][:] = numpy.squeeze(proj_result[2])

			# the 3D monopole
			for islab in range(A0_1.shape[0]):
				A0_1[islab,...] = A0_1[islab]*A0_2[islab].conj()

			# the 1D monopole
			proj_result, _ = project_to_basis(A0_1, edges)
			result['power_0'][:] = numpy.squeeze(proj_result[2])

		# save the number of modes and k
		result['k'][:] = numpy.squeeze(proj_result[0])
		result['modes'][:] = numpy.squeeze(proj_result[-1])

		# copy over any painting meta data
		if self.first is self.second:
			copy_meta(self.attrs, meta1)
		else:
			copy_meta(self.attrs, meta1, prefix='first')
			copy_meta(self.attrs, meta2, prefix='second')

		return result