Differentiable Stochastic Halo Occupation Distribution with Galaxy Intrinsic Alignments

TL;DR

diffHOD-IA delivers a fully differentiable implementation of the halo occupation distribution with intrinsic alignments, enabling end-to-end gradient-based inference for both HOD and IA parameters. By integrating differentiable sampling (including Dimroth–Watson misalignment) and differentiable estimators for IA and clustering statistics, it supports gradient-driven optimization and HMC-based inference directly at the catalog level. The approach is validated against the halotools-IA reference on Bolshoi-Planck, showing excellent agreement in one- and two-point statistics and accurate gradient recovery via autodiff. Applications demonstrate successful recovery of IA parameters from mock data and substantial speedups for Bayesian inference compared to non-differentiable pipelines or emulator-based methods. The framework, implemented in JAX, is poised for integration into differentiable cosmological pipelines and extension to broader IA statistics and higher-order analyses.

Abstract

We present diffHOD-IA, a fully differentiable implementation of a halo occupation distribution (HOD) model that incorporates galaxy intrinsic alignments (IA). Motivated by the diffHOD framework, we create a new implementation that extends differentiable galaxy population modeling to include orientation-dependent statistics crucial for weak gravitational lensing analyses. Our implementation combines this HOD formulation with an IA model, enabling end-to-end automatic differentiation from HOD and IA parameters through to the galaxy field. We additionally extend this framework to differentiably model two-point correlation functions, including galaxy clustering and IA statistics. We validate diffHOD-IA against the reference halotools-IA implementation using the Bolshoi-Planck simulation, demonstrating excellent agreement across both one-point and two-point statistics. We verify the accuracy of gradients computed via automatic differentiation by comparison with finite-difference estimates for both HOD and IA parameters. We present science use cases leveraging gradients in the simulations to recover the IA parameters of a galaxy field representative of the TNG300 simulation. Finally, we apply diffHOD-IA in a Hamiltonian Monte Carlo analysis and compare its performance with halotools-IA and a neural-network-based emulator, IAEmu. Unlike emulator-based approaches, diffHOD-IA provides differentiability at the catalog level, enabling integration into field-level inference pipelines and extension to arbitrary summary statistics for next-generation weak-lensing analyses. Our code is publicly available.

Figures (8)

• Figure 1: Validation of diffHOD-IA against the reference halotools-IA implementation for the tng300 fiducial HOD. Top left and center: Projected galaxy density fields across the simulation volume along the line of sight. Both implementations produce visually indistinguishable large-scale structure. Top right: Distribution of galaxy number counts $N_\text{gal}$ across 100 realizations using identical random seeds. Both implementations produce consistent galaxy number densities with similar scatter. Bottom left: Galaxy position-position correlation function $\xi(r)$ averaged over 100 realizations, with error bars indicating the standard deviation across realizations. The two implementations show excellent agreement across all scales. Bottom center and right: Galaxy position-orientation correlation function $\omega(r)$ and orientation-orientation correlation function $\eta(r)$. These correlations show strong agreement between the two implementations across all scales. $\eta(r)$ exhibits larger statistical noise and error bars due to the effects of galaxy shape noise. Despite the noise, the two implementations remain consistent within uncertainties.
• Figure 2: Gradients of the halo occupation distribution functions with respect to HOD parameters as a function of halo mass. Left panel: Gradients of the mean central galaxy occupation $\langle N_{\rm cen} \rangle$ with respect to $\log M_{\rm min}$ (blue) and $\sigma_{\log M}$ (orange). The gradients are largest in the transition region around $\log_{10} M/M_\odot \approx 12$ where the occupation probability transitions from 0 to 1, and vanish at high masses where $\langle N_{\rm cen} \rangle$ saturates to unity. Right panel: Gradients of the mean satellite galaxy occupation $\langle N_{\rm sat} \rangle$ with respect to all five HOD parameters: $\log M_{\rm min}$, $\sigma_{\log M}$, $\log M_0$, $\log M_1$, and $\alpha$. In both panels, solid lines show gradients computed via automatic differentiation and dotted points show finite difference estimates, demonstrating excellent agreement. The IA parameters $\mu_\text{cen}\xspace$ and $\mu_\text{sat}\xspace$ do not affect galaxy number counts and have zero gradient everywhere. These gradients enable efficient gradient-based inference of HOD parameters from galaxy clustering observations.
• Figure 3: Validation of differentiable sampling from the Dimroth--Watson distribution for galaxy-halo misalignment angles. Left panel: Probability distribution $P(\cos\theta)$ of misalignment angles for varying alignment strength $\mu$. The top axes show the corresponding misalignment angle $\theta$ in degrees. Solid lines show histograms from samples drawn using our differentiable inverse-CDF sampler; dashed lines show the analytic Dimroth--Watson PDF. The close agreement validates the differentiable sampling implementation. Positive $\mu$ (magenta) produces alignment with probability concentrated at $\cos\theta = \pm 1$, while negative $\mu$ (cyan) produces anti-alignment peaked at $\cos\theta = 0$. Right panel: Gradient of the probability distribution with respect to the alignment parameter, $\partial P / \partial \mu$. Dashed lines show analytic gradients derived from the PDF formula; scatter points show finite-difference gradients of the analytic Dimroth--Watson PDF; solid lines show gradients computed via automatic differentiation through the sampling procedure, where a Gaussian kernel density estimate is used to obtain a smooth density from discrete samples.
• Figure 4: Gradient-based recovery of IA parameters from 50 random initializations using moment-matching optimization. The target parameters ($\mu_{\rm cen} = 0.7905$, $\mu_{\rm sat} = 0.307$, pink star) represent the best-fit $\mu$ values to the tng300 HOD configuration, with the empirical uncertainty shown as pink contours. Optimization trajectories are shown as gray lines connecting initial positions to converged solutions. Brown circles denote optimizations using a single HOD realization seed per gradient step, while green squares show results when averaging over three seeds. The inset panel (lower left) shows a zoomed view of the convergence region, revealing tight clustering of final parameter estimates around the true values. Both single-seed and three-seed strategies successfully recover the target parameters across diverse initializations.
• Figure 5: Gradient-based recovery of IA parameters from 50 random initializations using correlation function matching optimization. The target parameters ($\mu_{\rm cen} = 0.79$, $\mu_{\rm sat} = 0.30$, pink star) represent the fiducial tng300 HOD configuration, with $1\sigma$ and $2\sigma$ uncertainty regions estimated from the scatter of maximum-likelihood fits across 50 independent HOD realizations shown in pink. Optimization trajectories are shown as gray lines connecting initial positions to converged solutions. Gray circles denote converged values. The inset panel (lower left) shows a zoomed view of the convergence region, revealing tight clustering of final parameter estimates around the true values.