Integrating the Sympy-based Boltzmann solver from PyCosmo

DifferentiableUniverseInitiative / jax_cosmo

A differentiable cosmology library in JAX

MIT License

172 stars 36 forks source link

Integrating the Sympy-based Boltzmann solver from PyCosmo #87

Open EiffL opened 2 years ago

EiffL commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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

moserb94 commented 2 years ago

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.