halox: Differentiable Halo-Model Library

• halox is a JAX-native library that provides differentiable computations of dark matter halo properties using analytic NFW profiles and standard fitting functions.
• It integrates halo mass functions and bias relations from Tinker et al., enabling gradient-based inference and machine-learning workflows in cosmology.
• Leveraging JAX-cosmo, halox delivers automatic differentiation and GPU acceleration, thereby streamlining large-scale structure modeling and cosmological parameter sensitivity analyses.

Searching arXiv for the cited halox paper and directly related JAX cosmology context. halox is a JAX-powered Python library for differentiable and accelerated computations of key properties of dark matter halos and of the halo mass function. It is designed as a JAX-native halo-model layer on top of differentiable cosmology software, wrapping standard halo-model ingredients—Navarro–Frenk–White profiles, halo mass functions, halo bias relations, and overdensity conversions—so that they can be inserted directly into gradient-based inference, Hamiltonian Monte Carlo, and machine learning workflows (Kéruzoré, 26 Sep 2025).

1. Scientific role and scope

halox is organized around three closely related scientific tasks. The first is the modeling of individual dark matter halos through analytic radial quantities for halos described by the Navarro–Frenk–White (NFW) profile. The second is the modeling of halo statistics in the Universe, especially the halo mass function (HMF) and halo bias as functions of mass, redshift, and cosmology using the fitting functions of Tinker et al. (2008, 2010). The third is the connection between halo statistics and large-scale structure through the calculation of the matter variance, $\sigma(R)$ or $\sigma(M)$, from the linear matter power spectrum supplied by JAX-cosmo.

These tasks place halox in a specific computational niche. It is intended for applications in which one needs a halo-model layer on top of differentiable cosmology codes, including predictions of halo abundances and clustering across cosmologies, analytic or semi-analytic large-scale-structure models, and gradient-based cosmological parameter inference with respect to parameters such as $\Omega_m$, $\sigma_8$, and $h$. The same design also supports differentiable physics and machine-learning use cases, including neural networks that embed analytic halo profiles and probabilistic workflows based on HMC or variational inference.

Domain	Quantity or model	Stated role
Individual halos	NFW radial profiles	Density, mass, potential, velocities, projected density
Halo statistics	HMF and halo bias	Dependence on mass, redshift, and cosmology
Large-scale structure	$\sigma(R)$ and $\sigma(M)$	Inputs to halo-model and HMF calculations

Within the broader JAX cosmology ecosystem, halox is described as the component that provides halo properties and halo statistics in a fully differentiable, GPU-friendly form. This positioning is important because it indicates that halox is not introduced as an alternative cosmological backend; rather, it depends on JAX-cosmo for cosmological calculations and specializes in the halo layer.

2. Mathematical content and physical ingredients

The library combines cosmological background quantities, spherical-overdensity halo definitions, analytic NFW expressions, and calibrated fitting functions for HMF and bias. Through JAX-cosmo it uses the critical density

$\rho_c(z) = \frac{3 H^2(z)}{8\pi G},$

and the differential comoving volume element

$$\frac{{\rm d}V_c}{{\rm d}\Omega\,{\rm d}z} = \frac{c}{H(z)}\chi^2(z).$$

These quantities enter number-count calculations and conversions between comoving densities and observables.

For halo structure, halox implements analytic NFW quantities following Łokas & Mamon (2001). The density profile is

$$\rho(r) = \frac{\rho_s}{(r/r_s)\left(1 + r/r_s\right)^2},$$

and the enclosed mass is

$\sigma(M)$0

The package also provides the gravitational potential $\sigma(M)$1, circular velocity,

$\sigma(M)$2

velocity dispersion $\sigma(M)$3 from the spherical Jeans equation plus the NFW density, and the projected surface density $\sigma(M)$4 as a function of projected radius. In practice these are parameterized by a mass definition such as $\sigma(M)$5 and a concentration parameter $\sigma(M)$6, with

