halox 2.0.0

Dark matter halo properties and halo mass functions using JAX

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.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.

All properties support cosmology dependence using jax-cosmo. More information on the modules available can be found in the documentation pages.

Testing

All functions available in halox are validated against existing, non-JAX-based software. 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.

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},
}