def test_solver_convergence(comm):
from mpi4py import MPI
pm = ParticleMesh(BoxSize=128., Nmesh=[8, 8, 8], comm=comm)
pm_ref = ParticleMesh(BoxSize=128., Nmesh=[8, 8, 8], comm=MPI.COMM_SELF)
Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')
Q = pm_ref.generate_uniform_particle_grid(shift=0)
def solve(pm):
solver = Solver(pm, Planck15, B=1)
q = numpy.array_split(Q, pm.comm.size)[pm.comm.rank]
wn = solver.whitenoise(1234)
dlin = solver.linear(wn, lambda k: Plin(k))
state = solver.lpt(dlin, q, a=0.1, order=2)
state = solver.nbody(state, leapfrog([0.1, 0.5, 1.0]))
d = {}
for key in 'X', 'P', 'F':
d[key] = numpy.concatenate(pm.comm.allgather(getattr(state, key)),
axis=0)
return d
s = solve(pm)
r = solve(pm_ref)
assert_allclose(s['X'], r['X'])
assert_allclose(s['P'], r['P'])
assert_allclose(s['F'], r['F'])
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_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_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 getPow(data1, data2=None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den1 = pm.paint(q + data1.reshape([-1, 3]))
if (data2 is not None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
den2 = pm.paint(q + data2.reshape([-1, 3]))
temp = FFTPower(first=den1, second=den2, mode='1d', BoxSize=128, dk=dk)
k, powspec = temp.power['k'], temp.power['power']
else:
temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk)
k, powspec = temp.power['k'], temp.power['power']
return [k, powspec.real]
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 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
class Test_makemap2(BaseScalarTest):
to_scalar = staticmethod(vmadfastpm.to_scalar)
pm = ParticleMesh(Nmesh=[4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
#x = pm.generate_uniform_particle_grid(shift=0.5)
#x = x[:1,:]
w = np.array([[3., 3.]],
dtype=np.float64) # particle cannot be on grid point
xx = pm.generate_whitenoise(seed=300, unitary=True, type='real')
y = NotImplemented
x = np.ones(len(w))
x_ = create_bases(x)
epsilon = 1e-3
def model(self, x):
compensation = self.pm.resampler.get_compensation()
layout = vmadfastpm.decompose(self.w, self.pm)
map = vmadfastpm.paint(self.w, x, layout, self.pm)
y = map + self.xx # bias needed to avoid zero derivative
# compensation for cic window
c = vmadfastpm.r2c(y)
c = vmadfastpm.apply_transfer(c,
lambda k: compensation(k, 1.0),
kind='circular')
map = vmadfastpm.c2r(c)
return map
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 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))
bs, nc, B = self['boxsize'], self['nc'], self['pm_nc_factor']
ncf = int(nc * B)
self['f_config'] = {
'boxsize':
bs,
'nc':
int(ncf),
'kvec':
fftk(shape=(ncf, ncf, ncf),
boxsize=bs,
symmetric=False,
dtype=self['dtype'])
}
def test_hermitian_weights(comm):
numpy.random.seed(42)
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
cfield = ComplexField(pm)
data = numpy.random.random(size=cfield.shape)
data = data[:] + 1j * data[:]
cfield[...] = data[:]
x = cfield.x
# iterate over symmetry axis
for i, slab in enumerate(SlabIterator(x, axis=2, symmetry_axis=2)):
# nonsingular weights give indices of positive frequencies
nonsig = slab.nonsingular
weights = slab.hermitian_weights
# weights == 2 when iterating frequency is positive
if numpy.float(slab.coords(2)) > 0.:
assert weights > 1
assert numpy.all(nonsig == True)
else:
assert weights == 1.0
assert numpy.all(nonsig == False)
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 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_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_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_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 make_lagfields(nc,
seed,
bs=1536,
T=40,
B=2,
simpath=sc_simpath,
outpath=sc_outpath,
Rsm=0):
fname = get_filename(nc, seed, T=T, B=B)
spath = simpath + fname
opath = outpath + fname
try:
os.makedirs(opath)
except:
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
rank = pm.comm.rank
lin = BigFileMesh(spath + '/linear', 'LinearDensityK').paint()
lin -= lin.cmean()
print(rank, 'lin field read')
header = '1,b1,b2,bg,bk'
names = header.split(',')
lag_fields = tools.getlagfields(
pm, lin, R=Rsm) # use the linear field at the desired redshift
print(rank, 'lag field created')
for i, ff in enumerate(lag_fields):
x = FieldMesh(ff)
x.save(opath + 'lag', dataset=names[i], mode='real')
def plotH1pk(fname, nc=256, fsize=15):
'''plot the power spectrum of halos and H1 on 'nc' grid'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
fig, ax = plt.subplots(1, 2, figsize=(9, 4)) #
for i, aa in enumerate(aafiles):
zz = zzfiles[i]
print(zz)
path = project + sim + '/fastpm_%0.4f/' % aa
k, pkm = np.loadtxt(path + 'pkm.txt').T[0], np.loadtxt(path +
'pkm.txt').T[1]
pkh = np.loadtxt(path + 'pkhmass.txt').T[1]
pkhp = np.loadtxt(path + 'pkhpos.txt').T[1]
pkh1 = np.loadtxt(path + 'pkh1mass.txt').T[1]
pkhm = np.loadtxt(path + 'pkhmassxm.txt').T[1]
pkh1m = np.loadtxt(path + 'pkh1massxm.txt').T[1]
pkhpm = np.loadtxt(path + 'pkhposxm.txt').T[1]
ax[0].plot(k, pkh, label=zz, color='C%d' % i)
ax[0].plot(k, pkhp, label=zz, ls="--", color='C%d' % i)
ax[1].plot(k, pkh1, label=zz, color='C%d' % i) #
ax[1].legend(ncol=2, fontsize=13)
ax[0].set_ylabel('P(k) Halos (Mass/Position)', fontsize=fsize)
ax[1].set_ylabel('P(k) H1 Mass', fontsize=fsize)
for axis in ax:
axis.set_xlim(2e-2, 1)
axis.set_ylim(9e1, 5e4)
axis.loglog()
#axis.grid(which='both', lw=0.5, color='gray')
fig.tight_layout()
fig.savefig(figpath + fname)
def fiddlebias(aa, saveb=False, bname='h1bias.txt', ofolder=None):
#Fiddling bias for halos. Deprecated.
print('Read in catalogs')
halocat = readincatalog(aa, matter=False)
hpos = halocat['Position']
hmass = halocat['Mass']
h1mass = dohod.HI_mass(hmass, aa)
dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint()
if ofolder is None: ofolder = project + '/%s/fastpm_%0.4f/'%(sim, aa)
#measure power
pm = ParticleMesh(BoxSize = bs, Nmesh = [nc, nc, nc])
pkm = FFTPower(dm/dm.cmean(), mode='1d').power
k, pkm = pkm['k'], pkm['power']
##H1
print("H1")
h1mesh = pm.paint(hpos, 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
b1h1 = pkh1m/pkm
b1h1sq = pkh1/pkm
if saveb:
np.savetxt(ofolder+bname, np.stack((k, b1h1, b1h1sq**0.5), axis=1),
header='k, pkh1xm/pkm, pkh1/pkm^0.5')
return k, b1h1, b1h1sq
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_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_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_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_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 createmesh(self, bs, nc, positions, weights):
'''use this to create mesh of HI
'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
mesh = pm.create(mode='real', value=0)
comm = pm.comm
rankweight = sum([wt.sum() for wt in weights])
totweight = comm.allreduce(rankweight)
for wt in weights:
wt /= totweight / float(nc)**3
lay = pm.decompose(positions[0])
mesh.paint(positions[0], mass=weights[0], layout=lay, hold=True)
if len(positions) > 1:
for i in range(self.nsat):
shift = np.random.normal(0, positions[1])
pos = positions[0] + shift
lay = pm.decompose(pos)
mesh.paint(pos,
mass=weights[1] / self.nsat,
layout=lay,
hold=True)
return mesh
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_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_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_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)
