rubix.cosmology package

rubix.cosmology package#

Submodules#

rubix.cosmology.base module#

class rubix.cosmology.base.BaseCosmology(Om0: float, w0: float, wa: float, h: float)[source]#

Bases: Module

Handle cosmological calculations using JAX-backed implementations.

The implementations follow the DSPS flat_wcdm module (ArgonneCPAC/dsps.) but are wrapped in a dataclass-style container so that the parameters are stored as jax arrays and can be reused safely in other JAX workflows.

Once initialized with the cosmological parameters, the class can be used to calculate various cosmological quantities.

Parameters:

  • Om0 (float) – The present day matter density.

  • w0 (float) – The present day dark energy equation of state.

  • wa (float) – The dark energy equation of state parameter.

  • h (float) – The dimensionless Hubble constant.

Om0#

Stored matter density parameter.

Type:

jnp.float32

w0#

Dark energy equation of state today.

Type:

jnp.float32

wa#

Evolution of the dark energy equation of state.

Type:

jnp.float32

h#

Dimensionless Hubble constant.

Type:

jnp.float32

Example

>>> # Create Planck15 cosmology
>>> from rubix.cosmology import COSMOLOGY
>>> cosmo = COSMOLOGY(0.3089, -1.0, 0.0, 0.6774)
age_at_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Return the age of the universe at the provided redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Age in seconds.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Calculate the age of the universe at redshift 0.5
>>> cosmo.age_at_z(0.5)
age_at_z0() → Float[Array, '...'][source]#

The function calculates the age of the universe at redshift 0.

Returns:

The age of the universe at redshift 0 (float).

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Calculate the age of the universe at redshift 0
>>> cosmo.age_at_z0()
angular_diameter_distance_to_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Compute the angular diameter distance for the given redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Angular diameter distance in Mpc.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Compute the angular diameter distance to redshift 0.5
>>> cosmo.angular_diameter_distance_to_z(0.5)
angular_scale(z: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Angular scale in kpc/arcsec at redshift z.

Parameters:

z (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Angular scale in kpc/arcsec.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Calculate the angular scale at redshift 0.5
>>> cosmo.angular_scale(0.5)
comoving_distance_to_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Calculate the comoving distance to the input redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Comoving distance in Mpc.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Calculate comoving distance to redshift 0.5
>>> cosmo.comoving_distance_to_z(0.5)
distance_modulus_to_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Compute the distance modulus for the requested redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Distance modulus.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Compute the distance modulus to redshift 0.5
>>> cosmo.distance_modulus_to_z(0.5)
lookback_to_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Calculate the lookback time to the requested redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Lookback time in seconds.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Calculate the lookback time to redshift 0.5
>>> cosmo.lookback_to_z(0.5)
luminosity_distance_to_z(redshift: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Compute the luminosity distance at the requested redshift.

Parameters:

redshift (Union[Float[Array, "..."], float]) – Redshift value(s).

Returns:

Luminosity distance in Mpc.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Compute the luminosity distance to redshift 0.5
>>> cosmo.luminosity_distance_to_z(0.5)
scale_factor_to_redshift(a: Float[Array, '...'] | float) → Float[Array, '...'][source]#

Convert the scale factor to redshift.

Parameters:

a (Union[Float[Array, "..."], float]) – Scale factor.

Returns:

Redshift 1/a - 1.

Return type:

Float[Array, “…”]

Example

>>> from rubix.cosmology import PLANCK15 as cosmo
>>> # Convert scale factor 0.5 to redshift
>>> cosmo.scale_factor_to_redshift(jnp.array(0.5))

rubix.cosmology.utils module#

rubix.cosmology.utils.trapz(xarr: Array | Float[Array, '...'], yarr: Array | Float[Array, '...']) → Array[source]#

Perform trapezoidal integration using _cumtrapz_scan_func.

Parameters:

  • xarr (Union[jnp.ndarray, Float[Array, "..."]]) – The x-coordinates.

  • yarr (Union[jnp.ndarray, Float[Array, "..."]]) – The y-values.

Returns:

Scalar results collected from the scan.

Return type:

jnp.ndarray

Example

>>> from rubix.cosmology.utils import trapz
>>> import jax.numpy as jnp

>>> x = jnp.array([0, 1, 2, 3])
>>> y = jnp.array([0, 1, 4, 9])
>>> print(trapz(x, y))

Module contents#

Contents