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 get_lpt(pm, z, cosmology, seed):
"""Evolve the linear power using a 2LPT solver,
so we get a good model of the density structure at the reionization redshift."""
a = 1 / (1 + z)
Plin = LinearPower(cosmology, redshift=0, transfer='EisensteinHu')
solver = Solver(pm, cosmology, B=1)
Q = pm.generate_uniform_particle_grid()
wn = solver.whitenoise(seed)
dlin = solver.linear(wn, Plin)
state = solver.lpt(dlin, Q, a=a, order=2)
return state
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 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_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 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 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_lpt():
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))
state1 = solver.lpt(dlin, Q, a=0.01, order=1)
state2 = solver.lpt(dlin, Q, a=1.0, order=1)
pt = MatterDominated(Planck15.Om0, a=[0.01, 1.0], a_normalize=1.0)
# print((state2.P[...] / state1.P[...]))
print((state2.P[...] - state1.P[...]) / state1.F[...])
fac = 1 / (0.01**2 * pt.E(0.01)) * (pt.Gf(1.0) - pt.Gf(0.01)) / pt.gf(0.01)
assert_allclose(state2.P[...], state1.P[...] + fac * state1.F[...])
from nbodykit.algorithms.fof import FOF
from nbodykit.algorithms.fftpower import FFTPower
from nbodykit.source import ArrayCatalog
from pmesh.pm import ParticleMesh
import numpy
pm = ParticleMesh(BoxSize=512,
Nmesh=[256, 256, 256],
dtype='f8',
resampler='tsc')
Q = pm.generate_uniform_particle_grid()
stages = numpy.linspace(0.1, 1.0, 20, endpoint=True)
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)
from nbodykit.algorithms.fftpower import FFTPower
from nbodykit.source import ArrayCatalog, MultipleSpeciesCatalog
from fastpm.multi import get_species_transfer_function_from_class
from pmesh.pm import ParticleMesh
import numpy
Planck15 = Planck15.clone(gauge='newtonian')
pm = ParticleMesh(BoxSize=512, Nmesh=[64, 64, 64], dtype='f4', resampler='tsc')
Q = pm.generate_uniform_particle_grid()
stages = numpy.linspace(0.1, 1.0, 10, endpoint=True)
#stages = [1.]
solver = Solver(pm, Planck15, B=2)
solver_multi = SolverMulti(pm, Planck15, B=2)
wn = solver.whitenoise(400, unitary=True)
dlin = solver.linear(wn, LinearPower(Planck15, 0))
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:
k[mask] = 1.
delta = np.interp(np.log10(k), logk, logpower)
power = norm * (10**delta) * (2 * np.pi**2) / k**3
power[mask] = 0.
k[mask] = 0.
return power
#time steps in scale factor
stages = np.linspace(0.1, 1., args.Nstep, endpoint=True)
a_output = 1. / (np.array(args.output_redshift) + 1.)
#IC 2LPT
t = time.time()
solver_IC = Solver(pm_IC, Planck15, B=2)
wn = solver_IC.whitenoise(2695896)
dlin = solver_IC.linear(wn, Power)
Q = pm.generate_uniform_particle_grid(shift=0)
solver = Solver(pm, Planck15, B=2)
state = solver.lpt(dlin, Q, stages[0], order=2)
if pm.comm.rank == 0:
print('Finish generating initial conditions with 2LPT. Time:',
time.time() - t)
X0 = state.X
V0 = state.V
a0 = np.array(stages[0])
RG = args.RG
Ng = args.Ng
Lbox = args.Lbox
print("Seed = ", Rdm_seed)
print("Lbox = %.1f" % Lbox, "Ng = %d" % Ng, "RG = %.2f" % RG)
# Choose cosmology
# -------------------------------------------
cosmology = nbcosmos.WMAP9
mycosmo = Cosmos(FLRW=True, obj=cosmology)
# generate linear density field at z=0
# -------------------------------------------
pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng, Ng, Ng])
Q = pm.generate_uniform_particle_grid(shift=0)
solver = Solver(pm, cosmology, B=1)
wn = solver.whitenoise(seed=Rdm_seed)
dlin = solver.linear(wn, lambda k: mycosmo.Pk_lin(k))
dx_field = dlin.c2r(
).value # dx_field is the density contrast field centered at 0
# initialize gsCR object, build xij^{-1} matrix for full 18 constraints
# -------------------------------------------
fg = gsCR(mycosmo, Lbox=Lbox, Nmesh=Ng, RG=RG, CONS=['full'])
fg.build_Xij_inv_matrix()
print("xij^{-1}:")
print(fg.xij_tensor_inv)
print("*********************************************")
# set peak position
def saveIC_bt2(IC_path, dx_field, Lbox, Ng, cosmology, redshift=99):
"""
Use Zel-dovich approximation to back-scale the linear density field to initial redshift,
paint baryon and dm particles from the grid,
and save initial condition in MP-Gadget/bluetides-ii format
Parameters
---------
: dx_field : linear density field at z=0
: Lbox : BoxSize, in Mpc/h, will be converted to kpc/h in the IC output
: Ng : Number of grid on which to paint the particle
: cosmology : nbodykit.cosmology, or astropy.cosmology.FLRW
: redshift : redshift of the initial condition
"""
mesh = ArrayMesh(dx_field,
BoxSize=Lbox) # density contrast field centered at zero
dk_field = mesh.compute(mode='complex')
shift_gas = -0.5 * (cosmology.Om0 - cosmology.Ob0) / cosmology.Om0
shift_dm = 0.5 * cosmology.Ob0 / cosmology.Om0
pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng, Ng, Ng])
solver = Solver(pm, cosmology, B=1)
Q_gas = pm.generate_uniform_particle_grid(shift=shift_gas)
Q_dm = pm.generate_uniform_particle_grid(shift=shift_dm)
scale_a = 1. / (1 + redshift)
state_gas = solver.lpt(dk_field, Q_gas, a=scale_a,
order=1) # order = 1 : Zel-dovich
state_dm = solver.lpt(dk_field, Q_dm, a=scale_a, order=1)
state = state_dm
with FileMPI(state.pm.comm, IC_path, create=True) as ff:
m0 = state.cosmology.rho_crit(0) * state.cosmology.Omega0_b * (
Lbox**3) / state.csize
m1 = state.cosmology.rho_crit(0) * (
state.cosmology.Om0 - state.cosmology.Omega0_b) * (Lbox**
3) / state.csize
with ff.create('Header') as bb:
bb.attrs['BoxSize'] = Lbox * 1000
bb.attrs['HubbleParam'] = state.cosmology.h
bb.attrs['MassTable'] = [m0, m1, 0, 0, 0, 0]
bb.attrs['OmegaM'] = state.cosmology.Om0
bb.attrs['OmegaB'] = state.cosmology.Omega0_b
bb.attrs['OmegaL'] = state.cosmology.Omega0_lambda
bb.attrs['Time'] = state.a['S']
bb.attrs['TotNumPart'] = [state.csize, state.csize, 0, 0, 0, 0]
ff.create_from_array('1/Position',
1000 * periodic_wrap(state.X, Lbox)) # in kpc/h
ff.create_from_array(
'1/Velocity',
state.V / np.sqrt(state.a['S'])) # old gadget convention for IC
dmID = np.arange(state.csize)
ff.create_from_array('1/ID', dmID)
#######################################################
ff.create_from_array('0/Position',
1000 * periodic_wrap(state_gas.X, Lbox))
ff.create_from_array('0/Velocity', state_gas.V / np.sqrt(state.a['S']))
gasID = np.arange(state.csize, 2 * state.csize)
ff.create_from_array('0/ID', gasID)
print("IC generated!")
print("*********************************************")
return state
wnk = fastpm.as_complex_field(x, pm)
rhok = fastpm.induce_correlation(wnk, powerspectrum, pm)
dx1, dx2 = fastpm.lpt(rhok, q, pm)
dx, p, f = fastpm.nbody(rhok, q, [0.1, 0.6, 1.0], Planck15, pm)
dx0, p0, f0 = fastpm.nbody(rhok, q, [0.1], Planck15, pm)
model.output(dx1=dx1, dx2=dx2, dx=dx, p=p, dx0=dx0, p0=p0, f0=f0, f=f)
wn = pm.generate_whitenoise(555, unitary=True)
x = wn[...]
x = numpy.stack([x.real, x.imag], -1)
from fastpm.core import Solver, leapfrog
solver = Solver(pm, Planck15, B=1)
linear = solver.linear(wn, powerspectrum)
dx1_f = lpt1(linear, q)
dx2_f = lpt1(lpt2source(linear), q)
lpt = solver.lpt(linear, Q=q, a=0.1)
print('comparing lpt order by order')
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])
Lbox = args.Lbox
print ("Seed = ",Rdm_seed)
print ("Lbox = %.1f"%Lbox,"Ng = %d"%Ng,"RG = %.2f"%RG)
# Choose cosmology
# -------------------------------------------
cosmology = nbcosmos.WMAP9
mycosmo = Cosmos(FLRW=True,obj=cosmology)
# generate linear density field at z=0
# -------------------------------------------
pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng,Ng,Ng])
Q = pm.generate_uniform_particle_grid(shift=0)
solver = Solver(pm,cosmology,B=1)
wn = solver.whitenoise(seed = Rdm_seed)
dlin = solver.linear(wn, lambda k: mycosmo.Pk_lin(k))
dx_field = dlin.c2r().value # dx_field is the density contrast field centered at 0
# initialize gsCR object, build xij^{-1} matrix for full 18 constraints
# -------------------------------------------
fg = gsCR(mycosmo,Lbox=Lbox,Nmesh=Ng,RG=RG,CONS=['full'])
fg.build_Xij_inv_matrix()
print ("xij^{-1}:")
print (fg.xij_tensor_inv)
print ("*********************************************")
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
