Implementation of the Magnitude Kernel and the RSD Kernel in jax_cosmo

[bhorowitz]

I implemented the new kernels to the best of my knowledge, based on the formulas here: https://arxiv.org/pdf/1812.05995.pdf

One thing I had to do for the RSD version is change the kernel to take a z1 variable (corresponding to z_{ell+1}).

NOTE! I wasn't able to get any of these functions working in CCL so I don't really have anything to compare directly against. In particular the s(z) function for the mag_bias I have no idea what a reasonable version would look like...

[bhorowitz]

It would be nice to not bother calculating the z1 component unless has_RSD is enabled, but idk how to check that from the angular cl function.

[EiffL]

So I was looking into the RSD part first, I see the issue with z1, and I managed to compare to CCL quickly, here are clustering cls with and without RSD:

So not exactly that ^^' but I suspect it might just be a scaling factor.

I see what you have implmented, and by curiosity I went to see how CCL does it, and they explicitly keep track of bessel derivatives in the integral: https://github.com/LSSTDESC/CCL/blob/60704b5245689e178a5e32fc0504156d8bb7ac1c/pyccl/tracers.py#L324
So, may have to adopt a similar mechanism somehow, so that we don't have to carry this z1 variable when not needed.
I'm thinking more about this ^^' Will check the mag kernel first

[bhorowitz]

EDIT: I fixed a few more bugs, and RSD now has very minimal effects on this example. Can you send me the pyccl output and I can play around some more?

[EiffL]

@all-contributors please add @bhorowitz for code

[allcontributors]

@EiffL

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

[dkirkby]

What is the status of this PR @EiffL?