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_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_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 __init__(self, cosmo, redshift):
from astropy.cosmology import FLRW
# convert astropy
if isinstance(cosmo, FLRW):
from nbodykit.cosmology import Cosmology
cosmo = Cosmology.from_astropy(cosmo)
# internal cosmology clone with nonlinear enabled
self.cosmo = cosmo.clone(nonlinear=True)
self.redshift = redshift
self._sigma8 = self.cosmo.sigma8
# store meta-data
self._attrs = {}
self._attrs['cosmo'] = dict(cosmo)
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 __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)
@author: Denise
"""
import numpy as np
import flowpm
from flowpm.tfbackground import dchioverda, rad_comoving_distance, a_of_chi as a_of_chi_tf, transverse_comoving_distance as trans_comoving_distance, angular_diameter_distance as ang_diameter_distance
from numpy.testing import assert_allclose
from scipy import interpolate
from nbodykit.cosmology import Cosmology
from astropy.cosmology import Planck15
import astropy.units as u
# Create a simple Planck15 cosmology without neutrinos, and makes sure sigma8
# is matched
ref_cosmo = Cosmology.from_astropy(Planck15.clone(m_nu=0 * u.eV))
ref_cosmo = ref_cosmo.match(sigma8=flowpm.cosmology.Planck15().sigma8.numpy())
def test_radial_comoving_distance():
""" This function tests the function computing the radial comoving distance.
"""
cosmo_tf = flowpm.cosmology.Planck15()
a = np.logspace(-2, 0.0)
z = 1 / a - 1
radial = rad_comoving_distance(cosmo_tf, a)
radial_astr = ref_cosmo.comoving_distance(z)
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)
def __init__(self, halos, seed=None, use_cache=False, comm=None, **params):
from halotools.empirical_models import model_defaults
# check type
from halotools.sim_manager import UserSuppliedHaloCatalog
if not isinstance(halos, UserSuppliedHaloCatalog):
raise TypeError("input halos catalog for HOD should be halotools.sim_manager.UserSuppliedHaloCatalog")
# try to extract meta-data from catalog
mdef = getattr(halos, 'mdef', 'vir')
cosmo = getattr(halos, 'cosmology', None)
Lbox = getattr(halos, 'Lbox', None)
redshift = getattr(halos, 'redshift', None)
# fail if we are missing any of these
for name, attr in zip(['cosmology', 'Lbox', 'redshift'], [cosmo, Lbox, redshift]):
if attr is None:
raise AttributeError("input UserSuppliedHaloCatalog must have '%s attribute" %name)
# promote astropy cosmology to nbodykit's Cosmology
from astropy.cosmology import FLRW
if isinstance(cosmo, FLRW):
from nbodykit.cosmology import Cosmology
cosmo = Cosmology.from_astropy(cosmo)
# store the halotools catalog
self._halos = halos
self.cosmo = cosmo
self.comm = comm
# set the seed randomly if it is None
if seed is None:
if self.comm.rank == 0:
seed = numpy.random.randint(0, 4294967295)
seed = self.comm.bcast(seed)
# grab the BoxSize from the halotools catalog
self.attrs['BoxSize'] = numpy.empty(3)
self.attrs['BoxSize'][:] = Lbox
# store mass and radius keys
self.mass = model_defaults.get_halo_mass_key(mdef)
self.radius = model_defaults.get_halo_boundary_key(mdef)
# check for any missing columns
needed = ['halo_%s' %s for s in ['x','y','z','vx','vy','vz', 'id', 'upid']] + [self.mass, self.radius]
missing = set(needed) - set(halos.halo_table.colnames)
if len(missing):
raise ValueError("missing columns from halotools UserSuppliedHaloCatalog: %s" %str(missing))
# this ensures that all columns in the halo catalog will propagate to the galaxy catalog
model_defaults.default_haloprop_list_inherited_by_mock = halos.halo_table.colnames
# store the attributes
self.attrs['mdef'] = mdef
self.attrs['redshift'] = redshift
self.attrs['seed'] = seed
self.attrs.update(params)
# add cosmo attributes
self.attrs['cosmo'] = dict(cosmo)
# make the model!
self._model = self.__makemodel__()
# set the HOD params
for param in self._model.param_dict:
if param not in self.attrs:
raise ValueError("missing '%s' parameter when initializing HOD" %param)
self._model.param_dict[param] = self.attrs[param]
# make the actual source
ArrayCatalog.__init__(self, self.__makesource__(), comm=comm, use_cache=use_cache)
# crash with no particles!
if self.csize == 0:
raise ValueError("no particles in catalog after populating HOD")
def fit_nbody(cosmo, state, stages, nc, pm_nc_factor=1, name="NBody"):
"""
Integrate the evolution of the state across the givent stages
Parameters:
-----------
cosmo: cosmology
Cosmological parameter object
state: tensor (3, batch_size, npart, 3)
Input state
stages: array
Array of scale factors
nc: int, or list of ints
Number of cells
pm_nc_factor: int
Upsampling factor for computing
pgdparams: array
list of pgdparameters [alpha, kl, ks] of size len(stages) - 1
Returns
-------
state: tensor (3, batch_size, npart, 3), or list of states
Integrated state to final condition, or list of intermediate steps
"""
with tf.name_scope(name):
state = tf.convert_to_tensor(state, name="state")
# Create a simple Planck15 cosmology without neutrinos, and makes sure sigma8
# is matched
nbdykit_cosmo = Cosmology.from_astropy(Planck15.clone(m_nu=0 * u.eV))
nbdykit_cosmo = nbdykit_cosmo.match(sigma8=cosmo.sigma8.numpy())
if isinstance(nc, int):
nc = [nc, nc, nc]
# Unrolling leapfrog integration to make tf Autograph happy
if len(stages) == 0:
return state
ai = stages[0]
# first force calculation for jump starting
state = tfpm.force(cosmo, state, nc, pm_nc_factor=pm_nc_factor)
k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3),
boxsize=np.array([FLAGS.box_size] * 3),
kmin=np.pi / FLAGS.box_size,
dk=2 * np.pi / FLAGS.box_size)
params = tf.Variable([FLAGS.alpha0, FLAGS.kl0, FLAGS.ks0],
dtype=tf.float32)
optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate)
x, p, f = ai, ai, ai
pgdparams = []
scale_factors = []
# Loop through the stages
for i in range(len(stages) - 1):
a0 = stages[i]
a1 = stages[i + 1]
ah = (a0 * a1)**0.5
# Kick step
state = tfpm.kick(cosmo, state, p, f, ah)
p = ah
# Drift step
state = tfpm.drift(cosmo, state, x, p, a1)
# Let's compute the target power spectrum at that scale factor
target_pk = HalofitPower(nbdykit_cosmo,
1. / a1 - 1.)(k).astype('float32')
for j in range(FLAGS.niter if i == 0 else FLAGS.niter_refine):
optimizer.minimize(partial(pgd_loss, params, state, target_pk),
params)
if j % 10 == 0:
loss, pk = pgd_loss(params,
state,
target_pk,
return_pk=True)
if j == 0:
pk0 = pk
print("step %d, loss:" % j, loss)
pgdparams.append(params.numpy())
scale_factors.append(a1)
print("Sim step %d, fitted params (alpha, kl, ks)" % i,
pgdparams[-1])
plt.loglog(k, target_pk, "k")
plt.loglog(k, pk0, ':', label='starting')
plt.loglog(k, pk, '--', label='after n steps')
plt.grid(which='both')
plt.savefig('PGD_fit_%0.2f.png' % a1)
plt.close()
# Optional PGD correction step
state = tfpm.pgd(state, params, nc, pm_nc_factor=pm_nc_factor)
x = a1
# Force
state = tfpm.force(cosmo, state, nc, pm_nc_factor=pm_nc_factor)
f = a1
# Kick again
state = tfpm.kick(cosmo, state, p, f, a1)
p = a1
return state, scale_factors, pgdparams
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)
def main(_):
cosmology = flowpm.cosmology.Planck15()
# Create a simple Planck15 cosmology without neutrinos, and makes sure sigma8
# is matched
nbdykit_cosmo = Cosmology.from_astropy(Planck15.clone(m_nu=0 * u.eV))
nbdykit_cosmo = nbdykit_cosmo.match(sigma8=cosmology.sigma8.numpy())
# Compute the k vectora that will be needed in the PGD fit
k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3),
boxsize=np.array([FLAGS.box_size] * 3),
kmin=np.pi / FLAGS.box_size,
dk=2 * np.pi / FLAGS.box_size)
# Create some initial conditions
klin = tf.constant(np.logspace(-4, 1, 512), dtype=tf.float32)
pk = linear_matter_power(cosmology, klin)
pk_fun = lambda x: tf.cast(
tf.reshape(
interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]),
klin, pk), x.shape), tf.complex64)
initial_conditions = flowpm.linear_field(
[FLAGS.nc, FLAGS.nc, FLAGS.nc],
[FLAGS.box_size, FLAGS.box_size, FLAGS.box_size],
pk_fun,
batch_size=FLAGS.batch_size)
initial_state = flowpm.lpt_init(cosmology, initial_conditions,
FLAGS.a_init)
stages = np.linspace(FLAGS.a_init, 1., FLAGS.nsteps, endpoint=True)
print('Starting simulation')
# Run the Nbody
states = flowpm.nbody(cosmology,
initial_state,
stages, [FLAGS.nc, FLAGS.nc, FLAGS.nc],
pm_nc_factor=FLAGS.B,
return_intermediate_states=True)
print('Simulation done')
# Initialize PGD params
alpha = tf.Variable([FLAGS.alpha0], dtype=tf.float32)
scales = tf.Variable([FLAGS.kl0, FLAGS.ks0], dtype=tf.float32)
optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate)
params = []
scale_factors = []
# We begin by fitting the last time step
for j, (a, state) in enumerate(states[::-1]):
# Let's compute the target power spectrum at that scale factor
target_pk = HalofitPower(nbdykit_cosmo,
1. / a - 1.)(k).astype('float32')
for i in range(FLAGS.niter if j == 0 else FLAGS.niter_refine):
optimizer.minimize(
partial(pgd_loss, alpha, scales, state, target_pk), [alpha] if
(FLAGS.fix_scales and j > 0) else [alpha, scales])
if i % 10 == 0:
loss, pk = pgd_loss(alpha,
scales,
state,
target_pk,
return_pk=True)
if i == 0:
pk0 = pk
print("step %d, loss:" % i, loss)
params.append(np.concatenate([alpha.numpy(), scales.numpy()]))
scale_factors.append(a)
print("Fitted params (alpha, kl, ks)", params[-1])
plt.loglog(k, target_pk, "k")
plt.loglog(k, pk0, ':', label='starting')
plt.loglog(k, pk, '--', label='after n steps')
plt.grid(which='both')
plt.savefig('PGD_fit_%0.2f.png' % a)
plt.close()
pickle.dump(
{
'B': FLAGS.B,
'nsteps': FLAGS.nsteps,
'params': params,
'scale_factors': scale_factors,
'cosmology': cosmology.to_dict(),
'boxsize': FLAGS.box_size,
'nc': FLAGS.nc
}, open(FLAGS.filename, "wb"))
