def paint_halos(pm, halocatalogue, BoxSize, m0, massmin, massmax):
if pm.comm.rank == 0:
hf = files.HaloFile(halocatalogue)
nhalo = hf.nhalo
halopos = numpy.float32(hf.read_pos())
halomass = numpy.float32(hf.read_mass() * m0)
logmass = numpy.log10(halomass)
mask = logmass > massmin
mask &= logmass < massmax
print logmass
halopos = halopos[mask]
logging.info("total number of halos in mass range is %d" % mask.sum())
else:
halopos = numpy.empty((0, 3), dtype='f4')
halomass = numpy.empty(0, dtype='f4')
P = {}
P['Position'] = halopos
Ntot = len(halopos)
Ntot = pm.comm.bcast(Ntot)
P['Position'] *= BoxSize
layout = pm.decompose(P['Position'])
tpos = layout.exchange(P['Position'])
pm.paint(tpos)
npaint = pm.comm.allreduce(len(tpos), op=MPI.SUM)
return Ntot
def main():
comm = MPI.COMM_WORLD
if comm.rank == 0:
h = files.HaloFile(ns.halocatalogue)
#print h.read_pos().shape()
N = h.read('Mass')
halo_pos = h.read('Position')
else:
N = None
halo_pos = None
N = comm.bcast(N)
halo_pos = comm.bcast(halo_pos)
halo_mass = N * ns.m0
if comm.rank == 0:
snapfile = files.Snapshot(ns.snapfilename, files.TPMSnapshotFile)
labelfile = files.Snapshot(ns.halolabel, files.HaloLabelFile)
npart = snapfile.npart
output = files.Snapshot.create(ns.output, files.TPMSnapshotFile, npart)
comm.bcast((snapfile, labelfile, output))
else:
snapfile, labelfile, output = comm.bcast(None)
comm.barrier()
Ntot = sum(snapfile.npart)
for i in range(0, Ntot, ns.bunchsize * comm.size):
if comm.rank == 0:
print i, Ntot
start, end, junk = slice(i, i + ns.bunchsize * comm.size).indices(Ntot)
mystart = start + (end - start) * comm.rank // comm.size
myend = start + (end - start) * (comm.rank + 1) // comm.size
position = snapfile.read('Position', mystart, myend)
velocity = snapfile.read('Velocity', mystart, myend)
ID = snapfile.read('ID', mystart, myend)
label = labelfile.read('Label', mystart, myend)
mask = label != 0
label = label[mask]
position2 = position[mask]
center = halo_pos[label]
mass = halo_mass[label]
dist = distp(center, position2, 1.0)
rwrong = numpy.einsum('ij,ij->i', dist, dist)**0.5
rwrong *= ns.boxsize
rfact = model(rwrong, mass)
dist *= rfact[:, None]
position2 = (center + dist) % 1.0
position[mask] = position2
print 'writing at', mystart
output.write('Position', mystart, position)
output.write('Velocity', mystart, velocity)
output.write('ID', mystart, ID)
def main():
sim = files.HaloFile(ns.sim_halo)
nbody = files.HaloFile(ns.Nbody_halo)
sim_mass = sim.read('Mass')
nbody_mass = nbody.read('Mass')
sim_pos = sim.read('Position')
sim_vel = sim.read('Velocity')
sim_link_l = sim.linking_length
sim_nhalo = sim.nhalo
nbody_nhalo = nbody.nhalo
minlength = min(nbody_nhalo, sim_nhalo)
print "starts to replace masses of halos."
with open(ns.output + '.halo.rpbmass', 'w') as ff:
numpy.int32(minlength).tofile(ff)
numpy.float32(sim_link_l).tofile(ff)
numpy.int32(nbody_mass[:minlength]).tofile(ff)
numpy.float32(sim_pos[:minlength]).tofile(ff)
numpy.float32(sim_vel[:minlength]).tofile(ff)
def parallel_read(self, columns, full=False):
if comm.rank == 0:
hf = files.HaloFile(self.pathhalo)
nhalo = hf.nhalo
halopos = numpy.float32(hf.read_pos())
halomass = numpy.float32(hf.read_mass() * self.m0)
logmass = numpy.log10(halomass)
halomask = logmass > self.logMmin
halomask &= logmass < self.logMmax
self.logger.info("total number of halos in mass range is %d" %
halomask.sum())
else:
halopos = None
halomask = None
halopos = comm.bcast(halopos)
halomask = comm.bcast(halomask)
for round, (P, PL) in enumerate(
izip(
files.read(comm,
self.pathmatter,
files.TPMSnapshotFile,
columns=['Position', 'Velocity'],
bunchsize=ns.bunchsize),
files.read(comm,
self.pathlabel,
files.HaloLabelFile,
columns=['Label'],
bunchsize=ns.bunchsize),
)):
mask = PL['Label'] != 0
mask &= halomask[PL['Label']]
self.logger.info("Number of particles in halos is %d" % mask.sum())
P['Position'][mask] = halopos[PL['Label'][mask]]
if self.rsd is not None:
dir = 'xyz'.index(self.rsd)
P['Position'][:, dir] += P['Velocity'][:, dir]
P['Position'][:,
dir] %= 1.0 # enforce periodic boundary conditions
P['Position'] *= self.BoxSize
P['Velocity'] *= self.BoxSize
if comm.rank == 0:
self.logger.info('round %d, nread %d' % (round, nread))
yield [P.get(key, None) for key in columns]
def paint_halos(Nmesh, halocatalogue, BoxSize, m0, massmin, massmax):
canvas = numpy.zeros((Nmesh, Nmesh, Nmesh), dtype='f4')
hf = files.HaloFile(halocatalogue)
nhalo = hf.nhalo
halopos = numpy.float32(hf.read_pos())
halomass = numpy.float32(hf.read_mass() * m0)
logmass = numpy.log10(halomass)
mask = logmass > massmin
mask &= logmass < massmax
halopos = halopos[mask]
logging.info("total number of halos in mass range is %d" % mask.sum())
halopos *= Nmesh
cic.paint(halopos, canvas, period=Nmesh)
return canvas
def main():
comm = MPI.COMM_WORLD
SNAP, LABEL = None, None
if comm.rank == 0:
SNAP = files.Snapshot(ns.snapfilename, files.TPMSnapshotFile)
LABEL = files.Snapshot(ns.halolabel, files.HaloLabelFile)
SNAP = comm.bcast(SNAP)
LABEL = comm.bcast(LABEL)
Ntot = sum(SNAP.npart)
assert Ntot == sum(LABEL.npart)
h = files.HaloFile(ns.halocatalogue)
N = h.read_mass()
N0 = Ntot - sum(N[1:])
# halos are assigned to ranks 0, 1, 2, 3 ...
halorank = numpy.arange(len(N)) % comm.size
# but non halos are special we will fix it later.
halorank[0] = -1
NonhaloStart = comm.rank * int(N0) // comm.size
NonhaloEnd = (comm.rank + 1) * int(N0) // comm.size
myNtotal = numpy.sum(N[halorank == comm.rank],
dtype='i8') + (NonhaloEnd - NonhaloStart)
print("Rank %d NonhaloStart %d NonhaloEnd %d myNtotal %d" %
(comm.rank, NonhaloStart, NonhaloEnd, myNtotal))
data = numpy.empty(myNtotal,
dtype=[
('Position', ('f4', 3)),
('Label', ('i4')),
('Rank', ('i4')),
])
allNtotal = comm.allgather(myNtotal)
start = sum(allNtotal[:comm.rank])
end = sum(allNtotal[:comm.rank + 1])
data['Position'] = SNAP.read("Position", start, end)
data['Label'] = LABEL.read("Label", start, end)
data['Rank'] = halorank[data['Label']]
# now assign ranks to nonhalo particles
nonhalomask = (data['Label'] == 0)
nonhalocount = comm.allgather(nonhalomask.sum())
data['Rank'][nonhalomask] = (sum(nonhalocount[:comm.rank]) +
numpy.arange(nonhalomask.sum())) % comm.size
mpsort.sort(data, orderby='Rank')
arg = data['Label'].argsort()
data = data[arg]
ul = numpy.unique(data['Label'])
bins = correlate.RBinning(40. / ns.boxsize, Nbins=ns.Nmesh)
sum1 = numpy.zeros(len(bins.centers))
for l in ul:
if l == 0: continue
start = data['Label'].searchsorted(l, side='left')
end = data['Label'].searchsorted(l, side='right')
pos = data['Position'][start:end]
dataset = correlate.points(pos, boxsize=1.0)
result = correlate.paircount(dataset, dataset, bins, np=0)
sum1 += result.sum1
if l % 1000 == 0:
print l
sum1 = comm.allreduce(sum1, MPI.SUM)
Ntot = sum(SNAP.npart)
RR = 4. / 3 * numpy.pi * numpy.diff(bins.edges**3) * (1.0 * Ntot * Ntot)
k = numpy.arange(ns.Nmesh // 2) * 2 * numpy.pi / ns.boxsize
# asymtotically zero at r. The mean doesn't matter as
# we don't use zero k mode anyways.
k, p = corrfrompower(bins.centers * ns.boxsize, sum1 / RR, R=k)
# inverse FT factor
p *= (2 * numpy.pi)**3
if comm.rank == 0:
if ns.output != '-':
ff = open(ns.output, 'w')
ff2 = open(ns.output + '.xi', 'w')
with ff2:
numpy.savetxt(ff2, zip(bins.centers, sum1 / RR - 1.0))
else:
ff = stdout
with ff:
# numpy.savetxt(ff, zip(bins.centers, sum1 / RR - 1.0))
numpy.savetxt(ff, zip(k, p))
def main():
comm = MPI.COMM_WORLD
if comm.rank == 0:
h = files.HaloFile(ns.halocatalogue)
#print h.read_pos().shape()
N = h.read('Mass')
halo_pos = h.read('Position')
else:
N = None
halo_pos = None
N = comm.bcast(N)
halo_pos = comm.bcast(halo_pos)
halo_mass = N * ns.m0
halo_ind = numpy.digitize(numpy.log10(halo_mass), ns.logMedges)
Nhalo_per_massbin = numpy.bincount(halo_ind,
minlength=len(ns.logMedges) + 1)
redges = numpy.linspace(0,
ns.Rmax / ns.boxsize,
ns.Rbins + 2,
endpoint=True)[1:]
den_prof = numpy.zeros((ns.Rbins, len(ns.logMedges) - 1))
count_prof = numpy.zeros((ns.Rbins + 2, len(ns.logMedges) + 1))
Ntot = 0
for round, (P, PL) in enumerate(
zip(
files.read(comm,
ns.snapfilename,
files.TPMSnapshotFile,
columns=['Position', 'Velocity'],
bunchsize=ns.bunchsize),
files.read(comm,
ns.halolabel,
files.HaloLabelFile,
columns=['Label'],
bunchsize=ns.bunchsize),
)):
m_ind = halo_ind[PL['Label']]
center = halo_pos[PL['Label']]
dist = distp(center, P['Position'], 1.0)
d_ind = numpy.digitize(dist, redges)
ind = numpy.ravel_multi_index((d_ind, m_ind), count_prof.shape)
count_prof.flat += numpy.bincount(ind, minlength=count_prof.size)
Ntot += len(PL['Label'])
Ntot = comm.allreduce(Ntot)
count_prof = comm.allreduce(count_prof)
count_prof = numpy.cumsum(count_prof, axis=0)
vol = 4 * numpy.pi / 3. * redges**3
# this is over density averaged over all halos in this mass bin
den_prof = count_prof[:-1, 1:-1] / vol[:, None] / \
Nhalo_per_massbin[None, 1:-1] / Ntot - 1
if comm.rank == 0:
if ns.output != '-':
ff = open(ns.output, 'w')
else:
ff = stdout
with ff:
ff.write('# nhalo_per_massbin: %s\n' %
' '.join([str(n) for n in Nhalo_per_massbin]))
ff.write('# edges of the mass bins %s\n' %
' '.join([str(m) for m in ns.logMedges]))
ff.write('# r cumulative_dens_prof\n')
numpy.savetxt(
ff,
numpy.concatenate((redges[:, None] * ns.boxsize, den_prof),
axis=-1))
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))