def test_deprecated_ehpower():
c = Cosmology()
with pytest.warns(FutureWarning):
Plin1 = EHPower(c, redshift=0)
Plin2 = LinearPower(c, 0., transfer='EisensteinHu')
assert_allclose(Plin1(0.1), Plin2(0.1))
with pytest.warns(FutureWarning):
Plin1 = NoWiggleEHPower(c, redshift=0)
Plin2 = LinearPower(c, 0., transfer='NoWiggleEisensteinHu')
assert_allclose(Plin1(0.1), Plin2(0.1))
def test_linear():
# linear power
c = Cosmology()
Plin = LinearPower(c, redshift=0)
# desired separation (in Mpc/h)
r = numpy.logspace(0, numpy.log10(150), 500)
# linear correlation
CF = CorrelationFunction(Plin)
CF.sigma8 = 0.82
xi1 = CF(100.)
assert_allclose(CF.redshift, CF.attrs['redshift'])
assert_allclose(CF.sigma8, CF.attrs['sigma8'])
# change sigma8
CF.sigma8 = 0.75
xi2 = CF(100.)
assert_allclose(xi1/xi2, (0.82/0.75)**2, rtol=1e-2)
# change redshift
CF.redshift = 0.55
xi3 = CF(100.)
D2 = CF.cosmo.scale_independent_growth_factor(0.)
D3 = c.scale_independent_growth_factor(0.55)
assert_allclose(xi2.max()/xi3.max(), (D2/D3)**2, rtol=1e-2)
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 __init__(self,FLRW=True,obj=None,H0=None,Om0=None,Ob0=None,Pk_lin=None):
if FLRW == True:
if obj is None:
raise ValueError(("give obj from astropy.cosmology ornbodykit.cosmology if FLRW is True"))
try:
from nbodykit.cosmology import LinearPower
except:
raise ImportError(("Need nbodykit.cosmology.LinearPower"))
self.H0 = obj.h * 100 # km/s/Mpc
self.Om0 = obj.Om0
self.Ob0 = obj.Ob0
self.F = (obj.Om0)**(0.6)
k_space = np.logspace(-4,3,500)
Plin = LinearPower(obj, redshift=0, transfer='EisensteinHu')
Pk_space = Plin(k_space)
self.Pk_lin = interpolate.interp1d(k_space,Pk_space,fill_value="extrapolate")
else:
if any(x is None for x in [H0,Om0,Ob0,Pk_lin]):
raise Exception("Set H0,Om0,Ob0,Pk_lin by hand if FLRW is not True")
self.H0 = H0
self.Om0 = Om0
self.Ob0 = Ob0
self.F = (Om0)**(0.6)
self.Pk_lin = Pk_lin
def save_powerspec(
omega0: float = 0.288,
omegab: float = 0.0472,
hubble: float = 0.7,
scalar_amp: float = 2.427e-9,
ns: float = 0.97,
outfile: str = "my_pk_linear.txt",
):
"""Generate linear powerspec and save into file"""
omegacdm = (omega0 - omegab
) # dark matter density = total mass density - baryon density
MYPlanck = Cosmology(
m_ncdm=[], # no massive neutrino
Omega0_b=omegab,
Omega0_cdm=omegacdm,
h=hubble,
n_s=ns,
A_s=scalar_amp,
) # \
# .match(sigma8=0.8159)
pklin0 = LinearPower(MYPlanck, redshift=0.0)
k = numpy.logspace(-3, 2, 10000, endpoint=True)
numpy.savetxt(outfile, list(zip(k, pklin0(k))))
def test_linear_norm():
# initialize the power
c = Cosmology().match(sigma8=0.82)
P = LinearPower(c, redshift=0, transfer='CLASS')
# compute for array
k = numpy.logspace(-3, numpy.log10(0.99 * c.P_k_max), 100)
Pk1 = P(k)
# change sigma8
P.sigma8 = 0.75
Pk2 = P(k)
assert_allclose(Pk1.max() / Pk2.max(), (0.82 / 0.75)**2, rtol=1e-2)
# change redshift
P.redshift = 0.55
Pk3 = P(k)
D2 = c.scale_independent_growth_factor(0.)
D3 = c.scale_independent_growth_factor(0.55)
assert_allclose(Pk2.max() / Pk3.max(), (D2 / D3)**2, rtol=1e-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=2)
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=0.3, order=2)
dnonlin = solver.nbody(state, leapfrog([0.3, 0.35]))
dnonlin.save('nonlin')
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 test_mcfit():
c = Cosmology()
Plin = LinearPower(c, redshift=0)
# do Pk to CF
k = numpy.logspace(-4, 2, 1024)
CF = pk_to_xi(k, Plin(k))
# do CF to Pk
r = numpy.logspace(-3, 3, 1024)
Pk2 = xi_to_pk(r, CF(r))(k)
idx = (k>1e-2)&(k<10.)
assert_allclose(Pk2[idx], Plin(k[idx]), rtol=1e-2)
def test_large_scales():
c = Cosmology()
k = numpy.logspace(-5, -2, 100)
# linear power
Plin = LinearPower(c, redshift=0)
# nonlinear power
Pnl = HalofitPower(c, redshift=0)
# zeldovich power
Pzel = ZeldovichPower(c, redshift=0)
assert_allclose(Plin(k), Pnl(k), rtol=1e-2)
assert_allclose(Plin(k), Pzel(k), rtol=1e-2)
def test_mcfit():
c = Cosmology()
Plin = LinearPower(c, redshift=0)
for ell in [0, 2, 4]:
# do Pk to CF; use Plin for ell>0 just for testing
k = numpy.logspace(-4, 2, 1024)
CF = pk_to_xi(k, Plin(k), ell=ell)
# do CF to Pk; use Plin for ell>0 just for testing
r = numpy.logspace(-2, 4, 1024)
Pk2 = xi_to_pk(r, CF(r), ell=ell)(k)
idx = (k > 1e-2) & (k < 10.)
assert_allclose(Pk2[idx], Plin(k[idx]), rtol=1e-2)
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[...])
def test_power(output, ref, ref_z, final=8):
"""Check the initial power against linear theory and a linearly grown IC power"""
print('testing', output)
pksim, pkgas, z = compute_power(output)
#Check types have the same power
fig = Figure(figsize=(5, 5), dpi=200)
canvas = FigureCanvasAgg(fig)
ax = fig.add_subplot(111)
pklin = LinearPower(WMAP9, redshift=z)
D = WMAP9.scale_independent_growth_factor(
z) / WMAP9.scale_independent_growth_factor(ref_z)
ax.plot(pksim['k'][1:],
pksim['power'][1:].real / pklin(pksim['k'][1:]),
label=output + " DM")
ax.plot(pkgas['k'][1:],
pkgas['power'][1:].real / pklin(pkgas['k'][1:]),
ls="-.",
label=output + " gas")
ax.plot(ref['k'][1:],
ref['power'][1:].real * D**2 / pklin(ref['k'][1:]),
ls="--",
label='ref, linearly grown')
ax.set_xscale('log')
ax.set_ylim(0.8, 1.2)
ax.legend()
fig.savefig(os.path.basename(output) + '.png')
# asserting the linear growth is 1% accurate
assert_allclose(pksim['power'][2:final],
ref['power'][2:final] * D**2,
rtol=0.012,
atol=0.0)
# Assert we agree with linear theory
assert_allclose(pksim['power'][2:final],
pklin(ref['k'])[2:final],
rtol=0.025 * D**2,
atol=0.0)
#Assert gas and DM are similar
assert_allclose(pkgas['power'][1:final],
pksim['power'][1:final],
rtol=0.01 * D**2,
atol=0.)
def data(self, seed=None):
"""
Return a LogNormalCatalog using the specific configuration of this
sample. The catalog also includes (RA, DEC, Z) coordinates in
addition to Position.
"""
redshift = 0.
nbar = self.N/self.BoxSize.prod()
# lognormal catalog
Plin = LinearPower(self.cosmo, redshift=self.redshift, transfer='EisensteinHu')
cat = LogNormalCatalog(Plin=Plin, nbar=nbar, BoxSize=self.BoxSize, Nmesh=512, seed=seed)
# add sky coordinates too
cat['RA'], cat['DEC'], cat['Z'] = CartesianToSky(cat['Position'], self.cosmo, observer=0.5*self.BoxSize)
# and n(z)
cat['NZ'] = nbar
return cat
def test_linear():
# initialize the power
c = Cosmology().match(sigma8=0.82)
P = LinearPower(c, redshift=0, transfer='CLASS')
# check velocity dispersion
assert_allclose(P.velocity_dispersion(), 5.898, rtol=1e-3)
# test sigma8
assert_allclose(P.sigma_r(8.), c.sigma8, rtol=1e-5)
# change sigma8
P.sigma8 = 0.80
c = c.match(sigma8=0.80)
assert_allclose(P.sigma_r(8.), P.sigma8, rtol=1e-5)
# change redshift and test sigma8(z)
P.redshift = 0.55
assert_allclose(P.sigma_r(8.), c.sigma8_z(P.redshift), rtol=1e-5)
# desired wavenumbers (in h/Mpc)
k = numpy.logspace(-3, 2, 500)
# initialize EH power
P1 = LinearPower(c, redshift=0., transfer="CLASS")
P2 = LinearPower(c, redshift=0., transfer='EisensteinHu')
P3 = LinearPower(c, redshift=0., transfer='NoWiggleEisensteinHu')
# check different transfers (very roughly)
Pk1 = P1(k)
Pk2 = P2(k)
Pk3 = P3(k)
assert_allclose(Pk1 / Pk1.max(), Pk2 / Pk2.max(), rtol=0.1)
assert_allclose(Pk1 / Pk1.max(), Pk3 / Pk3.max(), rtol=0.1)
# also try scalar
Pk = P(0.1)
def test_bad_transfer():
with pytest.raises(ValueError):
Plin = LinearPower(Cosmology(), redshift=0., transfer="BAD")
def test_from_astropy():
from astropy.cosmology import Planck15
Plin = LinearPower(Planck15, redshift=0)
Pk = Plin(0.1)
from vmad import Builder
from fastpm.force.lpt import lpt1, lpt2source
from vmad.lib import fastpm
import numpy
from nbodykit.cosmology import Planck15, LinearPower
pm = fastpm.ParticleMesh([32, 32, 32], BoxSize=128.)
powerspectrum = LinearPower(Planck15, 0)
q = pm.generate_uniform_particle_grid()
with Builder() as model:
x = model.input('x')
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
from nbodykit.lab import *
from nbodykit.cosmology import WMAP9, LinearPower, Cosmology
import numpy
MYPlanck = Cosmology(m_ncdm=[],
Omega0_b=0.0223/0.6774**2,
Omega0_cdm=0.1188/0.6774**2,
h=0.6774)\
.match(sigma8=0.8159)
pklin0 = LinearPower(MYPlanck, redshift=0.0)
k = numpy.logspace(-3, 2, 10000, endpoint=True)
numpy.savetxt('myplanck-z0.txt', list(zip(k, pklin0(k))))
from nbodykit.lab import *
from nbodykit.cosmology import WMAP9, LinearPower
import numpy
pklin0 = LinearPower(WMAP9, redshift=0.0)
pklin9 = LinearPower(WMAP9, redshift=9.0)
k = numpy.logspace(-3, 2, 10000, endpoint=True)
numpy.savetxt('wmap9-z0.txt', list(zip(k, pklin0(k))))
numpy.savetxt('wmap9-z9.txt', list(zip(k, pklin9(k))))
# Make catalogue measurements.
measurements = \
make_catalogue_measurements(kmin=progrc.kmin, kmax=progrc.kmax)
np.savez(
output_dir/output_filename.replace("likelihood", "pk"),
**measurements
)
# Make model predictions.
cosmo = BaseModel(cosmo_dir/progrc.cosmology_file)
mode_modifications = float(progrc.NG) \
* scale_dependence_modification(cosmo, Z)(measurements['k'])
mode_powers = LinearPower(cosmo, Z)(measurements['k'])
gaussianised_data_vector = measurements['pk'] ** (1./3.)
gaussianised_expectation_factor = np.array([
np.exp(loggamma(n + 1./3.) - loggamma(n)) / n ** (1./3.)
for n in measurements['nk']
])
gaussianised_variance_factor = np.array([
np.exp(loggamma(n + 2./3.) - loggamma(n)) / n ** (2./3.)
- np.exp(loggamma(n + 1./3.) - loggamma(n)) ** 2 / n ** (2./3.)
for n in measurements['nk']
])
# Evaluate likelihood and export results.
bias_start, bias_end, ninterval = progrc.bias
# asserting the linear growth is 1% accurate
assert_allclose(pksim['power'][2:final],
ref['power'][2:final] * D**2,
rtol=0.012,
atol=0.0)
# Assert we agree with linear theory
assert_allclose(pksim['power'][2:final],
pklin(ref['k'])[2:final],
rtol=0.025 * D**2,
atol=0.0)
#Assert gas and DM are similar
assert_allclose(pkgas['power'][1:final],
pksim['power'][1:final],
rtol=0.01 * D**2,
atol=0.)
# asserting the initial power spectrum is 1% accurate
print("testing IC power")
ref, refgas, ref_z = compute_power('output/IC')
pklin = LinearPower(WMAP9, redshift=ref_z)
#This checks that the power spectrum loading and rescaling code is working.
genpk = numpy.loadtxt("output/inputspec_IC.txt")
assert_allclose(pklin(genpk[:, 0]), genpk[:, 1], rtol=5e-4, atol=0.0)
#This checks that the output is working
test_power('output/IC', ref, ref_z)
test_power('output/PART_000', ref, ref_z)
test_power('output/PART_001', ref, ref_z)
test_power('output/PART_002', ref, ref_z)
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:
print(state.a['S'], state.a['P'], state.a['F'], state['1'].S[0],
state['1'].P[0], action, ai, ac, af)
def print(*args, **kwargs):
comm = pm.comm
from builtins import print
if comm.rank == 0:
print(*args, **kwargs)
def pprint(*args, **kwargs):
comm = pm.comm
from pprint import pprint
if comm.rank == 0:
pprint(*args, **kwargs)
Pss = LinearPower(Planck15, 0)
Pnn = lambda k: 1.0
ForwardModelHyperParameters = dict(q=pm.generate_uniform_particle_grid(),
stages=[0.5, 1.0],
cosmology=Planck15,
pm=pm)
ForwardOperator = cosmo4d.FastPMOperator.bind(**ForwardModelHyperParameters)
def monitor(state):
#problem.save('/tmp/bar-%04d' % state['nit'], state)
z = state.g
print(abs(z.c2r().r2c() - z)[...].max() / z.cnorm())
print(state)