$\sigma(M)$7

Here $\sigma(M)$8 is either the critical density or the mean matter density, depending on the spherical-overdensity convention.

For large-scale structure, halox computes the variance of the matter field from the linear matter power spectrum:

$\sigma(M)$9

with top-hat window

$\Omega_m$0

It also defines $\Omega_m$1 through the Lagrangian radius

$\Omega_m$2

so that $\Omega_m$3. This quantity is the key ingredient in the universal halo-mass-function and bias formulations.

The halo mass function implemented in halox follows Tinker et al. (2008) and is expressed as

$\Omega_m$4

The fitting function is

$\Omega_m$5

with parameters $\Omega_m$6 depending on overdensity definition and redshift. Halo bias follows Tinker et al. (2010), where the linear bias is written in the linear regime as

$\Omega_m$7

and the fitting function is expressed in terms of peak height,

$\Omega_m$8

A central practical feature is conversion among spherical-overdensity mass definitions. Halo masses satisfy

$\Omega_m$9

and halox provides conversions from one critical overdensity $\sigma_8$0 to another, as well as between critical-overdensity and mean-matter-overdensity conventions. Because these mappings are performed under an assumed NFW profile, the physical content of the conversion is tied to the NFW assumption rather than to a profile-independent identity.

3. Computational architecture and automatic differentiation

halox is written entirely in JAX and uses JAX-cosmo for cosmology. This architecture gives it two defining computational properties: JIT compilation to XLA and end-to-end automatic differentiation. Functions can be wrapped by jax.jit, which allows execution on CPUs, GPUs, and TPUs subject to JAX support; the HMF, halo bias, NFW profile evaluations, and $\sigma_8$1 calculations can then be vectorized over large mass and redshift arrays (Kéruzoré, 26 Sep 2025).

Differentiability is presented as a core capability rather than a secondary convenience. The library’s functions are automatically differentiable with respect to halo quantities such as $\sigma_8$2, $\sigma_8$3, concentration, and overdensity threshold, and also with respect to cosmological parameters passed through JAX-cosmo, including $\sigma_8$4, $\sigma_8$5, $\sigma_8$6, $\sigma_8$7, and $\sigma_8$8. Typical derivative operators include jax.grad, jax.jacobian, and jax.vjp/jax.jvp.

This design is especially relevant for HMC, where the log-posterior requires gradients through halo predictions, and for machine-learning settings in which halo quantities appear inside the computational graph. The same property is useful for sensitivity analyses, including differentiation of halo statistics with respect to cosmological parameters to probe degeneracies or Fisher information. The paper does not present detailed benchmarks, but it explicitly emphasizes that halox functions are compiled and GPU-accelerated and that the package adopts a lighter-weight halo-model approach than some more complex differentiable HMF codes.

Algorithmically, the HMF is evaluated numerically with JAX. The computation of $\sigma_8$9 proceeds through the power-spectrum integral, and the derivative $h$0 can be obtained by automatic differentiation of $h$1 or through analytic differentiation if the implementation is coded accordingly. The significance of this point is methodological: halox avoids explicit finite differencing and instead uses JAX-native derivative propagation, with accuracy limited by machine precision and the approximations involved in numerical integration.

4. Workflow integration and interface model

The paper is explicit about capabilities but brief about API surface. It does not spell out all class or function names. What it does specify is the intended pattern of use: a cosmology is defined via JAX-cosmo, halox functions are called with that cosmology and with mass or redshift arrays, and JAX autodiff is then applied to obtain derivatives. This gives halox the character of a composable backend component rather than a monolithic application.

In probabilistic inference, the intended use is to express a likelihood in terms of HMF or halo-bias predictions and allow a framework such as NumPyro to evaluate HMC or NUTS trajectories while gradients are computed automatically. In machine-learning settings, the same structure supports analytic halo quantities inside differentiable simulations or surrogate models. The paper also situates halox alongside JAX-based inference tools such as CosmoPower-JAX and candl, reinforcing that its role is interoperability within a broader differentiable ecosystem.

