cloelib: A Flexible Python Library for Computing Cosmological Observables in the Euclid Era

Technical Summary of cloelib: A Flexible Python Library for Computing Cosmological Observables in the Euclid Era

Problem Statement
The field of observational cosmology is entering a precision era driven by Stage IV galaxy surveys, including the European Space Agency's Euclid mission, DESI, the Vera C. Rubin Observatory, and the Nancy Grace Roman Space Telescope. These missions generate vast datasets requiring theoretical predictions that are both accurate and computationally efficient to constrain cosmological models. While existing frameworks (e.g., CCL, CosmoSIS, CAMB, CLASS) provide robust tools, they often lack the flexibility to seamlessly integrate diverse pre-existing components or to explore broad theoretical models alongside comprehensive systematic effects. Specifically, the original CLOE (Cosmology Likelihood for Observables in Euclid) software, while central to Euclid analyses, suffered from structural limitations: extending it for new observables or modeling choices required intrusive modifications, leading to technical debt and reduced maintainability. Furthermore, the framework was not originally designed to support the independent development and integration of new theoretical models or systematics within a unified pipeline, nor did it natively support modern gradient-based inference techniques.

Methodology and Software Design
cloelib is a fully Pythonic library designed to operationalize the flexibility and efficiency required for modern cosmological inference. It represents a structural evolution of the CLOE software, introducing a modular architecture based on Python protocols (PEP 544) to enforce a "plug-and-play" approach.

The library organizes functionality into four distinct layers:

1. Cosmological Backgrounds: Handling expansion history and distances.
2. Perturbation Theory: Managing linear and non-linear matter power spectra.
3. Observables: Computing cosmic shear, photometric/spectroscopic clustering, and CMB lensing.
4. Summary Statistics: Generating angular power spectra, correlation functions, and Legendre multipoles.

Key technical features include:

• Protocol-Based Interfaces: Components interact through well-defined interfaces (e.g., Background, Perturbations, Tracer, SpectroPower) rather than rigid inheritance hierarchies. This allows researchers to switch between different theoretical backends (e.g., Boltzmann solvers like CAMB and CLASS, or emulators like CosmoPower, BACCOemu, and EuclidEmulator2) without modifying the analysis pipeline.
• JAX Integration: A core feature is the native integration with JAX, enabling automatic differentiation, efficient gradient-based computation, and just-in-time (JIT) compilation. This transforms conventional pipelines into fully differentiable programs, facilitating advanced inference methods such as Hamiltonian Monte Carlo and neural network-based modeling.
• Systematic Modeling: The library includes implementations for survey-specific systematics, such as Alcock–Paczynski corrections, shear and photometric redshift calibration, and spectroscopic purity.
• Performance Optimization: Critical sections utilize JIT compilation, and a caching system avoids redundant evaluations. The library supports both linear and non-linear perturbation theories, interfacing with state-of-the-art emulators (e.g., HMCode2020Emu, FlamingoBaryonResponseEmulator) and spectroscopic codes (PBJ, comet-emu).

Key Contributions and Capabilities

• Unified Interface: cloelib is the first large-scale structure code to implement a unified interface to multiple cosmological codes using Python protocols, enabling seamless comparison of different theoretical implementations within a consistent setup.
• Differentiable Cosmology: By leveraging JAX, the library enables the training of neural network emulators and the use of gradient-based samplers, which are often difficult to implement efficiently in traditional frameworks.
• End-to-End Analysis: Integrated with the companion likelihood module cloelike, the library supports the full chain from observable computation to likelihood evaluation and posterior sampling for parameters.
• Modular Extensibility: The use of Python mixins and protocols allows users to extend functionality (e.g., adding baryonic effects) without altering underlying implementations.

Results and Performance
The paper reports computational performance benchmarks on an Apple M4 chip:

• Initialization: Initializing shear and position tracer classes takes approximately 0.4 seconds; subsequent calls are effectively instantaneous due to caching.
• Photometric Probes: Computing angular power spectra for a 3x2-pt analysis (3000 multipoles, 512 k-values, 500 redshifts, 6 bins) using the Limber approximation takes approximately 0.06 seconds per call after a one-second initialization phase. For a Euclid-like configuration with 13 bins (351 spectra), the time is approximately 0.1 seconds per call.
• Spectroscopic Probes: Computing full-shape spectroscopic Legendre multipoles for a single redshift bin takes approximately 3 milliseconds.
• Comparison: These results represent a significant speedup compared to the former CLOE software, while maintaining percent-level accuracy through the use of state-of-the-art emulators.

Significance
The paper positions cloelib as the reference theory computation infrastructure for Euclid's first cosmological release. Its significance lies in bridging traditional numerical cosmology with modern optimization techniques and emerging machine learning approaches. By providing a flexible, extensible, and differentiable platform, cloelib addresses the structural limitations of previous tools, enabling faster and more streamlined statistical analyses. It is designed to meet the demands of the next generation of precision cosmology experiments, facilitating systematic studies, model comparisons, and robust cross-validation of cosmological results within the Euclid Consortium and the wider community.