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_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"))