Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.
And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.
Integrating the Sympy-based Boltzmann solver from PyCosmo
See original GitHub issueFollowing 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.odeintto integrate in time the system - Profit!
And boom! You got yourself a diffable Boltzman solver!
Issue Analytics
- State:
- Created 2 years ago
- Comments:5 (3 by maintainers)
Top Results From Across the Web
Top Related Medium Post
Top Related StackOverflow Question
Troubleshoot Live Code
Top Related Reddit Thread
Top Related Hackernoon Post
Top Related Tweet
Top Related Dev.to Post
Top Related Hashnode Post
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.
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.