コード例 #1
0
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    Ntot = paint_darkmatter(pm, ns.filename, TPMSnapshotFile)

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize ** 3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if MPI.COMM_WORLD.rank == 0:
        print 'measure power done'

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()
コード例 #2
0
ファイル: measurepower.py プロジェクト: mjvakili/nbodykit
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    Ntot = paint_darkmatter(pm, ns.filename, TPMSnapshotFile)

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize**3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if MPI.COMM_WORLD.rank == 0:
        print 'measure power done'

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()
コード例 #3
0
ファイル: power.py プロジェクト: energy-berkeley/nbodykit
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    # setup the particle mesh object
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    # paint first input
    Ntot1 = ns.inputs[0].paint(ns, pm)

    # painting
    if MPI.COMM_WORLD.rank == 0:
        print 'painting done'
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'

    # do the cross power
    do_cross = len(ns.inputs) > 1 and ns.inputs[0] != ns.inputs[1]
    if do_cross:
        complex = pm.complex.copy()

        Ntot2 = ns.inputs[1].paint(ns, pm)
        if MPI.COMM_WORLD.rank == 0:
            print 'painting 2 done'
        pm.r2c()
        if MPI.COMM_WORLD.rank == 0:
            print 'r2c 2 done'
            
        # power in cross case: c1.real*c2.real + c1.imag*c2.imag
        complex.real *= pm.complex.real
        complex.imag *= pm.complex.imag

        if MPI.COMM_WORLD.rank == 0:
            print 'cross done'
    # do the auto power
    else:
        complex = pm.complex
        complex.real **= 2
        complex.imag **= 2
        Ntot2 = Ntot1 

    if ns.remove_shotnoise and not do_cross:
        shotnoise = pm.BoxSize ** 3 / (1.0*Ntot1)
    else:
        shotnoise = 0
 
    # call the appropriate function for 1d/2d cases
    if ns.mode == "1d":
        do1d(pm, complex, ns, shotnoise)

    if ns.mode == "2d":
        meta = {'box_size':pm.BoxSize, 'N1':Ntot1, 'N2':Ntot2, 'shot_noise':shotnoise}
        do2d(pm, complex, ns, shotnoise, **meta)
コード例 #4
0
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    Ntot = paint_darkmatter(pm, ns.filename1, TPMSnapshotFile)

    if MPI.COMM_WORLD.rank == 0:
        print 'painting done'
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'
    if ns.filename1 != ns.filename2:
        # cross power 
        complex = pm.complex.copy()
        numpy.conjugate(complex, out=complex)

        Ntot = paint_darkmatter(pm, ns.filename2, TPMSnapshotFile)

        if MPI.COMM_WORLD.rank == 0:
            print 'painting 2 done'
        pm.r2c()
        if MPI.COMM_WORLD.rank == 0:
            print 'r2c 2 done'
        complex *= pm.complex
        complex **= 0.5

        if MPI.COMM_WORLD.rank == 0:
            print 'cross done'
    else:
        # auto power 
        complex = pm.complex

    k, mu, p, N, edges = measure2Dpower(pm, complex, ns.binshift, ns.remove_cic, 0, ns.Nmu)
  
    if MPI.COMM_WORLD.rank == 0:
        print 'measure'

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k.flat, mu.flat, p.flat, N.flat), '%0.7g')
        myout.flush()
