Capse.jl

Capse.jl is a Julia package designed to emulate the computation of the CMB Angular Power Spectrum, with a speedup of several orders of magnitude compared to standard codes.

Installation

In order to install Capse.jl, run on the Julia REPL

using Pkg, Pkg.add(url="https://github.com/CosmologicalEmulators/Capse.jl")

Usage

In order to be able to use Capse.jl, there two major steps that need to be performed:

  • Instantiating the emulators, e.g. initializing the Neural Network, its weight and biases and the quantities employed in pre and post-processing
  • Use the instantiated emulators to retrieve the spectra

In the reminder of this section we are showing how to do this.

Instantiation

The most direct way to instantiate an official trained emulators is given by the following one-liner

Cℓ_emu = Capse.load_emulator(weights_folder);

where weights_folder is the path to the folder containing the files required to build up the network. Some of the trained emulators can be found on Zenodo and we plan to release more of them there in the future.

It is possible to pass an additional argument to the previous function, which is used to choose between the two NN backend now available:

  • SimpleChains, which is taylored for small NN running on a CPU
  • Lux, which can run both on CPUs and GPUs

SimpleChains.jl is faster expecially for small NN on the CPU. If you wanna use something running on a GPU, you should use Lux.jl, which can be done adding an additional argument to the load_emulator function, Capse.LuxEmulator

Cℓ_emu = Capse.load_emulator(weights_folder, Capse.LuxEmulator);

Each trained emulator should be shipped with a description within the JSON file. In order to print the description, just runs:

Capse.get_emulator_description(Cℓ_emu)
The parameters the model has been trained are, in the following order: ln10As, ns, H0, ωb, ωc, τ.
The emulator has been trained by Marco Bonici.
Marco Bonici email is bonici.marco@gmail.com.
The emulator has been trained on the high-precision-settings prediction as computed by the CAMB Boltzmann solver.
Warning

Cosmological parameters must be fed to Capse.jl with arrays. It is the user responsability to check the right ordering, by reading the output of the get_emulator_description method.

After loading a trained emulator, feed it some input parameters x in order to get the emulated $C_\ell$'s

x = rand(6) # generate some random input
Capse.get_Cℓ(x, Cℓ_emu) #compute the Cℓ's

Using SimpleChains.jl, we obtain a mean execution time of 45 microseconds

BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (minmax):  37.816 μs 2.830 ms   GC (min … max): 0.00% … 90.50%
 Time  (median):     44.850 μs               GC (median):    0.00%
 Time  (mean ± σ):   47.868 μs ± 66.928 μs   GC (mean ± σ):  3.51% ±  2.49%

        ▂▅█▃▂                                                 
  ▂▂▂▂▃▆█████▇▄▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▂▂▂▂▁▁▂ ▃
  37.8 μs         Histogram: frequency by time        82.3 μs <

 Memory estimate: 78.77 KiB, allocs estimate: 11.

Using Lux.jl, with the same architecture and weights, we obtain

BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (minmax):  69.908 μs 3.145 ms   GC (min … max): 0.00% … 89.92%
 Time  (median):     85.475 μs               GC (median):    0.00%
 Time  (mean ± σ):   89.488 μs ± 96.979 μs   GC (mean ± σ):  3.53% ±  3.19%

                           ▁▂▁▁▄█                      
  ▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▃▃▄▆▅▆████████▇▇▆▅▄▃▃▃▃▃▃▂▃▂▂▂▂▂▂▂▂▂▂▂ ▃
  69.9 μs         Histogram: frequency by time        99.8 μs <

 Memory estimate: 166.09 KiB, allocs estimate: 36.

SimpleChains.jl is about 2 times faster than Lux.jl and they give the same result up to floating point precision.

These benchmarks have been performed locally, with a 12th Gen Intel® Core™ i7-1260P.

Considering that a high-precision settings calculation performed with CAMB on the same machine requires around 60 seconds, Capse.jl is 5-6 order of magnitudes faster.

Authors

  • Marco Bonici, PostDoctoral researcher at Waterloo Center for Astrophysics
  • Federico Bianchini, PostDoctoral researcher at Stanford
  • Jaime Ruiz-Zapatero, PhD Student at Oxford
  • Marius Millea, Researcher at UC Davis and Berkeley Center for Cosmological Physics

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

Please make sure to update tests as appropriate.

License

Capse.jl is licensed under the MIT "Expat" license; see LICENSE for the full license text.