Source code for pynbody.analysis.cosmology

"""

cosmology
=========

A set of functions for common cosmological calculations. 

"""

import math
import numpy as np
numpy = np  # alias the alias
from .. import units
from ..configuration import config_parser

_interp_points = int(config_parser.get('general','cosmo-interpolation-points'))


def _a_dot(a, h0, om_m, om_l):
    om_k = 1.0 - om_m - om_l
    return h0 * a * np.sqrt(om_m * (a ** -3) + om_k * (a ** -2) + om_l)


def _a_dot_recip(*args):
    return 1. / _a_dot(*args)


[docs]def hzoverh0(a, omegam0): """ returns: H(a) / H0 = [omegam/a**3 + (1-omegam)]**0.5 """ return numpy.sqrt(omegam0 * numpy.power(a, -3) + (1. - omegam0))
def _lingrowthintegrand(a, omegam0): """ (e.g. eq. 8 in lukic et al. 2008) returns: da / [a*H(a)/H0]**3 """ return numpy.power((a * hzoverh0(a, omegam0)), -3) def _lingrowthfac(red, omegam0, omegal0, return_norm=False): """ returns: linear growth factor, b(a) normalized to 1 at z=0, good for flat lambda only a = 1/1+z b(a) = Delta(a) / Delta(a=1) [ so that b(z=0) = 1 ] (and b(a) [Einstein de Sitter, omegam=1] = a) Delta(a) = 5 omegam / 2 H(a) / H(0) * integral[0:a] [da / [a H(a) H0]**3] equation from peebles 1980 (or e.g. eq. 8 in lukic et al. 2008) """ # need to add w ~= , nonflat, -1 functionality import scipy.integrate if (abs(omegam0 + omegal0 - 1.) > 1.e-4): raise RuntimeError("Linear growth factors can only be calculated for flat cosmologies") a = 1 / (1. + red) # 1st calc. for z=z lingrowth = scipy.integrate.quad(_lingrowthintegrand, 0., a, (omegam0))[0] lingrowth *= 5. / 2. * omegam0 * hzoverh0(a, omegam0) # then calc. for z=0 (for normalization) a0 = 1. lingrowtha0 = scipy.integrate.quad( _lingrowthintegrand, 0., a0, (omegam0))[0] lingrowtha0 *= 5. / 2. * omegam0 * hzoverh0(a0, omegam0) lingrowthfactor = lingrowth / lingrowtha0 if return_norm: return lingrowthfactor, lingrowtha0 else: return lingrowthfactor
[docs]def linear_growth_factor(f, z=None): """Calculate the linear growth factor b(a), normalized to 1 at z=0, for the cosmology of snapshot f. The output is dimensionless. If a redshift z is specified, it is used in place of the redshift in output f. """ if z is None: z = f.properties['z'] omegam0 = f.properties['omegaM0'] omegal0 = f.properties['omegaL0'] return _lingrowthfac(z, omegam0, omegal0)
[docs]def rate_linear_growth(f, z=None, unit='h Gyr^-1'): """Calculate the linear growth rate b'(a), normalized to 1 at z=0, for the cosmology of snapshot f. The output is in 'h Gyr^-1' by default. If a redshift z is specified, it is used in place of the redshift in output f.""" if z is None: z = f.properties['z'] a = 1. / (1. + z) omegam0 = f.properties['omegaM0'] omegal0 = f.properties['omegaL0'] b, X = _lingrowthfac(z, omegam0, omegal0, return_norm=True) I = _lingrowthintegrand(a, omegam0) term1 = -(1.5 * omegam0 * a ** -3) * b / \ math.sqrt(1. - omegam0 + omegam0 * a ** -3) term2 = (2.5 * omegam0) * hzoverh0(a, omegam0) ** 2 * a * I / X res = units.h * (term1 + term2) * 100. * units.Unit("km s^-1 Mpc^-1") return res.in_units(unit, **f.conversion_context())
def _test_rate_linear_growth(f, z=None, unit='h Gyr^-1'): # coded up by AP to test linear growth *rate* equation above if z is None: z = f.properties['z'] a0 = 1. / (1. + z) a1 = a0 * 0.999 z0 = 1. / a0 - 1 z1 = 1. / a1 - 1 b0 = linear_growth_factor(f, z0) b1 = linear_growth_factor(f, z1) db = b1 - b0 unit = units.Unit(unit) dt = age(f, z1, unit ** -1) - age(f, z0, unit ** -1) return db / dt
[docs]def age(f, z=None, unit='Gyr'): """ Calculate the age of the universe in the snapshot f by integrating the Friedmann equation. The output is given in the specified units. If a redshift z is specified, it is used in place of the redshift in the output f. **Input**: *f*: SimSnap **Optional Keywords**: *z (None)*: desired redshift. Can be a single number, a list, or a numpy.ndarray. *unit ('Gyr')*: desired units for age output """ import scipy import scipy.integrate from scipy.interpolate import interp1d from ..array import SimArray if z is None: z = f.properties['z'] h0 = f.properties['h'] omM = f.properties['omegaM0'] omL = f.properties['omegaL0'] conv = units.Unit("0.01 s Mpc km^-1").ratio(unit, **f.conversion_context()) def get_age(x): x = 1.0 / (1.0 + x) return scipy.integrate.quad(_a_dot_recip, 0, x, (h0, omM, omL))[0] * conv if isinstance(z, np.ndarray) or isinstance(z, list): if len(z) > _interp_points: a_vals = np.logspace(-3,0, _interp_points) z_vals = 1./a_vals-1. log_age_vals = np.log(age(f, z_vals)) interp = interp1d(np.log(a_vals), log_age_vals, bounds_error=False) log_a_input = np.log(1./(1.+z)) results = np.exp(interp(log_a_input)) else: results = np.array(list(map(get_age, z))) results = results.view(SimArray) results.units = unit return results else: return get_age(z)
[docs]@units.takes_arg_in_units((1, "Gyr"), context_arg=0) def redshift(f, time): """ Calculate the redshift given a snapshot and a time since Big Bang in Gyr. Uses scipy.optimize.newton to do the root finding if number of elements in the time array is less than 1000, otherwise uses a linear interpolation. **Input**: *f*: SimSnap with cosmological parameters defined *time*: time since the Big Bang in Gyr for which a redshift should be returned. float, list, or numpy.ndarray """ from scipy.optimize import newton from scipy.interpolate import interp1d from .. import array def func(x, sim, time): return age(sim, x) - time if isinstance(time, list) or isinstance(time, np.ndarray): if len(time) > _interp_points: zs = np.logspace(3, -10, _interp_points) ages = age(f, zs) i = interp1d(ages, zs) return i(time) else: return np.array([newton(func, 1, args=(f, x)) for x in time]) else: return newton(func, 1, args=(f, time))
[docs]def rho_crit(f, z=None, unit=None): """Calculate the critical density of the universe in the snapshot f. z specifies the redshift. If z is none, the redshift of the provided snapshot is used. unit specifies the units of the returned density. If unit is None, the returned density will be in the units of f["mass"].units/f["pos"].units**3. If that unit cannot be calculated, the returned units are Msol kpc^-3 comoving. Note that you can get slightly confusing results if your simulation is in comoving units and you specify a different redshift z. Specifically, the physical density for the redshift you specify is calulated, but expressed as a comoving density *at the redshift of the snapshot*. This is intentional behaviour.""" if z is None: z = f.properties['z'] if unit is None: try: unit = f.dm["mass"].units / f.dm["pos"].units ** 3 except units.UnitsException: unit = units.NoUnit() if hasattr(unit, "_no_unit"): unit = units.Unit("Msol kpc^-3 a^-3") omM = f.properties['omegaM0'] omL = f.properties['omegaL0'] h0 = f.properties['h'] a = 1.0 / (1.0 + z) H_z = _a_dot(a, h0, omM, omL) / a H_z = units.Unit("100 km s^-1 Mpc^-1") * H_z rho_crit = (3 * H_z ** 2) / (8 * math.pi * units.G) return rho_crit.ratio(unit, **f.conversion_context())
[docs]def rho_M(f, z=None, unit=None): """Calculate the matter density of the universe in snapshot f. unit and z are used if not None, as by rho_crit. See also the note in rho_crit about confusion over comoving units in this case.""" if z is None: z = f.properties['z'] return f.properties['omegaM0'] * rho_crit(f, 0, unit) * (1.0 + z) ** 3
[docs]def H(f): """Calculate the Hubble parameter of the universe in snapshot f""" return f.properties['h'] * hzoverh0(f.properties['a'], f.properties['omegaM0']) * units.Unit("100 km s^-1 Mpc^-1")
[docs]def add_hubble(f): """Add the hubble flow to velocities in snapshot f""" f['vel'] += f['pos'] * H(f)
[docs]def comoving_to_physical(ar): """Given an array, modify it to be in physical units (remove any dependencies on a or aform).""" a_power = ar.units._power_of("a") aform_power = ar.units._power_of("aform") if a_power != 0: a = ar.sim.properties['a'] ar *= a ** a_power ar /= units.Unit("a") ** a_power if aform_power != 0: aform = ar.sim['aform'] ar *= aform ** aform_power ar /= units.Unit("aform")