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')
real = pm.generate_whitenoise(seed=123, mode='real')
complex = pm.generate_whitenoise(seed=123, mode='complex')
assert_array_equal(real, complex.c2r())
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_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_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_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 test_kdk(comm):
# Or use an object from astropy.
cosmo = lambda : None
cosmo.Om0 = 0.3
cosmo.Ode0 = 0.7
cosmo.Ok0 = 0.0
pm = ParticleMesh(BoxSize=128.0, Nmesh=(4,4,4), comm=comm, dtype='f8')
vm = fastpm.Evolution(pm, B=1, shift=0.5)
dlink = pm.generate_whitenoise(12345, mode='complex', unitary=True)
dlink, ampl = _addampl(dlink)
data = dlink.c2r()
data[...] = 0
sigma = data.copy()
sigma[...] = 1.0
code = vm.code()
code.Decompress(C='dlin_k')
code.KDKSimulation(dlin_k='dlin_k', mesh='mesh', cosmo=cosmo, astart=0.1, aend=1.0, Nsteps=5)
code.Decompose()
code.Paint()
code.Residual(data_x=data, sigma_x=sigma)
code.Chi2(variable='residual')
_test_model(code, dlink, ampl)
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_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 test_nody():
""" Checking end to end nbody
"""
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)
stages = np.linspace(0.1, 1.0, 10, endpoint=True)
# 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)
finalstate = solver.nbody(statelpt, leapfrog(stages))
final_cube = pm.paint(finalstate.X)
# Same thing with flowpm
tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
state = tfpm.lpt_init(tlinear, a0, order=1)
state = tfpm.nbody(state, stages, nc)
tfread = pmutils.cic_paint(tf.zeros_like(tlinear), state[0]).numpy()
assert_allclose(final_cube, tfread[0], atol=1.2)
def test_gravity(comm):
from pmesh.pm import ParticleMesh
import fastpm.operators as operators
pm = ParticleMesh(BoxSize=4.0, Nmesh=(4, 4), comm=comm, method='cic', dtype='f8')
vm = fastpm.Evolution(pm, shift=0.5)
dlink = pm.generate_whitenoise(12345, 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.ForcePaint()
code.Force(factor=0.1)
code.Chi2(variable='f')
print(code)
# FIXME: without the shift some particles have near zero dx1.
# or near 1 dx1.
# the gradient is not well approximated by the numerical if
# any of the left or right value shifts beyond the support of
# the window.
#
dlink, ampl = _addampl(dlink)
_test_model(code, dlink, ampl)
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_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_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_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_grad(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, mode='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_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_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_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_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_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_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_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 make_data(bs, nc, seed, nsteps, path='', z=0):
#initiate
pm = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8')
ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep/'%(z*10, bs, nc, seed, nsteps)
#
if pm.comm.rank == 0:
print('Make data for seed = %d'%seed)
print('Data to be saved at path - %s'%ofolder)
klin, plin = numpy.loadtxt('ics_matterpow_0.dat', unpack = True)
pk = interpolate(klin, plin)
cosmo = Planck15.clone(Omega_cdm = 0.2685, h = 0.6711, Omega_b = 0.049)
#pk = EHPower(Planck15, redshift=0)
#cosmo = Planck15
s_truth = pm.generate_whitenoise(seed, type='complex')\
.apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
.c2r()
#dynamics
aa = 1.0/(1+z)
if pm.comm.rank == 0:
print('Evolve to redshift = %0.1f, scale factor = %0.2f'%(z, aa))
stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)
start = time()
dynamic_model = NBodyModel(cosmo, pm, B=2, steps=stages)
#dynamic_model = LPTModel(cosmo, pm, B=2, steps=stages)
#Save data
X, V, final = dynamic_model.get_code().compute(['X', 'V', 'final'], init={'parameters':s_truth})
end = time()
print('Time taken = ', end-start)
save_map(s_truth, ofolder + 'mesh', 's')
save_map(final, ofolder + 'mesh', 'd')
if pm.comm.rank == 0:
print('X, V computed')
cat = ArrayCatalog({'Position': X, 'Velocity' : V}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
kdd = KDDensity(cat).density
cat['KDDensity'] = kdd
cat['InitPosition'] = dynamic_model.get_code().engine.q
cat.save(ofolder + 'dynamic/1', ('InitPosition', 'Position', 'Velocity', 'KDDensity'))
if pm.comm.rank == 0:
print('dynamic model created')
#FOF
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.find_features(peakcolumn='KDDensity')
fofcat['Mass'] = fofcat['Length'] * cosmo.Om0 * 27.7455 * pm.BoxSize.prod() / pm.Nmesh.prod()
fofcat.save(ofolder+'FOF', ('CMPosition', 'CMVelocity', 'PeakPosition', 'PeakVelocity', 'Length', 'Mass'))
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_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_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_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_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.Prior(dlin_k='dlin_k', powerspectrum=lambda k : 1.0)
code.CopyVariable(x='prior', y='chi2')
dlink, ampl = _addampl(dlink)
_test_model(code, dlink, ampl)
def simulate(comm, ns):
pm = ParticleMesh(BoxSize=ns.BoxSize,
Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh],
dtype='f8',
comm=comm)
gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)
Q = pm.generate_uniform_particle_grid(shift=0)
print(Q.min(axis=0), Q.max(axis=0))
def convolve(k, v):
kmag = sum(ki**2 for ki in k)**0.5
ampl = (PowerSpectrum(kmag) / v.BoxSize.prod())**0.5
return v * ampl
dlinear = gaussian.apply(convolve)
DX1 = numpy.zeros_like(Q)
layout = pm.decompose(Q)
# Fill it in one dimension at a time.
for d in range(pm.ndim):
DX1[..., d] = dlinear \
.apply(dx1_transfer(d)) \
.c2r().readout(Q, layout=layout)
a0 = time_steps[0]
# 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
S = DX1 * pt.D1(a=a0)
V = S * a0**2 * pt.f1(a0) * pt.E(a0)
state = State(Q, S, V)
fpm = ParticleMesh(BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh * ns.boost,
resampler='tsc',
dtype='f8')
ns.scheme(fpm, state, time_steps, ns.factors)
r = Result()
r.Q = Q
r.DX1 = DX1
r.S = S
r.V = V
r.dlinear = dlinear
return r
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 test_paint(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
real = pm.generate_whitenoise(mode='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 test_lpt(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.Chi2(variable='s')
dlink, ampl = _addampl(dlink)
_test_model(code, dlink, ampl)
def simulate(comm, ns):
pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm)
gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)
Q = pm.generate_uniform_particle_grid(shift=0)
print(Q.min(axis=0), Q.max(axis=0))
def convolve(k, v):
kmag = sum(ki**2 for ki in k) ** 0.5
ampl = (PowerSpectrum(kmag) / v.BoxSize.prod()) ** 0.5
return v * ampl
dlinear = gaussian.apply(convolve)
DX1 = numpy.zeros_like(Q)
layout = pm.decompose(Q)
# Fill it in one dimension at a time.
for d in range(pm.ndim):
DX1[..., d] = dlinear \
.apply(dx1_transfer(d)) \
.c2r().readout(Q, layout=layout)
a0 = time_steps[0]
# 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
S = DX1 * pt.D1(a=a0)
V = S * a0 ** 2 * pt.f1(a0) * pt.E(a0)
state = State(Q, S, V)
fpm = ParticleMesh(BoxSize=pm.BoxSize, Nmesh=pm.Nmesh * ns.boost, resampler='tsc', dtype='f8')
ns.scheme(fpm, state, time_steps, ns.factors)
r = Result()
r.Q = Q
r.DX1 = DX1
r.S = S
r.V = V
r.dlinear = dlinear
return r
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 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_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_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_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_transpose(comm):
pm = ParticleMesh(BoxSize=[8.0, 16.0, 32.0], Nmesh=[4, 6, 8], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='real')
comp1t = comp1.ctranspose([0, 1, 2])
assert_array_equal(comp1t.Nmesh, comp1.Nmesh)
assert_array_equal(comp1t.BoxSize, comp1.BoxSize)
assert_array_equal(comp1t.cnorm(), comp1.cnorm())
comp1t = comp1.ctranspose([1, 2, 0])
assert_array_equal(comp1t.Nmesh, comp1.Nmesh[[1, 2, 0]])
assert_array_equal(comp1t.BoxSize, comp1.BoxSize[[1, 2, 0]])
comp1tt = comp1t.ctranspose([1, 2, 0])
comp1ttt = comp1tt.ctranspose([1, 2, 0])
assert_allclose(comp1ttt, comp1)
def test_transpose(comm):
pm = ParticleMesh(BoxSize=[8.0, 16.0, 32.0],
Nmesh=[4, 6, 8],
comm=comm,
dtype='f8')
comp1 = pm.generate_whitenoise(1234, type='real')
comp1t = comp1.ctranspose([0, 1, 2])
assert_array_equal(comp1t.Nmesh, comp1.Nmesh)
assert_array_equal(comp1t.BoxSize, comp1.BoxSize)
assert_array_equal(comp1t.cnorm(), comp1.cnorm())
comp1t = comp1.ctranspose([1, 2, 0])
assert_array_equal(comp1t.Nmesh, comp1.Nmesh[[1, 2, 0]])
assert_array_equal(comp1t.BoxSize, comp1.BoxSize[[1, 2, 0]])
comp1tt = comp1t.ctranspose([1, 2, 0])
comp1ttt = comp1tt.ctranspose([1, 2, 0])
assert_allclose(comp1ttt, comp1)
def test_cdot_cnorm(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
comp1 = pm.generate_whitenoise(1234, mode='complex')
norm1 = comp1.cdot(comp1)
norm2 = comp1.cnorm()
norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value)))**2).sum()
assert_allclose(norm2, norm3)
# p-2 norm of the full complex matrix (norm2) is
# NORM = Self + Lower + Upper
# CDOT = Self + Lower
# Upper = Lower
self_conj_sum = 0
for ind1 in numpy.ndindex(*(list(comp1.cshape))):
c = [(4 - i) % 4 for i, n in zip(ind1, comp1.cshape)]
ci = numpy.ravel_multi_index(c, comp1.Nmesh)
ii = numpy.ravel_multi_index(ind1, comp1.Nmesh)
v = comp1.cgetitem(ind1)
if ci == ii:
self_conj_sum += abs(v)**2
assert_allclose(norm1 * 2, norm2 + self_conj_sum)
def test_1d(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8], comm=comm, dtype='f8')
real = pm.generate_whitenoise(seed=123, type='real')
complex = pm.generate_whitenoise(seed=123, type='complex')
assert_array_equal(real, complex.c2r())
def test_readout_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, type='real')
def objective(real, pos, layout):
value = real.readout(pos, layout=layout)
obj = (value ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(real, pos, layout, v_real=None, v_pos=None):
value = real.readout(pos, layout=layout)
v_value = real.readout_jvp(pos, v_self=v_real, v_pos=v_pos, layout=layout)
return comm.allreduce((v_value * value * 2).sum())
def backward_gradient(real, pos, layout):
value = real.readout(pos, layout=layout)
return real.readout_vjp(pos, v=value * 2, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += 0.5
pos *= pm.BoxSize / pm.Nmesh
layout = pm.decompose(pos)
obj = objective(real, pos, layout)
grad_real, grad_pos = backward_gradient(real, pos, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(*grad_real.cshape):
dx1, r1 = perturb(real, ind1, dx)
ng1 = (objective(r1, pos, layout) - obj)
bag1 = grad_real.cgetitem(ind1) * dx1
fag1 = forward_gradient(real, pos, layout, v_real=r1 - real)
# print (dx1, dx, ind1, ag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(real, pos1, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(real, pos, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'a', ag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
def test_paint_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, type='real')
def objective(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
obj = (real[...] ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(pos, mass, layout, v_pos=None, v_mass=None):
real = pm.paint(pos, mass=mass, layout=layout)
jvp = pm.paint_jvp(pos, mass=mass, v_mass=v_mass, v_pos=v_pos, layout=layout)
return comm.allreduce((jvp * real * 2)[...].sum())
def backward_gradient(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
return pm.paint_vjp(real * 2, pos, mass=mass, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
numpy.random.seed(9999)
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += (numpy.random.uniform(size=pos.shape)) * 0.8 + 0.1
pos *= pm.BoxSize / pm.Nmesh
mass = numpy.ones(len(pos)) * 2
layout = pm.decompose(pos)
obj = objective(pos, mass, layout)
grad_pos, grad_mass = backward_gradient(pos, mass, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(real.csize):
dx1, mass1, old = perturb_mass(mass, ind1[0], dx, comm)
ng1 = (objective(pos, mass1, layout) - obj)
bag1 = get_mass(grad_mass, ind1[0], comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_mass=mass1 - mass)
# print (dx1, dx, ind1, fag1, bag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(pos1, mass, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'v_pos', (pos1 - pos)[ind1[0]], 'f*g', fag1, 'bag', bag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
def test_1d(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8], comm=comm, dtype='f8')
real = pm.generate_whitenoise(seed=123, type='real')
complex = pm.generate_whitenoise(seed=123, type='complex')
assert_array_equal(real, complex.c2r())