cosmicjellyfish.cosmology.jaxcosmo_theory#

Pure-JAX theory layer of the jaxcosmo backend.

Everything in this module is a pure function of the cosmological parameters, written with jax.numpy only, so that it can be traced by jax.jit / jax.jacfwd. No NumPy/SciPy objects (splines, np.vectorize …) are used here; the conversion to the spline-like objects expected by cosmicjellyfish.cosmology.cosmology.cosmo_functions happens at the outer boundary, in cosmicjellyfish.cosmology.jaxcosmo_provider.

Model#

  • Flat universe, massless neutrinos only (mnu = 0), no radiation in the background (this is the jax_cosmo convention: E^2 = Om a^-3 + Ode f(a)).

  • Linear power: Eisenstein & Hu transfer function, normalised to sigma8.

  • sigma8 is derived from As with the JAX port of the symbolic_pofk fit (cosmicjellyfish.cosmology.jaxcosmo_amplitude), inside the traced function, so that the derivative with respect to As is exact with respect to the fit.

  • Non-linear power: Halofit (takahashi2012 or smith2003) as implemented in jax_cosmo.

Units#

Public functions of this module return the CosmicJellyfish conventions: distances in Mpc, H/c in 1/Mpc, wavenumbers in h/Mpc and power spectra in Mpc^3. Internally jax_cosmo works in Mpc/h, h/Mpc and (Mpc/h)^3; the conversions are done in compute_tables() only.

Notes

jax_cosmo caches tabulated background quantities in cosmo._workspace (mutable state). A fresh jax_cosmo.Cosmology is therefore created inside every traced function and the caches are primed before any vmap, so no tracer can leak out of a transformation.

jax must run in double precision for Fisher work, call enable_x64() before creating any array.

Module attributes

PARAM_ORDER

Order of the cosmological parameters used by the vector interface of the autodiff layer.

FREE_NAME_MAP

Mapping from CosmicJellyfish parameter names to the canonical names of this backend.

Functions

compute_tables

Evaluate all theory quantities on a (z, k) grid, in CosmicJellyfish units.

enable_x64

Switch JAX to double precision (required for Fisher matrices).

make_cosmology

Create a fresh jax_cosmo.Cosmology from CosmicJellyfish-style parameters.

params_from_vector

Unpack a parameter vector (ordered as names) into the dict used by this module.

Classes

JaxCosmoTheory

jax_cosmo implementation of cosmicjellyfish.cosmology.jax_theory.JaxTheory.

TheoryConfig

Static (non-differentiable) numerical choices of the JAX theory.