コード例 #5
0
ファイル: qpm.py プロジェクト: sashmit/pmesh
    def run(self, P, aout=[]):
        logaout = numpy.log(aout)
        logaout.sort()

        pm = ParticleMesh(self.BoxSize, self.Nmesh, verbose=False, dtype='f4')
        self.pm = pm
        #SaveSnapshot(pm.comm, 'gridic-256', P)

        dloga = 0.1
        timesteps = list(numpy.arange(numpy.log(self.a0), 0.0, dloga))
        if timesteps[-1] < 0.0:
            timesteps.append(timesteps[-1] + dloga)

        for istep in range(len(timesteps)):

            # force at x(n+1)
            self.Accel(pm, P)

            # do the remaining KickB of last step
            if istep > 0:
                # KickB vel from n+1/2 to n+1
                self.Kick(P, 0.5 * (loga1 + loga2), loga2)

            loga1 = timesteps[istep]

            if istep == len(timesteps) - 1:
                # no more steps
                break
            if loga1 > logaout.max():
                # no more output
                break

            # now vel and pos are both at n+1, notify the caller! 
            yield self.PM_STEP_DONE, numpy.exp(loga1)

            loga2 = timesteps[istep + 1]

            # kickA
            # vel n -> n+1/2
            self.Kick(P, loga1, 0.5 * (loga1 + loga2))

            # drift
            # pos n -> n + 1

            # detect snapshot times
            left = logaout.searchsorted(loga1, side='left')
            right = logaout.searchsorted(loga2, side='right')

            if left != right:
                self.Drift(P, loga1, logaout[left])
                yield self.WRITE_SNAPSHOT, numpy.exp(logaout[left])
                for i in range(left + 1, right):
                    self.Drift(P, logaout[i-1], logaout[i])
                    yield self.WRITE_SNAPSHOT, numpy.exp(logaout[i])
                self.Drift(P, logaout[right - 1], loga2)
            else:
                self.Drift(P, loga1, loga2)

        yield self.FINISHED, numpy.exp(loga1)
コード例 #6
0
def main():

    if MPI.COMM_WORLD.rank == 0:
        print "importing done"

    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype="f4")

    Ntot = paint_darkmatter(pm, ns.filename1, TPMSnapshotFile)

    if MPI.COMM_WORLD.rank == 0:
        print "painting done"
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print "r2c done"

    complex = pm.complex.copy()
    numpy.conjugate(complex, out=complex)

    Ntot = paint_darkmatter(pm, ns.filename2, TPMSnapshotFile)
    if MPI.COMM_WORLD.rank == 0:
        print "painting 2 done"
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print "r2c 2 done"
    complex *= pm.complex
    complex **= 0.5

    if MPI.COMM_WORLD.rank == 0:
        print "cross done"
    k, p = measurepower(pm, complex, ns.binshift, ns.remove_cic, 0)

    if MPI.COMM_WORLD.rank == 0:
        print "measure"

    if pm.comm.rank == 0:
        if ns.output != "-":
            myout = open(ns.output, "w")
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), "%0.7g")
        myout.flush()
コード例 #7
0
def main():
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh)

    Ntot = paint_halos(pm, ns.halocatalogue, ns.BoxSize, ns.m0, ns.massmin, ns.massmax)

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize ** 3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()
コード例 #8
0
def overdensity(Pos, Mass, Nmesh, BoxSize, smoothing):
    """ Pos and smoothing is given in the same unit as BoxSize """
    Ndim = Pos.shape[1]

    assert Ndim == 3

    # first convert to Nmesh units:

    smoothing = smoothing * (1.0 * Nmesh / BoxSize)

    pm = ParticleMesh(BoxSize, Nmesh, verbose=False)

    layout = pm.decompose(Pos)
    tpos = layout.exchange(Pos)

    if numpy.isscalar(P.Mass):
        tmass = P.Mass
    else:
        tmass = layout.exchange(P.Mass)

    pm.r2c(tpos, tmass)

    D = numpy.empty(len(Pos), dtype='f8')

    tmp = pm.c2r(
        tpos, 
        TransferFunction.Inspect('K0', (0, 0, 0)),
#        TransferFunction.NormalizeDC,
        TransferFunction.RemoveDC,
        TransferFunction.Trilinear,
        TransferFunction.Gaussian(smoothing), 
        TransferFunction.Trilinear,
        )
    D[:] = layout.gather(tmp, mode='sum')
    return D
