Efficiently emulating distribution functions in gigaparsec volumes for varying cosmological parameters

1. Problem Statement

Cosmological surveys (e.g., Euclid, Rubin, Roman) are approaching statistical precision where theoretical systematics, particularly the Halo Mass Function (HMF), become a limiting factor. Accurately predicting the HMF across a wide range of cosmological parameters and halo masses requires:

• Large Volumes: To capture rare, high-mass clusters (the exponential tail of the HMF) and minimize cosmic variance.
• High Resolution: To resolve low-mass haloes and accurately model baryonic physics (hydrodynamics).
• Parameter Diversity: To map the HMF across the full $\Lambda$CDM parameter space for simulation-based inference (SBI).

The Trade-off: Traditional periodic $N$-body simulations cannot simultaneously achieve gigaparsec-scale volumes, high resolution (to resolve low-mass haloes), and broad parameter coverage due to prohibitive computational costs. Hydrodynamic simulations are even more expensive, often restricted to small volumes ($\sim 25-100$ Mpc$^3$) or single cosmologies.

2. Methodology

The authors propose a framework that combines targeted "zoom" simulations (small regions extracted from large parent volumes) with differentiable neural density estimators to emulate global distribution functions.

A. Data Source: The Quijote BSQ Suite

• Parent Simulations: The study utilizes the Quijote BSQ (Big Sobol Sequence) suite, consisting of 32,768 gravity-only $N$-body simulations.
• Parameters: Each simulation covers a $(1000 \, h^{-1}\text{Mpc})^3$ volume with $512^3$ particles, varying five cosmological parameters ($\Omega_m, \Omega_b, h, n_s, \sigma_8$) using a Sobol low-discrepancy sequence.
• Region Selection: Instead of using the full periodic boxes, the authors select small sub-volumes (regions) of size $V \sim (50 \, h^{-1}\text{Mpc})^3$ from these parent simulations. These regions span a range of local overdensities ($\delta$).

B. The Emulator: Conditional Normalizing Flow

The core of the method is a Conditional Normalizing Flow trained to predict the HMF in a specific region based on:

1. Global Parameters ($\theta$): The cosmological parameters of the parent simulation.
2. Local Environment: The overdensity percentile ($p_S$) of the selected region relative to the parent distribution (rather than the raw overdensity value $\delta$, to avoid modeling the complex PDF of $\delta$ explicitly).
3. Mass Bin: The target halo mass.

Architecture Details:

• Decomposition: The HMF is decomposed into a global baseline (mean HMF of all training regions) and a residual ($r_{Sj}$). The flow predicts the residual.
• Input Vector: $x_{Sj} = [p_S, \theta, m_j]$.
• Output: The predicted residual $r_{Sj}$, which modifies the baseline to produce the region-specific HMF.
• Likelihood: The model treats halo counts as a Poisson process. The training objective minimizes the negative log-likelihood (NLL) averaged over Monte Carlo samples of the latent space to handle the intractable marginalization of the Poisson rate.

C. Global Reconstruction via Integration

To recover the global HMF for a given cosmology:

1. The emulator predicts the HMF for a specific region given a random overdensity percentile.
2. The global distribution is reconstructed by integrating (summing) the predictions over the distribution of overdensities.
3. This is implemented via a Monte Carlo approach: tiling the global volume with $N_{eff}$ independent regions, assigning each a random percentile $p_S \sim U(0, 100)$, and summing the predicted halo counts.

3. Key Contributions

1. Cost Efficiency: The method achieves global HMF emulation using only $\sim 0.026\%$ of the original simulation volume. A suite of zoom simulations can emulate a gigaparsec-scale periodic simulation at a fraction of the computational cost.
2. Dynamic Range Extension: By using zoom regions, the method pushes the dynamic range to much lower halo masses (down to $\sim 10^{10} M_\odot$ in potential hydrodynamic applications) than possible in periodic simulations of the same particle count, while still capturing high-mass clusters via the integration of rare overdense regions.
3. Differentiability: The model is built in JAX, making it fully differentiable. This allows direct integration into gradient-based inference frameworks (e.g., BlackJax) for cosmological parameter estimation.
4. Percentile Conditioning: The novel use of overdensity percentiles as a conditioning variable simplifies the integration step, avoiding the need to explicitly model the complex, cosmology-dependent PDF of physical overdensities.

4. Results

• Accuracy:
  • Region Level: The emulator residuals are consistent with Poisson scatter. The signed deviance residuals show no systematic bias, and zero-count probabilities are well-calibrated.
  • Global Level: The reconstructed global HMF achieves a median fractional error of $\sim 2\%$ for halo masses below $10^{14.5} M_\odot$. For 95% of test cosmologies, the error remains below 5% in this mass range.
  • High Mass: Errors at the high-mass end converge toward the shot-noise limit, as expected for rare objects.
• Parameter Sensitivity: The emulator correctly captures the strong dependence of HMF normalization on $\Omega_m$ and the sensitivity of the high-mass tail to $\sigma_8$ and $n_s$.
• Training Requirements:
  • Volume: Regions of $\sim (50-70 \, h^{-1}\text{Mpc})^3$ are sufficient to suppress cosmic variance; larger volumes yield diminishing returns.
  • Sample Size: Training on 4,096 unique parent simulations (selecting 1-2 regions each) is sufficient to achieve median absolute fractional errors $< 5.5\%$.
  • Diversity: Sampling a wide range of overdensities at fixed cosmology is more critical for reducing error than sampling a wider range of cosmologies at fixed overdensity.

5. Significance and Future Implications

• Baryonic Physics: The framework is designed to be applied to hydrodynamic zoom simulations (e.g., FLARES, CAMELS). This would allow for the emulation of baryonic effects on the HMF across gigaparsec volumes and wide parameter spaces, a task currently impossible with full hydrodynamic periodic simulations.
• Survey Interpretation: The method provides the precise, parameter-dependent HMF predictions required for upcoming wide-field surveys (Euclid, Roman, Rubin), enabling better constraints on dark matter and dark energy.
• Beyond HMF: The approach is generalizable to other distribution functions (e.g., mass-concentration relations) and higher-order statistics (e.g., two-point correlation functions) by combining zoom simulations with field-level emulators.
• Active Learning: The authors suggest that future iterations could use Active Learning to optimize the selection of training regions, potentially further reducing the number of required simulations.

In summary, this work demonstrates that targeted zoom simulations combined with neural density estimators offer a computationally feasible path to generating high-fidelity, large-volume cosmological predictions across broad parameter spaces, effectively breaking the traditional volume-resolution trade-off.