jaxcapse

Repo containing the jaxcapse emulator.

Documentation

• Stable Documentation - Latest release documentation
• Development Documentation - Latest development version documentation

Installation and usage

In order to install jaxcapse, you can just run

pip install jaxcapse

If you prefer to use the latest version from the repository, you can clone it, enter it, and run

pip install .

In order to use the emulators, you have to import jaxcapse and load a trained emulator

import jaxcapse
import jax.numpy as np
trained_emu = jaxcapse.load_emulator("/path/to/emu/")

Then you are good to! You have to create an input array and retrieve your calculation result

input_array = np.array([...]) #write in the relevant numbers
result = trained_emu.get_Cl(input_array)

# The returned grid always matches the prediction. Subsampled emulators with
# at most 2048 knots are transparently interpolated with a cubic spline.
ell = trained_emu.get_ell_grid()
ell_training = trained_emu.get_training_ell_grid()

For automatically interpolated grids, source bounds within 0.1 of an integer are snapped to that integer. Bounds farther away are moved inward.

CAMB Mnu-w0-wa-CDM emulators

The published CAMB + CosmoRec artifact contains TT, TE, EE, BB, and PP models. jaxcapse downloads them into a separate model-specific cache so they cannot be confused with the existing LCDM models:

import jax.numpy as jnp
import jaxcapse

params = jnp.array([3.044, 0.965, 0.054, 67.4, 0.02237, 0.120, 0.06, -1.0, 0.0])
tt = jaxcapse.trained_emulators["camb_mnuw0wacdm"]["TT"]
D_ell_TT = tt.get_Cl(params)
ell = tt.get_ell_grid()  # 2:9500

The input order is ln10As, ns, tau, H0, omega_b, omega_c, Mnu, w0, wa, restricted to w0 + wa < -0.5. Exact Mnu = 0 was not in the training set. Despite the method name, TT/TE/EE/BB outputs are lensed D_ell in microK^2; PP is [ell*(ell+1)]^2 C_ell^phiphi/(2*pi) (dimensionless).

For a more detailed explanation, check the tutorial in the notebooks folder, which also shows a comparison with the standard CAMB Boltzmann solver.

Citing

Free usage of the software in this repository is provided, given that you cite our release paper.

M. Bonici, F. Bianchini, J. Ruiz-Zapatero, Capse: efficient and auto-differentiable CMB power spectra emulation