def to_real_field(self):
"""
Return the RealField stored on disk.
.. note::
The mesh stored on disk must be stored with ``mode=real``
Returns
-------
real : pmesh.pm.RealField
an array-like object holding the mesh loaded from disk in
configuration space
"""
if self.isfourier:
return NotImplemented
# the real field to paint to
pmread = self.pm
with FileMPI(comm=self.comm, filename=self.path)[self.dataset] as ds:
if self.comm.rank == 0:
self.logger.info("reading real field from %s" % self.path)
real2 = RealField(pmread)
start = numpy.sum(self.comm.allgather(real2.size)[:self.comm.rank], dtype='intp')
end = start + real2.size
real2.unravel(ds[start:end])
return real2
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)
def to_real_field(self):
"""
Return the RealField stored on disk.
.. note::
The mesh stored on disk must be stored with ``mode=real``
Returns
-------
real : pmesh.pm.RealField
an array-like object holding the mesh loaded from disk in
configuration space
"""
if self.isfourier:
return NotImplemented
# the real field to paint to
pmread = self.pm
with BigFileMPI(comm=self.comm,
filename=self.path)[self.dataset] as ds:
if self.comm.rank == 0:
self.logger.info("reading real field from %s" % self.path)
real2 = RealField(pmread)
start = sum(self.comm.allgather(real2.size)[:self.comm.rank])
end = start + real2.size
real2.unsort(ds[start:end])
return real2
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)
def test_inplace_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(
-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
complex2 = real.r2c(out=Ellipsis)
assert real._base in complex2._base
assert_almost_equal(numpy.asarray(complex),
numpy.asarray(complex2),
decimal=7)
real = complex2.c2r()
real2 = complex2.c2r(out=Ellipsis)
assert real2._base in complex2._base
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
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)
def test_inplace_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
complex2 = real.r2c(out=Ellipsis)
assert real._base in complex2._base
assert_almost_equal(numpy.asarray(complex), numpy.asarray(complex2), decimal=7)
real = complex2.c2r()
real2 = complex2.c2r(out=Ellipsis)
assert real2._base in complex2._base
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_real_apply(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
def filter(x, v):
return x[0] * 10 + x[1]
real.apply(filter, out=Ellipsis)
for i, x, slab in zip(real.slabs.i, real.slabs.x, real.slabs):
assert_array_equal(slab, x[0] * 10 + x[1])
def test_real_apply(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
def filter(x, v):
xnormp = x.normp()
assert_allclose(xnormp, sum(xi ** 2 for xi in x))
return x[0] * 10 + x[1]
real.apply(filter, out=Ellipsis)
for i, x, slab in zip(real.slabs.i, real.slabs.x, real.slabs):
assert_array_equal(slab, x[0] * 10 + x[1])
def test_asarray(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
a = numpy.asarray(real)
assert a is real.value
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f4')
real = RealField(pm)
a = numpy.asarray(real)
assert a is real.value
real = RealField(pm)
a = numpy.array(real, copy=False)
assert a is real.value
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
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[...])
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[...])
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)
def test_real_iter(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
for i, x, slab in zip(real.slabs.i, real.slabs.x, real.slabs):
assert_same_base(slab, real.value)
assert_almost_equal(slab.shape,
sum(x[d]**2 for d in range(len(pm.Nmesh))).shape)
for a, b in zip(slab.x, x):
assert_array_equal(a, b)
for a, b in zip(slab.i, i):
assert_array_equal(a, b)
def test_wrong_ndim(comm):
numpy.random.seed(42)
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
rfield = RealField(pm)
data = numpy.random.random(size=rfield.shape)
rfield[...] = data[:]
# SlabIterator only works for 2D or 3D coordinate meshes
with pytest.raises(NotImplementedError):
for slab in SlabIterator([rfield.x[0]], axis=0, symmetry_axis=None):
pass
def paint(self, pm, datasource):
"""
Paint the ``DataSource`` specified by ``input`` onto the
``ParticleMesh`` specified by ``pm``
Parameters
----------
pm : ``ParticleMesh``
particle mesh object that does the painting
datasource : ``DataSource``
the data source object representing the field to paint onto the mesh
Returns
-------
stats : dict
dictionary of statistics, usually only containing `Ntot`
"""
real = RealField(pm)
real[:] = 0
stats = {}
Nlocal = 0
# open the datasource stream (with no defaults)
with datasource.open() as stream:
# just paint density as usual
if self.moment == 0:
for [position] in stream.read(['Position']):
self.basepaint(real, position, paintbrush=self.paintbrush)
Nlocal += len(position)
# paint density-weighted velocity moments
else:
for position, velocity in stream.read(['Position',
'Velocity']):
self.basepaint(
real,
position,
weight=velocity[:, self._comp_index]**self.moment,
paintbrush=self.paintbrush)
Nlocal += len(position)
# total N
stats['Ntot'] = self.comm.allreduce(Nlocal)
# normalize config-space velocity field by mean number density
# this is (Nmesh**3 / V) / (Ntot / V)
norm = pm.Nmesh.prod() / (stats['Ntot'])
real[:] *= norm
return real, stats
def test_readout(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
real.value[...] = 1.0
pos = numpy.ones((1, 2))
out = numpy.empty((1), dtype='f8')
real.readout(pos, out=out)
pos = numpy.ones((1, 2), dtype='f4')
out = numpy.empty((1), dtype='f8')
real.readout(pos, out=out)
pos = numpy.ones((1, 2), dtype='f4')
out = numpy.empty((1), dtype='f4')
real.readout(pos, out=out)
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())
def test_wrong_coords_shape(comm):
numpy.random.seed(42)
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
rfield = RealField(pm)
data = numpy.random.random(size=rfield.shape)
rfield[...] = data[:]
# coords arrays should not be squeezed
x = [numpy.squeeze(xx) for xx in rfield.x]
with pytest.raises(ValueError):
for slab in SlabIterator(x, axis=0, symmetry_axis=None):
pass
def test_2d_slab(comm):
numpy.random.seed(42)
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
rfield = RealField(pm)
data = numpy.random.random(size=rfield.shape)
rfield[...] = data[:]
x = rfield.x
for i, slab in enumerate(SlabIterator(x, axis=0, symmetry_axis=None)):
assert slab.__str__() == slab.__repr__()
assert slab.shape == (8, )
assert_array_equal(slab.hermitian_weights, numpy.ones(slab.shape))
assert_array_equal(rfield[slab.index], data[i])
def test_readout(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
real.value[...] = 1.0
pos = numpy.ones ((1, 2))
out = numpy.empty((1), dtype='f8')
real.readout(pos, out=out)
pos = numpy.ones ((1, 2), dtype='f4')
out = numpy.empty((1), dtype='f8')
real.readout(pos, out=out)
pos = numpy.ones ((1, 2), dtype='f4')
out = numpy.empty((1), dtype='f4')
real.readout(pos, out=out)
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)
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())
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)
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)
def paint(self, pm, paintbrush='cic'):
"""
Paint the FKP weighted density field: ``data - alpha*randoms`` using
the input `ParticleMesh`
The are two different weights that enter into the painting procedure:
1. **completeness weights**: these weight each number density
field for data and randoms seperately
2. **FKP weights**: these weight the total FKP density field, i.e.,
they weight ``data - alpha*randoms``
Parameters
----------
pm : ParticleMesh
the particle mesh instance to paint the density field to
Returns
-------
stats : dict
a dictionary of FKP statistics, including total number, normalization,
and shot noise parameters (see equations 13-15 of Beutler et al. 2013)
"""
from pmesh.pm import RealField
if self.closed:
raise ValueError("'paint' operation on a closed FKPCatalog")
# setup
columns = ['Position', 'Nbar', 'FKPWeight', 'Weight']
stats = {}
A_ran = A_data = 0.
S_ran = S_data = 0.
N_ran = N_data = 0.
W_ran = W_data = 0.
# clear the density mesh
real = RealField(pm)
real[:] = 0
# alpha determined from completeness weights
alpha = self.alpha
# paint -1.0*alpha*N_randoms
for [position, nbar, fkp_weight,
comp_weight] in self.read('randoms', columns):
# total weight each galaxy gets is the product of FKP
# weight and completeness weight
weight = fkp_weight * comp_weight
self.painter.basepaint(real,
position,
weight=weight,
paintbrush=paintbrush)
A_ran += (nbar * comp_weight * fkp_weight**2).sum()
N_ran += len(position)
S_ran += (weight**2).sum()
W_ran += comp_weight.sum()
# randoms get -alpha factor
real[:] *= -alpha
A_ran = self.comm.allreduce(A_ran)
N_ran = self.comm.allreduce(N_ran)
S_ran = self.comm.allreduce(S_ran)
W_ran = self.comm.allreduce(W_ran)
# paint the data
for [position, nbar, fkp_weight,
comp_weight] in self.read('data', columns):
# total weight each galaxy gets is the product of FKP
# weight and completeness weight
weight = fkp_weight * comp_weight
self.painter.basepaint(real,
position,
weight=weight,
paintbrush=paintbrush)
A_data += (nbar * comp_weight * fkp_weight**2).sum()
N_data += len(position)
S_data += (weight**2).sum()
W_data += comp_weight.sum()
A_data = self.comm.allreduce(A_data)
N_data = self.comm.allreduce(N_data)
S_data = self.comm.allreduce(S_data)
W_data = self.comm.allreduce(W_data)
# store the stats (see equations 13-15 of Beutler et al 2013)
# see equations 13-15 of Beutler et al 2013
stats['W_data'] = W_data
stats['W_ran'] = W_ran
stats['N_data'] = N_data
stats['N_ran'] = N_ran
stats['A_data'] = A_data
stats['A_ran'] = A_ran
stats['S_data'] = S_data
stats['S_ran'] = S_ran
stats['alpha'] = alpha
stats['A_ran'] *= alpha
stats['S_ran'] *= alpha**2
stats['shot_noise'] = (stats['S_ran'] + stats['S_data']) / stats[
'A_ran'] # the final shot noise estimate for monopole
# go from number to number density
real[:] *= numpy.product(pm.Nmesh / pm.BoxSize)
return real, stats
def to_real_field(self, out=None, normalize=True):
"""
Paint the density field, by interpolating the position column
on to the mesh.
This computes the following meta-data attributes in the process of
painting, returned in the :attr:`attrs` attributes of the returned
RealField object:
- N : int
the (unweighted) total number of objects painted to the mesh
- W : float
the weighted number of total objects, equal to the collective
sum of the 'weight' column
- shotnoise : float
the Poisson shot noise, equal to the volume divided by ``N``
- num_per_cell : float
the mean number of weighted objects per cell
.. note::
The density field on the mesh is normalized as :math:`1+\delta`,
such that the collective mean of the field is unity.
See the :ref:`documentation <painting-mesh>` on painting for more
details on painting catalogs to a mesh.
Returns
-------
real : :class:`pmesh.pm.RealField`
the painted real field; this has a ``attrs`` dict storing meta-data
"""
# check for 'Position' column
if self.position not in self.source:
msg = "in order to paint a CatalogSource to a RealField, add a "
msg += "column named '%s', representing the particle positions" % self.position
raise ValueError(msg)
pm = self.pm
Nlocal = 0 # (unweighted) number of particles read on local rank
Wlocal = 0 # (weighted) number of particles read on local rank
# the paint brush window
paintbrush = window.methods[self.window]
# initialize the RealField to return
if out is not None:
assert isinstance(
out, RealField), "output of to_real_field must be a RealField"
numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh)
toret = out
else:
toret = RealField(pm)
toret[:] = 0
# for interlacing, we need two empty meshes if out was provided
# since out may have non-zero elements, messing up our interlacing sum
if self.interlaced:
real1 = RealField(pm)
real1[:] = 0
# the second, shifted mesh (always needed)
real2 = RealField(pm)
real2[:] = 0
# read the necessary data (as dask arrays)
columns = [self.position, self.weight, self.value, self.selection]
Position, Weight, Value, Selection = self.source.read(columns)
# ensure the slices are synced, since decomposition is collective
Nlocalmax = max(pm.comm.allgather(len(Position)))
# paint data in chunks on each rank;
# we do this by chunk 8 million is pretty big anyways.
chunksize = _global_options['paint_chunk_size']
for i in range(0, Nlocalmax, chunksize):
s = slice(i, i + chunksize)
if len(Position) != 0:
# selection has to be computed many times when data is `large`.
sel = self.source.compute(Selection[s])
# be sure to use the source to compute
position, weight, value = \
self.source.compute(Position[s], Weight[s], Value[s])
# FIXME: investigate if move selection before compute
# speeds up IO.
position = position[sel]
weight = weight[sel]
value = value[sel]
else:
# workaround a potential dask issue on empty dask arrays
position = numpy.empty((0, 3), dtype=Position.dtype)
weight = None
value = None
selection = None
if weight is None:
weight = numpy.ones(len(position))
if value is None:
value = numpy.ones(len(position))
# track total (selected) number and sum of weights
Nlocal += len(position)
Wlocal += weight.sum()
# no interlacing
if not self.interlaced:
lay = pm.decompose(position,
smoothing=0.5 * paintbrush.support)
p = lay.exchange(position)
w = lay.exchange(weight)
v = lay.exchange(value)
pm.paint(p,
mass=w * v,
resampler=paintbrush,
hold=True,
out=toret)
# interlacing: use 2 meshes separated by 1/2 cell size
else:
lay = pm.decompose(position,
smoothing=1.0 * paintbrush.support)
p = lay.exchange(position)
w = lay.exchange(weight)
v = lay.exchange(value)
H = pm.BoxSize / pm.Nmesh
# in mesh units
shifted = pm.affine.shift(0.5)
# paint to two shifted meshes
pm.paint(p,
mass=w * v,
resampler=paintbrush,
hold=True,
out=real1)
pm.paint(p,
mass=w * v,
resampler=paintbrush,
transform=shifted,
hold=True,
out=real2)
Nglobal = pm.comm.allreduce(Nlocal)
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh" %
(Nglobal, self.source.csize))
# now the loop over particles is done
if not self.interlaced:
# nothing to do, toret is already filled.
pass
else:
# compose the two interlaced fields into the final result.
c1 = real1.r2c()
c2 = real2.r2c()
# and then combine
for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs):
kH = sum(k[i] * H[i] for i in range(3))
s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(
0.5 * 1j * kH)
# FFT back to real-space
# NOTE: cannot use "toret" here in case user supplied "out"
c1.c2r(real1)
# need to add to the returned mesh if user supplied "out"
toret[:] += real1[:]
# unweighted number of objects
N = pm.comm.allreduce(Nlocal)
# weighted number of objects
W = pm.comm.allreduce(Wlocal)
# weighted number density (objs/cell)
nbar = 1. * W / numpy.prod(pm.Nmesh)
# make sure we painted something or nbar is nan; in which case
# we set the density to uniform everywhere.
if N == 0:
warnings.warn(("trying to paint particle source to mesh, "
"but no particles were found!"), RuntimeWarning)
# shot noise is volume / un-weighted number
shotnoise = numpy.prod(pm.BoxSize) / N
# save some meta-data
toret.attrs = {}
toret.attrs['shotnoise'] = shotnoise
toret.attrs['N'] = N
toret.attrs['W'] = W
toret.attrs['num_per_cell'] = nbar
csum = toret.csum()
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh" %
(N, self.source.csize))
self.logger.info("mean particles per cell is %g", nbar)
self.logger.info("sum is %g ", csum)
self.logger.info("normalized the convention to 1 + delta")
if normalize:
if nbar > 0:
toret[...] /= nbar
else:
toret[...] = 1
return toret
def test_shape_real(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
assert (tuple(real.cshape) == (8, 8))
def to_real_field(self, out=None, normalize=True):
"""
Paint the density field, by interpolating the position column
on to the mesh.
This computes the following meta-data attributes in the process of
painting, returned in the :attr:`attrs` attributes of the returned
RealField object:
- N : int
the (unweighted) total number of objects painted to the mesh
- W : float
the weighted number of total objects, equal to the collective
sum of the 'weight' column
- shotnoise : float
the Poisson shot noise, equal to the volume divided by ``N``
- num_per_cell : float
the mean number of weighted objects per cell
.. note::
The density field on the mesh is normalized as :math:`1+\delta`,
such that the collective mean of the field is unity.
See the :ref:`documentation <painting-mesh>` on painting for more
details on painting catalogs to a mesh.
Returns
-------
real : :class:`pmesh.pm.RealField`
the painted real field; this has a ``attrs`` dict storing meta-data
"""
pm = self.pm
Nlocal = 0 # (unweighted) number of particles read on local rank
Wlocal = 0 # (weighted) number of particles read on local rank
W2local = 0 # sum of weight square. This is used to estimate shotnoise.
# the paint brush window
resampler = window.methods[self.resampler]
# initialize the RealField to return
if out is not None:
assert isinstance(out, RealField), "output of to_real_field must be a RealField"
numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh)
toret = out
else:
toret = RealField(pm)
toret[:] = 0
# for interlacing, we need two empty meshes if out was provided
# since out may have non-zero elements, messing up our interlacing sum
if self.interlaced:
real1 = RealField(pm)
real1[:] = 0
# the second, shifted mesh (always needed)
real2 = RealField(pm)
real2[:] = 0
Position = self.Position
Weight = self.Weight
Value = self.Value
Selection = self.Selection
# ensure the slices are synced, since decomposition is collective
Nlocalmax = max(pm.comm.allgather(len(Position)))
H = pm.BoxSize / pm.Nmesh
# paint data in chunks on each rank;
# we do this by chunk 8 million is pretty big anyways.
max_chunksize = _global_options['paint_chunk_size']
# use a local scope to avoid having two copies of data in memory
def dochunk(s):
if len(Position) != 0:
# selection has to be computed many times when data is `large`.
columns = [Position[s]]
if Weight is not None:
columns.append(Weight[s])
if Value is not None:
columns.append(Value[s])
if Selection is not None:
columns.append(Selection[s])
# be sure to use the source to compute
data = self.source.compute(columns)
sel = Ellipsis if Selection is None else data.pop()
value = None if Value is None else data.pop()[sel]
weight = None if Weight is None else data.pop()[sel]
position = data.pop()[sel]
else:
# workaround a potential dask issue on empty dask arrays
position = numpy.empty((0, 3), dtype=Position.dtype)
weight = None
value = None
if weight is None:
weight = numpy.ones(len(position))
if value is None:
value = numpy.ones(len(position))
# track total (selected) number and sum of weights
Nlocal = len(position)
Wlocal = weight.sum()
W2local = (weight ** 2).sum()
# no interlacing
if not self.interlaced:
lay = pm.decompose(position, smoothing=0.5 * resampler.support)
else:
lay = pm.decompose(position, smoothing=1.0 * resampler.support)
# if we are receiving too many particles, abort and retry with a smaller chunksize
newlengths = pm.comm.allgather(lay.newlength)
if any([newlength > 2 * max_chunksize for newlength in newlengths]):
if pm.comm.rank == 0:
self.logger.info("Throttling chunksize as some ranks will receive too many particles. (%d > %d)" % (max(newlengths), max_chunksize * 2))
raise StopIteration
p = lay.exchange(position)
w = lay.exchange(weight)
v = lay.exchange(value)
if not self.interlaced:
pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=toret)
# interlacing: use 2 meshes separated by 1/2 cell size
else:
# in mesh units
shifted = pm.affine.shift(0.5)
# paint to two shifted meshes
pm.paint(p, mass=w * v, resampler=resampler, hold=True, out=real1)
pm.paint(p, mass=w * v, resampler=resampler, transform=shifted, hold=True, out=real2)
return Nlocal, Wlocal, W2local
import gc
i = 0
chunksize = max_chunksize
while i < Nlocalmax:
s = slice(i, i + chunksize)
if pm.comm.rank == 0:
self.logger.info("Chunk %d ~ %d / %d " % (i, i + chunksize, Nlocalmax))
try:
Nlocal1, Wlocal1, W2local1 = dochunk(s)
chunksize = min(max_chunksize, int(chunksize * 1.5))
except StopIteration:
chunksize = chunksize // 2
if chunksize < 1:
raise RuntimeError("Cannot find a chunksize that fits into memory.")
continue
finally:
# collect unfreed items
gc.collect()
Nlocal += Nlocal1
Wlocal += Wlocal1
W2local += W2local1
Nglobal = pm.comm.allreduce(Nlocal)
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh"
% (Nglobal, self.source.csize))
i = i + chunksize
# now the loop over particles is done
if not self.interlaced:
# nothing to do, toret is already filled.
pass
else:
# compose the two interlaced fields into the final result.
c1 = real1.r2c()
c2 = real2.r2c()
# and then combine
for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs):
kH = sum(k[i] * H[i] for i in range(3))
s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(0.5 * 1j * kH)
# FFT back to real-space
# NOTE: cannot use "toret" here in case user supplied "out"
c1.c2r(real1)
# need to add to the returned mesh if user supplied "out"
toret[:] += real1[:]
# unweighted number of objects
N = pm.comm.allreduce(Nlocal)
# weighted number of objects
W = pm.comm.allreduce(Wlocal)
# weighted number of objects
W2 = pm.comm.allreduce(W2local)
# weighted number density (objs/cell)
nbar = 1. * W / numpy.prod(pm.Nmesh)
# make sure we painted something or nbar is nan; in which case
# we set the density to uniform everywhere.
if N == 0:
warnings.warn(("trying to paint particle source to mesh, "
"but no particles were found!"),
RuntimeWarning
)
# shot noise is volume / un-weighted number
shotnoise = numpy.prod(pm.BoxSize) * W2 / W ** 2
# save some meta-data
toret.attrs = {}
toret.attrs['shotnoise'] = shotnoise
toret.attrs['N'] = N
toret.attrs['W'] = W
toret.attrs['W2'] = W
toret.attrs['num_per_cell'] = nbar
csum = toret.csum()
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh" %(N, self.source.csize))
self.logger.info("mean particles per cell is %g", nbar)
self.logger.info("sum is %g ", csum)
if normalize:
if nbar > 0:
toret[...] /= nbar
else:
toret[...] = 1
if pm.comm.rank == 0:
self.logger.info("normalized the convention to 1 + delta")
return toret
def to_real_field(self, out=None, normalize=True):
"""
Paint the density field, by interpolating the position column
on to the mesh.
This computes the following meta-data attributes in the process of
painting, returned in the :attr:`attrs` attributes of the returned
RealField object:
- N : int
the (unweighted) total number of objects painted to the mesh
- W : float
the weighted number of total objects, equal to the collective
sum of the 'weight' column
- shotnoise : float
the Poisson shot noise, equal to the volume divided by ``N``
- num_per_cell : float
the mean number of weighted objects per cell
.. note::
The density field on the mesh is normalized as :math:`1+\delta`,
such that the collective mean of the field is unity.
See the :ref:`documentation <painting-mesh>` on painting for more
details on painting catalogs to a mesh.
Returns
-------
real : :class:`pmesh.pm.RealField`
the painted real field; this has a ``attrs`` dict storing meta-data
"""
pm = self.pm
Nlocal = 0 # (unweighted) number of particles read on local rank
Wlocal = 0 # (weighted) number of particles read on local rank
W2local = 0 # sum of weight square. This is used to estimate shotnoise.
# the paint brush window
resampler = window.methods[self.resampler]
# initialize the RealField to return
if out is not None:
assert isinstance(
out, RealField), "output of to_real_field must be a RealField"
numpy.testing.assert_array_equal(out.pm.Nmesh, pm.Nmesh)
toret = out
else:
toret = RealField(pm)
toret[:] = 0
# for interlacing, we need two empty meshes if out was provided
# since out may have non-zero elements, messing up our interlacing sum
if self.interlaced:
real1 = RealField(pm)
real1[:] = 0
# the second, shifted mesh (always needed)
real2 = RealField(pm)
real2[:] = 0
Position = self.Position
Weight = self.Weight
Value = self.Value
Selection = self.Selection
# ensure the slices are synced, since decomposition is collective
Nlocalmax = max(pm.comm.allgather(len(Position)))
H = pm.BoxSize / pm.Nmesh
# paint data in chunks on each rank;
# we do this by chunk 8 million is pretty big anyways.
max_chunksize = _global_options['paint_chunk_size']
# use a local scope to avoid having two copies of data in memory
def dochunk(s):
if len(Position) != 0:
# selection has to be computed many times when data is `large`.
columns = [Position[s]]
if Weight is not None:
columns.append(Weight[s])
if Value is not None:
columns.append(Value[s])
if Selection is not None:
columns.append(Selection[s])
# be sure to use the source to compute
data = self.source.compute(columns)
sel = Ellipsis if Selection is None else data.pop()
value = None if Value is None else data.pop()[sel]
weight = None if Weight is None else data.pop()[sel]
position = data.pop()[sel]
else:
# workaround a potential dask issue on empty dask arrays
position = numpy.empty((0, 3), dtype=Position.dtype)
weight = None
value = None
if weight is None:
weight = numpy.ones(len(position))
if value is None:
value = numpy.ones(len(position))
# track total (selected) number and sum of weights
Nlocal = len(position)
Wlocal = weight.sum()
W2local = (weight**2).sum()
# no interlacing
if not self.interlaced:
lay = pm.decompose(position, smoothing=0.5 * resampler.support)
else:
lay = pm.decompose(position, smoothing=1.0 * resampler.support)
# if we are receiving too many particles, abort and retry with a smaller chunksize
recvlengths = pm.comm.allgather(lay.recvlength)
if any(
[recvlength > 2 * max_chunksize
for recvlength in recvlengths]):
if pm.comm.rank == 0:
self.logger.info(
"Throttling chunksize as some ranks will receive too many particles. (%d > %d)"
% (max(recvlengths), max_chunksize * 2))
raise StopIteration
p = lay.exchange(position)
w = lay.exchange(weight)
v = lay.exchange(value)
if not self.interlaced:
pm.paint(p,
mass=w * v,
resampler=resampler,
hold=True,
out=toret)
# interlacing: use 2 meshes separated by 1/2 cell size
else:
# in mesh units
shifted = pm.affine.shift(0.5)
# paint to two shifted meshes
pm.paint(p,
mass=w * v,
resampler=resampler,
hold=True,
out=real1)
pm.paint(p,
mass=w * v,
resampler=resampler,
transform=shifted,
hold=True,
out=real2)
return Nlocal, Wlocal, W2local
import gc
i = 0
chunksize = max_chunksize
while i < Nlocalmax:
s = slice(i, i + chunksize)
if pm.comm.rank == 0:
self.logger.info("Chunk %d ~ %d / %d " %
(i, i + chunksize, Nlocalmax))
try:
Nlocal1, Wlocal1, W2local1 = dochunk(s)
except StopIteration:
chunksize = chunksize // 2
if chunksize < 1:
raise RuntimeError(
"Cannot find a chunksize that fits into memory.")
continue
finally:
# collect unfreed items
gc.collect()
Nlocal += Nlocal1
Wlocal += Wlocal1
W2local += W2local1
Nglobal = pm.comm.allreduce(Nlocal)
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh" %
(Nglobal, self.source.csize))
i = i + chunksize
chunksize = min(max_chunksize, int(chunksize * 1.5))
# now the loop over particles is done
if not self.interlaced:
# nothing to do, toret is already filled.
pass
else:
# compose the two interlaced fields into the final result.
c1 = real1.r2c()
c2 = real2.r2c()
# and then combine
for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs):
kH = sum(k[i] * H[i] for i in range(3))
s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(
0.5 * 1j * kH)
# FFT back to real-space
# NOTE: cannot use "toret" here in case user supplied "out"
c1.c2r(real1)
# need to add to the returned mesh if user supplied "out"
toret[:] += real1[:]
# unweighted number of objects
N = pm.comm.allreduce(Nlocal)
# weighted number of objects
W = pm.comm.allreduce(Wlocal)
# weighted number of objects
W2 = pm.comm.allreduce(W2local)
# weighted number density (objs/cell)
nbar = 1. * W / numpy.prod(pm.Nmesh)
# make sure we painted something or nbar is nan; in which case
# we set the density to uniform everywhere.
if N == 0:
warnings.warn(("trying to paint particle source to mesh, "
"but no particles were found!"), RuntimeWarning)
# shot noise is volume / un-weighted number
shotnoise = numpy.prod(pm.BoxSize) * W2 / W**2
# save some meta-data
toret.attrs = {}
toret.attrs['shotnoise'] = shotnoise
toret.attrs['N'] = N
toret.attrs['W'] = W
toret.attrs['W2'] = W
toret.attrs['num_per_cell'] = nbar
csum = toret.csum()
if pm.comm.rank == 0:
self.logger.info("painted %d out of %d objects to mesh" %
(N, self.source.csize))
self.logger.info("mean particles per cell is %g", nbar)
self.logger.info("sum is %g ", csum)
if normalize:
if nbar > 0:
toret[...] /= nbar
else:
toret[...] = 1
if pm.comm.rank == 0:
self.logger.info("normalized the convention to 1 + delta")
return toret
def paint(self, pm, datasource):
"""
Paint the ``DataSource`` specified by ``input`` onto the
``ParticleMesh`` specified by ``pm``
Parameters
----------
pm : ``ParticleMesh``
particle mesh object that does the painting
datasource : ``DataSource``
the data source object representing the field to paint onto the mesh
Returns
-------
stats : dict
dictionary of statistics, usually only containing `Ntot`
"""
stats = {}
real = RealField(pm)
real[:] = 0
if isinstance(datasource, DataSource):
# open the datasource stream (with no defaults)
if self.interlaced:
real2 = RealField(pm)
real2[...] = 0
with datasource.open() as stream:
Nlocal = 0
if self.weight is None:
for [position] in stream.read(['Position']):
if not self.interlaced:
self.basepaint(real, position, paintbrush=self.paintbrush)
else:
self.shiftedpaint(real, real2, position, paintbrush=self.paintbrush)
Nlocal += len(position)
else:
for position, weight in stream.read(['Position', self.weight]):
if not self.interlaced:
self.basepaint(real, position, weight=weight, paintbrush=self.paintbrush)
else:
self.shiftedpaint(real, real2, position, weight=weight, paintbrush=self.paintbrush)
Nlocal += len(position)
if self.interlaced:
c1 = real.r2c()
c2 = real2.r2c()
H = pm.BoxSize / pm.Nmesh
for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs):
kH = sum(k[i] * H[i] for i in range(3))
s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(0.5 * 1j * kH)
c1.c2r(real)
stats['Ntot'] = self.comm.allreduce(Nlocal)
elif isinstance(datasource, GridSource):
datasource.read(real)
stats['Ntot'] = datasource.Ntot
# apply the filters.
mean = self.comm.allreduce(real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("Mean = %g" % mean)
if self.normalize:
real[...] *= 1. / mean
mean = self.comm.allreduce(real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("Renormalized mean = %g" % mean)
if self.setMean is not None:
real[...] += (self.setMean - mean)
if self.fk:
if self.comm.rank == 0:
self.logger.info("applying transformation fk %s" % self.fk)
def function(k, kx, ky, kz):
from numpy import exp, sin, cos
return eval(self.fk)
complex = real.r2c()
for kk, slab in zip(complex.slabs.x, complex.slabs):
k = sum([k ** 2 for k in kk]) ** 0.5
slab[...] *= function(k, kk[0], kk[1], kk[2])
complex.c2r(real)
mean = self.comm.allreduce(real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("after fk, mean = %g" % mean)
if self.frho:
if self.comm.rank == 0:
self.logger.info("applying transformation frho %s" % self.frho)
def function(rho):
return eval(self.frho)
if self.comm.rank == 0:
self.logger.info("example value before frho %g" % real.flat[0])
for slab in real.slabs:
slab[...] = function(slab)
if self.comm.rank == 0:
self.logger.info("example value after frho %g" % real.flat[0])
mean = self.comm.allreduce(real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("after frho, mean = %g" % mean)
return real, stats
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)
def paint(self, pm, datasource):
"""
Paint the ``DataSource`` specified by ``input`` onto the
``ParticleMesh`` specified by ``pm``
Parameters
----------
pm : ``ParticleMesh``
particle mesh object that does the painting
datasource : ``DataSource``
the data source object representing the field to paint onto the mesh
Returns
-------
stats : dict
dictionary of statistics, usually only containing `Ntot`
"""
stats = {}
real = RealField(pm)
real[:] = 0
if isinstance(datasource, DataSource):
# open the datasource stream (with no defaults)
if self.interlaced:
real2 = RealField(pm)
real2[...] = 0
with datasource.open() as stream:
Nlocal = 0
if self.weight is None:
for [position] in stream.read(['Position']):
if not self.interlaced:
self.basepaint(real,
position,
paintbrush=self.paintbrush)
else:
self.shiftedpaint(real,
real2,
position,
paintbrush=self.paintbrush)
Nlocal += len(position)
else:
for position, weight in stream.read(
['Position', self.weight]):
if not self.interlaced:
self.basepaint(real,
position,
weight=weight,
paintbrush=self.paintbrush)
else:
self.shiftedpaint(real,
real2,
position,
weight=weight,
paintbrush=self.paintbrush)
Nlocal += len(position)
if self.interlaced:
c1 = real.r2c()
c2 = real2.r2c()
H = pm.BoxSize / pm.Nmesh
for k, s1, s2 in zip(c1.slabs.x, c1.slabs, c2.slabs):
kH = sum(k[i] * H[i] for i in range(3))
s1[...] = s1[...] * 0.5 + s2[...] * 0.5 * numpy.exp(
0.5 * 1j * kH)
c1.c2r(real)
stats['Ntot'] = self.comm.allreduce(Nlocal)
elif isinstance(datasource, GridSource):
datasource.read(real)
stats['Ntot'] = datasource.Ntot
# apply the filters.
mean = self.comm.allreduce(real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("Mean = %g" % mean)
if self.normalize:
real[...] *= 1. / mean
mean = self.comm.allreduce(
real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("Renormalized mean = %g" % mean)
if self.setMean is not None:
real[...] += (self.setMean - mean)
if self.fk:
if self.comm.rank == 0:
self.logger.info("applying transformation fk %s" % self.fk)
def function(k, kx, ky, kz):
from numpy import exp, sin, cos
return eval(self.fk)
complex = real.r2c()
for kk, slab in zip(complex.slabs.x, complex.slabs):
k = sum([k**2 for k in kk])**0.5
slab[...] *= function(k, kk[0], kk[1], kk[2])
complex.c2r(real)
mean = self.comm.allreduce(
real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("after fk, mean = %g" % mean)
if self.frho:
if self.comm.rank == 0:
self.logger.info("applying transformation frho %s" % self.frho)
def function(rho):
return eval(self.frho)
if self.comm.rank == 0:
self.logger.info("example value before frho %g" % real.flat[0])
for slab in real.slabs:
slab[...] = function(slab)
if self.comm.rank == 0:
self.logger.info("example value after frho %g" % real.flat[0])
mean = self.comm.allreduce(
real.sum(dtype='f8')) / real.Nmesh.prod()
if self.comm.rank == 0:
self.logger.info("after frho, mean = %g" % mean)
return real, stats
def main():
ns = ap.parse_args()
comm = MPI.COMM_WORLD
ff = bigfile.BigFileMPI(comm, ns.fastpm)
with ff['.'] as bb:
BoxSize = bb.attrs['BoxSize'][0]
Redshift = 1 / bb.attrs['ScalingFactor'][0] - 1
Nmesh = int(BoxSize / ns.resolution * 2)
# round it to 8.
Nmesh -= Nmesh % 8
if comm.rank == 0:
logger.info("source = %s", ns.fastpm)
logger.info("output = %s", ns.output)
logger.info("BoxSize = %g", BoxSize)
logger.info("Redshift = %g", Redshift)
logger.info("Nmesh = %g", Nmesh)
pm = ParticleMesh([Nmesh, Nmesh, Nmesh], BoxSize, comm=comm)
real = RealField(pm)
real[...] = 0
with ff['Position'] as ds:
logger.info(ds.size)
for i in range(0, ds.size, ns.chunksize):
sl = slice(i, i + ns.chunksize)
pos = ds[sl]
layout = pm.decompose(pos)
lpos = layout.exchange(pos)
real.paint(lpos, hold=True)
mean = real.cmean()
if comm.rank == 0:
logger.info("mean particle per cell = %s", mean)
real[...] /= mean
real[...] -= 1
complex = real.r2c()
for k, i, slab in zip(complex.slabs.x, complex.slabs.i, complex.slabs):
k2 = sum(kd**2 for kd in k)
# tophat
f = tophat(ns.filtersize, k2**0.5)
slab[...] *= f
# zreion
slab[...] *= Bk(k2**0.5)
slab[...] *= (1 + Redshift)
real = complex.c2r()
real[...] += Redshift
mean = real.cmean()
if comm.rank == 0:
logger.info("zreion.mean = %s", mean)
buffer = numpy.empty(real.size, real.dtype)
real.sort(out=buffer)
if comm.rank == 0:
logger.info("sorted for output")
with bigfile.BigFileMPI(comm, ns.output, create=True) as ff:
with ff.create_from_array(ns.dataset, buffer) as bb:
bb.attrs['BoxSize'] = BoxSize
bb.attrs['Redshift'] = Redshift
bb.attrs['TopHatFilterSize'] = ns.filtersize
bb.attrs['Nmesh'] = Nmesh
#
# hack: compatible with current MPGadget. This is not really needed
# we'll remove the bins later, since BoxSize and Nmesh are known.
with ff.create("XYZ_bins", dtype='f8', size=Nmesh) as bb:
if comm.rank == 0:
bins = numpy.linspace(0, BoxSize * 1000., Nmesh, dtype='f8')
bb.write(0, bins)
if comm.rank == 0:
logger.info("done. written at %s", ns.output)