コード例 #9
0
def main():
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh)

    Ntot = paint_halos(pm, ns.halocatalogue, ns.BoxSize, ns.m0, ns.massmin,
                       ns.massmax)

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize**3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()
コード例 #10
0
def main():
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')
    if pm.comm.rank == 0:
        
        hf = files.HaloFile(ns.halocatalogue)
        nhalo = hf.nhalo
        halopos = hf.read_pos()
        halopos = pm.comm.bcast(halopos)
    else:
        halopos = pm.comm.bcast(None)

    Ntot = 0
    for round, (P, PL) in enumerate(izip(
                files.read(pm.comm, ns.filename, files.TPMSnapshotFile, 
                    columns=['Position'], bunchsize=ns.bunchsize),
                files.read(pm.comm, ns.halolabel, files.HaloLabelFile, 
                    columns=['Label'], bunchsize=ns.bunchsize),
                )):
        mask = PL['Label'] != 0
        logging.info("Number of particles in halos is %d" % mask.sum())
        P['Position'][mask] = halopos[PL['Label'][mask]]

        P['Position'] *= ns.BoxSize
        layout = pm.decompose(P['Position'])
        tpos = layout.exchange(P['Position'])
        #print tpos.shape
        pm.paint(tpos)
        npaint = pm.comm.allreduce(len(tpos), op=MPI.SUM) 
        nread = pm.comm.allreduce(len(P['Position']), op=MPI.SUM) 
        if pm.comm.rank == 0:
            logging.info('round %d, npaint %d, nread %d' % (round, npaint, nread))
        Ntot = Ntot + nread

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize ** 3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()

        hf = files.HaloFile(ns.halocatalogue)
        mass = hf.read_mass()
        M1 = ((mass[1:] * 1.0)** 2).sum(dtype='f8') / (1.0 * Ntot) ** 2 * ns.BoxSize ** 3
        logging.info("p[-inf] = %g, M1 = %g" % (p[-1], M1))
コード例 #11
0
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    Ntot = paint_darkmatter(pm, ns.filename1, TPMSnapshotFile)

    if MPI.COMM_WORLD.rank == 0:
        print 'painting done'
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'
    if ns.filename1 != ns.filename2:
        # cross power
        complex = pm.complex.copy()
        numpy.conjugate(complex, out=complex)

        Ntot = paint_darkmatter(pm, ns.filename2, TPMSnapshotFile)

        if MPI.COMM_WORLD.rank == 0:
            print 'painting 2 done'
        pm.r2c()
        if MPI.COMM_WORLD.rank == 0:
            print 'r2c 2 done'
        complex *= pm.complex
        complex **= 0.5

        if MPI.COMM_WORLD.rank == 0:
            print 'cross done'
    else:
        # auto power
        complex = pm.complex

    k, mu, p, N, edges = measure2Dpower(pm, complex, ns.binshift,
                                        ns.remove_cic, 0, ns.Nmu)

    if MPI.COMM_WORLD.rank == 0:
        print 'measure'

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k.flat, mu.flat, p.flat, N.flat), '%0.7g')
        myout.flush()
コード例 #12
0
def strain_tensor(Pos, Mass, Nmesh, BoxSize, smoothing):
    """ Pos and smoothing is given in the same unit as BoxSize """
    Ndim = Pos.shape[1]

    assert Ndim == 3

    # first convert to Nmesh units:

    smoothing = smoothing * (1.0 * Nmesh / BoxSize)

    pm = ParticleMesh(BoxSize, Nmesh, verbose=False)

    layout = pm.decompose(Pos)
    tpos = layout.exchange(Pos)

    if numpy.isscalar(P.Mass):
        tmass = P.Mass
    else:
        tmass = layout.exchange(P.Mass)

    pm.r2c(tpos, tmass)

    S = numpy.empty((len(Pos), 
                Ndim, Ndim), dtype='f8')

    for i, j in numpy.ndindex(Ndim, Ndim):
        if i > j: continue
        tmp = pm.c2r(
            tpos, 
            TransferFunction.RemoveDC,
            TransferFunction.Trilinear,
            TransferFunction.Gaussian(smoothing), 
            TransferFunction.Poisson, 
            TransferFunction.Constant(4 * numpy.pi * G),
            TransferFunction.Constant(Nmesh ** -2 * BoxSize ** 2),
            TransferFunction.Trilinear,
            TransferFunction.SuperLanzcos(i), 
            TransferFunction.SuperLanzcos(j), 
            TransferFunction.Constant(Nmesh ** 1 * BoxSize ** -1),
            TransferFunction.Constant(Nmesh ** 1 * BoxSize ** -1),
            )

        tmp = layout.gather(tmp, mode='sum')
        # symmetric!
        S[..., i, j] = tmp
        S[..., j, i] = tmp
    return S
