[ANN] SymBoltz.jl: a symbolic-numeric, approximation-free and differentiable linear Einstein-Boltzmann solver for cosmology

Package Announcements
, ,
1

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 parameter H_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 :star:

31 Likes
2

Really cool stuff, I had no idea what you were actually doing with MTK but now this might be the coolest application :sweat_smile:

7 Likes
3

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!

1 Like
4

Hi @hersle , cool stuff! Nice to see that what you showed me back in Rome is now released:)
I have a couple of questions:

  1. 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.
  2. 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).
  3. CMB lensing. Have you implemented the lensing of the CMB spectra or for the moment you just have the unlensed CMB spectra?
2 Likes
5

It’s probably more dependent on me? MTK integration with Gauss adjoint will just work out of the box, but

We need to do those two in order to make it generally stable. It’s near the top of my list. Then all MTK codes should be fine.

2 Likes
6

@ChrisRackauckas , this is with Enzyme, right? Will you include support for other AD systems as well, like MoonCake?

7

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 :sweat_smile:), but need to get this to support the implicit methods as well (currently only the explicit methods).

2 Likes
8

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.

9

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!

10

Hi and thanks, @marcobonici and @schiarenza!

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:

11

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.

1 Like
12

I have a few questions/comments regarding getting started. I was looking at SymBoltz.jl/scripts/cmb_source.jl, and the script does not work for me.
In line 15 I get an error at

τs = sol1[M.τ]

M has no member τ. Is it supposed to be M.t?

Also, on line 19:

println("max(abs(S2-S1)) = ", maximum(abs.((S2 .- S1)[1])))

I have to change this line for the code to run.

On line 31

Θls = SymBoltz.los_integrate(S, ls, ks, τs, us, u′s)

gives an error:

ERROR: MethodError: no method matching los_integrate(::Tuple{…}, ::StepRange{…}, ::StepRangeLen{…}, ::Vector{…}, ::StepRangeLen{…}, ::Vector{…})
The function `los_integrate` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  los_integrate(::AbstractArray{T}, ::AbstractArray{T}, ::AbstractArray, ::AbstractRange, ::AbstractArray, ::AbstractRange, ::AbstractArray; integrator, verbose) where T<:Real
   @ SymBoltz ~/.julia/packages/SymBoltz/kRJhS/src/spectra.jl:158

Stacktrace:
 [1] top-level scope
   @ REPL[36]:1
Some type information was truncated. Use `show(err)` to see complete types.

I am running the code using julia 1.11.0 version. Is it possible that the documentation is outdated or written with an older julia version?

I really appreciate your effort in writing this package. Thanks a lot!!!

Edit: M does have tau if I install julia version 1.12.0, but then SymBoltz does not have a member named cmb_ks. Strange!

13

Hi @Shashank. The scripts/ directory just contains some dirty scripts I have used for debugging during development of earlier versions. It is not updated. I should probably remove it.

Please refer to the documentation pages for up-to-date examples. And do let me know if you want to do something not covered there :slight_smile:

14

Thanks a lot. The scripts on the documentation page work well for me, but only if I run them on version 1.12.0, but do not work on 1.11.0. The members of M are different depending on which version of julia I am using. There may be some problem with my 1.11.0 installation.

15

Hi @hersle! Thanks! I am not expert in the field. I would need these WDM solutions for studying stellar stream sub-halo perturbations. In particular I am interested in the DM family named Relativistic Fermionic model (https://arxiv.org/pdf/1908.10806). I will write DM to discuss further. Best.

1 Like
16

In “getting started”

ls = 10:10:1000
jl = SphericalBesselCache(ls)
Dls = spectrum_cmb(:TT, prob, jl; normalization = :Dl, unit = u"μK")
plot(ls, Dls; xlabel = "l", ylabel = "l (l+1) Cₗ / 2π", label = nothing)

does not work. I get an error

ERROR: UndefVarError: `SphericalBesselCache` not defined in `Main`
Suggestion: check for spelling errors or missing imports.
Stacktrace:
 [1] top-level scope
   @ REPL[3]:1

Is SphericalBesselCache supposed to be imported from an external module or defined by the user?

Edit: I resolved the issue. I installed SphericalFunctions.jl then uninstalled Symboltz and then reinstalled SymBoltz. Now SphericalBesselCache is there.

17

Apologies if my questions are somewhat naive, I’m out of the GR game for some years.

So the way I understand this, you solve some FLRW-style ODE for homogeneous isotropic background geometry, and then on top of that compute stability in direction of various inhomogeneities, and have some machinery to match to observational data?

Something I’d find very interesting is to look at homogeneous anisotropic background, i.e. Bianchi, especially type 8 and 9 with the fun chaotic BKL / Mixmaster attractor + particle horizons near the initial singularity / big bang. I could talk at lengths about Mixmaster, but never got to look at how its particle/causal horizons interact with inhomogenous pertubations (I’m more of an ODE than GR person and I would fail miserably if you asked me to write down equations for that).

Mixmaster may look numerically intractable on a naive glance (the heteroclinic connections can mess you up!), but is actually very tractable if you locally view it as pertubations of Bianchi 1 and 2 (and kinda 6_0 and 7_0 and LRS / Taub). You just need to give up “have one set of equations on one set of coordinates” and stitch together different methods in different regimes.

Any intentions on looking in that direction, ie anisotropic background geometries?

1 Like
18

@Shashank I added SphericalBesselCache in a recent update. The reason for this error was simply that you were using the documentation for a newer version of SymBoltz than you had installed. The same is probably true what you encountered with different members of M above. Typing using Pkg; Pkg.update() would update SymBoltz and fix the problem, and reinstalling it like you did accomplishes the same. SphericalFunctions.jl has nothing to do with it, though. :slight_smile:

19

@foobar_lv2 Yes, it solves ODEs for a linearly perturbed FLRW metric, i.e. first-order corrections to a homogeneous and isotropic universe. SymBoltz produces theoretical predictions for a given model, which can be compared to observations (observational data is outside SymBoltz’ scope).

I don’t have any hands-on familiarity with these models, but it sounds interesting. As long as they have a “background+perturbations”-structure that is not too different from (say) ΛCDM, I think they should be possible to implement in SymBoltz’ framework. Up until now my main priority has been to get the code to work and agree well with existing standard codes for close-to standard models.

2 Likes