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

This paper presents ABCMB (Autodifferentiable Boltzmann solver for the CMB), a new differentiable Einstein-Boltzmann solver designed for cosmic microwave background (CMB) analysis. As precision cosmology moves toward higher-dimensional parameter spaces and more complex models, the ability to perform efficient, gradient-based parameter estimation is essential to overcome the computational expense of traditional Markov-Chain Monte Carlo methods.

Architecture and philosophy

ABCMB utilizes a significant departure from predecessor codes by employing an object-oriented architecture refactored for improved extensibility. The code is built using JAX, which allows it to be Just-In-Time (JIT) compiled for faster run times and makes it backend-agnostic, meaning it can be trivially accelerated on GPUs. The core philosophy is to minimize the need for users to modify source files; instead, new physics can be introduced by defining new fluid modules that are passed into ABCMB at initialization.

The computation is managed by a centralized model that dispatches several specialized modules:

• Background modules for computing thermodynamic quantities and managing recombination.

• Perturbation modules for solving the evolution of coupled ordinary differential equations.

• Spectrum modules for integrating transfer functions to compute power spectra.

Key capabilities and physics

The package is a complete code capturing important effects to linear order in CDM cosmology. It provides a state-of-the-art Einstein-Boltzmann solver that achieves subpercent agreement with established codes like CLASS. The primary breakthroughs of the package include:

1. Ease of use and extensibility, allowing new physics to be added without ever needing to open a source file.

2. A high precision recombination calculation provided by the companion code HyRex, which is a differentiable version of HYREC-2.

3. A fast, state-of-the-art BBN calculation via the inclusion of LINX, enabling precise theoretical predictions rather than interpolating pre-tabulated results.

The solver currently includes effects such as lensing, E-mode polarization, and massive neutrinos.

Differentiability and gradients

A major advantage of ABCMB is that it is fully differentiable, meaning functional gradients are computed using automatic differentiation (AD) rather than numerical derivatives. This makes the code particularly powerful for Fisher forecasting and enables the use of efficient, gradient-based sampling algorithms such as Hamiltonian Monte Carlo (HMC). Because ABCMB uses forward AD, it is well-suited for its large number of outputs, including thousands of C and P(k) values.

The differentiability extends across the entire pipeline; when used with LINX, the entire pipeline from BBN through recombination and the CMB power spectrum is fully differentiable. This allows researchers to explore complex interdependencies between BBN abundances, baryon density, and effective neutrino numbers.

Performance and validation

ABCMB has been validated against HYREC-2 and CLASS, showing subpercent accuracy in the matter power spectrum and CMB temperature/E-mode spectra. In terms of performance, the code is highly efficient on GPU hardware. While CLASS uses various approximation schemes to achieve speed, ABCMB's JAX-based implementation allows it to remain competitive or even faster in certain regimes.

The paper highlights several performance characteristics:

• On NVIDIA GPUs, ABCMB can outperform the CLASS run time when considering modes up to = 4000.

• The runtime scales more favorably with the number of modes considered than codes like CLASS do.

• The code is backend-agnostic, allowing it to be run on both CPU and GPU.

1. Integration of Fully Differentiable Black Box Simulators via Forward Automatic Differentiation (AD)

• What the improved AI system can do: Instead of treating external scientific or engineering simulators as non-differentiable black boxes that require slow, error-prone numerical finite-difference methods, the system can integrate them directly into its optimization loop. By utilizing forward AD, the AI can compute stable, high-precision functional gradients through complex, non-linear differential equations. This enables the use of hyper-efficient gradient-based sampling (e.g., Hamiltonian Monte Carlo) for real-time parameter estimation and design optimization in high-dimensional scientific datasets.

2. Object-Oriented Modular Extensibility for Autonomous Knowledge Injection

• What the improved AI system can do: The system can move away from monolithic architectures toward a template inheritance model. When the AI encounters new domain knowledge (e.g., a new set of chemical reaction rates, physical laws, or mathematical axioms), it can automatically wrap that knowledge into an independent, object-oriented module that inherits from existing base classes. This allows the AI to integrate and reason with entirely new species of logic or physics within its core engine without requiring model retraining or manual modification of its underlying source code.

3. High-Dimensional Output Optimization via Forward-Mode Gradient Computation

• What the improved AI system can do: In specialized engineering or generative design tasks where the output dimensionality (e.g., thousands of physical properties, structural points, or molecular states) is significantly larger than the input parameter space, the system can switch from standard reverse-mode backpropagation to forward-mode differentiation (jacfwd). This allows for more efficient and computationally stable gradient calculations across massive output vectors, facilitating hyper-accurate optimization in complex design landscapes.

4. JIT-Compiled, Backend-Agnostic Reasoning Kernels

• What the improved AI system can do: The system can implement Just-In-Time (JIT) compilation for its mathematical and logical reasoning subroutines. This allows the AI to dynamically compile its own reasoning chains into highly optimized machine code that is backend-agnostic, meaning the same logic can transition seamlessly between CPU, GPU, and TPU environments. This ensures maximum computational throughput and minimal latency when scaling from simple logical tasks to heavy numerical simulations.

We present ABCMB, a differentiable Einstein-Boltzmann solver for the cosmic microwave background (CMB). ABCMB is an analysis-ready code package capturing important effects to linear order in Λ CDM cosmology. It computes the CMB power spectrum and includes effects like lensing, polarization, massive neutrinos, and a state-of-the-art treatment of Big Bang Nucleosynthesis 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, owing to ABCMB's favorable run time scaling with maximum multipole moment. 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.