コード例 #13
0
ファイル: power-fancy.py プロジェクト: DonRegan/nbodykit
def main():

    if MPI.COMM_WORLD.rank == 0:
        print 'importing done'

    chain = [
        TransferFunction.NormalizeDC,
        TransferFunction.RemoveDC,
        TransferFunction.Gaussian(1),
    ]

    if ns.remove_cic == 'anisotropic':
        chain.append(AnisotropicCIC)
    if ns.remove_cic == 'isotropic':
        chain.append(IsotropicCIC)

    # setup the particle mesh object
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')

    # paint first input
    Ntot1 = ns.inputs[0].paint(ns, pm)

    r1 = pm.real.copy()
    # painting
    if MPI.COMM_WORLD.rank == 0:
        print 'painting done'
    pm.r2c()
    if MPI.COMM_WORLD.rank == 0:
        print 'r2c done'

    # filter the field 
    pm.transfer(chain)

    # do the cross power
    do_cross = len(ns.inputs) > 1

    if do_cross:
        c1 = pm.complex.copy()

        Ntot2 = ns.inputs[1].paint(ns, pm)

        if MPI.COMM_WORLD.rank == 0:
            print 'painting 2 done'

        pm.r2c()
        if MPI.COMM_WORLD.rank == 0:
            print 'r2c 2 done'

        # filter the field 
        pm.transfer(chain)
        c2 = pm.complex.copy()

        if MPI.COMM_WORLD.rank == 0:
            print 'convolution done'

        dot(pm, c2, c2) 
        c2 = pm.complex

    # do the auto power
    else:
        c1 = pm.complex
        c2 = pm.complex

        Ntot2 = Ntot1 

    if ns.remove_shotnoise and not do_cross:
        shotnoise = pm.BoxSize ** 3 / (1.0*Ntot1)
    else:
        shotnoise = 0
 
    # only need one mu bin if 1d case is requested
    if ns.mode == "1d": ns.Nmu = 1 

    # do the calculation
    meta = {'box_size':pm.BoxSize, 'N1':Ntot1, 
            'N2':Ntot2, 'shot_noise': shotnoise}
    result = measurepower(pm, c1, c2, ns.Nmu, binshift=ns.binshift, 
                            shotnoise=shotnoise, los=ns.los)
    
    # format the output appropriately
    if ns.mode == "1d":
        # this writes out 0 -> mean k, 2 -> mean power, 3 -> number of modes
        meta['edges'] = result[-1][0] # write out kedges as metadata
        result = map(numpy.ravel, (result[i] for i in [0, 2, 3]))
    elif ns.mode == "2d":
        result = dict(zip(['k','mu','power','modes','edges'], result))
        
    if MPI.COMM_WORLD.rank == 0:
        print 'measure'
        storage = plugins.PowerSpectrumStorage.get(ns.mode, ns.output)
        storage.write(result, **meta)
