🔭 astrophysics

halox: Dark matter halo properties and large-scale structure calculations using JAX

The paper introduces **halox**, a JAX-powered Python library that enables differentiable and accelerated computations of dark matter halo properties and the halo mass function, facilitating their integration into gradient-based workflows such as Hamiltonian Monte Carlo sampling and machine learning applications.

Original authors: Florian Kéruzoré, Lance A. Moreau

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Florian Kéruzoré, Lance A. Moreau

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the universe as a giant, invisible scaffolding made of "dark matter." You can't see this scaffolding, but it holds everything together. The heavy, clumpy parts of this scaffolding are called dark matter halos. These halos are like the invisible foundations of a city; they are the gravitational "wells" where galaxies (our cities) and clusters of galaxies (our metropolises) eventually form and settle.

To understand how the universe grew, scientists need to measure these invisible foundations very precisely. They need to know their size, their density, and how many of them exist at different times in the universe's history.

The Problem: The Old Tools Were Too Slow
For a long time, scientists used software tools (like a toolbox called colossus) to calculate these properties. Think of these old tools as a very skilled, but slow, human accountant. They could do the math, but if you asked them to calculate the properties of millions of halos at once, or if you needed to tweak the numbers slightly to see how the results changed (a process called "gradient-based analysis"), the human accountant would take forever. They were also stuck working on a single desk (a standard computer processor) and couldn't easily use the super-fast power of modern graphics cards (GPUs).

The Solution: Enter halox
The authors of this paper built a new tool called halox. If the old tools were a human accountant, halox is like a team of super-fast robots working in a high-tech factory, all coordinated by a smart brain called JAX.

Here is what makes halox special, using simple analogies:

The "Smart Brain" (Differentiability):
In the past, if a scientist wanted to see how a tiny change in the universe's rules would affect the halos, they had to manually re-calculate everything from scratch. halox is "differentiable." Imagine driving a car where the dashboard instantly tells you exactly how much your speed will change if you press the gas pedal just a tiny bit. halox does this for the universe. It can instantly tell scientists how the results change if they tweak the input numbers. This is crucial for advanced methods like "Hamiltonian Monte Carlo," which is a fancy way of saying "finding the best answer by exploring many possibilities efficiently."
The "Super-Fast Factory" (Hardware Acceleration):
While the old tools worked on a single lane road, halox is built to drive on a massive highway. It is designed to run on GPUs (the powerful chips usually found in video game computers). The paper shows that when halox runs on these GPUs, it can calculate the properties of dark matter halos 20 to 65 times faster than the old tools, depending on the task.
The "Cheat Sheet" (Neural Network Emulators):
Some of the math required to understand these halos involves complex integrals (summing up infinite tiny pieces), which is like trying to count every grain of sand on a beach one by one. It takes a long time.
halox includes a "neural network emulator." Think of this as a highly trained student who has studied the beach so thoroughly that they can guess the number of grains of sand almost instantly with 99% accuracy. Instead of doing the slow, heavy math every time, halox can use this "cheat sheet" to get the answer 95 times faster than the baseline method, while still being incredibly accurate.

What Does halox Actually Do?
The paper lists specific "ingredients" the tool can calculate:

Cosmology: It calculates the basic rules of the universe (like how dense space is).
Halo Shapes: It describes the shape of the dark matter clouds (using models like NFW and Einasto), telling us how dense they are at the center versus the edge.
Concentration: It figures out how "clumped" the matter is in a halo.
The "Mass Function": It predicts how many halos of a certain size exist at a specific time in the universe.
Bias: It calculates how likely these halos are to cluster together.

The Results
The authors tested halox against the old standard (colossus).

On a standard computer (CPU), the old tool was still slightly faster for simple tasks because it uses very optimized, traditional code.
However, on a GPU, halox crushed the competition, running 20 to 65 times faster.
When using the "cheat sheet" (emulator), it was 34 times faster than the old tool even on the GPU.

Who is Using It?
The paper mentions that halox is already being used in upcoming scientific projects, such as creating realistic simulations of the sky and analyzing data for the LSST Dark Energy Science Collaboration (a major project studying the mysterious force pushing the universe apart).

In Summary
halox is a new, super-fast, and "smart" calculator for the invisible foundations of the universe. It allows scientists to run complex simulations on powerful graphics cards and instantly see how small changes in their theories affect the results, speeding up the discovery process for understanding how our universe is built.

Technical Summary of halox: Dark matter halo properties and large-scale structure calculations using JAX

Problem Statement

