Module Ymir.Cosmology

A cosmology is one record of tensors: flat and curved ΛCDM, wCDM and w0waCDM, with or without radiation and massive neutrinos, are values of it. Every parameter is a leaf, so one compiled program and one batch serve every model, and a derivative or a sampler moves Ω_k across 0 with no branch.

let planck = Cosmology.planck2018 ~codata:Codata.v2022 Nx.float64
let z = Nx.create Nx.float64 [| 4 |] [| 0.; 1.; 3.; 1100. |]
let gyr = Unit.giga Units.julian_year
let ages = Cosmology.age planck z |> Quantity.value gyr
(* 13.7872 5.85157 2.14385 0.000366537 *)

Terms. The scale factor is a = 1/(1 + z). The Hubble rate H(a) is the expansion rate, and E(a) = H(a)/H₀. A component is a part of the energy budget: cold matter (baryons and cold dark matter), photons, neutrinos, dark energy and curvature; its density parameter Ω_i(a) is its density over the critical density at a. The Hubble distance is D_H = c/H₀, and χ is the line-of-sight comoving distance. A cosmology's lanes are the broadcast shape of its leaves and of a call's redshifts.

The expansion rate. With Ω_r(a) the radiation's part,

a^4 E^2(a) = Ω_cb a + Ω_k a^2 + Ω_de a^4 exp(-3(1 + w0 + wa) ln a - 3 wa (1 - a))
             + Ω_r(a)
Ω_r(a)     = Ω_γ (1 + 7/8 (4/11)^(4/3) N_ur) + Σ_i Ω_γ 7/8 T_ncdm^4 F(y_i a)
Ω_γ        = 32 π G σ T_CMB^4 / (3 c^3 H0^2)
y_i        = m_i c^2 / (k_B T_ncdm T_CMB),   T_ncdm = 0.71611
N_ur       = N_eff - k (T_ncdm / (4/11)^(1/3))^4
F(y)       = (120 / 7 π^4) ∫_0^∞ x^2 sqrt(x^2 + y^2) / (e^x + 1) dx
Ω_de       = 1 - Ω_cb - Ω_k - Ω_r(1)

for k massive species. Dark energy closes the budget. The massive species sit at T_ncdm T_CMB, and each counts as (T_ncdm/(4/11)^1/3)⁴ massless species at zero mass, so every result is continuous in each mass at 0. G is the Codata release's, rounded once to the payload's dtype. ΛCDM is w0 = -1, wa = 0, where the dark energy's factor is 1 exactly; flat is Ω_k = 0, exactly; T_CMB = 0 removes the photons and every neutrino, massive ones included.

Evaluation. Every distance and time is one Gauss–Legendre sum of a⁴E², in a variable where its integrand is smooth from today to the big bang, over nodes fixed per dtype: 64 for distances and 56 for ages in float64, 24 and 20 in float32. The age is integrated from the big bang. Each function is a formula of nx operations: batched, compiled and differentiable in every parameter and in z, its derivative the sum's. Rune.grad (Nx.Ptree.instantiate (module Cosmology.Structure)) returns a cosmology whose fields are the derivatives, each per the unit its field holds.

Accuracy. Against the exact values of the same model, each distance, volume and time, Cosmology.hubble and Cosmology.critical_density has a relative error below 2^-46 in float64 and 2^-20 in float32, and each density parameter an error below the same bound as a fraction of the budget, on the validation box: Ω_cb ∈ [0.01, 3], Ω_k ∈ [-4.5, 2], w0 ∈ [-3, 1], wa ∈ [-3, 2] with w0 + wa < 0, h ∈ [0.4, 1], T_CMB ∈ [0, 3] K, N_eff ∈ [0, 5], masses ∈ [0, 1] eV, redshifts in [0, 1100] for distances, volumes and lookback times and [0, 10⁶] for ages, and a⁴E² at least 0.3 of the sum of its terms' magnitudes on the range each function integrates. Outside the box the same rule runs with no promise: where the terms of a⁴E² nearly cancel, as in a universe whose expansion stalls and resumes, the integrand has a near-singularity no fixed rule resolves.

Batching. A cosmology is an operand of an elementwise function. Its leaves, m_nu's without its last axis, broadcast with the redshifts from the right, and every result has the lanes' shape: leaves of shape [chains; 1] against redshifts [n] give [chains; n]. Leaves of shape [k; 1] give what vmap over k lanes of scalar leaves gives, row for row.

The domain. A distance, volume or time is NaN, with a zero derivative, where z ≤ -1, H₀ ≤ 0 or a⁴E² ≤ 0 at a node of its rule; Cosmology.hubble, Cosmology.density_parameter and Cosmology.critical_density are NaN where z ≤ -1, H₀ ≤ 0 or E²(z) ≤ 0. Cosmology.transverse is NaN where H₀ ≤ 0. A NaN parameter gives NaN with a NaN derivative. Every other lane is untouched, and nothing raises on a tensor's values.

Errors. Each function raises Invalid_argument, eagerly or at trace, on static data: a dtype other than float32 and float64, as in "Cosmology.hubble: float16 has no rule; the cosmology computes in float32 and float64"; a unit that does not convert, as in "Cosmology.age: h0 is in 1e3 m s^-1, which does not convert to s^-1: their quotient keeps m"; a scalar m_nu, as in "Cosmology.comoving_distance: m_nu is a scalar; its last axis lists the massive species ([1] for one, [0] for none)"; and shapes that do not broadcast, as in "Cosmology.distance_modulus: the cosmology's lanes [4] and z [1590] do not broadcast; for every redshift in each of 4 lanes, give leaves of shape [4; 1] or map with Rune.vmap".

Realisations

A realisation is a paper's flat ΛCDM fit, a starting point for record update: { planck with omega_k; w0 }. It transcribes the paper's decimals and rounds each field once to the dtype: Ω_cb is (ω_b + ω_c)/h² from the physical densities the paper samples, computed exactly. Its Ω_k, w0 and wa are 0, -1 and 0. The Planck realisations hold T_CMB = 2.7255 K (Fixsen 2009), N_eff = 3.046 and one massive species of 0.06 eV; the WMAP ones T_CMB = 2.725 K, N_eff = 3.04 and none. Each raises Invalid_argument for a dtype other than float32 and float64.

Expansion

Each function takes a cosmology and redshifts and returns the lanes' shape.

Distances

A distance is in c/u for h0 in u. For H₀ in km s^-1 Mpc^-1 that unit is exactly 299792.458 Mpc, and Quantity.value mpc applies it as one rounded factor.

Volumes

A volume is in (c/u)³ sr^-1 for h0 in u.

Times

A time is in 1/u for h0 in u.