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()
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()
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)
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()
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)
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()
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()
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
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()
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))
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()
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
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)
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,
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))
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
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)