halox

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

halox provides differentiable and GPU-accelerated implementations of widely used dark matter halo and large-scale structure calculations. All functions support jax.jit, jax.vmap, and jax.grad; halo masses, redshifts, and cosmological parameters are all valid differentiation targets, enabling efficient gradient-based workflows such as Hamiltonian Monte Carlo sampling or machine learning pipelines.

Features

• NFW Profiles and Einasto Profiles: Radial profiles of dark matter halos (density, enclosed mass, gravitational potential, circular velocity, velocity dispersion, projected surface density)

• Mass-concentration relations: Concentration–mass relations

• Large-Scale Structure quantities: Large-scale structure (σ(R), σ(M))

• Halo Mass Function: Halo mass function (Tinker et al. 2008)

• Halo Bias: Halo bias (Tinker et al. 2010)

• Using the σ(M) emulator: Neural-network emulator for σ(M), enabling up to 95× speedup over the analytic calculation

All quantities support cosmology dependence via jax-cosmo.

Installation

halox can be installed via pip:

pip install halox

See Installation for source installation instructions.

Quick start

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

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

cosmo = cosmology.Planck18()

# Halo mass function at z = 0.5
M = jnp.logspace(13, 15, 100)  # h⁻¹ M☉
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)  # h⁻¹ Mpc
rho = h.density(r)

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

Learn more

Tutorials

• Installation

• Automatic differentiation and JAX features

• The σ(M) emulator

• Comparison with colossus

• Other utilities

Physics modules gallery

• NFW profiles

• Einasto profiles

• Concentration-mass relations

• Large-scale structure

• Halo mass function

• Halo bias

Reference

• Cosmology

• NFW halos

• Einasto halos

• Concentration-mass relations

• Large-scale structure

• Emulators

• Halo mass functions

• Halo bias

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