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_cnorm_log(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0)
norm2 = comp1.cnorm(norm = lambda x: numpy.log(x.real ** 2 + x.imag ** 2))
norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2)).sum()
assert_allclose(norm2, norm3)
def test_whitenoise_mean(comm):
# the whitenoise shall preserve the large scale.
pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
complex1 = pm0.generate_whitenoise(seed=8, unitary=True, mean=1.0)
assert_allclose(complex1.c2r().cmean(), 1.0)
def test_paint(comm):
from pmesh.pm import ParticleMesh
# initialize a random white noise field in real space
pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
real = pm.generate_whitenoise(type='real', seed=3333) # a RealField
# FFT to a ComplexField
complex = real.r2c()
# gather to all ranks --> shape is now (8,8,8) on all ranks
data = numpy.concatenate(comm.allgather(numpy.array(real.ravel())))\
.reshape(real.cshape)
# mesh from data array
realmesh = ArrayMesh(data, BoxSize=128., comm=comm)
# paint() must yield the RealField
assert_array_equal(real, realmesh.to_field(mode='real'))
# gather complex to all ranks --> shape is (8,8,5) on all ranks
cdata = numpy.concatenate(comm.allgather(numpy.array(complex.ravel())))\
.reshape(complex.cshape)
# mesh from complex array
cmesh = ArrayMesh(cdata, BoxSize=128., comm=comm)
# paint() yields the ComplexField
assert_allclose(complex, cmesh.to_field(mode='complex'), atol=1e-8)
def test_cnorm_grad(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex')
def objective(comp1):
return comp1.cnorm()
grad_comp1 = comp1 * 2
grad_comp1.decompress_vjp(grad_comp1)
ng = []
ag = []
ind = []
dx = 1e-7
for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
dx1, c1 = perturb(comp1, ind1, dx)
ng1 = (objective(c1) - objective(comp1)) / dx
ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx
if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
print (ind1, 'a', ag1, 'n', ng1)
comm.barrier()
ng.append(ng1)
ag.append(ag1)
ind.append(ind1)
assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
def test_decompose(comm):
pm = ParticleMesh(BoxSize=4.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 1000
else:
Npar = 0
pos = 4.0 * (numpy.random.uniform(size=(Npar, 3)))
pos = pm.generate_uniform_particle_grid(shift=0.5)
all_pos = numpy.concatenate(comm.allgather(pos), axis=0)
for resampler in ['cic', 'tsc', 'db12']:
def test(resampler):
print(resampler)
truth = numpy.zeros(pm.Nmesh, dtype='f8')
affine = window.Affine(ndim=3, period=4)
window.FindResampler(resampler).paint(truth, all_pos,
transform=affine)
truth = comm.bcast(truth)
layout = pm.decompose(pos, smoothing=resampler)
npos = layout.exchange(pos)
real = pm.paint(npos, resampler=resampler)
full = numpy.zeros(pm.Nmesh, dtype='f8')
full[real.slices] = real
full = comm.allreduce(full)
#print(full.sum(), pm.Nmesh.prod())
#print(truth.sum(), pm.Nmesh.prod())
#print(comm.rank, npos, real.slices)
assert_almost_equal(full, truth)
# can't yield!
test(resampler)
def test_c2c(comm):
# this test requires pfft-python 0.1.16.
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128')
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()
real2 = complex.c2r()
assert numpy.iscomplexobj(real)
assert numpy.iscomplexobj(real2)
assert numpy.iscomplexobj(complex)
assert_array_equal(complex.cshape, pm.Nmesh)
assert_array_equal(real2.cshape, pm.Nmesh)
assert_array_equal(real.cshape, pm.Nmesh)
real.readout(npos)
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
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_operators(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type='real', value=0)
complex = pm.create(type='complex', value=0)
real = real + 1
real = 1 + real
real = real + real.value
real = real * real.value
real = real * real
assert isinstance(real, RealField)
complex = 1 + complex
assert isinstance(complex, ComplexField)
complex = complex + 1
assert isinstance(complex, ComplexField)
complex = complex + complex.value
assert isinstance(complex, ComplexField)
complex = numpy.conj(complex) * complex
assert isinstance(complex, ComplexField)
assert (real == real).dtype == numpy.dtype('?')
assert not isinstance(real == real, RealField)
assert not isinstance(complex == complex, ComplexField)
assert not isinstance(numpy.sum(real), RealField)
complex = numpy.conj(complex)
def test_preview(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='real')
preview = comp1.preview(axes=(0, 1, 2))
preview = comp1.preview(Nmesh=4, axes=(0, 1, 2))
for ind1 in numpy.ndindex(*(list(comp1.cshape))):
assert_allclose(preview[ind1], comp1.cgetitem(ind1))
preview1 = comp1.preview(Nmesh=4, axes=(0, 1))
previewsum1 = preview.sum(axis=2)
assert_allclose(preview1, previewsum1)
preview2 = comp1.preview(Nmesh=4, axes=(1, 2))
previewsum2 = preview.sum(axis=0)
assert_allclose(preview2, previewsum2)
preview3 = comp1.preview(Nmesh=4, axes=(0, 2))
previewsum3 = preview.sum(axis=1)
assert_allclose(preview3, previewsum3)
preview4 = comp1.preview(Nmesh=4, axes=(2, 0))
previewsum4 = preview.sum(axis=1).T
assert_allclose(preview4, previewsum4)
preview5 = comp1.preview(Nmesh=4, axes=(0,))
previewsum5 = preview.sum(axis=(1, 2))
assert_allclose(preview5, previewsum5)
preview6 = comp1.preview(Nmesh=8, axes=(0,))
def test_negnyquist(comm):
# the nyquist mode wave number in the hermitian complex field must be negative.
# nbodykit depends on this behavior.
# see https://github.com/bccp/nbodykit/pull/459
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
c = pm.create(type='complex')
assert (c.x[-1][0][-1] < 0).all()
assert (c.x[-1][0][:-1] >= 0).all()
def test_indices_c2c(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='c16')
comp = pm.create(type='complex')
real = pm.create(type='real')
assert_almost_equal(comp.x[0], [[0], [0.785], [-1.571], [-0.785]], decimal=3)
assert_almost_equal(comp.x[1], [[0, 0.785, -1.571, -0.785]], decimal=3)
assert_almost_equal(real.x[0], [[0], [2], [-4], [-2]], decimal=3)
assert_almost_equal(real.x[1], [[0, 2, -4, -2]], decimal=3)
def test_c2c_r2c_edges(comm):
pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='c16')
pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='f8')
real1 = pm1.create(type='real')
real2 = pm2.create(type='real')
assert_allclose(real1.x[0], real2.x[0])
assert_allclose(real1.x[1], real2.x[1])
assert_allclose(real1.x[2], real2.x[2])
def test_cdot_cnorm(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex')
norm1 = comp1.cdot(comp1)
norm2 = comp1.cnorm()
norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2).sum()
assert_allclose(norm2, norm3)
assert_allclose(norm2, norm1)
def test_grid(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid, id = pm.generate_uniform_particle_grid(shift=0.5, return_id=True)
assert len(id) == len(grid)
allid = numpy.concatenate(comm.allgather(id), axis=0)
# must be all unique
assert len(numpy.unique(allid)) == len(allid)
assert numpy.max(allid) == len(allid) - 1
assert numpy.min(allid) == 0
def test_create_typenames(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type=RealField, value=0)
from pmesh.pm import _typestr_to_type
real = pm.create(type=RealField, value=0)
real = pm.create(type=_typestr_to_type('real'), value=0)
real.cast(type=_typestr_to_type('real'))
real.cast(type=RealField)
real.cast(type=RealField)
def test_lic(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(Nmesh=[8, 8], comm=comm, BoxSize=8.0)
vx = pm.create(type='real')
vy = pm.create(type='real')
vx = vx.apply(lambda r, v: r[0])
vy = vy.apply(lambda r, v: 1.0)
# FIXME: This only tests if the code 'runs', but no correctness
# is tested.
r = lic([vx, vy], kernel=lambda s: 1.0, length=2.0, ds=0.1)
def test_cdot_grad(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex', mean=1)
comp2 = pm.generate_whitenoise(1235, type='complex', mean=1)
def objective(comp1, comp2):
return comp1.cdot(comp2).real
y = objective(comp1, comp2)
grad_comp2 = comp1.cdot_vjp(1.0)
grad_comp1 = comp2.cdot_vjp(1.0)
grad_comp1.decompress_vjp(grad_comp1)
grad_comp2.decompress_vjp(grad_comp2)
print(grad_comp1)
print("comp1")
ng = []
ag = []
ind = []
dx = 1e-7
for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
dx1, c1 = perturb(comp1, ind1, dx)
ng1 = (objective(c1, comp2) - objective(comp1, comp2)) / dx
ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx
if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
print (ind1, 'a', ag1, 'n', ng1)
comm.barrier()
ng.append(ng1)
ag.append(ag1)
ind.append(ind1)
assert_allclose(ng, ag, rtol=1e-5)
print("comp2")
ng = []
ag = []
ind = []
dx = 1e-7
for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
dx1, c2 = perturb(comp2, ind1, dx)
ng1 = (objective(comp1, c2) - objective(comp1, comp2)) / dx
ag1 = grad_comp2.cgetitem(ind1) * dx1 / dx
if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
print (ind1, 'a', ag1, 'n', ng1)
comm.barrier()
ng.append(ng1)
ag.append(ag1)
ind.append(ind1)
assert_allclose(ng, ag, rtol=1e-5)
def test_fof_parallel_no_merge(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm)
Q = pm.generate_uniform_particle_grid()
cat = ArrayCatalog({'Position' : Q}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm)
fof = FOF(cat, linking_length=0.9, nmin=0)
labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0)
# one particle per group
assert max(labels) == cat.csize - 1
def test_2d_2d(comm):
import pfft
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8')
pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8')
assert pm._use_padded == False
real = pm.generate_whitenoise(seed=123, type='real')
complex = pm.generate_whitenoise(seed=123, type='complex')
assert_array_equal(real, complex.c2r())
real2 = pm2.generate_whitenoise(seed=123, type='real')
assert real2.shape[:2] == real.shape
def test_cdot_types(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex')
comp2 = pm.generate_whitenoise(1239, type='untransposedcomplex')
with pytest.raises(TypeError):
norm1 = comp1.cdot(comp2)
norm2 = comp2.cdot(comp1)
# this should work though fragile, because only on the Nmesh and 1 rank
# the shape of values is the same.
norm1 = comp1.cdot(comp2.value)
norm2 = comp2.cdot(comp1.value)
def test_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type='real', value=0)
raw = real._base.view_raw()
real[...] = 2
real[::2, ::2] = -2
real3 = real.copy()
complex = real.r2c()
assert_almost_equal(numpy.asarray(real), numpy.asarray(real3), decimal=7)
real2 = complex.c2r()
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_real_resample(comm):
from functools import reduce
pmh = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
pml = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
reall = pml.create(type='real')
reall.apply(lambda i, v: (i[0] % 2) * (i[1] %2 ), kind='index', out=Ellipsis)
for resampler in ['nearest', 'cic', 'tsc', 'cubic']:
realh = pmh.upsample(reall, resampler=resampler, keep_mean=False)
reall2 = pml.downsample(realh, resampler=resampler)
# print(resampler, comm.rank, comm.size, reall, realh)
assert_almost_equal(reall.csum(), realh.csum())
assert_almost_equal(reall.csum(), reall2.csum())
def test_cdot(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='complex')
comp2 = pm.generate_whitenoise(1239, type='complex')
norm1 = comp1.cdot(comp2)
norm2 = comp2.cdot(comp1)
norm_r = comp1.c2r().cdot(comp2.c2r()) / pm.Nmesh.prod()
assert_allclose(norm2.real, norm_r)
assert_allclose(norm1.real, norm_r)
assert_allclose(norm1.real, norm2.real)
assert_allclose(norm1.imag, -norm2.imag)
def test_paint(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
real = pm.generate_whitenoise(type='real', seed=3333)
complex = real.r2c()
realmesh = FieldMesh(real)
complexmesh = FieldMesh(complex)
assert_array_equal(real, realmesh.to_field(mode='real'))
assert_array_equal(complex, complexmesh.to_field(mode='complex'))
assert_allclose(complex, realmesh.to_field(mode='complex'))
assert_allclose(real, complexmesh.to_field(mode='real'))
def __init__(self, comm, Nmesh, BoxSize, dtype):
# ensure self.comm is set, though usually already set by the child.
self.comm = comm
self.dtype = dtype
if Nmesh is None or BoxSize is None:
raise ValueError("both Nmesh and BoxSize must not be None to initialize ParticleMesh")
Nmesh = numpy.array(Nmesh)
if Nmesh.ndim == 0:
ndim = 3
else:
ndim = len(Nmesh)
_Nmesh = numpy.empty(ndim, dtype='i8')
_Nmesh[:] = Nmesh
self.pm = ParticleMesh(BoxSize=BoxSize, Nmesh=_Nmesh,
dtype=self.dtype, comm=self.comm)
self.attrs['BoxSize'] = self.pm.BoxSize.copy()
self.attrs['Nmesh'] = self.pm.Nmesh.copy()
# modify the underlying method
# actions may have been overriden!
self._actions = []
self.base = None
def test_whitenoise(comm):
# the whitenoise shall preserve the large scale.
pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8')
pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8')
complex1_down = ComplexField(pm0)
complex2_down = ComplexField(pm0)
complex1 = pm1.generate_whitenoise(seed=8, unitary=True)
complex2 = pm2.generate_whitenoise(seed=8, unitary=True)
complex1.resample(complex1_down)
complex2.resample(complex2_down)
mask1 = complex1_down.value != complex2_down.value
assert_array_equal(complex1_down.value, complex2_down.value)
def main():
from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("--ndim", type=int, choices=[2, 3], default=2, help="Number of dimensions, 2 or 3")
ap.add_argument("--nmesh", type=int, default=256, help="Size of FFT mesh")
ns = ap.parse_args()
pm = ParticleMesh(BoxSize=32.0, Nmesh=[ns.nmesh] * ns.ndim, comm=MPI.COMM_WORLD)
u = pm.create(mode='real')
def transfer(i, v):
r = [(ii - 0.5 * ni) * (Li / ni) for ii, ni, Li in zip(i, v.Nmesh, v.BoxSize)]
r2 = sum(ri ** 2 for ri in r)
return 4.0 * numpy.arctan(numpy.exp(3 - r2))
u = u.apply(transfer, kind='index')
du = pm.create(mode='real')
du[...] = 0
steps = numpy.linspace(0, 16, 321, endpoint=True)
tmonitor = [0, 4, 8, 11.5, 15]
def monitor(t, dt, u_k, dv_k):
norm = u_k.cnorm()
if pm.comm.rank == 0:
print("---- timestep %5.3f, step size %5.4f" % (t, dt))
print("norm of u_k is %g." % norm)
for tm in tmonitor.copy():
if abs(t - tm) > dt * 0.5: continue
preview = u_k.c2r().preview(Nmesh=min([512, ns.nmesh]), axes=(0, 1))
if pm.comm.rank == 0:
print("writing a snapshot")
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg
fig = Figure(figsize=(8, 8))
ax = fig.add_subplot(111)
ax.imshow(preview.T, origin='lower', extent=(0, pm.BoxSize[0], 0, pm.BoxSize[1]))
canvas = FigureCanvasAgg(fig)
fig.savefig('klein-gordon-result-%05.3f.png' % t, dpi=128)
tmonitor.remove(tm)
kgsolver(steps, u, du, lambda u : numpy.sin(u), monitor=monitor)
def __init__(self, field, points, Nmesh, smoothing=None):
from pmesh.pm import ParticleMesh
self.field = field
self.points = points
self.Nmesh = Nmesh
self.smoothing = smoothing
self.pm = ParticleMesh(BoxSize=self.field.BoxSize, Nmesh=[self.Nmesh]*3, dtype='f4', comm=self.comm)
def test_fof_parallel_merge(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm)
Q = pm.generate_uniform_particle_grid(shift=0)
Q1 = Q.copy()
Q1[:] += 0.01
Q2 = Q.copy()
Q2[:] -= 0.01
Q3 = Q.copy()
Q3[:] += 0.02
cat = ArrayCatalog({'Position' :
numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm)
fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True)
labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0)
assert max(labels) == pm.Nmesh.prod() - 1
assert all(numpy.bincount(labels) == 4)
def test_complex_iter(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
complex = ComplexField(pm)
for x, slab in zip(complex.slabs.x, complex.slabs):
assert_array_equal(slab.shape,
sum(x[d]**2 for d in range(len(pm.Nmesh))).shape)
for a, b in zip(slab.x, x):
assert_almost_equal(a, b)
def test_fof_parallel_merge(comm):
CurrentMPIComm.set(comm)
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=[32, 32, 32], Nmesh=[32, 32, 32], comm=comm)
Q = pm.generate_uniform_particle_grid(shift=0)
Q1 = Q.copy()
Q1[:] += 0.01
Q2 = Q.copy()
Q2[:] -= 0.01
Q3 = Q.copy()
Q3[:] += 0.02
cat = ArrayCatalog({'Position' :
numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True)
labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0)
assert max(labels) == pm.Nmesh.prod() - 1
assert all(numpy.bincount(labels) == 4)
def __init__(self, pm, cosmology, B=1):
"""
"""
if not isinstance(cosmology, Cosmology):
raise TypeError("only nbodykit.cosmology object is supported")
fpm = ParticleMesh(Nmesh=pm.Nmesh * B, BoxSize=pm.BoxSize, dtype=pm.dtype, comm=pm.comm, resampler=pm.resampler)
self.pm = pm
self.fpm = fpm
self.cosmology = cosmology
def test_resample(comm):
if comm.size == 4:
comm.barrier()
from pmesh.pm import ParticleMesh
# testing the gradient of down sampling, which we use the most.
pm = ParticleMesh(BoxSize=128.0, Nmesh=(8, 8), comm=comm)
vm = fastpm.Evolution(pm, shift=0.5)
dlink = pm.generate_whitenoise(1234, mode='complex')
dlink, ampl = _addampl(dlink)
code = vm.code()
code.Decompress(C='dlin_k')
code.C2R(C='dlin_k', R='mesh')
code.Resample(Neff=4)
code.Chi2(variable='mesh')
_test_model(code, dlink, ampl)
def test_whitenoise_untransposed(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
f1 = pm.generate_whitenoise(seed=3333, type='untransposedcomplex')
f2 = pm.generate_whitenoise(seed=3333, type='transposedcomplex')
f1r = numpy.concatenate(comm.allgather(numpy.array(f1.ravel())))
f2r = numpy.concatenate(comm.allgather(numpy.array(f2.ravel())))
assert_array_equal(f1r, f2r)
# this should have asserted r2c transforms as well.
r1 = f1.c2r()
r2 = f2.c2r()
r1r = numpy.concatenate(comm.allgather(numpy.array(r1.ravel())))
r2r = numpy.concatenate(comm.allgather(numpy.array(r2.ravel())))
assert_array_equal(r1r, r2r)
def test_field_compressed(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='c16')
comp = pm.create(type='complex')
real = pm.create(type='real')
assert comp.compressed == False
assert real.compressed == False
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
comp = pm.create(type='complex')
real = pm.create(type='real')
assert comp.compressed == True
assert real.compressed == False
def test_lpt2():
""" Checking lpt2_source, this also checks the laplace and gradient kernels
"""
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
source = fpmops.lpt2source(lineark).c2r()
# Same thing from tensorflow
tfsource = tfpm.lpt2_source(
pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)))
tfread = pmutils.c2r3d(tfsource).numpy()
assert_allclose(source, tfread[0], atol=1e-5)
def _makesource(self, BoxSize, Nmesh):
from nbodykit import mockmaker
from pmesh.pm import ParticleMesh
# the particle mesh for gridding purposes
_Nmesh = numpy.empty(3, dtype='i8')
_Nmesh[:] = Nmesh
pm = ParticleMesh(BoxSize=BoxSize, Nmesh=_Nmesh, dtype='f4', comm=self.comm)
# growth rate to do RSD in the Zel'dovich approx
f = self.cosmo.scale_independent_growth_rate(self.attrs['redshift'])
if self.comm.rank == 0:
self.logger.info("Growth Rate is %g" % f)
# compute the linear overdensity and displacement fields
delta, disp = mockmaker.gaussian_real_fields(pm, self.Plin, self.attrs['seed'],
unitary_amplitude=self.attrs['unitary_amplitude'],
inverted_phase=self.attrs['inverted_phase'],
compute_displacement=True,
logger=self.logger)
if self.comm.rank == 0:
self.logger.info("gaussian field is generated")
# poisson sample to points
# this returns position and velocity offsets
kws = {'bias':self.attrs['bias'], 'seed':self.attrs['seed'], 'logger' : self.logger}
pos, disp = mockmaker.poisson_sample_to_points(delta, disp, pm, self.attrs['nbar'], **kws)
if self.comm.rank == 0:
self.logger.info("poisson sampling is generated")
# move particles from initial position based on the Zeldovich displacement
pos[:] = (pos + disp) % BoxSize
# velocity from displacement (assuming Mpc/h)
# this is f * H(z) * a / h = f 100 E(z) a --> converts from Mpc/h to km/s
z = self.attrs['redshift']
velocity_norm = f * 100 * self.cosmo.efunc(z) / (1+z)
vel = velocity_norm * disp
# return data
dtype = numpy.dtype([
('Position', ('f4', 3)),
('Velocity', ('f4', 3)),
('VelocityOffset', ('f4', 3))
])
source = numpy.empty(len(pos), dtype)
source['Position'][:] = pos[:] # in Mpc/h
source['Velocity'][:] = vel[:] # in km/s
source['VelocityOffset'][:] = f*disp[:] # in Mpc/h
return source, pm
def __init__(self,
datasource,
Nmesh,
seed=12345,
ratio=0.01,
smoothing=None,
format='hdf5'):
from pmesh.pm import ParticleMesh
self.datasource = datasource
self.Nmesh = Nmesh
self.seed = seed
self.ratio = ratio
self.smoothing = smoothing
self.format = format
self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize,
Nmesh=[self.Nmesh] * 3,
dtype='f4',
comm=self.comm)
def test_cic_readout():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
base = 100 * np.random.random(nc**3).reshape(nc, nc, nc).astype(np.float32)
pmmesh = pm.create(mode='real', value=base)
pmread = pmmesh.readout(pos)
with tf.Session() as sess:
mesh = cic_readout(
tf.constant(base.reshape((1, nc, nc, nc)), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)))
sess.run(tf.global_variables_initializer())
tfread = sess.run(mesh)
assert_allclose(pmread, tfread[0], rtol=1e-06)
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_complex_apply(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
complex = ComplexField(pm)
def filter(k, v):
return k[0] + k[1] * 1j
complex.apply(filter, out=Ellipsis)
for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs):
assert_array_equal(slab, x[0] + x[1] * 1j)
def test_cic_paint():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
wts = np.random.random(nparticle).astype(np.float32)
# Painting with pmesg
pmmesh = pm.paint(pos, mass=wts)
with tf.Session() as sess:
mesh = cic_paint(tf.zeros((1, nc, nc, nc), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)),
weight=wts.reshape(1, nparticle))
sess.run(tf.global_variables_initializer())
tfmesh = sess.run(mesh)
assert_allclose(pmmesh, tfmesh[0], atol=1e-06)
def test_coords(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid_x = pm.create_coords('real')
assert len(grid_x) == 3
assert grid_x[0].dtype == pm.dtype
grid_i = pm.create_coords('real', return_indices=True)
assert len(grid_i) == 3
grid_x = pm.create_coords('complex')
grid_i = pm.create_coords('complex', return_indices=True)
assert len(grid_x) == 3
assert grid_x[0].dtype == pm.dtype
assert len(grid_i) == 3
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f4')
grid_x = pm.create_coords('transposedcomplex')
grid_i = pm.create_coords('transposedcomplex', return_indices=True)
assert len(grid_x) == 3
assert grid_x[0].dtype == pm.dtype
assert len(grid_i) == 3
def __init__(self, pm, B=1, shift=0, dtype='f8'):
self.fpm = ParticleMesh(Nmesh=pm.Nmesh * B,
BoxSize=pm.BoxSize,
dtype=dtype,
comm=pm.comm)
q = operators.create_grid(pm, shift=shift, dtype=dtype)
N = pm.comm.allreduce(len(q))
self.mean_number_count = 1.0 * N / (1.0 * self.fpm.Nmesh.prod())
ParticleMeshVM.__init__(self, pm, q)
MPINumeric.__init__(self, pm.comm)
def test_lpt1():
""" Checking lpt1, this also checks the laplace and gradient kernels
"""
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
lpt = fpmops.lpt1(lineark, grid)
# Same thing from tensorflow
with tf.Session() as sess:
state = tfpm.lpt1(
pmutils.r2c3d(tf.expand_dims(tf.constant(lineark.c2r()), axis=0)),
grid.reshape((1, -1, 3)) * nc / bs)
tfread = sess.run(state)
assert_allclose(lpt, tfread[0] * bs / nc, atol=1e-5)
def test_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
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()
real2 = complex.c2r()
real.readout(npos)
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_lpt_prior(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm)
vm = fastpm.Evolution(pm, shift=0.5)
dlink = pm.generate_whitenoise(1234, mode='complex')
code = vm.code()
code.Decompress(C='dlin_k')
code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0)
code.Decompose()
code.Paint()
code.Chi2(variable='mesh')
code.Prior(powerspectrum=lambda k : 1.0)
code.Multiply(a='prior', b='prior', f=0.1)
code.Add(a='prior', b='chi2', c='chi2')
dlink, ampl = _addampl(dlink)
_test_model(code, dlink, ampl)
def test_lpt_init():
"""
Checking lpt init
"""
a0 = 0.1
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
solver = Solver(pm, Planck15, B=1)
# Generate initial state with fastpm
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
statelpt = solver.lpt(lineark, grid, a0, order=1)
# Same thing with flowpm
tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
tfread = tfpm.lpt_init(tlinear, a0, order=1).numpy()
assert_allclose(statelpt.X, tfread[0, 0] * bs / nc, rtol=1e-2)
def test_lpt1_64():
""" Checking lpt1, this also checks the laplace and gradient kernels
This variant of the test checks that it works for cubes of size 64
"""
nc = 64
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
# Compute lpt1 from fastpm with matching kernel order
lpt = fpmops.lpt1(lineark, grid)
# Same thing from tensorflow
tfread = tfpm.lpt1(
pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)),
grid.reshape((1, -1, 3)) * nc / bs).numpy()
assert_allclose(lpt, tfread[0] * bs / nc, atol=5e-5)
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 main(ns):
comm = MPI.COMM_WORLD
result = simulate(comm, ns)
pm = ParticleMesh(BoxSize=ns.BoxSize,
Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh],
dtype='f8',
comm=comm)
report = analyze(pm, result)
if comm.rank == 0:
write_report(ns.output, report)
def test_fdownsample(comm):
""" fourier space resample, deprecated """
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])
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(pm2)
tmp = ComplexField(pm2)
complex1.resample(tmp)
assert_almost_equal(complex2[...], tmp[...], decimal=5)
complex1.c2r().resample(tmp)
assert_almost_equal(complex2[...], tmp[...], decimal=5)
complex1.resample(tmpr)
assert_almost_equal(tmpr.r2c(), tmp[...])
complex1.c2r().resample(tmpr)
assert_almost_equal(tmpr.r2c(), tmp[...])
def test_untransposed_complex_apply(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
complex = UntransposedComplexField(pm)
def filter(k, v):
knormp = k.normp()
assert_allclose(knormp, sum(ki**2 for ki in k))
return k[0] + k[1] * 1j + k[2]
complex = complex.apply(filter)
for i, x, slab in zip(complex.slabs.i, complex.slabs.x, complex.slabs):
assert_array_equal(slab, x[0] + x[1] * 1j + x[2])
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 fiddlebiasgal(aa, suff, nc=nc, mcfv=[1.], saveb=False, bname='h1bias', ofolder=None):
'''Fiddle bias for galaxies'''
if ofolder is None: ofolder = project + '/%s/fastpm_%0.4f/'%(sim, aa)
pm = ParticleMesh(BoxSize = bs, Nmesh = [nc, nc, nc])
print('Read in catalogs')
cencat = BigFileCatalog(project + sim + '/fastpm_%0.4f/cencat'%aa)
satcat = BigFileCatalog(project + sim + '/fastpm_%0.4f/satcat'%aa+suff)
cpos, spos = cencat['Position'], satcat['Position']
cmass, smass = cencat['Mass'], satcat['Mass']
pos = np.concatenate((cpos, spos), axis=0)
dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint()
pkm = FFTPower(dm/dm.cmean(), mode='1d').power
k, pkm = pkm['k'], pkm['power']
b1, b1sq = np.zeros((k.size, len(mcfv))), np.zeros((k.size, len(mcfv)))
for imc, mcf in enumerate(mcfv):
print(mcf)
ch1mass = HI_masscutfiddle(cmass, aa, mcutf=mcf)
sh1mass = HI_masscutfiddle(smass, aa, mcutf=mcf)
h1mass = np.concatenate((ch1mass, sh1mass), axis=0)
#
h1mesh = pm.paint(pos, mass=h1mass)
pkh1 = FFTPower(h1mesh/h1mesh.cmean(), mode='1d').power['power']
pkh1m = FFTPower(h1mesh/h1mesh.cmean(), second=dm/dm.cmean(), mode='1d').power['power']
#Bias
b1[:, imc] = pkh1m/pkm
b1sq[:, imc] = pkh1/pkm
np.savetxt(ofolder+bname+'auto'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1sq**0.5), axis=1),
header='mcut factors = %s\nk, pkh1xm/pkm, pkh1/pkm^0.5'%mcfv)
np.savetxt(ofolder+bname+'cross'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1), axis=1),
header='mcut factors = %s\nk, pkh1xm/pkm, pkh1mx/pkm'%mcfv)
return k, b1, b1sq
def test_decompose(comm):
pm = ParticleMesh(BoxSize=4.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 1000
else:
Npar = 0
pos = 4.0 * (numpy.random.uniform(size=(Npar, 3)))
pos = pm.generate_uniform_particle_grid(shift=0.5)
all_pos = numpy.concatenate(comm.allgather(pos), axis=0)
for resampler in ['cic', 'tsc', 'db12']:
def test(resampler):
print(resampler)
truth = numpy.zeros(pm.Nmesh, dtype='f8')
affine = window.Affine(ndim=3, period=4)
window.FindResampler(resampler).paint(truth,
all_pos,
transform=affine)
truth = comm.bcast(truth)
layout = pm.decompose(pos, smoothing=resampler)
npos = layout.exchange(pos)
real = pm.paint(npos, resampler=resampler)
full = numpy.zeros(pm.Nmesh, dtype='f8')
full[real.slices] = real
full = comm.allreduce(full)
#print(full.sum(), pm.Nmesh.prod())
#print(truth.sum(), pm.Nmesh.prod())
#print(comm.rank, npos, real.slices)
assert_almost_equal(full, truth)
# can't yield!
test(resampler)
def finalize(self):
self['aout'] = numpy.array(self['aout'])
self.pm = ParticleMesh(BoxSize=self['boxsize'],
Nmesh=[self['nc']] * self['ndim'],
resampler=self['resampler'],
dtype='f4')
mask = numpy.array([a not in self['stages'] for a in self['aout']],
dtype='?')
missing_stages = self['aout'][mask]
if len(missing_stages):
raise ValueError(
'Some stages are requested for output but missing: %s' %
str(missing_stages))
def set_up():
params = get_params()
pm = ParticleMesh(Nmesh=params['Nmesh'],
BoxSize=params['BoxSize'],
comm=MPI.COMM_WORLD,
resampler='cic')
# generate initial conditions
cosmo = Planck15.clone(P_k_max=30)
x = pm.generate_uniform_particle_grid(shift=0.1)
BoxSize2D = [deg / 180. * np.pi for deg in params['BoxSize2D']]
logger = None
sim = lightcone.WLSimulation(stages=numpy.linspace(0.1,
1.0,
params['N_steps'],
endpoint=True),
cosmology=cosmo,
pm=pm,
boxsize2D=BoxSize2D,
params=params,
logger=None)
kmaps = [sim.mappm.create('real', value=0.) for ii in range(1)]
return pm, cosmo, x, kmaps, sim.DriftFactor, sim.mappm, sim
def test_c2c_r2c_edges(comm):
pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='c16')
pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='f8')
real1 = pm1.create(type='real')
real2 = pm2.create(type='real')
assert_allclose(real1.x[0], real2.x[0])
assert_allclose(real1.x[1], real2.x[1])
assert_allclose(real1.x[2], real2.x[2])
