CCL vs JaxCosmo comparaison nb: NaNs for some parameters

Hi,
if you use in jax_cosmo/notebooks/CCL_comparison.ipynb

# We first define equivalent CCL and jax_cosmo cosmologies

Omega_c_fidu =  0.2589
sigma8_fidu  = 0.8159
Omega_b_fidu = 0.0486 
h_fidu       = 0.6774 
n_s_fidu     =  0.9667 
w0_fidu      = -1.0

params_fidu = [Omega_c_fidu,sigma8_fidu,Omega_b_fidu,h_fidu,n_s_fidu,w0_fidu]


params_pathos =[[ 0.24986424,  0.55773633,  0.04630792,  0.61299905,  0.99486902,
        -1.44603776],
       [ 0.17611082,  0.71919114,  0.05874876,  0.67659724,  0.95329858,
        -1.41323164],
       [ 0.11939933,  0.44445573,  0.06162388,  0.79803975,  1.00682757,
        -1.77138837],
       [ 0.21723086,  0.4209687 ,  0.03940861,  0.73820567,  1.05830809,
        -1.77511559],
       [ 0.45726321,  0.60605765,  0.05444171,  0.62133009,  0.88678424,
        -1.83988558],
       [ 0.11030471,  0.49406485,  0.05778253,  0.88668811,  0.97388076,
        -1.94438598],
       [ 0.23417783,  0.6072679 ,  0.03107563,  0.72839096,  0.91472034,
        -0.94570266],
       [ 0.31449189,  0.48404425,  0.05613155,  0.56604649,  0.91932335,
        -1.11252456],
       [ 0.42029479,  0.40920462,  0.03177297,  0.76163707,  0.93350105,
        -1.78184122],
       [ 0.35075444,  0.43376266,  0.03103901,  0.65861007,  0.99913718,
        -1.70636958]]

params = params_fidu  # Ok

params = params_pathos[0]


Omega_c,sigma8,Omega_b,h,n_s,w0 = params

cosmo_ccl = ccl.Cosmology(
    Omega_c=Omega_c, Omega_b=Omega_b, h=h, sigma8 = sigma8, n_s=n_s, Neff=0,
    transfer_function='eisenstein_hu', matter_power_spectrum='halofit')

cosmo_jax = Cosmology(Omega_c=Omega_c, Omega_b=Omega_b, h=h, sigma8=sigma8, n_s=n_s,
                      Omega_k=0., w0=w0, wa=0.)

There are differences in many places and NaN for Cl in all tests (WL, gg clustering, cross)

Many thanks @jecampagne for finding these issues >.< this is super helpful.

To document our private communications, you have managed to track down the issue to a very specific place in the computation of the non-linear halofit correction.

Specifically, the issue comes from here:

jax_cosmo/jax_cosmo/power.py

Line 111 in c2f8b59

root = interp(np.atleast_1d(1.0), sigma, r)

When a is small enough, the non linear scale becomes really small, and sometimes goes outside of the range of scales
r = np.logspace(-3, 1, 256) considered in this function. If that happens, the return value of R_nl may become negative, and then all hell breaks loose further down in the halofit computation, as we would be taking roots of negative numbers, causing either complex or nan numbers to pop up.
The fix here is easy enough, we just want to prevent the non linear scale to reach negative numbers (or even zero for that matter).

To resolve an issue like this, here would be the recommended course of action:

: write a test that fails in the situation encountered
: implement a fix
: check that the test passes with the new fix
: open a PR

All good, fix is implemented in this PR #78
it simply computes the root in log space, then takes an exponential, to ensure positivity. A small clipping at 1e-6 ensures that we never actually reach r=0. This should no longer result in any NaNs.

@all-contributors please add @jecampagne for bug reports

@EiffL

I've put up a pull request to add @jecampagne! 🎉