In cosmology and astrophysics, modeling dark matter halos is essential for understanding the large-scale structure of the Universe. While numerous toolkits exist for halo modeling (e.g., halofit, halotools, colossus, pyCCL), there is a lack of publicly available software that combines standard halo property calculations with differentiability and hardware acceleration. The increasing complexity of cosmological data and the rise of AI-driven workflows (such as Hamiltonian Monte Carlo sampling and machine learning) necessitate frameworks where physical models are not only fast but also differentiable with respect to input parameters. Existing libraries generally lack these specific capabilities, creating a bottleneck for gradient-based inference and optimization tasks.

Methodology

The halox library addresses this gap by implementing a suite of dark matter halo and large-scale structure (LSS) calculations using JAX, a high-performance numerical computing framework.

Core Implementation

JAX Backend: All functions are written in JAX and rely on JAX-cosmo for cosmology-dependent calculations. This allows for:
- Just-in-Time (JIT) Compilation: Functions are compiled for execution on CPUs and GPUs.
- Automatic Differentiation: All outputs are differentiable with respect to inputs, including halo mass, redshift, and cosmological parameters.
- Vectorization: Native support for parallelization via jax.vmap.
Physical Quantities Implemented:
- Cosmological Quantities: Critical density ( $\rho_c$ ) and differential comoving volume elements.
- Halo Profiles: Radially dependent properties for NFW and Einasto profiles (density, enclosed mass, gravitational potential, circular velocity, velocity dispersion, and projected surface density), based on analytical derivations from Łokas & Mamon (2001) and Retana-Montenegro et al. (2012).
- Concentration-Mass Relations: Implementations of relations from Duffy et al. (2008), Klypin et al. (2011), Prada et al. (2012), and Child et al. (2018).
- Large-Scale Structure: RMS variance of matter distribution ( $\sigma(R)$ ) and within the Lagrangian radius of a halo ( $\sigma(M)$ ).
- Halo Mass Function (HMF) & Bias: Models from Tinker et al. (2008, 2010).
- Overdensities: Support for spherical overdensity (SO) definitions with conversion utilities.

Emulation Strategy

To address the computational expense of numerical integration required for $\sigma(M)$ (a primary bottleneck for HMF and bias calculations), halox includes a neural network emulator:

Architecture: A multi-layer perceptron (MLP) with three hidden layers of width 64.
Training: Trained on a Sobol sample of the analytical halox $\sigma(M)$ implementation across log-mass, redshift, and cosmological parameters ( $\Omega_b, \Omega_c, h, n_s, \sigma_8$ ).
Integration: The emulator is optional and can be passed as an argument to standard functions, preserving the differentiable and hardware-accelerated nature of the workflow while providing approximate results.

Key Results and Performance

The paper presents benchmarks comparing halox against colossus (a widely used, highly optimized non-JAX library) and the halox analytical baseline.

CPU Performance: On CPU architectures, colossus currently outperforms halox by a factor of approximately 3 for analytical computations, attributed to colossus's use of highly efficient libraries like CAMB and optimized interpolation schemes.
GPU Acceleration: Utilizing GPUs significantly accelerates halox:
- Analytical computations speed up by a factor of >20.
- Emulated computations speed up by a factor of ~65.
Emulation Speedup: Using the neural network emulator for $\sigma(M)$ calculations yields a speedup of up to 95× compared to the analytical baseline and 34× compared to colossus on GPU architectures.
Accuracy:
- The emulator predicts $\sigma(M)$ and halo bias with accuracy within 1%.
- The HMF prediction is accurate within a few percent across the tested parameter space.
- Errors increase slightly in the regime of extremely massive halos ( $M_{200c} > 10^{15} h^{-1} M_\odot$ ), which are rare in simulations and observations.
Validation: All functions are validated against Astropy, colossus, and Gala across varying masses, redshifts, and cosmologies. These tests are integrated into an automated CI/CD pipeline.

Significance and Claims

The authors position halox as a critical tool for modern, AI-driven cosmological analysis. Its primary significance lies in enabling differentiable and hardware-accelerated workflows for halo properties, which were previously unavailable in a unified, public package.

Gradient-Based Workflows: The automatic differentiation capabilities allow halox to be used directly in sensitivity analyses, Hamiltonian Monte Carlo sampling for Bayesian inference, and machine learning applications.
Current Adoption: At the time of writing, the library is already being utilized in upcoming analyses, including the development of differentiable halo catalog generators and realistic simulations of the extragalactic millimeter-wave sky. It has been referenced in cosmology publications by the LSST Dark Energy Science Collaboration and others.
Modest Claims: The authors acknowledge that while halox offers superior performance on GPUs and unique differentiability, it does not yet surpass the raw CPU speed of established libraries like colossus for analytical calculations. The library is presented as a complementary tool that unlocks new computational paradigms rather than a direct replacement for all existing CPU-optimized codes.

The paper concludes that halox demonstrates strong potential for GPU-based cosmological analyses, delivering significant speedups alongside the critical capability of automatic differentiation.