halox

JAX-powered Python library for differentiable dark matter halo property and mass function calculations.

Installation

halox can be installed via pip:

pip install halox

For a manual installation, see the documentation pages.

Features

halox offers a JAX-powered differentiable and GPU-accelerated implementation of some widely used properties of dark matter halos and large-scale structure, including:

• halox.halo: Radial profiles of dark matter halos following Navarro-Frenk-White (NFW) and Einasto distributions;
• halox.cm: Mass-concentration relations of dark matter halos;
• halox.lss: Large-scale structure ($$\sigma(R)$$, $$\sigma(M)$$);
• halox.hmf: The halo mass function, quantifying the abundance of dark matter halos in mass and redshift and its dependence on cosmological parameters;
• halox.bias: The halo bias;
• halox.emus: A neural-network emulator for $$\sigma(M)$$, providing up to 95× speedup over the analytic calculation while remaining differentiable and GPU-accelerated.

All functions support jax.jit, jax.vmap, and jax.grad. Halo masses, redshifts, and cosmological parameters are all valid differentiation targets. All properties support cosmology dependence using jax-cosmo. More information on the modules available can be found in the documentation pages.

Quick start

import jax
import jax.numpy as jnp
from halox import cosmology, hmf, halo

jax.config.update("jax_enable_x64", True)  # optional

cosmo = cosmology.Planck18()

# Halo mass function at z = 0.5
M = jnp.logspace(13, 15, 100)
dn_dlnM = hmf.tinker08_mass_function(M, 0.5, cosmo)

# NFW profile
h = halo.nfw.NFWHalo(1e14, 5.0, 0.0, cosmo)
r = jnp.logspace(-1, 1, 50)
rho = h.density(r)

# Derivative w.r.t. halo mass
dhmf_dM = jax.grad(lambda M: hmf.tinker08_mass_function(M, 0.5, cosmo))(1e14)

Units

All input/output units for halox functions are reported in docstrings. Units are assumed to be in proper coordinates (not comoving) and include factors of little h.

Testing

All functions available in halox are validated against existing, non-JAX-based software, and a 100% coverage threshold is enforced. Cosmology calculations are validated against Astropy for varying cosmological parameters and redshifts. Other quantities are validated against Colossus or gala for varying halo masses, redshifts, critical overdensities, and cosmological parameters. These tests are included in the automatic CI/CD pipeline; a visual comparison is also included in the documentation.

Documentation

For more detail on the code and features, please visit our documentation pages.

Contributing

Contributions are welcome! See CONTRIBUTING.md for development setup, code conventions, and the pull-request process.

Citation

If you use halox for your research, please cite the original paper:

@ARTICLE{2025arXiv250922478K,
       author = {{K{\'e}ruzor{\'e}}, Florian},
        title = "{halox: Dark matter halo properties and large-scale structure calculations using JAX}",
      journal = {arXiv e-prints},
     keywords = {Instrumentation and Methods for Astrophysics, Cosmology and Nongalactic Astrophysics},
         year = 2025,
        month = sep,
          eid = {arXiv:2509.22478},
        pages = {arXiv:2509.22478},
          doi = {10.48550/arXiv.2509.22478},
archivePrefix = {arXiv},
       eprint = {2509.22478},
 primaryClass = {astro-ph.IM},
}