ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum

TL;DR

ABCMB is a differentiable, GPU-compatible Python+JAX package for computing the CMB power spectrum, incorporating advanced physical effects and enabling efficient gradient-based cosmological analyses.

Contribution

It introduces ABCMB, a novel, extensible, and GPU-accelerated Einstein-Boltzmann solver with accurate gradients for the first time in this context.

Findings

Sub-percent agreement with CLASS

Faster run times on GPU

Stable gradients for Fisher analyses

Abstract

We present ABCMB, a differentiable Einstein-Boltzmann solver for the cosmic microwave background (CMB). ABCMB is a complete code capturing important effects to linear order in $\Lambda{\rm CDM}$ cosmology. It computes the CMB power spectrum and includes effects like lensing, polarization, massive neutrinos, and a state-of-the-art treatment of BBN and recombination. ABCMB has sub-percent-level agreement with CLASS and can be run on a GPU with competitive, and sometimes even faster, run times. It is refactored compared to previous codes and takes advantage of object-oriented programming to improve extensibility, meaning new physics can be added to it without the need for modifying source files. ABCMB provides accurate and stable gradients to the user, making Fisher analyses straightforward, and enabling the use of efficient gradient-based sampling methods.