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 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_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_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_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_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_fof_parallel_no_merge(comm):
CurrentMPIComm.set(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)
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_fof_fully_connected(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(BoxSize=[4, 4, 4], Nmesh=[4, 4, 4], comm=comm)
Q = pm.generate_uniform_particle_grid(shift=0)
cat = ArrayCatalog({'Position': numpy.concatenate([Q], axis=0)},
BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh,
comm=comm)
fof = FOF(cat, linking_length=2, nmin=63, absolute=True)
labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0)
assert_array_equal(labels, 1)
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_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 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 getPow_dis(data1, data2=None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den1 = pm.paint(q)
power = 0
if (data2 is not None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den2 = pm.paint(q)
for ii in range(3):
den1[:] = data1[:, :, :, ii]
den2[:] = data2[:, :, :, ii]
temp = FFTPower(first=den1,
second=den2,
mode='1d',
BoxSize=128,
dk=dk)
k, power = temp.power['k'], power + temp.power['power']
else:
for ii in range(3):
den1[:] = data1[:, :, :, ii]
temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk)
k, power = temp.power['k'], power + temp.power['power']
return [k, power.real]
def test_grid_shifted(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
grid = grid + 4.0
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
g2 = layout.exchange(grid)
real = pm.paint(grid, layout=layout)
#print(real, g2 % 8, real.slices)
assert_allclose(real, 1.0)
grid = grid - 6.1
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
real = pm.paint(grid, layout=layout)
assert_allclose(real, 1.0)
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_grid_shifted(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
grid = grid + 4.0
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
g2 = layout.exchange(grid)
real = pm.paint(grid, layout=layout)
#print(real, g2 % 8, real.slices)
assert_allclose(real, 1.0)
grid = grid - 6.1
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
real = pm.paint(grid, layout=layout)
assert_allclose(real, 1.0)
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_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 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_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_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_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 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_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 Halos(bs, nc, seed, nstep, seed_hod, Omega_m, p_alpha, p_logMin, p_logM1, p_logM0, p_sigma_logM):
'''
'''
# setup initial conditions
Omegacdm = Omega_m-0.049,
cosmo = NBlab.cosmology.Planck15.clone(Omega_cdm=Omegacdm,
h = 0.6711, Omega_b = 0.049)
power = NBlab.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
# run FOF to construct Halo catalog
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.to_halos(particle_mass=cat.attrs['M0'], cosmo = cosmo, redshift = 0.0)
return fofcat
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 ("*********************************************")
from fastpm.core import Solver as Solver
from fastpm.core import leapfrog
from nbodykit.cosmology import Planck15, EHPower
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],
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])
def monitor(action, ai, ac, af, state, event):
if not state.synchronized: return
def test_grid(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
real = pm.paint(grid)
assert_array_equal(real, 1.0)
def main():
"""
Convert grids4plots 3d grids to 2d slices.
"""
#####################################
# PARSE COMMAND LINE ARGS
#####################################
ap = ArgumentParser()
ap.add_argument('--inbasepath',
type=str,
default='$SCRATCH/perr/grids4plots/',
help='Input base path.')
ap.add_argument('--outbasepath',
type=str,
default='$SCRATCH/perr/grids4plots/',
help='Output base path.')
ap.add_argument(
'--inpath',
type=str,
default=
'main_calc_Perr_2020_Sep_22_18:44:31_time1600800271.dill', # laptop
#default='main_calc_Perr_2020_Aug_26_02:49:57_time1598410197.dill', # cluster
help='Input path.')
ap.add_argument('--Rsmooth',
type=float,
default=2.0,
help='3D Gaussian smoothing applied to field.')
# min and max index included in output. inclusive.
ap.add_argument('--ixmin',
type=int,
default=5,
help='xmin of output. must be between 0 and Ngrid.')
ap.add_argument('--ixmax', type=int, default=5, help='xmax of output')
ap.add_argument('--iymin', type=int, default=0, help='ymin of output')
ap.add_argument('--iymax', type=int, default=-1, help='ymax of output')
ap.add_argument('--izmin', type=int, default=0, help='zmin of output')
ap.add_argument('--izmax', type=int, default=-1, help='zmax of output')
cmd_args = ap.parse_args()
verbose = True
#####################################
# START PROGRAM
#####################################
comm = CurrentMPIComm.get()
rank = comm.rank
path = os.path.join(os.path.expandvars(cmd_args.inbasepath),
cmd_args.inpath)
if rank == 0:
print('path: ', path)
initialized_slicecat = False
for fname in os.listdir(path):
if fname.startswith('SLICE'):
continue
full_fname = os.path.join(path, fname)
print('%d Reading %s' % (rank, full_fname))
inmesh = BigFileMesh(full_fname,
dataset='tmp4storage',
header='header')
Ngrid = inmesh.attrs['Ngrid']
boxsize = inmesh.attrs['boxsize']
# apply smoothing
mesh = apply_smoothing(inmesh, mode='Gaussian', R=cmd_args.Rsmooth)
del inmesh
# convert indices to modulo ngrid
ixmin = cmd_args.ixmin % Ngrid
ixmax = cmd_args.ixmax % Ngrid
iymin = cmd_args.iymin % Ngrid
iymax = cmd_args.iymax % Ngrid
izmin = cmd_args.izmin % Ngrid
izmax = cmd_args.izmax % Ngrid
# convert to boxsize units (Mpc/h)
xmin = float(ixmin) / float(Ngrid) * boxsize
xmax = float(ixmax + 1) / float(Ngrid) * boxsize
ymin = float(iymin) / float(Ngrid) * boxsize
ymax = float(iymax + 1) / float(Ngrid) * boxsize
zmin = float(izmin) / float(Ngrid) * boxsize
zmax = float(izmax + 1) / float(Ngrid) * boxsize
if not initialized_slicecat:
# Generate catalog with positions of slice points, to readout mesh there.
# First generate all 3D points. Then keep only subset in slice.
# THen readout mesh at those points.
# use pmesh generate_uniform_particle_grid
# http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight=
# readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid
partmesh = ParticleMesh(BoxSize=boxsize,
Nmesh=[Ngrid, Ngrid, Ngrid])
ptcles = partmesh.generate_uniform_particle_grid(shift=0.0,
dtype='f8')
#print("type ptcles", type(ptcles), ptcles.shape)
#print("head ptcles:", ptcles[:5,:])
dtype = np.dtype([('Position', ('f8', 3))])
# number of rows is given by number of ptcles on this rank
uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype)
uni_cat_array['Position'] = ptcles
uni_cat = ArrayCatalog(uni_cat_array,
comm=None,
BoxSize=boxsize * np.ones(3),
Nmesh=[Ngrid, Ngrid, Ngrid])
del ptcles
del uni_cat_array
print("%d: Before cut: local Nptcles=%d, global Nptcles=%d" %
(comm.rank, uni_cat.size, uni_cat.csize))
# only keep points in the slice
uni_cat = uni_cat[(uni_cat['Position'][:, 0] >= xmin)
& (uni_cat['Position'][:, 0] < xmax)
& (uni_cat['Position'][:, 1] >= ymin)
& (uni_cat['Position'][:, 1] < ymax)
& (uni_cat['Position'][:, 2] >= zmin)
& (uni_cat['Position'][:, 2] < zmax)]
print("%d: After cut: local Nptcles=%d, global Nptcles=%d" %
(comm.rank, uni_cat.size, uni_cat.csize))
initialized_slicecat = True
# read out full 3D mesh at catalog positions. this is a numpy array
slicecat = readout_mesh_at_cat_pos(mesh=mesh,
cat=uni_cat,
readout_window='nearest')
if rank == 0:
print('slicecat type:', type(slicecat))
slicecat = GatherArray(slicecat, comm, root=0)
if rank == 0:
if not slicecat.shape == ((ixmax - ixmin + 1) *
(iymax - iymin + 1) *
(izmax - izmin + 1), ):
raise Exception(
'Unexpected shape of particles read out on slice: %s' %
str(slicecat.shape))
slicecat = slicecat.reshape(
(ixmax - ixmin + 1, iymax - iymin + 1, izmax - izmin + 1))
print('slicecat shape:', slicecat.shape)
if verbose:
print('slicecat:', slicecat)
# convert to a mesh. assume full numpy array sits on rank 0.
Lx = xmax - xmin
Ly = ymax - ymin
Lz = zmax - zmin
if Lx == 0.: Lx = boxsize / float(Ngrid)
if Ly == 0.: Ly = boxsize / float(Ngrid)
if Lz == 0.: Lz = boxsize / float(Ngrid)
BoxSize_slice = np.array([Lx, Ly, Lz])
slicemesh = ArrayMesh(slicecat, BoxSize=BoxSize_slice, root=0)
outshape = slicemesh.compute(mode='real').shape
if verbose:
print('slicemesh: ', slicemesh.compute(mode='real'))
# write to disk
outpath = os.path.join(
os.path.expandvars(cmd_args.outbasepath), cmd_args.inpath,
'SLICE_R%g_%d-%d_%d-%d_%d-%d/' %
(cmd_args.Rsmooth, ixmin, ixmax, iymin, iymax, izmin, izmax))
if rank == 0:
if not os.path.exists(outpath):
os.makedirs(outpath)
full_outfname = os.path.join(outpath, fname)
if rank == 0:
print('Writing %s' % full_outfname)
slicemesh.save(full_outfname)
if rank == 0:
print('Wrote %s' % full_outfname)
def weigh_and_shift_uni_cat(delta_for_weights,
displacements,
Nptcles_per_dim,
out_Ngrid,
boxsize,
internal_scale_factor_for_weights=None,
out_scale_factor=None,
cosmo_params=None,
weighted_CIC_mode=None,
uni_cat_generator='pmesh',
plot_slices=False,
verbose=False,
return_catalog=False):
"""
Make uniform catalog, weigh by delta_for_weights, displace by displacements
and interpolate to grid which we return.
Parameters
----------
delta_for_weights : None or pmesh.pm.RealField object
Particles are weighted by 1+delta_for_weights. If None, use weight=1.
displacements : list
[Psi_x, Psi_y, Psi_z] where Psi_i are pmesh.pm.RealField objects,
holding the displacement field in different directions (on the grid).
If None, do not shift.
uni_cat_generator : string
If 'pmesh', use pmesh for generating uniform catalog.
If 'manual', use an old serial code.
Returns
-------
delta_shifted : FieldMesh object
Density delta_shifted of shifted weighted particles (normalized to mean
of 1 i.e. returning 1+delta). Returned if return_catalog=False.
attrs : meshsource attrs
Attrs of delta_shifted. Returned if return_catalog=False.
catalog_shifted : ArrayCatalog object
Shifted catalog, returned if return_catalog=True.
"""
comm = CurrentMPIComm.get()
# ######################################################################
# Generate uniform catalog with Nptcles_per_dim^3 particles on regular
# grid
# ######################################################################
if uni_cat_generator == 'pmesh':
# use pmesh generate_uniform_particle_grid
# http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight=
# readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid
pmesh = ParticleMesh(
BoxSize=boxsize,
Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim])
ptcles = pmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8')
#print("type ptcles", type(ptcles), ptcles.shape)
#print("head ptcles:", ptcles[:5,:])
dtype = np.dtype([('Position', ('f8', 3))])
# number of rows is given by number of ptcles on this rank
uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype)
uni_cat_array['Position'] = ptcles
uni_cat = ArrayCatalog(
uni_cat_array,
comm=None,
BoxSize=boxsize * np.ones(3),
Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim])
print("%d: local Nptcles=%d, global Nptcles=%d" %
(comm.rank, uni_cat.size, uni_cat.csize))
del ptcles
del uni_cat_array
elif uni_cat_generator == 'manual':
# Old serial code to generate regular grid and scatter across ranks.
if comm.rank == 0:
# Code copied from do_rec_v1.py and adopted
# Note that nbkit UniformCatalog is random catalog, but we want a catalog
# where each ptcle sits at grid points of a regular grid.
# This is what we call 'regular uniform' catalog.
Np = Nptcles_per_dim
dtype = np.dtype([('Position', ('f8', 3))])
# Have Np**3 particles, and each particle has position x,y,z and weight 'Weight'
uni_cat_array = np.empty((Np**3, ), dtype=dtype)
# x components in units such that box ranges from 0 to 1. Note dx=1/Np.
#x_components_1d = np.linspace(0.0, (Np-1)*(L/float(Np)), num=Np, endpoint=True)/L
x_components_1d = np.linspace(0.0, (Np - 1) / float(Np),
num=Np,
endpoint=True)
ones_1d = np.ones(x_components_1d.shape)
# Put particles on the regular grid
print("%d: Fill regular uniform catalog" % comm.rank)
uni_cat_array['Position'][:,
0] = np.einsum('a,b,c->abc',
x_components_1d, ones_1d,
ones_1d).reshape(
(Np**3, ))
uni_cat_array['Position'][:,
1] = np.einsum('a,b,c->abc', ones_1d,
x_components_1d,
ones_1d).reshape(
(Np**3, ))
uni_cat_array['Position'][:,
2] = np.einsum('a,b,c->abc', ones_1d,
ones_1d,
x_components_1d).reshape(
(Np**3, ))
print("%d: Done filling regular uniform catalog" % comm.rank)
# in nbkit0.3 units must be in Mpc/h
uni_cat_array['Position'] *= boxsize
else:
uni_cat_array = None
# Scatter across all ranks
print("%d: Scatter array" % comm.rank)
from nbodykit.utils import ScatterArray
uni_cat_array = ScatterArray(uni_cat_array, comm, root=0, counts=None)
print("%d: Scatter array done. Shape: %s" %
(comm.rank, str(uni_cat_array.shape)))
# Save in ArrayCatalog object
uni_cat = ArrayCatalog(uni_cat_array)
uni_cat.attrs['BoxSize'] = np.ones(3) * boxsize
uni_cat.attrs['Nmesh'] = np.ones(3) * Nptcles_per_dim
uni_cat.attrs['Nmesh_internal'] = np.ones(3) * Nmesh_orig
else:
raise Exception('Invalid uni_cat_generator %s' % uni_cat_generator)
########################################################################
# Set weight of particles in uni_cat to delta (interpolated to ptcle
# positions)
########################################################################
if delta_for_weights is None:
# set all weights to 1
uni_cat['Mass'] = np.ones(uni_cat['Position'].shape[0])
else:
# weight by delta_for_weights
nbkit03_utils.interpolate_pm_rfield_to_catalog(
delta_for_weights, uni_cat, catalog_column_to_save_to='Mass')
print("%d: rms Mass: %g" %
(comm.rank, np.sqrt(np.mean(np.array(uni_cat['Mass'])**2))))
# optionally plot weighted uniform cat before shifting
if plot_slices:
# paint the original uni_cat to a grid and plot slice
import matplotlib.pyplot as plt
tmp_meshsource = uni_cat.to_mesh(Nmesh=out_Ngrid,
value='Mass',
window='cic',
compensated=False,
interlaced=False)
# paint to get delta(a_internal)
tmp_outfield = tmp_meshsource.paint(mode='real')
# linear rescale factor from internal_scale_factor_for_weights to
# out_scale_factor
rescalefac = nbkit03_utils.linear_rescale_fac(
internal_scale_factor_for_weights,
out_scale_factor,
cosmo_params=cosmo_params)
tmp_outfield = 1.0 + rescalefac * (tmp_outfield - 1.0)
tmp_mesh = FieldMesh(tmp_outfield)
plt.imshow(tmp_mesh.preview(Nmesh=32, axes=(0, 1)))
if comm.rank == 0:
plt_fname = 'inmesh_Np%d_Nm%d_Ng%d.pdf' % (Nptcles_per_dim,
Nmesh_orig, out_Ngrid)
plt.savefig(plt_fname)
print("Made %s" % plt_fname)
del tmp_meshsource, rescalefac, tmp_outfield, tmp_mesh
# ######################################################################
# Shift uniform catalog particles by Psi (changes uni_cat)
# ######################################################################
nbkit03_utils.shift_catalog_by_psi_grid(
cat=uni_cat,
in_displacement_rfields=displacements,
pos_column='Position',
pos_units='Mpc/h',
displacement_units='Mpc/h',
boxsize=boxsize,
verbose=verbose)
#del Psi_rfields
if return_catalog:
# return shifted catalog
return uni_cat
else:
# return density of shifted catalog, delta_shifted
# ######################################################################
# paint shifted catalog to grid, using field_to_shift as weights
# ######################################################################
print("%d: paint shifted catalog to grid using mass weights" %
comm.rank)
# this gets 1+delta
if weighted_CIC_mode == 'sum':
delta_shifted, attrs = paint_utils.weighted_paint_cat_to_delta(
uni_cat,
weight='Mass',
Nmesh=out_Ngrid,
weighted_paint_mode=weighted_CIC_mode,
normalize=True, # compute 1+delta
verbose=verbose,
to_mesh_kwargs={
'window': 'cic',
'compensated': False,
'interlaced': False
})
# this get rho
elif weighted_CIC_mode == 'avg':
delta_shifted, attrs = paint_utils.mass_avg_weighted_paint_cat_to_rho(
uni_cat,
weight='Mass',
Nmesh=out_Ngrid,
verbose=verbose,
to_mesh_kwargs={
'window': 'cic',
'compensated': False,
'interlaced': False
})
else:
raise Exception('Invalid weighted_CIC_mode %s' % weighted_CIC_mode)
# ######################################################################
# rescale to output redshift
# ######################################################################
if internal_scale_factor_for_weights != out_scale_factor:
# linear rescale factor from internal_scale_factor_for_weights to
# out_scale_factor
rescalefac = nbkit03_utils.linear_rescale_fac(
internal_scale_factor_for_weights,
out_scale_factor,
cosmo_params=cosmo_params)
delta_shifted *= rescalefac
# print some info:
if comm.rank == 0:
print(
"%d: Linear rescalefac from a=%g to a=%g, rescalefac=%g" %
(comm.rank, internal_scale_factor_for_weights,
out_scale_factor, rescalefac))
raise Exception(
'Check if rescaling of delta_shifted is correct. Looks like 1+delta.'
)
if verbose:
print("%d: delta_shifted: min, mean, max, rms(x-1):" % comm.rank,
np.min(delta_shifted), np.mean(delta_shifted),
np.max(delta_shifted),
np.mean((delta_shifted - 1.)**2)**0.5)
# get 1+deta mesh from field
#outmesh = FieldMesh(1 + out_delta)
# print some info: this makes code never finish (race condition maybe?)
#nbkit03_utils.rfield_print_info(outfield, comm, 'outfield: ')
return delta_shifted, attrs
def test_grid(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
real = pm.paint(grid)
assert_array_equal(real, 1.0)
class Data_Generation():
log=classlogger
def __init__(self, parameters_box, parameters_igm, use_baryonic=False, transfer="EisensteinHu", path_transfer=None, column=None, cosmo=None):
self.redshift= parameters_box['redshift']
self.width=parameters_box['width']
chosen_cosmo=parameters_box['cosmology']
#Use the baryonic power spectrum instead
#Then the resulting overdensity field does not correspond to cdm-density, but to baryonic density
#Jeans length needed [h**-1 Mpc]
#Note that Jeans length scales with redshift
#For proper values look: Choudhury, Sraianda ... 2001
self.use_baryonic=use_baryonic
self.Jeans_length=parameters_box['jeans_length']
#specifies specific cosmological model
if chosen_cosmo=='Planck15':
self.cosmo = cosmology.Planck15
elif chosen_cosmo=='UserDefined':
self.cosmo = cosmo
else:
print(chosen_cosmo, 'not implemented right now')
print('Continue with Planck15 cosmology instead')
self.cosmo = cosmology.Planck15
self.comoving_distance=self.cosmo.comoving_distance(self.redshift)
#linear power spectrum
Plin = LinearPower(self.cosmo, self.redshift, transfer, path_transfer, column)
if self.use_baryonic:
self.Plin=lambda k: 1/(1+self.Jeans_length**2*k**2)**2*Plin(k)
self.Plin.redshift=self.redshift
self.Plin.sigma8=self.cosmo.sigma8
self.Plin.cosmo=self.cosmo
if self.log.isEnabledFor(logging.INFO):
self.log.info('Start computation of baryonic density field')
else:
self.Plin=Plin
if self.log.isEnabledFor(logging.INFO):
self.log.info('Start computation of cdm density field')
#Parameters describing the box
self.background_field_x=parameters_box['background_box'][0]
self.background_field_y=parameters_box['background_box'][1]
self.background_sampling_x=parameters_box['background_sampling'][0]
self.background_sampling_y=parameters_box['background_sampling'][1]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Background has been specified:')
self.log.info('-->Comoving length: {}'.format([self.background_field_x, self.background_field_y]))
self.log.info('-->sampled:{}'.format([self.background_sampling_x, self.background_sampling_y]))
#initializes pseudo-gaussian random generator
self.seed=parameters_box['seed']
#if True, the seed gaussian has unitary_amplitude
self.unitary_amplitude=True
self.inverted_phase=False
#Also compute displacement
self.compute_displacement=True
self.BoxSize=[self.background_field_x, self.background_field_y, self.width]
self.N_pix=parameters_box['nr_pix']
self.bias=2
#The following values describe the state of the IGM and have to match whatever specified in the inversion
self.beta=parameters_igm['beta']
self.alpha=2-1.4*self.beta
self.tildeC=parameters_igm['tilde_c']
if self.log.isEnabledFor(logging.INFO):
self.log.info('Generator initialized with parameters')
self.log.info('-->Cosmology: {}'.format(chosen_cosmo))
self.log.info('-->Box size: {}'.format(self.BoxSize))
self.log.info('-->Number pixels: {}'.format(self.N_pix))
self.log.info('-->Redshift: {}'.format(self.redshift))
self.log.info('-->seed: {}'.format(self.seed))
self.log.info('-->Jeans: {}'.format(self.Jeans_length))
# the particle mesh for gridding purposes
_Nmesh = np.empty(3, dtype='i8')
_Nmesh[0] = self.background_sampling_x
_Nmesh[1] = self.background_sampling_y
_Nmesh[2] = self.N_pix
self.pm = ParticleMesh(BoxSize=self.BoxSize, Nmesh=_Nmesh, dtype='f4')
return
#Compute Density Field. The function strongly matches the implementation in nbodykit.
#Return:
#delta: Density Perturbation (rho/rho_med)
#comoving: Comoving distance to each point at line of sight [h^-1 Mpc]
#vel: Peculiar velocity field [km/s]
def Compute_Density_Field(self):
# growth rate to do RSD in the Zel'dovich approximation
f = self.cosmo.scale_independent_growth_rate(self.redshift)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Growth rate is {}'.format(f))
#Lineat density and displacement field
delta, disp = gaussian_real_fields(self.pm, self.Plin, self.seed,
unitary_amplitude=self.unitary_amplitude,
inverted_phase=self.inverted_phase,
compute_displacement=self.compute_displacement)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Gaussian field generated')
# apply the lognormal transformation to the initial conditions density
# this creates a positive-definite delta (necessary for Poisson sampling)
lagrangian_bias = self.bias - 1
delta = lognormal_transform(delta, bias=lagrangian_bias)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Density field projected to lognormal')
#Compute peculiar velocity field from linear displacement field
if self.compute_displacement:
self.displacement=disp
velocity_norm = f * 100 * self.cosmo.efunc(self.redshift) / (1+self.redshift)
vel_x = velocity_norm * disp[0]
vel_y = velocity_norm * disp[1]
vel_z = velocity_norm * disp[2]
vel=[vel_x.value, vel_y.value, vel_z.value]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Velocity field computed')
#Computes comoving distances along line of sight
comoving=self.cosmo.comoving_distance(self.redshift)-self.width+self.width/self.N_pix*np.linspace(1, self.N_pix, self.N_pix)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Density field succesfully computed')
return [delta, vel, comoving]
#Computes the Covariance matrix, depending on the linear autocorrelation
#Correlation may look different in nonlinear and mildly nonlinear regime (due to lognormal transform)
#In case of non biased gaussian fields (expectation value equal to zero) the output matches the covariance matrix
#diff: The difference value between two points in real space
#Autocorrelation is given by the inverse Fourier transform of power spectrum (Wiener-Khinchin-Theorem)
#comoving: vector of comoving coordinates on which the density field is evaluated
#WARNING: Due to rebinning the comoving field could have changed and is not equal binned anymore
#As Covariance matrix is symmetric, only the upper half is computed, then both merged together by transposition
def Compute_Linear_Covariance(self, comoving):
if self.log.isEnabledFor(logging.INFO):
self.log.info('Computation of prior covariance started')
Covariance=np.zeros((np.size(comoving), np.size(comoving)))
k = np.linspace(10**(-5), 10, 10**6)
size = len(k)//2
Pk = self.Plin(k)
fourier_coeff = np.abs(np.fft.fftn(Pk)[0:size+1])
frqs = np.linspace(0, 0.1*size, size+1)
cf_lin = Spline(frqs, fourier_coeff)
diff=np.zeros((np.size(comoving), np.size(comoving)))
for i in range(np.size(comoving)):
for j in range(np.size(comoving)):
diff[i, j]=np.abs(comoving[i]-comoving[j])
Covariance=cf_lin(diff)
Covariance /= Covariance[0, 0]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Prior Covariance computed')
return Covariance
#This Function computes the selection of a single line of sight from the delta sample
#i, j: indices of background sources at maximal redshift
#Return: Density-field along one line-of-sight, evaluated at the pixels of comoving
def Select_LOS(self, delta, vel, index):
i=index[0].astype(int)
j=index[1].astype(int)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Line of sight selected: {}'.format(index))
return [delta[i, j, :], vel[i, j, :]]
#Poisson sample to overdensity field,
#matches the nbodykit routine
def PoissonSample(self, delta, parameters_sampling):
nbar=parameters_sampling['nbar']
seed1=parameters_sampling['seed1']
seed2=parameters_sampling['seed2']
comm = self.pm.comm
# mean number of objects per cell
H = self.BoxSize / self.pm.Nmesh
overallmean = H.prod() * nbar
# number of objects in each cell (per rank, as a RealField)
cellmean = delta * overallmean
# create a random state with the input seed
rng = MPIRandomState(seed=seed1, comm=comm, size=delta.size)
# generate poissons. Note that we use ravel/unravel to
# maintain MPI invariane.
Nravel = rng.poisson(lam=cellmean.ravel())
N = self.pm.create(type='real')
N.unravel(Nravel)
Ntot = N.csum()
if self.log.isEnabledFor(logging.INFO):
self.log.info('Poisson sampling done, total number of objects is {}'.format(Ntot))
pos_mesh = self.pm.generate_uniform_particle_grid(shift=0.0)
disp_mesh = np.empty_like(pos_mesh)
# no need to do decompose because pos_mesh is strictly within the
# local volume of the RealField.
N_per_cell = N.readout(pos_mesh, resampler='nnb')
for i in range(N.ndim):
disp_mesh[:, i] = self.displacement[i].readout(pos_mesh, resampler='nnb')
# fight round off errors, if any
N_per_cell = np.int64(N_per_cell + 0.5)
pos = pos_mesh.repeat(N_per_cell, axis=0)
disp = disp_mesh.repeat(N_per_cell, axis=0)
del pos_mesh
del disp_mesh
if self.log.isEnabledFor(logging.INFO):
self.log.info("Catalog produced. Assigning in cell shift.")
# FIXME: after pmesh update, remove this
orderby = np.int64(pos[:, 0] / H[0] + 0.5)
for i in range(1, delta.ndim):
orderby[...] *= self.pm.Nmesh[i]
orderby[...] += np.int64(pos[:, i] / H[i] + 0.5)
# sort by ID to maintain MPI invariance.
pos = mpsort.sort(pos, orderby=orderby, comm=comm)
disp = mpsort.sort(disp, orderby=orderby, comm=comm)
if self.log.isEnabledFor(logging.INFO):
self.log.info("Sorting done")
rng_shift = MPIRandomState(seed=seed2, comm=comm, size=len(pos))
in_cell_shift = rng_shift.uniform(0, H[i], itemshape=(delta.ndim,))
pos[...] += in_cell_shift
pos[...] %= self.pm.BoxSize
if self.log.isEnabledFor(logging.INFO):
self.log.info("Catalog shifted.")
#Catalog needs to be shifted in z-coordinate, such that pos and comoving match
pos[...,0]+=(self.comoving_distance-self.width)*np.ones(pos.shape[0])
return pos, disp
#Projects baryonic density field to neutral hydrogen overdensity field
def Find_Neutral_Hydrogen_Fraction(self, density, redshift):
density_h1=self.tildeC*(1+redshift)**6*density**(self.alpha)
if self.log.isEnabledFor(logging.INFO):
self.log.info("Density field projected to neutral hydrogen density field")
return density_h1
def Project_Velocity(self, vel, direction):
return vel[direction]
def look_den_slice(net1, net2, net3, s, title):
matplotlib.rc('xtick', labelsize=12)
matplotlib.rc('ytick', labelsize=12)
matplotlib.rc('font', size=12)
cmap = 'coolwarm'
assert net1.shape[-1] == 3
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
net1 = pm.paint(q + net1.reshape([-1, 3]))
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
net2 = pm.paint(q + net2.reshape([-1, 3]))
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
net3 = pm.paint(q + net3.reshape([-1, 3]))
fig = plt.figure(figsize=(6, 4))
amp_low = min(np.percentile(net1[:, :, s], 5),
np.percentile(net2[:, :, s], 5),
np.percentile(net3[:, :, s], 5))
amp_high = max(np.percentile(net1[:, :, s], 95),
np.percentile(net2[:, :, s], 95),
np.percentile(net3[:, :, s], 95))
amp_low1 = min(np.percentile((net1[:, :, s] - net2[:, :, s]), 5),
np.percentile((net1[:, :, s] - net3[:, :, s]), 5))
amp_high1 = max(np.percentile((net1[:, :, s] - net2[:, :, s]), 95),
np.percentile((net1[:, :, s] - net3[:, :, s]), 95))
amp = max(np.abs(amp_low1), amp_high1)
gs = GridSpec(2, 4, height_ratios=[1, 1], width_ratios=[1, 1, 1, 0.1])
plt.subplot(gs[0, 0])
im = plt.imshow(net1[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('fastPM')
plt.subplot(gs[0, 1])
plt.imshow(net2[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('2LPT')
plt.subplot(gs[0, 2])
plt.imshow(net3[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('U-Net')
cbax = plt.subplot(gs[0, 3])
cbar = fig.colorbar(mappable=im,
cax=cbax,
orientation='vertical',
ticklocation='right')
cbar.ax.tick_params(labelsize=12)
plt.subplot(gs[1, 1])
im = plt.imshow((net1[:, :, s] - net2[:, :, s]),
cmap=cmap,
vmin=-amp,
vmax=amp)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title(r'fastPM $-$ 2LPT')
plt.subplot(gs[1, 2])
plt.imshow((net1[:, :, s] - net3[:, :, s]), cmap=cmap, vmin=-amp, vmax=amp)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title(r'fastPM $-$ U-Net')
cbax = plt.subplot(gs[1, 3])
cbar = fig.colorbar(mappable=im,
cax=cbax,
orientation='vertical',
ticklocation='right')
cbar.ax.tick_params(labelsize=12)
fig.subplots_adjust(hspace=.2)
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
rank = pm.comm.rank
wsize = pm.comm.size
comm = pm.comm
if rank == 0:
print(args)
print(outfolder)
cube_size = nc / ncube
shift = cube_size
cube_length = cube_size * bs / nc
pmsmall = ParticleMesh(BoxSize=bs / ncube,
Nmesh=[cube_size, cube_size, cube_size],
dtype=np.float32,
comm=MPI.COMM_SELF)
gridsmall = pmsmall.generate_uniform_particle_grid(shift=0)
for zz in [3.5, 4.0]:
aa = 1 / (1 + zz)
cats, meshes = setupuvmesh(zz,
suff=suff,
sim=sim,
pm=pm,
profile=profile,
stellar=stellar)
cencat, satcat = cats
h1meshfid, h1mesh, lmesh, uvmesh = meshes
if ncsim == 10240:
dm = BigFileMesh(scratchyf+sim+'/fastpm_%0.4f/'%aa+\
'/1-mesh/N%04d'%nc,'').paint()
else: dm = BigFileMesh(project+sim+'/fastpm_%0.4f/'%aa+\
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
