Diffhalos: A Generative Model of Cosmological Lightcones of Dark Matter Halos

Technical Summary of "Diffhalos: A Generative Model of Cosmological Lightcones of Dark Matter Halos"

Problem Statement
Modern cosmological analyses of galaxy surveys require synthetic lightcones—simulated catalogs of galaxies spanning cosmic time—to test theoretical models and interpret observational data. While traditional methods rely on running massive gravitational N-body simulations or interpolating between output snapshots to generate these lightcones, these approaches are computationally expensive and often lack differentiability. Furthermore, generating physically realistic synthetic populations requires incorporating Mass Assembly Histories (MAHs), which describe how halos grow through mergers and accretion. Existing models for MAHs, such as the extended Press-Schechter formalism or standard fitting functions, often lack the flexibility to capture the full diversity of assembly paths or the computational efficiency required for gradient-based inference. There is a need for a generative framework that can produce lightcones of dark matter halos, subhalos, and their MAHs with statistical accuracy, while supporting automatic differentiation for efficient parameter optimization.

Methodology
The authors introduce Diffhalos, a generative model implemented in JAX to leverage automatic differentiation and GPU acceleration. The framework decomposes the generation of a cosmological lightcone into three distinct, differentiable stages:

1. Host Halo Generation:
   The model generates a population of host halos within a lightcone volume by sampling from the Halo Mass Function (HMF), $n_{\text{halo}}(m_{\text{halo}}, z)$. The authors explore two differentiable implementations:

   • Halox-based: A JAX implementation of the analytical halo model (using the Tinker et al. 2008 fitting formula) which computes the HMF via the linear matter power spectrum. This allows for direct computation of gradients with respect to cosmological parameters.
   • Emulator-based: To reduce memory overhead associated with numerical integration in the analytical approach, the authors train a Multilayer Perceptron (MLP) to emulate the cumulative HMF, $N_{\text{halo}}(> m_{\text{halo}}, z)$. The differential HMF is then derived via jax.grad.
     Sampling is performed using inverse transformation sampling on a grid of redshifts and masses. The authors also implement a Quasi-Monte Carlo (QMC) generator that uses weighted grids to approximate the population, significantly reducing memory footprint while maintaining statistical consistency.

2. Subhalo Generation:
   Once host halos are established, the model populates them with subhalos using a differentiable Conditional Subhalo Mass Function (CSHMF). The CSHMF is parameterized as a cumulative distribution $N_{\text{sub}}(> \mu | M_{\text{host}})$, where $\mu = M_{\text{sub}}/M_{\text{host}}$. The authors utilize a "sig-slope" kernel (a power law with a sigmoid-varying slope) calibrated against the Discovery simulations (HACC code) and Jiang & van den Bosch (2016) fitting functions. Similar to the host generation, both Monte Carlo and QMC methods are available for subhalo sampling.

3. Mass Assembly History (MAH) Assignment:
   The final stage assigns a MAH to every host and subhalo. The MAH is approximated using the Diffmah functional form, parameterized by a set of parameters $\theta_{\text{MAH}}$. To map these parameters to the simulated population, the authors introduce DiffmahNet, a normalizing flow architecture trained on merger trees from the Small MultiDark Planck (SMDPL) N-body simulation. DiffmahNet learns the probability distribution $P(\theta_{\text{MAH}} | z, M_{\text{halo}}(z))$, enabling the generation of diverse, physically plausible assembly histories conditioned on the halo's mass and redshift at the time of observation.

Key Contributions

• Diffhalos Framework: A fully differentiable, JAX-based library for generating cosmological lightcones of dark matter halos, subhalos, and their MAHs.
• Differentiable HMF and CSHMF: The implementation of differentiable halo mass functions (via Halox and custom emulators) and conditional subhalo mass functions, enabling gradient-based optimization with respect to cosmological and astrophysical parameters.
• DiffmahNet: A normalizing flow model that generalizes previous parametric MAH models (Hearin et al. 2021) to generate MAHs for entire populations of halos and subhalos, capturing both mean growth trends and population variance.
• Efficiency Tools: The introduction of Quasi-Monte Carlo (QMC) generators for both host and subhalo populations, offering orders-of-magnitude improvements in memory efficiency compared to standard Monte Carlo realizations.
• Public Code: The release of the Diffhalos codebase to support the development and calibration of galaxy-halo connection models.

Results
The authors validate the Diffhalos framework against N-body simulation data:

• HMF Accuracy: The Halox implementation shows percent-level agreement with the Colossus library across a wide range of masses and redshifts ($z=0$ to $5$) for different cosmologies. The MLP-based emulator achieves $\le 10\%$ accuracy, which is comparable to the intrinsic accuracy of the Tinker et al. (2008) fitting function itself.
• Subhalo Statistics: The subhalo generator reproduces the total mass function of host halos and subhalos from the Discovery simulations with reasonable fidelity for host masses $10^{10} < M_{\text{host}}/M_\odot < 10^{14}$ and $0 < z < 4$.
• MAH Fidelity: DiffmahNet accurately reproduces the mean mass growth and the variance of MAHs found in the SMDPL simulation for halos identified at different epochs ($t_{\text{obs}} = 6.42$ and $13.8$ Gyr).
• Gradient Calculation: The framework successfully computes gradients of the HMF with respect to cosmological parameters ($\Omega_m, \sigma_8$), demonstrating the utility of the differentiable approach for inference.

Significance and Claims
The paper positions Diffhalos as a "proof-of-concept" tool designed primarily to support the development and calibration of galaxy-halo connection models. The authors emphasize that the differentiable nature of the framework allows for efficient gradient-based applications, such as Hamiltonian Monte Carlo and gradient descent, which are critical for optimizing semi-analytic models (SAMs) or empirical models against observational data.

The authors explicitly state that the current implementation is a proof-of-concept and modestly notes that the cosmology dependence of the subhalo mass function and MAHs is not yet fully calibrated across a suite of cosmologies (currently trained only on the SMDPL cosmology). They identify future work as necessary to integrate state-of-the-art emulators, calibrate cosmology dependence for subhalos and MAHs, and extend the framework to include spatial clustering (sky coordinates) and baryonic effects. The primary value proposition is providing a computationally efficient, differentiable alternative to N-body simulations for generating synthetic halo populations, thereby accelerating the optimization of galaxy formation models.