cosmicjellyfish.cosmology.jaxcosmo_amplitude#
JAX port of the SYREN-NEW As <-> sigma8 conversion.
jax-cosmo normalises its linear power spectrum with sigma8 and has no
primordial amplitude A_s input. CosmicJellyfish parameterises cosmologies
with 10^9 A_s, so the jaxcosmo backend needs a differentiable map
(10^9 A_s, Omega_m, Omega_b, h, n_s, m_nu, w0, w_a) -> sigma8 that is
evaluated inside the traced function. Treating sigma8 as a constant
would silently give wrong derivatives with respect to A_s and to every
parameter on which the conversion depends.
The expressions below are the “max precision” symbolic fits of
symbolic_pofk.linear_new (Bartlett et al.), ported verbatim to
jax.numpy so that they can be traced, jitted and differentiated. The
coefficients are copied from the upstream symbolic_pofk package (the
revision used for development is locked in uv.lock); tests should
compare against that package.
All functions are pure and accept scalars or arrays. As is 10^9 A_s.
The fit is valid for the symbolic_pofk training domain; callers are
responsible for staying inside it.
Functions
Emulated |
|
Emulated |