A notable aspect of the interface model is the distinction between what halox provides directly and what it expects from adjacent packages or user input. Cosmology is delegated to JAX-cosmo. Concentration–mass specifics are not developed in the paper and are left to external differentiable packages such as Diffprof or to user-supplied concentrations. This division of labor suggests a modular design philosophy in which halo structure, halo statistics, and cosmological background calculations remain separable but differentiably connected.

5. Validation, assumptions, and limitations

The package is described as validated against established tools in two layers. Cosmology calculations, including distances and critical density, are validated against Astropy for a variety of cosmologies and redshifts. Halo quantities—NFW profiles, halo mass function, and halo bias—are validated against colossus across different halo masses, redshifts, critical overdensities, and cosmological parameters. These tests are included in continuous integration on the GitHub repository and are also visualized in the online documentation (Kéruzoré, 26 Sep 2025).

The physical assumptions are narrow and explicit. halox relies on JAX-cosmo and is therefore geared toward $h$2CDM cosmologies. It assumes the NFW profile for dark matter halos and uses the Tinker et al. (2008, 2010) fitting functions for HMF and bias, with validity inherited from those original calibrations over their specified mass ranges, redshift ranges—typically $h$3–3—and overdensity definitions such as $h$4 and $h$5. It models dark matter–only halos and does not explicitly include baryonic effects such as gas profiles or feedback.

The main limitations follow directly from those assumptions. halox offers a lighter, halo model-based approach to the HMF than some more complex differentiable HMF tools, so applications that require very high-fidelity halo mass predictions with sophisticated corrections may require other software. Concentration–mass relations are also not elaborated in the paper; the package implements NFW profiles, but concentration modeling may rely on external codes or user input. A plausible implication is that halox prioritizes differentiable interoperability and direct analytic halo modeling over maximal phenomenological breadth.

A common misconception would be to treat halox as introducing a new empirical halo model. The description instead makes clear that its main physical ingredients are standard ones—NFW, Tinker HMF, Tinker bias, and spherical-overdensity conversions—implemented in a JAX-native, differentiable form. The novelty is therefore computational and integrative rather than a replacement of the underlying halo-model calibrations.

6. Place within the JAX cosmology ecosystem and prospective extensions

halox is positioned within an expanding ecosystem of JAX-based cosmology tools that includes JAX-cosmo, JAXPM, SHAMNet, CosmoPower-JAX, and candl. Within that ecosystem it serves as the halo-properties and halo-statistics layer, enabling analytic halo calculations to be inserted into differentiable cosmological pipelines. This role is consequential because halo abundances, bias, and radial profiles are standard building blocks in cluster counts, halo catalogs, and large-scale-structure modeling, yet these ingredients often become bottlenecks when gradient-based methods are required.

The future directions indicated in the paper are extensions rather than changes of principle. These include support for additional halo profiles such as Einasto or cored profiles, updated HMF and bias fitting functions, baryonic extensions through links to gas-halo models such as picasso, and extended mass or redshift coverage tailored to upcoming surveys (Kéruzoré, 26 Sep 2025). Since halox is presented as part of a rapidly expanding JAX cosmology ecosystem, a plausible implication is that later releases may broaden the range of halo ingredients while preserving the same differentiable, hardware-accelerated computational model.

In that sense, halox occupies a specific methodological position. It formalizes standard dark-matter-halo calculations in a framework where vectorization, JIT compilation, and exact autodiff through the implemented computation graph are first-class requirements. For research programs that already adopt JAX-based cosmological inference or differentiable simulation pipelines, halox provides the halo-model layer needed to propagate gradients through NFW structure, mass-function calculations, bias relations, and overdensity conversions without leaving the JAX stack.