def weighted_map(ipix, npix, weights, localsize, comm):
""" Make a map from particles, for quantities like
W(t) = \int dx delta(t, x) w
Parameters
----------
ipix: array_like
weights : array_like
Returns
-------
Wmap, Nmap; distributed maps
Wmap is the weighted map. Nmap is the number of objects
"""
ipix, labels = numpy.unique(ipix, return_inverse=True)
N = numpy.bincount(labels)
weights = numpy.bincount(labels, weights)
#print("shrink to %d from %d" % (len(ipix), len(labels)))
del labels
pairs = numpy.empty(len(ipix) + 1, dtype=[('ipix', 'i4'), ('N', 'i4'), ('weights', 'f8') ])
pairs['ipix'][:-1] = ipix
pairs['weights'][:-1] = weights
pairs['N'][:-1] = N
pairs['ipix'][-1] = npix - 1 # trick to make sure the final length is correct.
pairs['weights'][-1] = 0
pairs['N'][-1] = 0
disa = DistributedArray(pairs, comm=comm)
disa.sort('ipix')
w = disa['ipix'].bincount(weights=disa['weights'].local, local=False, shared_edges=False)
N = disa['ipix'].bincount(weights=disa['N'].local, local=False, shared_edges=False)
if npix - w.cshape[0] != 0:
if comm.rank == 0:
print('padding -- this shouldnt have occured ', npix, w.cshape)
# pad with zeros, since the last few bins can be empty.
ipadding = DistributedArray.cempty((npix - w.cshape[0],), dtype='i4', comm=comm)
fpadding = DistributedArray.cempty((npix - w.cshape[0],), dtype='f8', comm=comm)
fpadding.local[:] = 0
ipadding.local[:] = 0
w = DistributedArray.concat(w, fpadding)
N = DistributedArray.concat(N, ipadding)
w = DistributedArray.concat(w, localsize=localsize)
N = DistributedArray.concat(N, localsize=localsize)
return w.local, N.local
def test_distributed_array_concat(comm):
from nbodykit.utils import DistributedArray, EmptyRank
data = numpy.array(
comm.scatter([
numpy.array([
0,
1,
], 'i4'),
numpy.array([
2,
3,
], 'i4'),
numpy.array([], 'i4'),
numpy.array([
4,
], 'i4'),
]))
da = DistributedArray(data, comm)
assert da.cshape[0] == 5
assert_array_equal(comm.allgather(da.coffset), [0, 2, 4, 4])
cc = DistributedArray.concat(da, da)
assert_array_equal(numpy.concatenate(comm.allgather(cc.local)),
[0, 1, 2, 3, 4, 0, 1, 2, 3, 4])
def getinsat(self, mHIsat, satid, totalsize, localsize, comm):
da = DistributedArray(satid, comm)
mHI = da.bincount(mHIsat, shared_edges=False)
zerosize = totalsize - mHI.cshape[0]
zeros = DistributedArray.cempty(cshape=(zerosize, ), dtype=mHI.local.dtype, comm=comm)
zeros.local[...] = 0
mHItotal = DistributedArray.concat(mHI, zeros, localsize=localsize)
return mHItotal
def getinsat(self, mHIsat, satid, totalsize, localsize, comm):
#print(comm.rank, np.all(np.diff(satid) >=0))
#diff = np.diff(satid)
#if comm.rank == 260:
# print(satid[:-1][diff <0], satid[1:][diff < 0])
da = DistributedArray(satid, comm)
mHI = da.bincount(mHIsat, shared_edges=False)
zerosize = totalsize - mHI.cshape[0]
zeros = DistributedArray.cempty(cshape=(zerosize, ), dtype=mHI.local.dtype, comm=comm)
zeros.local[...] = 0
mHItotal = DistributedArray.concat(mHI, zeros, localsize=localsize)
return mHItotal
def test_distributed_array_concat(comm):
from nbodykit.utils import DistributedArray, EmptyRank
data = numpy.array(comm.scatter(
[numpy.array([0, 1, ], 'i4'),
numpy.array([2, 3, ], 'i4'),
numpy.array([], 'i4'),
numpy.array([4, ], 'i4'),
]))
da = DistributedArray(data, comm)
assert da.cshape[0] == 5
assert_array_equal(
comm.allgather(da.coffset),
[ 0, 2, 4, 4]
)
cc = DistributedArray.concat(da, da)
assert_array_equal(
numpy.concatenate(comm.allgather(cc.local)),
[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
)
N.local[:10])
print('rank, cen.csize, N.cshape : ', rank, cen.csize, N.cshape)
print('rank, cen.size, Nlocal.size : ', rank, cen.size, N.local.size)
print(cen.csize - N.cshape)
zerosize = (cen.csize - N.cshape[0])
#start = (zerosize *rank // wsize)
#end = (zerosize *(rank+1) // wsize)
#zeros = DistributedArray(np.zeros(end-start), comm=comm)
print(zerosize, N.local.dtype)
zeros = DistributedArray.cempty(cshape=(zerosize, ),
dtype=N.local.dtype,
comm=comm)
zeros.local[...] = 0
print(rank, cen.csize, N.cshape, zeros.cshape, N.cshape + zeros.cshape)
N2 = DistributedArray.concat(N, zeros, localsize=cnsat.size)
print('rank, shid, N : ', rank, shid[:10], N2.local[:10])
print('rank, chid, csat, N : ', rank, chid[:10], cnsat[:10],
N2.local[:10])
print('rank, cen.csize, N.cshape : ', rank, cen.csize, N2.cshape)
print('rank, cen.size, Nlocal.size : ', rank, cen.size, N2.local.size)
print(rank, cnsat - N2.local)
print(np.allclose(cnsat - N2.local, 0))
#Do for mass
M = da.bincount(smass, shared_edges=False)
zerosize = (cen.csize - N.cshape[0])
start = (zerosize * rank // wsize)