def test_cosmology_get_pk():
c = Cosmology()
p = c.get_pk(z=0, k=0.1)
p1 = c.Spectra.get_pk(z=0, k=0.1)
# ensure the dro did use get_pk of Spectra rather than that from Primordial
assert_allclose(p, p1)
def test_cosmology_a_max():
c = Cosmology(gauge='synchronous', a_max=2.0)
print(c.parameter_file)
assert c.a_max == 2.0
t = c.Omega_m(-0.1)
t = c.efunc(-0.1)
t = c.scale_independent_growth_factor(-0.1)
def test_cosmology_a_max():
c = Cosmology(gauge='synchronous', a_max=2.0)
print(c.parameter_file)
assert c.a_max == 2.0
t = c.Omega_m(-0.1)
t = c.efunc(-0.1)
t = c.scale_independent_growth_factor(-0.1)
def test_cosmology_get_pk():
c = Cosmology()
p = c.get_pk(z=0, k=0.1)
p1 = c.Spectra.get_pk(z=0, k=0.1)
# ensure the dro did use get_pk of Spectra rather than that from Primordial
assert_allclose(p, p1)
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_halofit():
# initialize the power
c = Cosmology().match(sigma8=0.82)
P = HalofitPower(c, redshift=0)
# k is out of range
with pytest.raises(ValueError):
Pk = P(2 * c.P_k_max)
# compute for scalar
Pk = P(0.1)
# 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_deprecated_init():
# all valid deprecated kwargs
with pytest.warns(FutureWarning):
c1 = Cosmology(H0=67.6, Om0=0.31, flat=True)
c2 = Cosmology(H0=67.6, Om0=0.31, Ode0=0.7, flat=False, w0=-0.9)
# missing valid and deprecated
with pytest.raises(Exception):
c = Cosmology(h=0.7, flat=True)
# passing arguments and mixing
with pytest.raises(Exception):
c = Cosmology(0.7, flat=True)
# parameter conflict
with pytest.raises(Exception):
c3 = Cosmology(H0=70., flat=True, h=0.7)
assert_allclose(c1.h, 0.676)
assert_allclose(c2.h, 0.676)
assert_allclose(c1.Om0, 0.31)
assert_allclose(c2.Om0, 0.31)
assert_allclose(c1.Ok0, 0.)
assert_allclose(c2.Ode0, 0.7)
assert_allclose(c2.w0_fld, -0.9)
def test_linear():
# linear power
c = Cosmology()
Plin = LinearPower(c, redshift=0, transfer='EisensteinHu')
# 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_sigma8_z():
z = numpy.linspace(0, 1, 100)
c = Cosmology()
s8_z = c.sigma8_z(z)
D_z = c.scale_independent_growth_factor(z)
assert_allclose(s8_z, D_z * c.sigma8, rtol=1e-3)
def test_set_Omega0_cb():
c = Cosmology().match(Omega0_cb=0.4)
assert_allclose(c.Omega0_cb, 0.4)
c = Cosmology().match(Omega0_m=0.4)
assert_allclose(c.Omega0_m, 0.4)
def test_sigma8_z():
z = numpy.linspace(0, 1, 100)
c = Cosmology()
s8_z = c.sigma8_z(z)
D_z = c.scale_independent_growth_factor(z)
assert_allclose(s8_z, D_z*c.sigma8, rtol=1e-3)
def test_immutable():
c = Cosmology()
with pytest.raises(ValueError):
c.A_s = 2e-9 # immutable
# can add non-CLASS attributes still
c.test = 'TEST'
def test_bad_input():
with pytest.raises(ValueError):
c = Cosmology(gauge='BAD')
# specifying w0_fld + Omega_Lambda is inconsistent
with pytest.raises(ValueError):
c = Cosmology(Omega_Lambda=0.7, w0_fld=-0.9)
def test_cosmology_clone():
c = Cosmology(gauge='synchronous')
c1 = Cosmology(gauge='newtonian')
assert 'newtonian' in c1.parameter_file
c2 = Cosmology(P_k_max=1.01234567)
assert '1.01234567' in c2.parameter_file
def test_immutable():
c = Cosmology()
with pytest.raises(ValueError):
c.A_s = 2e-9 # immutable
# can add non-CLASS attributes still
c.test = 'TEST'
def test_massive_neutrinos():
# single massive neutrino
c = Cosmology(m_ncdm=0.6)
assert c.N_ncdm == 1
# do not need 0 values
with pytest.raises(ValueError):
c = Cosmology(m_ncdm=[0.6, 0.])
def test_conflicts():
# h is the correct param
with pytest.raises(Exception):
c = Cosmology(h=0.7, H0=70)
with pytest.raises(Exception):
c = Cosmology(Omega0_b=0.04, Omega_b=0.04)
# Omega_b is the correct param
with pytest.raises(Exception):
c = Cosmology(Omega0_b=0.04, omega_b=0.02)
def test_cosmology_vect():
c = Cosmology(gauge='synchronous')
assert_allclose(c.Omega_cdm([0]), c.Omega0_cdm)
assert_array_equal(c.Omega_cdm([]).shape, [0])
assert_array_equal(c.Omega_cdm([0]).shape, [1])
assert_array_equal(c.Omega_cdm([[0]]).shape, [1, 1])
assert_array_equal(c.rho_k([[0]]).shape, [1, 1])
k, z = numpy.meshgrid([0, 1], [0.01, 0.05, 0.1, 0.5], sparse=True, indexing='ij')
pk = c.get_pk(z=z, k=k)
assert_array_equal(pk.shape, [2, 4])
def __init__(self, datar, randomr, nofz, universe_params=None):
#
# make sure catalogs have the desired format
#
# cosmology
if universe_params is None:
universe_params = universeparams
cosmo = Cosmology(Om0=universe_params['Om0'],
H0=universe_params['H0'],
flat=universe_params['flat'])
data = catit(datar)
randoms = catit(randomr)
data = nb.ArrayCatalog(data)
randoms = nb.ArrayCatalog(randoms)
data['Position'] = SkyToCartesian(data['RA'],
data['DEC'],
data['Z'],
cosmo=cosmo)
randoms['Position'] = SkyToCartesian(randoms['RA'],
randoms['DEC'],
randoms['Z'],
cosmo=cosmo)
#
self.data = data
self.randoms = randoms
self.nofz = nofz
def test_bad_no_Ob0():
from astropy.cosmology import FlatLambdaCDM
c = FlatLambdaCDM(Om0=0.3, H0=70) # no Ob0
with pytest.raises(ValueError):
c = Cosmology.from_astropy(c)
def test_bad_astropy_class():
from astropy.cosmology import w0wzCDM
c = w0wzCDM(Om0=0.3, H0=70, Ode0=0.7) # no Ob0
with pytest.raises(ValueError):
c = Cosmology.from_astropy(c)
def test_linear_nbody():
from nbodykit.cosmology import Cosmology
cosmo = Cosmology(Omega0_cdm=0.26,
Omega0_b=0.04,
m_ncdm=[0.6],
P_k_max=1e-2)
linearnbody = LinearNbody(cosmo, c_b=0, c_ncdm_1ev_z0=0)
a0 = 0.01
k, q0, p0 = linearnbody.seed_from_synchronous(cosmo, a0)
k, qt, pt = linearnbody.seed_from_synchronous(cosmo, 1.0)
a, q, p = linearnbody.integrate(k, q0, p0, [a0, 1.0])
# shall be within 5% to class solution
assert_allclose(q[-1] / qt, 1., rtol=0.1)
# need very high precision for reversibility.
a, q, p = linearnbody.integrate(k, qt, pt, [1.0, a0], rtol=1e-9)
a, q, p = linearnbody.integrate(k, q[-1], p[-1], [a0, 1.0], rtol=1e-9)
assert_allclose(q[-1], qt, rtol=1e-4)
assert_allclose(p[-1], pt, rtol=1e-4)
def test_bad_astropy_class():
from astropy.cosmology import w0wzCDM
c = w0wzCDM(Om0=0.3, H0=70, Ode0=0.7) # no Ob0
with pytest.raises(ValueError):
c = Cosmology.from_astropy(c)
def test_bad_no_Ob0():
from astropy.cosmology import FlatLambdaCDM
c = FlatLambdaCDM(Om0=0.3, H0=70) # no Ob0
with pytest.raises(ValueError):
c = Cosmology.from_astropy(c)
def test_set_sigma8():
# set sigma8 by adjusting A_s internally
c = Cosmology().match(sigma8=0.80)
# run CLASS and compute sigma8
assert_allclose(c.sigma8, 0.80)
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_astropy_compat():
c = Cosmology(gauge='synchronous', m_ncdm=[0.06])
assert_allclose(c.Odm(0), c.Odm0)
assert_allclose(c.Ogamma(0), c.Ogamma0)
assert_allclose(c.Ob(0), c.Ob0)
assert_allclose(c.Onu(0), c.Onu0)
assert_allclose(c.Ok(0), c.Ok0)
assert_allclose(c.Ode(0), c.Ode0)
assert_array_equal(c.has_massive_nu, True)
def test_cosmology_sane():
c = Cosmology(gauge='synchronous', verbose=True)
assert_allclose(c.Omega_cdm(0), c.Omega0_cdm)
assert_allclose(c.Omega_g(0), c.Omega0_g)
assert_allclose(c.Omega_b(0), c.Omega0_b)
assert_allclose(c.Omega_ncdm(0), c.Omega0_ncdm)
assert_allclose(c.Omega_ur(0), c.Omega0_ur)
assert_allclose(c.Omega_ncdm(0), c.Omega0_ncdm_tot)
assert_allclose(c.Omega_pncdm(0), c.Omega0_pncdm)
assert_allclose(c.Omega_m(0), c.Omega0_m)
assert_allclose(c.Omega_r(0), c.Omega0_r)
# total density in 10e10 Msun/h unit
assert_allclose(c.rho_tot(0), 27.754999)
# comoving distance to z=1.0 in Mpc/h unit.
assert_allclose(c.comoving_distance(1.0), 3396.157391 * c.h)
# conformal time in Mpc unit.
assert_allclose(c.tau(1.0), 3396.157391)
assert_allclose(c.efunc(0), 1.) # hubble in Mpc/h km/s unit
assert_allclose(c.efunc(0) - c.efunc(1 / 0.9999 - 1),
0.0001 * c.efunc_prime(0), rtol=1e-3)
def make_cosmo(param_names, param_vals):
"""
Frequently, the parameters of our input will not align with what we need.
So, we convert them to the appropriate format.
:param param_names:
:param param_vals:
:return:
"""
param_dict = dict([(pn, param_vals[pn]) for pn in param_names])
# TODO it'd be nice if this could be somehow generalized.
param_dict['N_ncdm'] = 3.0
param_dict['N_ur'] = param_dict['Neff'] - 3.0
del param_dict['Neff']
#param_dict['h'] = param_dict['H0']/100
#del param_dict['H0']
param_dict['h'] = param_dict['h0']
del param_dict['h0']
#param_dict['w0_fld'] = param_dict['w']
w = param_dict['w']
del param_dict['w']
param_dict['Omega_cdm'] = param_dict['Omega_m'] - param_dict['Omega_b']
del param_dict['Omega_b']
del param_dict['Omega_m']
param_dict['Omega_ncdm'] = [
param_dict['Omeganuh2'] / (param_dict['h']**2), 0.0, 0.0
]
#param_dict['m_ncdm']=None
#param_dict['omch2'] = (param_dict['Omega_m'] - param_dict['Omega_b'] - param_dict['Omega0_ncdm_tot'])*(param_dict['h']**2)
del param_dict['Omeganuh2']
#param_dict['A_s'] = 10**(np.log10(np.exp(param_dict['ln_1e10_A_s']))-10.0)
param_dict['A_s'] = 10**(-10) * np.exp(param_dict['ln(10^{10}A_s)'])
del param_dict['ln(10^{10}A_s)']
#print param_dict
# this is seriously what i have to do here.
C = Cosmology()
C2 = C.from_dict(param_dict)
C3 = C2.clone(w0_fld=w)
return C3 #Cosmology(**param_dict)
def test_from_astropy():
from astropy.cosmology import FlatLambdaCDM, LambdaCDM
from astropy.cosmology import FlatwCDM, wCDM
from astropy.cosmology import Flatw0waCDM, w0waCDM
# LambdaCDM
flat = {'H0': 70, 'Om0': 0.3, 'Ob0': 0.04, 'Tcmb0': 2.7}
for cls in [FlatLambdaCDM, LambdaCDM]:
if "Flat" in cls.__name__:
x = cls(**flat)
else:
x = cls(Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Omega0_fld, 0.) # Omega0_lambda is nonzero
assert_allclose(c.Odm0, x.Odm0)
# w0 CDM
for cls in [FlatwCDM, wCDM]:
if "Flat" in cls.__name__:
x = cls(w0=-0.9, **flat)
else:
x = cls(w0=-0.9, Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Odm0, x.Odm0)
assert_allclose(c.w0, x.w0)
assert_allclose(c.Omega0_lambda, 0.) # Omega_fld is nonzero
# w0,wa CDM
for cls in [Flatw0waCDM, w0waCDM]:
if "Flat" in cls.__name__:
x = cls(w0=-0.9, wa=0.01, **flat)
else:
x = cls(w0=-0.9, wa=0.01, Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Odm0, x.Odm0)
assert_allclose(c.w0, x.w0)
assert_allclose(c.wa, x.wa)
assert_allclose(c.Omega0_lambda, 0.) # Omega_fld is nonzero
def test_from_astropy():
from astropy.cosmology import FlatLambdaCDM, LambdaCDM
from astropy.cosmology import FlatwCDM, wCDM
from astropy.cosmology import Flatw0waCDM, w0waCDM
# LambdaCDM
flat = {'H0':70, 'Om0':0.3, 'Ob0':0.04, 'Tcmb0':2.7}
for cls in [FlatLambdaCDM, LambdaCDM]:
if "Flat" in cls.__name__:
x = cls(**flat)
else:
x = cls(Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Omega0_fld, 0.) # Omega0_lambda is nonzero
assert_allclose(c.Odm0, x.Odm0)
# w0 CDM
for cls in [FlatwCDM, wCDM]:
if "Flat" in cls.__name__:
x = cls(w0=-0.9, **flat)
else:
x = cls(w0=-0.9, Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Odm0, x.Odm0)
assert_allclose(c.w0, x.w0)
assert_allclose(c.Omega0_lambda, 0.) # Omega_fld is nonzero
# w0,wa CDM
for cls in [Flatw0waCDM, w0waCDM]:
if "Flat" in cls.__name__:
x = cls(w0=-0.9, wa=0.01, **flat)
else:
x = cls(w0=-0.9, wa=0.01, Ode0=0.75, **flat)
c = Cosmology.from_astropy(x)
assert_allclose(c.Ok0, x.Ok0)
assert_allclose(c.Odm0, x.Odm0)
assert_allclose(c.w0, x.w0)
assert_allclose(c.wa, x.wa)
assert_allclose(c.Omega0_lambda, 0.) # Omega_fld is nonzero
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_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_zeldovich():
# zeldovich power
Pzel = ZeldovichPower(Cosmology(), redshift=0)
# desired separation (in Mpc/h)
r = numpy.logspace(0, numpy.log10(150), 500)
# zeldovich correlation
CF = CorrelationFunction(Pzel)
xi = CF(r)
def test_halofit():
# nonlinear power
Pnl = HalofitPower(Cosmology(), redshift=0)
# desired separation (in Mpc/h)
r = numpy.logspace(0, numpy.log10(150), 500)
# nonlinear correlation
CF = CorrelationFunction(Pnl)
xi = CF(r)
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_old_Omega_syntax():
c1 = Cosmology(Omega_b=0.04)
c2 = Cosmology(Omega0_b=0.04)
assert c1.Omega0_b == c2.Omega0_b
c1 = Cosmology(T_cmb=2.7)
c2 = Cosmology(T0_cmb=2.7)
assert c1.T0_cmb == c2.T0_cmb
c1 = Cosmology(Omega0_k=0.05)
c2 = Cosmology(Omega_k=0.05)
assert c1.Omega0_k == c2.Omega0_k
c1 = Cosmology(Omega0_lambda=0.7)
c2 = Cosmology(Omega_lambda=0.7)
c3 = Cosmology(Omega0_Lambda=0.7)
assert c1.Omega0_lambda == c2.Omega0_lambda
assert c1.Omega0_lambda == c3.Omega0_lambda
def test_cosmology_vect():
c = Cosmology(gauge='synchronous')
assert_allclose(c.Omega_cdm([0]), c.Omega0_cdm)
assert_array_equal(c.Omega_cdm([]).shape, [0])
assert_array_equal(c.Omega_cdm([0]).shape, [1])
assert_array_equal(c.Omega_cdm([[0]]).shape, [1, 1])
assert_array_equal(c.rho_k([[0]]).shape, [1, 1])
k, z = numpy.meshgrid([0, 1], [0.01, 0.05, 0.1, 0.5],
sparse=True,
indexing='ij')
pk = c.get_pk(z=z, k=k)
assert_array_equal(pk.shape, [2, 4])
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 hook(value):
def fixdtype(dtype):
if isinstance(dtype, list):
true_dtype = []
for field in dtype:
if len(field) == 3:
true_dtype.append(
(str(field[0]), str(field[1]), field[2]))
if len(field) == 2:
true_dtype.append((str(field[0]), str(field[1])))
return true_dtype
return dtype
def fixdata(data, N, dtype):
if not isinstance(dtype, list):
return data
# for structured array,
# the last dimension shall be a tuple
if N > 0:
return [fixdata(i, N - 1, dtype) for i in data]
else:
assert len(data) == len(dtype)
return tuple(data)
d = None
if '__dtype__' in value:
dtype = fixdtype(value['__dtype__'])
shape = value['__shape__']
a = fixdata(value['__data__'], len(shape), dtype)
d = numpy.array(a, dtype=dtype)
if '__unit__' in value:
if d is None:
d = value['__data__']
d = Quantity(d, Unit(value['__unit__']))
if '__cosmo__' in value:
d = Cosmology.from_dict(value['__cosmo__'])
if d is not None:
return d
if '__complex__' in value:
real, imag = value['__complex__']
return real + 1j * imag
return value
def hook(value):
def fixdtype(dtype):
if isinstance(dtype, list):
true_dtype = []
for field in dtype:
if len(field) == 3:
true_dtype.append((str(field[0]), str(field[1]), field[2]))
if len(field) == 2:
true_dtype.append((str(field[0]), str(field[1])))
return true_dtype
return dtype
def fixdata(data, N, dtype):
if not isinstance(dtype, list):
return data
# for structured array,
# the last dimension shall be a tuple
if N > 0:
return [fixdata(i, N - 1, dtype) for i in data]
else:
assert len(data) == len(dtype)
return tuple(data)
d = None
if '__dtype__' in value:
dtype = fixdtype(value['__dtype__'])
shape = value['__shape__']
a = fixdata(value['__data__'], len(shape), dtype)
d = numpy.array(a, dtype=dtype)
if '__unit__' in value:
if d is None:
d = value['__data__']
d = Quantity(d, Unit(value['__unit__']))
if '__cosmo__' in value:
d = Cosmology.from_dict(value['__cosmo__'])
if d is not None:
return d
if '__complex__' in value:
real, imag = value['__complex__']
return real + 1j * imag
return value
def test_to_astropy():
from astropy.cosmology import FlatLambdaCDM, LambdaCDM
from astropy.cosmology import FlatwCDM, wCDM
from astropy.cosmology import Flatw0waCDM, w0waCDM
# lambda CDM
for cls in [FlatLambdaCDM, LambdaCDM]:
if "Flat" in cls.__name__:
c1 = Cosmology(Omega_k=0.)
else:
c1 = Cosmology(Omega_k=0.05)
c2 = c1.to_astropy()
assert isinstance(c2, cls)
assert_allclose(c2.Ok0, c1.Omega0_k, rtol=1e-3)
# w0 CDM
for cls in [FlatwCDM, wCDM]:
if "Flat" in cls.__name__:
c1 = Cosmology(w0_fld=-0.9, Omega_k=0.)
else:
c1 = Cosmology(w0_fld=-0.9, Omega_k=0.05)
c2 = c1.to_astropy()
assert isinstance(c2, cls)
assert_allclose(c2.Ok0, c1.Omega0_k, rtol=1e-3)
assert_allclose(c2.w0, c1.w0_fld)
# wa w0 CDM
for cls in [Flatw0waCDM, w0waCDM]:
if "Flat" in cls.__name__:
c1 = Cosmology(wa_fld=0.05, w0_fld=-0.9, Omega_k=0.)
else:
c1 = Cosmology(wa_fld=0.05, w0_fld=-0.9, Omega_k=0.05)
c2 = c1.to_astropy()
assert isinstance(c2, cls)
assert_allclose(c2.Ok0, c1.Omega0_k, rtol=1e-3)
assert_allclose(c2.w0, c1.w0_fld)
assert_allclose(c2.wa, c1.wa_fld)
def test_from_file():
import tempfile, pickle
with tempfile.NamedTemporaryFile(mode='w') as ff:
ff.write("H0=70\nomega_b = 0.0266691\nomega_cdm = 0.110616\nT_cmb=2.7255\n")
ff.seek(0)
# load from a file and check values
c = Cosmology.from_file(ff.name)
assert_allclose(c.Omega0_b*c.h**2, 0.0266691)
assert_allclose(c.Omega0_cdm*c.h**2, 0.110616)
# clone
c2 = c.clone(Omega0_b=0.04)
assert_allclose(c2.Omega0_b, 0.04)
# serialize and make sure we get the same
s = pickle.dumps(c)
c1 = pickle.loads(s)
assert_allclose(c.Omega0_cdm, c1.Omega0_cdm)
assert_allclose(c.Omega0_b, c1.Omega0_b)
def __init__(self, cosmo, redshift, transfer='CLASS'):
from astropy.cosmology import FLRW
# convert astropy
if isinstance(cosmo, FLRW):
from nbodykit.cosmology import Cosmology
cosmo = Cosmology.from_astropy(cosmo)
# store a copy of the cosmology
self.cosmo = cosmo.clone()
# set sigma8 to the cosmology value
self._sigma8 = self.cosmo.sigma8
# setup the transfers
if transfer not in transfers.available:
raise ValueError("'transfer' should be one of %s" %str(transfers.available))
self.transfer = transfer
# initialize internal transfers
c = self.cosmo.clone() # transfers get an internal copy
self._transfer = getattr(transfers, transfer)(c, redshift)
self._fallback = transfers.EisensteinHu(c, redshift) # fallback to analytic when out of range
# normalize to proper sigma8
self._norm = 1.
self.redshift = 0;
self._norm = (self._sigma8 / self.sigma_r(8.))**2 # sigma_r(z=0, r=8)
# set redshift
self.redshift = redshift
# store meta-data
self._attrs = {}
self._attrs['transfer'] = transfer
self._attrs['cosmo'] = dict(cosmo)
def test_cosmology_transfer():
c = Cosmology()
t = c.get_transfer(z=0)
assert 'h_prime' in t.dtype.names
assert 'k' in t.dtype.names
assert 'd_cdm' in t.dtype.names
def test_cosmology_density():
c = Cosmology(gauge='synchronous')
z = [0, 1, 2, 5, 9, 99]
assert_allclose(c.rho_cdm(z), c.Omega_cdm(z) * c.rho_tot(z))
assert_allclose(c.rho_g(z), c.Omega_g(z) * c.rho_tot(z))
assert_allclose(c.rho_ncdm(z), c.Omega_ncdm(z) * c.rho_tot(z))
assert_allclose(c.rho_b(z), c.Omega_b(z) * c.rho_tot(z))
assert_allclose(c.rho_m(z), c.Omega_m(z) * c.rho_tot(z))
assert_allclose(c.rho_r(z), c.Omega_r(z) * c.rho_tot(z))
assert_allclose(c.rho_ur(z), c.Omega_ur(z) * c.rho_tot(z))
def test_clone():
c = Cosmology(gauge='synchronous', tol_background_integration=1e-5)
c2 = c.clone(Omega0_b=0.04)
assert_allclose(c2.Omega0_b, 0.04)
c2 = c2.clone()
assert_allclose(c2.Omega0_b, 0.04)
def __init__(self, simname, halo_finder, redshift, comm=None):
from halotools.sim_manager import CachedHaloCatalog, DownloadManager
from halotools.sim_manager.supported_sims import supported_sim_dict
# do seme setup
self.comm = comm
meta_cols = ['Lbox', 'redshift', 'particle_mass']
# try to automatically load from the Halotools cache
exception = None
if self.comm.rank == 0:
kws = {'simname':simname, 'halo_finder':halo_finder, 'redshift':redshift}
try:
cached_halos = CachedHaloCatalog(dz_tol=0.1, **kws)
fname = cached_halos.fname # the filename to load
meta = {k:getattr(cached_halos, k) for k in meta_cols}
except Exception as e:
# try to download on the root rank
try:
# download
dl = DownloadManager()
dl.download_processed_halo_table(dz_tol=0.1, **kws)
# access the cached halo catalog and get fname attribute
# NOTE: this does not read the data
cached_halos = CachedHaloCatalog(dz_tol=0.1, **kws)
fname = cached_halos.fname
meta = {k:getattr(cached_halos, k) for k in meta_cols}
except Exception as e:
exception = e
else:
fname = None
meta = None
# re-raise a download error on all ranks if it occurred
exception = self.comm.bcast(exception, root=0)
if exception is not None:
raise exception
# broadcast the file we are loading
fname = self.comm.bcast(fname, root=0)
meta = self.comm.bcast(meta, root=0)
# initialize an HDF catalog and add Position/Velocity
cat = HDFCatalog(fname, comm=comm)
cat['Position'] = transform.StackColumns(cat['halo_x'], cat['halo_y'], cat['halo_z'])
cat['Velocity'] = transform.StackColumns(cat['halo_vx'], cat['halo_vy'], cat['halo_vz'])
# get the cosmology from Halotools
cosmo = supported_sim_dict[simname]().cosmology # this is astropy cosmology
cosmo = Cosmology.from_astropy(cosmo)
# initialize the HaloCatalog
HaloCatalog.__init__(self, cat, cosmo, meta['redshift'], mdef='vir', mass='halo_mvir')
# add some meta-data
# NOTE: all Halotools catalogs have to these attributes
self.attrs['BoxSize'] = meta['Lbox']
self.attrs['redshift'] = meta['redshift']
self.attrs['particle_mass'] = meta['particle_mass']
# save the cosmology
self.cosmo = cosmo
self.attrs['cosmo'] = dict(self.cosmo)
