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_ncdm(comm):
pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm)
Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')
solver = Solver(pm, Planck15, B=1)
Q = pm.generate_uniform_particle_grid(shift=0)
wn = solver.whitenoise(1234)
dlin = solver.linear(wn, lambda k: Plin(k))
state = solver.lpt(dlin, Q, a=1.0, order=2)
dnonlin = solver.nbody(state, leapfrog([0.1, 1.0]))
dnonlin.save('nonlin')
def test_solver():
Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')
solver = Solver(pm, Planck15, B=1)
Q = pm.generate_uniform_particle_grid(shift=0)
wn = solver.whitenoise(1234)
dlin = solver.linear(wn, lambda k: Plin(k))
state = solver.lpt(dlin, Q, a=1.0, order=2)
dnonlin = solver.nbody(state, leapfrog([1.0]))
dnonlin.save('nonlin')
def test_solver():
Plin = EHPower(Planck15, redshift=0)
solver = Solver(pm, Planck15, B=2)
Q = pm.generate_uniform_particle_grid()
wn = solver.whitenoise(1234)
dlin = solver.linear(wn, lambda k: Plin(k))
state = solver.lpt(dlin, Q, a=0.3, order=2)
dnonlin = solver.nbody(state, leapfrog([0.3, 0.35]))
dnonlin.save('nonlin')
def gorfpael(stages):
# Reversed Leap-Frog
stack = []
for action, ai, ac, af in list(leapfrog(stages))[::-1]:
if action == 'K':
# need to pop the F before K to ensure the correct force is used.
stack.append((action, af, ac, ai))
elif action == 'F':
yield action, af, ac, ai
for item in stack:
yield item
stack = []
else:
yield action, af, ac, ai
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
def test_solver(comm):
pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm)
solver = Solver(pm, Planck15, B=1)
P_prm = Planck15.Primordial.get_pkprim
tf = get_species_transfer_function_from_class(Planck15, 9)
Q = pm.generate_uniform_particle_grid(shift=0)
wn = solver.whitenoise(1234)
prm = solver.primordial(wn, P_prm)
ic = solver.lpt(prm, {
'0': (Baryon, tf['d_b'], tf['dd_b']),
'1': (CDM, tf['d_cdm'], tf['dd_cdm']),
'4': (NCDM, tf['d_ncdm[0]'], tf['dd_ncdm[0]']),
}, Q, a=0.1)
print('0', ic.species['0'].S[0], ic.species['0'].P[0], ic.species['0'].Q[0])
print('1', ic.species['1'].S[0], ic.species['1'].P[0], ic.species['1'].Q[0])
print('4', ic.species['4'].S[0], ic.species['4'].P[0], ic.species['4'].Q[0])
c2 = CoreSolver(pm, Planck15, B=1)
Pk = lambda k: Planck15.get_pk(k, z=0)
dlin = c2.linear(wn, Pk)
ic2 = c2.lpt(dlin, Q, 0.1, order=1)
print(ic2.S[0], ic2.P[0], ic2.Q[0])
final2 = c2.nbody(ic2, leapfrog([0.1, 1.0]))
final = solver.nbody(ic, leapfrog([0.1, 1.0]))
print('0', final.species['0'].F[0])
print('1', final.species['1'].F[0])
print('4', final.species['4'].F[0])
print(final2.F[0])
final.to_catalog()
def __init__(self,
linear,
astart=0.1,
aend=1.0,
boost=2,
Nsteps=5,
cosmo=None):
self.comm = linear.comm
if cosmo is None:
cosmo = linear.Plin.cosmo
self.cosmo = cosmo
# the linear density field mesh
self.linear = linear
self.attrs.update(linear.attrs)
asteps = numpy.linspace(astart, aend, Nsteps)
self.attrs['astart'] = astart
self.attrs['aend'] = aend
self.attrs['Nsteps'] = Nsteps
self.attrs['asteps'] = asteps
self.attrs['boost'] = boost
solver = Solver(self.linear.pm, cosmology=self.cosmo, B=boost)
Q = self.linear.pm.generate_uniform_particle_grid(shift=0.5)
self.linear = linear
dlin = self.linear.to_field(mode='complex')
state = solver.lpt(dlin, Q, a=astart, order=2)
state = solver.nbody(
state,
leapfrog(numpy.linspace(astart, aend, Nsteps + 1, endpoint=True)))
H0 = 100.
self.RSD = 1.0 / (H0 * aend * self.cosmo.efunc(1.0 / aend - 1))
self._size = len(Q)
CatalogSource.__init__(self, comm=linear.comm, use_cache=False)
self._csize = self.comm.allreduce(self._size)
self['Displacement'] = state.S
self['InitialPosition'] = state.Q
self['ConjugateMomentum'] = state.P # a ** 2 / H0 dx / dt
def gorfpael(stages):
# Reversed Leap-Frog, the verbose way
# this is identical to leapfrog(stages[::-1])
stack = []
# compute the force first as we didn't save it
yield ('F', stages[-1], stages[-1], stages[-1])
for action, ai, ac, af in list(leapfrog(stages))[::-1]:
if action == 'D':
yield action, af, ac, ai
for action, af, ac, ai in stack:
assert action == 'F'
yield action, af, ac, ai
stack = []
elif action == 'F':
# need to do F after D to ensure the time tag is right.
assert ac == af
stack.append((action, af, ai, ai))
else:
yield action, af, ac, ai
def func_gal_catalogue(bs, nc, seed, nstep, seed_hod, Omega_m, p_alpha,
p_logMin, p_logM1, p_logM0, p_sigma_logM):
folder = "L%04d_N%04d_S%04d_%02dstep" % (bs, nc, seed, nstep)
# setup initial conditions
Omegacdm = Omega_m - 0.049,
cosmo = cosmology.Planck15.clone(Omega_cdm=Omegacdm,
h=0.6711,
Omega_b=0.049)
power = cosmology.LinearPower(cosmo, 0)
klin = np.logspace(-4, 2, 1000)
plin = power(klin)
pkfunc = interpolate(klin, plin)
# run the simulation
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
Q = pm.generate_uniform_particle_grid()
stages = numpy.linspace(0.1, 1.0, nstep, endpoint=True)
solver = Solver(pm, cosmo, B=2)
wn = solver.whitenoise(seed)
dlin = solver.linear(wn, pkfunc)
state = solver.lpt(dlin, Q, stages[0])
state = solver.nbody(state, leapfrog(stages))
# create the catalogue
cat = ArrayCatalog(
{
'Position': state.X,
'Velocity': state.V,
'Displacement': state.S,
'Density': state.RHO
},
BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh,
M0=Omega_m * 27.75e10 * bs**3 / (nc / 2.0)**3)
cat['KDDensity'] = KDDensity(cat).density
cat.save('%s/Matter' % (folder),
('Position', 'Velocity', 'Density', 'KDDensity'))
# run FOF
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.to_halos(particle_mass=cat.attrs['M0'],
cosmo=cosmo,
redshift=0.0)
fofcat.save('%s/FOF' % (folder),
('Position', 'Velocity', 'Mass', 'Radius'))
# run HOD
params = {
'alpha': p_alpha,
'logMmin': p_logMin,
'logM1': p_logM1,
'logM0': p_logM0,
'sigma_logM': p_sigma_logM
}
halos = HaloCatalog(fofcat, cosmo=cosmo, redshift=0.0, mdef='vir')
halocat = halos.to_halotools(halos.attrs['BoxSize'])
hod = HODCatalog(halocat, seed=seed_hod, **params)
hod.save('%s/HOD' % (folder), ('Position', 'Velocity'))
return folder, cat, fofcat, hod
dx1, dx2 = model.compute(['dx1', 'dx2'], init=dict(x=x))
print('model', dx1.std(axis=0), dx2.std(axis=0))
print('fastpm', dx1_f.std(axis=0), dx2_f.std(axis=0))
print('model', dx1[0], dx2[0])
print('fastpm', dx1_f[0], dx2_f[0])
print('comparing lpt dx and p ')
dx0, p0 = model.compute(['dx0', 'p0'], init=dict(x=x))
print('model', dx0.std(axis=0), p0.std(axis=0))
print('fastpm', lpt.S.std(axis=0), lpt.P.std(axis=0))
print('comparing single step')
lpt = solver.lpt(linear, Q=q, a=0.1)
state = solver.nbody(lpt, leapfrog([0.1]))
dx0, p0, f0 = model.compute(['dx0', 'p0', 'f0'], init=dict(x=x))
print('model', dx0.std(axis=0), p0.std(axis=0), f0.std(axis=0))
print('fastpm', state.S.std(axis=0), state.P.std(axis=0), state.F.std(axis=0))
print('model', dx0[0], p0[0], f0[0])
print('fastpm', state.S[0], state.P[0], state.F[0])
print('comparing multi step')
dx, p, f = model.compute(['dx', 'p', 'f'], init=dict(x=x))
lpt = solver.lpt(linear, Q=q, a=0.1)
state = solver.nbody(lpt, leapfrog([0.1, 0.6, 1.0]))
print('model', dx.std(axis=0), p.std(axis=0), f.std(axis=0))
print('fastpm', state.S.std(axis=0), state.P.std(axis=0), state.F.std(axis=0))
print('model', dx[0], p[0], f[0])
solver = Solver(pm, Planck15, B=2)
solver_ncdm = SolverNCDM(pm, Planck15, B=2)
wn = solver.whitenoise(400)
dlin = solver.linear(wn, EHPower(Planck15, 0))
lpt = solver.lpt(dlin, Q, stages[0])
#lpt.S = numpy.float32(lpt.S)
def monitor(action, ai, ac, af, state, event):
if pm.comm.rank == 0:
print(state.a['S'], state.a['P'], state.a['F'], state.S[0], state.P[0],
action, ai, ac, af)
state1 = solver.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
state2 = solver_ncdm.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
if pm.comm.rank == 0:
print('----------------')
cat1 = ArrayCatalog({'Position': state1.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat2 = ArrayCatalog({'Position': state2.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
r1 = FFTPower(cat1, mode='1d')
r2 = FFTPower(cat2, mode='1d')
cat1.save('NCDM/core', ('Position', ))
cat2.save('NCDM/ncdm', ('Position', ))
r1.save('NCDM/core-power.json')
r2.save('NCDM/ncdm-power.json')
lpt = solver.lpt(dlin, Q, stages[0], order=2)
def monitor(action, ai, ac, af, state, event):
if pm.comm.rank == 0:
print(state.a['S'], state.a['P'], state.a['F'], state.S[0], state.P[0],
action, ai, ac, af)
def monitor_multi(action, ai, ac, af, state, event):
if pm.comm.rank == 0:
print(state.a['S'], state.a['P'], state.a['F'], state['1'].S[0],
state['1'].P[0], action, ai, ac, af)
state1 = solver.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
tf = get_species_transfer_function_from_class(Planck15, 1.0 / stages[0] - 1)
prm = solver_multi.primordial(wn, Planck15.Primordial.get_pk)
ic = solver_multi.lpt(prm, {
'0': (Baryon, tf['d_b'], tf['dd_b']),
'1': (CDM, tf['d_cdm'], tf['dd_cdm']),
'4': (NCDM, tf['d_ncdm[0]'], tf['dd_ncdm[0]']),
},
Q,
a=stages[0],
order=2)
state2 = solver_multi.nbody(ic.copy(), leapfrog(stages), monitor=monitor_multi)
for action, ai, ac, af in list(leapfrog(stages))[::-1]:
if action == 'D':
yield action, af, ac, ai
for action, af, ac, ai in stack:
assert action == 'F'
yield action, af, ac, ai
stack = []
elif action == 'F':
# need to do F after D to ensure the time tag is right.
assert ac == af
stack.append((action, af, ai, ai))
else:
yield action, af, ac, ai
print(list(leapfrog(stages)))
print('----')
print(list(gorfpael(stages)))
print('++++')
print(list(leapfrog(stages[::-1])))
state = solver.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
if pm.comm.rank == 0:
print('----------------')
reverse = solver.nbody(state, leapfrog(stages[::-1]), monitor=monitor)
#print((lpt.X - reverse.X).max())
assert_allclose(lpt.X, reverse.X)
cat1 = ArrayCatalog({'Position' : lpt.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat2 = ArrayCatalog({'Position' : reverse.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
def gorfpael(stages):
# Reversed Leap-Frog
stack = []
for action, ai, ac, af in list(leapfrog(stages))[::-1]:
if action == 'K':
# need to pop the F before K to ensure the correct force is used.
stack.append((action, af, ac, ai))
elif action == 'F':
yield action, af, ac, ai
for item in stack:
yield item
stack = []
else:
yield action, af, ac, ai
state = solver.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
if pm.comm.rank == 0:
print('----------------')
reverse = solver.nbody(state, gorfpael(stages), monitor=monitor)
#print((lpt.X - reverse.X).max())
assert_allclose(lpt.X, reverse.X)
cat1 = ArrayCatalog({'Position' : lpt.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat2 = ArrayCatalog({'Position' : reverse.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat1.save('Janus/Truth', ('Position',))
cat2.save('Janus/Reverse', ('Position',))
V = fac1 * V0 + fac2 * state.V
#save the snapshot
cat = ArrayCatalog({
'Position': X,
'Velocity': V
},
BoxSize=state.pm.BoxSize)
cat.save(
args.save + '/FastPM_Nmesh%d_Nstep%d_z%.2f' %
(args.Nmesh, args.Nstep, 1. / a - 1.),
('Position', 'Velocity'))
del X, V
if state.pm.comm.rank == 0:
print(
'Finish writing snapshot at redshift %.2f' % (1. / a - 1.),
'Time:',
time.time() - t)
a0 = np.array(af)
X0 = state.X
V0 = state.V
if state.pm.comm.rank == 0:
print('Finish redshift %.2f' % (1. / af - 1.), 'Time:',
time.time() - t)
#run FastPM
state = solver.nbody(state, leapfrog(stages), monitor=monitor)
def test_leapfrog():
l = list(leapfrog(numpy.linspace(0.1, 1.0, 2, endpoint=True)))
assert len(l) == 1 + 4 * (2 - 1)
l = list(leapfrog([1.0]))
assert len(l) == 1
