cosmicjellyfish.cosmology.jaxcosmo_halofit

Smooth Halofit for the jaxcosmo theory layer.

jax_cosmo.power.halofit locates the nonlinear scale by interpolating sigma(R) = 1 linearly on a fixed 256-node grid in log R. The resulting nonlinear spectrum is only piecewise smooth in the cosmological parameters: whenever the root moves between two grid nodes the derivative jumps. For a 1% parameter step those jumps reach ~13% in individual redshift rows, which is fatal both for autodiff (it differentiates one linear piece) and for finite differences (they average over several).

This module keeps the Halofit prescription and its coefficients - taken from jax_cosmo.power (MIT licence, Copyright (c) 2022 Differentiable Universe Initiative) - and only replaces the root find with a fixed-count Newton iteration on the smooth integral

\[\sigma^2(R) = \int d\ln k \, \frac{k^3 P_{\rm lin}(k, a)}{2\pi^2} e^{-k^2 R^2},\]

which converges to machine precision and is differentiable. Everything else (n_eff, C, the a_n ... nu_n polynomials and equation C2) follows jax-cosmo so that the two implementations agree to the accuracy of their respective root finds.

Functions

halofit	Nonlinear matter power spectrum P(k) in (Mpc/h)^3.
nonlinear_scale	(k_nl, n_eff, C) in h/Mpc and dimensionless, smooth in the parameters.
nonlinear_scale_residual	|growth^2 sigma^2(R_nl) - 1| at the Newton solution (~1e-15 when converged).