SymBoltz.jl is a new Julia package for solving the linear Einstein-Boltzmann equations in cosmology. It is a “Boltzmann code” similar to CAMB (in Fortran + Python) and CLASS (in C + Python), which are the two most popular codes in use today.
These codes integrate a number of ODEs that describe the evolution and expansion of the universe since the Big Bang under a given cosmological model. These ODEs are very stiff, and Boltzmann codes have traditionally resorted to several approximation schemes to remove this stiffness, which make them quite complicated. From the solutions of the ODEs, the codes produce theoretical predictions for (statistics of) the cosmic microwave background, large-scale structure and distances (e.g. to supernovae). The predictions are compared to observed data to constrain cosmological models, commonly with Bayesian inference and MCMC methods.
The current standard model of cosmology is the ΛCDM model. It assumes that General Relativity is the theory of gravity, and that the universe consists of the cosmological constant (Λ), cold dark matter (CDM), baryons (“normal matter”), photons (interacting with baryons) and neutrinos (with and/or without mass). However, more precise data, the ΛCDM model is under tension and looks inconsistent. For example, prediction-to-data fits give different values for the Hubble parameterH_0 using supernova distances (probing the late universe and giving H_0 \approx 73\,\text{km}/\text{s}\,\text{Mpc}) versus statistics of the cosmic microwave background (probing mostly the early universe and giving H_0 \approx 67\,\text{km}/\text{s}\,\text{Mpc}). A big part of cosmology today is therefore to explore modified cosmological models, so Boltzmann codes are often modified to implement alternative models.
Motivated by this context, these are the main features of SymBoltz:
Symbolic-numeric: Models are specified with symbolic equations in a simple, convenient and extensible high-level format, then automatically compiled to fast numerical functions that are integrated by ODE solvers. This makes it much easier to implement modified models: one only has to write down the new/modified equations in one place in simple symbolic form, and SymBoltz automatically does some transformations on the symbolic equations that takes care of many “chores” one must do manually when modifying traditional codes. For example, SymBoltz can specify the equations of the full ΛCDM model in only 277 lines of code, while the equivalent code in CLASS is spread over 10 files with 27721 lines of code (i.e. a 100x simplification). The symbolic-numeric interface is made possible by ModelingToolkit.jl (big thanks to @cryptic.ax and @ChrisRackauckas).
Approximation-free: The full (stiff) equations are solved at all times without tight-coupling, ultrarelativistic fluid and radiation-streaming approximations (TCA, UFA and RSA) using efficient implicit ODE integrators from OrdinaryDiffEq.jl (again big thanks to @ChrisRackauckas and others). The lack of approximation schemes found in traditional codes make the code internals much simpler. This makes it much easier to implement modified models, because one does not have to invent new approximation schemes, validate the approximations for the modified model, and so on. The user only needs to provide one full set of equations to solve, which is not always the case with other codes.
Differentiable: Get derivatives of any output (e.g. cosmic microwave background statistics) with respect to any input (e.g. cosmological parameters) using automatic differentiation. For now this works with ForwardDiff.jl, and performance is good in some cases but not others. Cosmology is currently stuck with the Metropolis-Hastings algorithm, while next-generation cosmological surveys will make MCMCs more demanding with larger parameter spaces. My wet dream is to support reverse-mode to take fast likelihood gradients in full prediction-to-data fits with differentiable MCMC samplers, like HMC and NUTS.
SymBoltz.jl is available with documentation containing several tutorials and examples. You can also read more about the code design in the recent arXiv paper. I am very much open to questions, feedback, suggestions, issues and contributions to help shape and spread the package! I hope it can help to drive the cosmology community in Julia. And it would make me very happy if you want to try it or star the Github repository
Hi @hersle ! Great package!
Do you think it could be applied with warm DM, like in CLASS?
In particular I am intersted in mildly warm (mc^2 ~ 50-300 keV) particles.
If it is not a feature yet, I would like to contribute.
Thanks!
Hi @hersle , cool stuff! Nice to see that what you showed me back in Rome is now released:)
I have a couple of questions:
Backward AD. You clearly state in the paper that this is not wokring at the moment. Do you have any idea on the time scale to have that working? with a student of mine, @schiarenza , we are working on the new version of her NonLimber code, Blast.jl, which will soon include more contributions compared to the first release (e.g. RSD, Magnification bias, fNL and CMB lensing) and backward AD. It would be extremely cool if we could use your code as the cosmology engine beside one of my emulators.
Extensions. Have you considered to include an extension like non clustering Dark Energy as in w0-wa? It shouldn’t be too difficult, as it is just a matter of including some contributions to the background (and would make a nice addition as this is the main extension to LCDM that is being analyzed).
CMB lensing. Have you implemented the lensing of the CMB spectra or for the moment you just have the unlensed CMB spectra?
Mooncake is already supported. This is about making sure the rules, which apply to all AD front ends (Mooncake, Enzyme, ReverseDiff, Tracker, and Zygote) fully support the MTK features, which is currently missing a bit of SciMLStructures.jl support to be complete. We also have some early demonstrations of Enzyme.jl and Mooncake.jl direct adjoints on the solvers, now with integration tests that those do not regress (tests in the AD libraries that force them to keep compatibility ), but need to get this to support the implicit methods as well (currently only the explicit methods).
Thanks Chris! Yep, I remember when I first discovered it, I thought ModelingToolkit.jl + OrdinaryDiffEq.jl would be a very powerful combination for this project. I really appreciate your help and responsiveness with bug fixes and new features.
Hi @martinmestre and thank you! Warm dark matter models are not implemented yet, but I think it would be very interesting to add them and to test the extensibility of the code. Do you have a particular model in mind or a reference paper with equations? I do not have much prior experience with these models, but I suppose they could be based off the implementation of massive neutrinos, which evolve momentum bins of their distribution function in time. Let me know and I will also help you!
Like Chris says, this is mostly dependent on the underlying libraries, at least for what concerns the ODE solutions of the background and perturbations. I rely on standard ModelingToolkit.jl and OrdinaryDiffEq.jl functionality, so I hope reverse-mode support will eventually come “for free”.
I would love for you to train your emulator with SymBoltz! Right now it may not output all the quantities you want, though. Do you (plan to) use the direct output from CAMB/CLASS for RSD, fNL, etc.? Or do you compute it yourself from simpler quantities?
@ChrisRackauckas Just FYI, the ODE solving in SymBoltz relies on first solving one “background” ODE, followed by several independent “perturbation” ODEs where the “background” solution is supplied through interpolating DataInterpolations.jl splines (as time-dependent parameters). I don’t know if that makes reverse-mode more or less complicated. It would be fantastic if that works.
Yes, w_0w_a dark energy is included (see also e.g. here or here). I am also interesting in implementing other models to properly test the extensibility of the code. Let me know if you have something in mind.
It is not implemented yet. For now I have focused on getting good agreement with CLASS for the background + perturbations + matter power spectrum + (unlensed) CMB power spectrum for ΛCDM, which is challenging on its own. I have not done the calculation before, but from what I can tell it looks like another line-of-sight integration with a different source function, which should be relatively easy to add to the current framework. I opened this issue, feel free to track it and drop any comments there:
That should just work. It would be better to use discrete adjoints and the interpolation of the solver if that’s what you’re trying to do though (discrete adjoints since that’s required for the solver interpolation). But DataInterpolations in this form should just work.
