Warm-started Chebyshev geometric multigrid is competitive with distributed FFTs for FastPM and enables a differentiable moving-mesh particle–mesh gravity solver in JAX.
Differentiable Stochastic Halo Occupation Distribution
1 Pith paper cite this work, alongside 5 external citations. Polarity classification is still indexing.
abstract
In this work, we demonstrate how differentiable stochastic sampling techniques developed in the context of deep Reinforcement Learning can be used to perform efficient parameter inference over stochastic, simulation-based, forward models. As a particular example, we focus on the problem of estimating parameters of Halo Occupancy Distribution (HOD) models which are used to connect galaxies with their dark matter halos. Using a combination of continuous relaxation and gradient parameterization techniques, we can obtain well-defined gradients with respect to HOD parameters through discrete galaxy catalogs realizations. Having access to these gradients allows us to leverage efficient sampling schemes, such as Hamiltonian Monte-Carlo, and greatly speed up parameter inference. We demonstrate our technique on a mock galaxy catalog generated from the Bolshoi simulation using the Zheng et al. 2007 HOD model and find near identical posteriors as standard Markov Chain Monte Carlo techniques with an increase of ~8x in convergence efficiency. Our differentiable HOD model also has broad applications in full forward model approaches to cosmic structure and cosmological analysis.
fields
astro-ph.IM 1years
2026 1verdicts
CONDITIONAL 1representative citing papers
citing papers explorer
-
Fast(er)PM and Moving Mesh: JAX-native Geometric Multigrid Methods
Warm-started Chebyshev geometric multigrid is competitive with distributed FFTs for FastPM and enables a differentiable moving-mesh particle–mesh gravity solver in JAX.