コード例 #14
0
ファイル: powerspectrum.py プロジェクト: sashmit/pmesh
    Ntot = file.open('1/ID').size

    myslice = slice(
            MPI.COMM_WORLD.rank * Ntot // MPI.COMM_WORLD.size,
            (MPI.COMM_WORLD.rank + 1) * Ntot // MPI.COMM_WORLD.size,
            )

    P = lambda : None

    P.Pos = file.open('1/Position')[myslice] 
    
    NumPart = len(P.Pos)
    if MPI.COMM_WORLD.rank == 0:
        print '#', Nmesh, BoxSize

    pm = ParticleMesh(BoxSize, Nmesh, verbose=False)
    if pm.comm.allreduce(P.Pos.max(), MPI.MAX) < BoxSize * 0.1:
        print '# boxsize seems to be 1.0, corrected'
        P.Pos *= BoxSize

    layout = pm.decompose(P.Pos)
    tpos = layout.exchange(P.Pos)
    pm.r2c(tpos)

    psout = numpy.empty(Nmesh)
    wout = numpy.empty(Nmesh)

    pm.transfer( 
            TransferFunction.Trilinear,
            TransferFunction.NormalizeDC,
            TransferFunction.RemoveDC,
コード例 #15
0
def main():
    pm = ParticleMesh(ns.BoxSize, ns.Nmesh, dtype='f4')
    if pm.comm.rank == 0:

        hf = files.HaloFile(ns.halocatalogue)
        nhalo = hf.nhalo
        halopos = hf.read_pos()
        halopos = pm.comm.bcast(halopos)
    else:
        halopos = pm.comm.bcast(None)

    Ntot = 0
    for round, (P, PL) in enumerate(
            izip(
                files.read(pm.comm,
                           ns.filename,
                           files.TPMSnapshotFile,
                           columns=['Position'],
                           bunchsize=ns.bunchsize),
                files.read(pm.comm,
                           ns.halolabel,
                           files.HaloLabelFile,
                           columns=['Label'],
                           bunchsize=ns.bunchsize),
            )):
        mask = PL['Label'] != 0
        logging.info("Number of particles in halos is %d" % mask.sum())
        P['Position'][mask] = halopos[PL['Label'][mask]]

        P['Position'] *= ns.BoxSize
        layout = pm.decompose(P['Position'])
        tpos = layout.exchange(P['Position'])
        #print tpos.shape
        pm.paint(tpos)
        npaint = pm.comm.allreduce(len(tpos), op=MPI.SUM)
        nread = pm.comm.allreduce(len(P['Position']), op=MPI.SUM)
        if pm.comm.rank == 0:
            logging.info('round %d, npaint %d, nread %d' %
                         (round, npaint, nread))
        Ntot = Ntot + nread

    if ns.remove_shotnoise:
        shotnoise = pm.BoxSize**3 / Ntot
    else:
        shotnoise = 0

    pm.r2c()

    k, p = measurepower(pm, pm.complex, ns.binshift, ns.remove_cic, shotnoise)

    if pm.comm.rank == 0:
        if ns.output != '-':
            myout = open(ns.output, 'w')
        else:
            myout = stdout
        numpy.savetxt(myout, zip(k, p), '%0.7g')
        myout.flush()

        hf = files.HaloFile(ns.halocatalogue)
        mass = hf.read_mass()
        M1 = ((mass[1:] * 1.0)**
              2).sum(dtype='f8') / (1.0 * Ntot)**2 * ns.BoxSize**3
        logging.info("p[-inf] = %g, M1 = %g" % (p[-1], M1))
コード例 #16
0
ファイル: genic.py プロジェクト: sashmit/pmesh
def GridIC(PowerSpectrum, BoxSize, Ngrid, order=3, preshift=False,
        shift=0.5, ZAonly=False, dtype='f8'):
    """ 2LPT IC from PowerSpectrum for particle grid of Ngrid
    
        CPARAM is a Cosmology object. We also need CPARAM.PowerSpectrum object.

        order is the force differentialtion kernel order. 0 or 3.
        This rather long code does 

        (http://arxiv.org/pdf/astro-ph/9711187v1.pdf)

        A few strange things to notice. 
        The real space gaussian field has an amplitude of 1.0. 
        In gaussian field it is 0.707 in amplitude. (grid **3 to adjust that)

        (Also see http://www.design.caltech.edu/erik/Misc/Gaussian.html)
        (And what FFTW really computes)

        After applying the phase each component is further reduced to 0.5.
        (thus FFT back from delta_k with unity power doesn't give us 
        The PowerSpectrum we use is Pk/(2pi)**3. This is the convention used in
        Gadget.

        The sign of terms.  We agree with the paper but shall pull out the - sign in D2
        in Formula D2; 
        The final result agrees with Martin's code(ic_2lpt_big). 
        The final result differ with 2LPTic by -1.

        Factor 3/7 is multiplied to 2LPT field, abs(D2) and abs(D1) shall be
        applied the ZA and 2LPT before shifting the particles and adding the
        velocity.

        Position of initial points. If set to the center of cells the small
        scale power is smoothed. 
        COLA does a global shift after the readout. This matters if one wants to
        evolve the position by 2LPT. We follow COLA, but give an option to do
        the preshift shift.
    """
    # convert to the internal vel units of Gadget a**2 xdot

    D1 = 1.0
    D2 = D1 ** 2 

    pm = ParticleMesh(BoxSize, Ngrid, verbose=False, dtype='f4')

    x0 = pm.partition.local_i_start
    ni = pm.partition.local_ni
    Nlocal = numpy.prod(ni)

    pos = numpy.empty((Nlocal, 3), dtype=dtype)
    ID = numpy.empty(Nlocal, dtype=('i8'))

    view = pos.reshape(list(ni) + [3])
    view[:, :, :, 0] = numpy.arange(ni[0])[:, None, None] + x0[0]
    view[:, :, :, 1] = numpy.arange(ni[1])[None, :, None] + x0[1]
    view[:, :, :, 2] = numpy.arange(ni[2])[None, None, :] + x0[2]

    view *= 1.0 * BoxSize / Ngrid
    if preshift:
        pos += shift * BoxSize / Ngrid

    # now set up the ranks
    Nlist = numpy.array(pm.comm.allgather(Nlocal), dtype='i8')
    offset = numpy.cumsum(Nlist)
    ID = numpy.arange(Nlocal)
    if pm.comm.rank > 0:
        ID += offset[pm.comm.rank - 1]
    P = dict()
    P['Position'] = pos
    P['ID'] = ID

    layout = pm.decompose(P['Position'])
    tpos = layout.exchange(P['Position'])

    GlobalRNG = numpy.random.RandomState(299995)
    seed = GlobalRNG.randint(999999999, size=pm.comm.size*11)[::11][pm.comm.rank]
    RNG = numpy.random.RandomState(seed)

    pm.real[:] = RNG.normal(scale=1.0, size=pm.real.shape)
#    realstd = pm.comm.allreduce((pm.real ** 2).sum(), MPI.SUM)
#    if pm.comm.rank == 0:
#        print 'realstd', (realstd / pm.Nmesh ** 3) ** 0.5

    pm.real *= Ngrid ** -1.5

    pm.r2c()
#    realstd = pm.comm.allreduce((pm.complex.real ** 2).sum(), MPI.SUM)
#    if pm.comm.rank == 0:
#        print 'complex std', (realstd / (1. + pm.Nmesh//2 +1) / pm.Nmesh ** 2) ** 0.5

    def Transfer(comm, complex, w):
        w2 = 0
        for wi in w:
            w2 = w2 + wi ** 2
        w2 **= 0.5
        w2 *= 1.0 * Ngrid / BoxSize
        wt = PowerSpectrum.PofK(w2)
        wt *= (2 * numpy.pi) ** 3 * (BoxSize) ** -3 * D1 ** 2
        wt **= 0.5 
        wt[w2 == 0] = 0
        # cut at nyquist
        wt[w2 >= numpy.pi / (BoxSize) * Ngrid] =0 
        complex[:] *= wt

    pm.transfer( [
            TransferFunction.RemoveDC,
            Transfer,
            TransferFunction.Poisson,
            TransferFunction.Constant((1.0 * Ngrid / BoxSize) ** -2),
            ])

    # now we have the 'potential' field in K-space

    # ZA displacements
    P['ZA'] = numpy.empty_like(pos)

    for dir in range(3):
        pm.c2r( [
                TransferFunction.SuperLanzcos(dir, order=order),
                TransferFunction.Constant(-1.0 * Ngrid / BoxSize),
                ])
        tmp = pm.readout(tpos)
        tmp = layout.gather(tmp, mode='sum')
        P['ZA'][:, dir] = tmp

    # additional source term for 2 lpt correction

    # diag terms
    diag = []
    for i, dir in enumerate([(0, 0), (1, 1), (2, 2)]):
        pm.c2r([
                TransferFunction.SuperLanzcos(dir[0], order=order),
                TransferFunction.SuperLanzcos(dir[1], order=order),
                TransferFunction.Constant((1.0 * Ngrid / BoxSize) ** 2),
               ])
        diag.append(pm.real.copy())

    field = diag[0] * diag[1]
    field += diag[1] * diag[2] 
    field += diag[2] * diag[0]
    diag = []

    # off terms
    for i, dir in enumerate([(0, 1), (0, 2), (1, 2)]):
        pm.c2r([
                TransferFunction.SuperLanzcos(dir[0], order=order),
                TransferFunction.SuperLanzcos(dir[1], order=order),
                TransferFunction.Constant((1.0 * Ngrid / BoxSize) ** 2),
               ])
        field -= pm.real ** 2
        
    field *= Ngrid ** -3.0
    pm.real[:] = field
    field = []

    pm.r2c()

    P['2LPT'] = numpy.empty_like(pos)

    tmp = pm.readout(tpos)
    P['digrad'] = layout.gather(tmp, mode='sum')

    for dir in range(3):
        pm.c2r([
                TransferFunction.Poisson,
                TransferFunction.SuperLanzcos(dir, order=0),
                TransferFunction.Constant((1.0 * Ngrid / BoxSize) ** -2),
                TransferFunction.Constant(-1.0 * Ngrid / BoxSize),
                ])
        tmp = pm.readout(tpos)
        tmp = layout.gather(tmp, mode='sum')
        P['2LPT'][:, dir] = tmp

    P['2LPT'] *= 3.0 / 7
    # std of displacements
    ZA2 = pm.comm.allreduce(numpy.einsum('ij,ij->', P['ZA'], P['ZA'],
        dtype='f8'), MPI.SUM)
    LPT2 = pm.comm.allreduce(numpy.einsum('ij,ij->', P['2LPT'], P['2LPT'],
        dtype='f8'), MPI.SUM)
    ZAM = pm.comm.allreduce(numpy.max(P['ZA']), MPI.MAX)
    LPTM = pm.comm.allreduce(numpy.max(P['2LPT']), MPI.MAX)

    # norm of the 3-vector!
    ZA2 /= Ngrid ** 3
    LPT2 /= Ngrid ** 3

    stats = dict(
            BoxSize=BoxSize,
            Ngrid=Ngrid,
            stdZA=ZA2 ** 0.5 / BoxSize * Ngrid,
            std2LPT= LPT2 ** 0.5 / BoxSize * Ngrid,
            maxZA= ZAM ** 0.5 / BoxSize * Ngrid,
            max2LPT= LPTM ** 0.5 / BoxSize * Ngrid,
            T=str(pm.T))
    if not preshift:
        P['Position'] += shift * BoxSize / Ngrid
    return P, stats
コード例 #17
0
ファイル: power.py プロジェクト: DonRegan/nbodykit
def compute_power(ns, comm=None):
    """
    Compute the power spectrum. Given a `Namespace`, this is the function,
    that computes and saves the power spectrum. It does all the work.
    
    Parameters
    ----------
    ns : argparse.Namespace
        the parser namespace corresponding to the ``initialize_power_parser``
        functions
    comm : MPI.Communicator
        the communicator to pass to the ``ParticleMesh`` object
    """
    rank = comm.rank if comm is not None else MPI.COMM_WORLD.rank
    
    # set logging level
    logger.setLevel(ns.log_level)
    
    if rank == 0:
        logger.info('importing done')

    chain = [TransferFunction.NormalizeDC, TransferFunction.RemoveDC]
    if ns.remove_cic == 'anisotropic':
        chain.append(AnisotropicCIC)
    if ns.remove_cic == 'isotropic':
        chain.append(IsotropicCIC)
        
    # setup the particle mesh object, taking BoxSize from the painters
    pm = ParticleMesh(ns.inputs[0].BoxSize, ns.Nmesh, dtype='f4', comm=comm)

    # paint first input
    Ntot1 = paint(ns.inputs[0], pm, ns)

    # painting
    if rank == 0:
        logger.info('painting done')
    pm.r2c()
    if rank == 0:
        logger.info('r2c done')

    # filter the field 
    pm.transfer(chain)

    # do the cross power
    do_cross = len(ns.inputs) > 1 and ns.inputs[0] != ns.inputs[1]

    if do_cross:
        
        # crash if box size isn't the same
        if not numpy.all(ns.inputs[0].BoxSize == ns.inputs[1].BoxSize):
            raise ValueError("mismatch in box sizes for cross power measurement")
        
        c1 = pm.complex.copy()
        Ntot2 = paint(ns.inputs[1], pm, ns)

        if rank == 0:
            logger.info('painting 2 done')

        pm.r2c()
        if rank == 0:
            logger.info('r2c 2 done')

        # filter the field 
        pm.transfer(chain)
        c2 = pm.complex
  
    # do the auto power
    else:
        c1 = pm.complex
        c2 = pm.complex
        Ntot2 = Ntot1 

    if ns.remove_shotnoise and not do_cross:
        shotnoise =  pm.BoxSize.prod() / (1.0*Ntot1)
    else:
        shotnoise = 0
 
    # only need one mu bin if 1d case is requested
    if ns.mode == "1d": ns.Nmu = 1 

    # do the calculation
    Lx, Ly, Lz = pm.BoxSize
    meta = {'Lx':Lx, 'Ly':Ly, 'Lz':Lz, 'volume':Lx*Ly*Lz, 
            'N1':Ntot1, 'N2':Ntot2, 'shot_noise': shotnoise}
    result = measurepower(pm, c1, c2, ns.Nmu, binshift=ns.binshift, 
                            shotnoise=shotnoise, los=ns.los, dk=ns.dk, 
                            kmin=ns.kmin, poles=ns.poles)
    
    # format the output appropriately
    if len(ns.poles):
        pole_result, pkmu_result, edges = result
        result = dict(zip(['k','mu','power','modes','edges'], pkmu_result+(edges,)))
    elif ns.mode == "1d":
        # this writes out 0 -> mean k, 2 -> mean power, 3 -> number of modes
        meta['edges'] = result[-1][0] # write out kedges as metadata
        result = map(numpy.ravel, (result[i] for i in [0, 2, 3]))
    elif ns.mode == "2d":
        result = dict(zip(['k','mu','power','modes','edges'], result))
        
    if rank == 0:
        # save the power
        logger.info('measurement done; saving power to %s' %ns.output)
        storage = plugins.PowerSpectrumStorage.new(ns.mode, ns.output)
        storage.write(result, **meta)
        
        # save the multipoles
        if len(ns.poles):
            if ns.pole_output is None:
                raise RuntimeError("you specified multipoles to compute, but did not provide an output file name")
            meta['edges'] = edges[0]
            
            # format is k pole_0, pole_1, ...., modes_1d
            logger.info('saving ell = %s multipoles to %s' %(",".join(map(str,ns.poles)), ns.pole_output))
            result = [x for x in numpy.vstack(pole_result)]
            storage = plugins.PowerSpectrumStorage.new('1d', ns.pole_output)
            storage.write(result, **meta)