Integrating the Sympy-based Boltzmann solver from PyCosmo

[EiffL]

Following a very nice talk from Beatrice Moser, I looked a little bit at how PyCosmo is implementing their own Boltzmann code and it's actually pretty nice and could be pretty easily ported to JAX I think.

The ODE is specificied using Sympy, which makes it pretty easy to write down the equations, and then they use their own sympy2c code to transform the python code into JIT compiled code to run the actual ODE solver.

So that's pretty cool :-) But we could make it way cooler by doing the following:

• Copy/Paste the Sympy equations from PyCosmo (which is GPL licensed)
• Use Miles' sympy2jax library https://github.com/MilesCranmer/sympy2jax to turn the equations into JAX functions
• Use jax.experimental.odeint to integrate in time the system
• Profit!

And boom! You got yourself a diffable Boltzman solver!

[EiffL]

And tagging @moserb94 on this issue.

After reading more of the PyCosmo code I realized that it still relied on an external RecFast implementation however :-( so not completely as easy as I was hoping. One would have to recode RecFast also in JAX, in the same way as Zack did it for Bolt.jl here: https://github.com/xzackli/Bolt.jl/blob/main/src/ionization/recfast.jl

The whole thing is still doable :-) But slightly more involved than I was hoping.

[EiffL]

Also I'm thinking it would be worthwhile to host a JAX Boltzmann solver as a separate project, as opposed to integrating it directly into jax-cosmo.

[moserb94]

Thank you for tagging me. It sounds pretty cool (even though I just started looking into JAX)! Yes, we rely on RecFast++ to compute the evolution of thermodynamic quantities and we are actually planning to move to HYREC-2 in the near future.

[EiffL]

interesting :-) By any chance.... would you be thinking of doing an implementation in sympy as well for hyrec-2 ?

[moserb94]

The idea was to use hyrec-2 the way we use RecFast++ (meaning a python wrapper and an interpolator), rather than rewriting the code in sympy... so that wouldn't help I guess.