---
title: 'diffHOD-IA: Differentiable HOD & IA Modeling'
url: https://www.emergentmind.com/topics/diffhod-ia
type: topic
---

# diffHOD-IA: Differentiable HOD & IA Modeling

diffHOD-IA is a fully differentiable implementation of the stochastic Halo Occupation Distribution (HOD) framework incorporating galaxy intrinsic alignment (IA) statistics, developed to address parameter inference and simulation-based modeling requirements for next-generation weak gravitational lensing cosmology. Built on differentiable sampling methods (notably Gumbel-Softmax relaxation and reparameterization), and implemented in JAX, diffHOD-IA enables automatic differentiation with respect to both HOD and IA parameters through the full pipeline, allowing the application of gradient-based inference algorithms such as Hamiltonian Monte Carlo (HMC) not only for galaxy number statistics but also for alignment-dependent, orientation-sensitive observables [2602.04977].

## 1. HOD and IA Model Structure

The base HOD logic follows the widely used Zheng et al. (2007) prescription, parameterizing the mean number of central and satellite galaxies per halo as a function of halo mass:
\[
\langle N_{\rm cen}(M)\rangle = \frac12\left[1 + \operatorname{erf}\left(\frac{\log M - \log M_{\min}}{\sigma_{\log M}}\right)\right]
\]
Central galaxy occupation is sampled via a (possibly relaxed) Bernoulli process. Satellites are assigned with mean
\[
\langle N_{\rm sat}(M)\rangle = \langle N_{\rm cen}(M)\rangle \left(\frac{M-M_0}{M_1}\right)^\alpha
\]
with samples drawn from a Poisson or (in the differentiable relaxation) a Binomial distribution. Central locations are set to halo centers; satellites are assigned to subhalos (ranked by $m_{\text{peak}}$), falling back to NFW halo profiles when necessary.

The IA component models the alignment of galaxy shapes with halo properties using a Dimroth–Watson (“spin-2 symmetric maximum-entropy”) distribution. The misalignment angle $\theta$ has probability density
\[
P(\theta,\phi)\,d\theta\,d\phi = \frac{B(\kappa)}{2\pi} \exp[-\kappa \cos^2\theta] \sin\theta\,d\theta\,d\phi
\]
where $\kappa=-\tan\left(\frac{\pi}{2}\mu\right)$ and $(\mu_{\rm cen}, \mu_{\rm sat})$ control the degree of alignment for centrals and satellites. These parameters adjust the orientation statistics that are central to weak-lensing contaminant modeling.

## 2. Differentiable Stochastic Sampling and Implementation

diffHOD-IA leverages continuous relaxations to render discrete galaxy sampling differentiable. For central galaxies, the relaxed Bernoulli (a.k.a. BinConcrete) reparameterization is
\[
N_{\rm cen}^{\rm relax} = \sigma\left(\frac{\log(p/(1-p)) + \epsilon}{\tau}\right), \;\; \epsilon\sim\operatorname{Logistic}(0,1),\;\; \tau=0.1
\]
For satellites, Poisson draws are replaced by binomial relaxation: set $N_{\max}$ satellite slots (fiducially $N_{\max}=48$), each assigned via an independent relaxed Bernoulli with $p=\langle N_{\rm sat}\rangle/N_{\max}$. This guarantees differentiability with respect to HOD parameters across all stochastic catalog realizations.

Sampling NFW radial positions and IA misalignment angles proceeds via differentiable inverse-CDF sampling using Newton's method. Assignment of satellites to subhalos employs a softmax over ranked subhalo lists:
\[
q_i = \frac{\exp(-\mathrm{rank}_i / t_{\mathrm{rank}})}{\sum_{j} \exp(-\mathrm{rank}_j / t_{\mathrm{rank}})}, \;\; t_{\mathrm{rank}}=0.5
\]
All computational components use JAX primitives with `jax.jit` and `vmap` for efficient vectorization and parallelization.

## 3. Correlation Functions and Summary Statistics

diffHOD-IA computes both standard galaxy clustering and IA summary statistics in a differentiable fashion, enabling gradient-based optimization or sampling of IA parameters with respect to observed two-point functions. Let $(i, j)$ index galaxy pairs with spatial separation $|\Delta\mathbf x|$ within a radial bin $r_k$. The following statistics are implemented:

- **Galaxy–galaxy autocorrelation:**
  \[
  \xi_{gg}(r_k) = \frac{DD(r_k)}{RR(r_k)} - 1
  \]
  with $DD(r_k)$ the weighted pair count from the simulated catalog and $RR(r_k)$ the analytic random pair count.

- **Position–orientation correlation (galaxy–shape, or $\omega$):**
  \[
  \omega(r_k) = \langle (\hat e_i\cdot\hat r_{ij})^2\rangle_{(i, j)\in \mathcal B_k} - \frac13
  \]

- **Shape–shape (orientation–orientation, or $\eta$):**
  \[
  \eta(r_k) = \langle (\hat e_i\cdot\hat e_j)^2\rangle_{(i, j)\in \mathcal B_k} - \frac13
  \]

All statistics are automatically differentiable with respect to the underlying HOD and IA parameters.

## 4. Gradient Validation and Inference

diffHOD-IA provides end-to-end differentiability not only through HOD parameterizations but also through IA orientation sampling, as validated by comparison of autodiff gradients with finite-difference estimates. For both $\partial \langle N_{\rm cen}\rangle/\partial\theta$ and IA parameter partials (e.g., $\partial P/\partial\mu$ for misalignment PDF), numerical (finite-difference) and autodiff (JAX) gradients were found to be in excellent agreement [2602.04977]. This supports the use of diffHOD-IA within sophisticated inference workflows.

Gradient-based inference workflows enabled include:

- **Moment-matching optimization** for IA parameters $(\mu_{\rm cen}, \mu_{\rm sat})$ to match target orientation distributions,
- **Two-point correlation function fitting** by minimizing loss over $\omega(r)$ with inverse-variance weighting,
- **Hamiltonian Monte Carlo (HMC)** using the NumPyro NUTS algorithm for posterior inference, leveraging JAX autodiff gradients for efficiency.

In application to mock catalogs (Bolshoi-Planck, tng300-matched HOD), posteriors for IA parameters were in close agreement with reference implementations and neural-network emulators, but at orders-of-magnitude lower computational cost.

| Method          |  $\mu_{\rm cen}$           |  $\mu_{\rm sat}$          |
|-----------------|---------------------------|---------------------------|
| halotools-IA    | $0.793 \pm 0.017$         | $0.294 \pm 0.056$         |
| IAEmu           | $0.799 \pm 0.022$         | $0.317 \pm 0.049$         |
| diffHOD-IA      | $0.802 \pm 0.016$         | $0.318 \pm 0.048$         |

Wall-clock performance for diffHOD-IA HMC on an NVIDIA A100 is $\sim$5 min (4 chains, 1,500 steps), compared to 1 day for halotools-IA MCMC on 150 CPU cores.

## 5. Comparison with Reference Implementations

Extensive validation was performed against halotools-IA (an industry-standard non-differentiable HOD+IA simulator) and a neural-network–based emulator, IAEmu. Across 100 realizations, diffHOD-IA reproduces not only the mean galaxy count but also one-point and two-point statistics ($\xi$, $\omega$, $\eta$) to within sample variance. Posterior inferences for IA parameters match across methods to within $0.4\sigma$ [2602.04977].

Unlike emulator-based approaches, diffHOD-IA enables differentiability of the full galaxy catalog realization, supporting integration with field-level inference pipelines and extension to arbitrary differentiable summary statistics.

## 6. Implementation, Applications, and Extensions

The diffHOD-IA package is open source (https://github.com/snehjp2/diffHOD-IA) and built in JAX for both research reproducibility and high-performance autodiff. Interfaces are included for

- Bolshoi-Planck and tng300 halo/galaxy catalogs,
- KD-tree–based neighbor search using SciPy,
- NumPyro for gradient-based HMC.

Potential extensions—many under active development—include generalization to assembly bias HODs, alternative IA prescriptions (including radial and shape alignment), 2D projected statistics, and higher-order or field-level statistics via differentiable gravity solvers and halo finders.

A plausible implication is that the diffHOD-IA framework provides a scalable and extensible foundation for simulation-based inference in cosmological analyses where both galaxy-halo connection and orientation-dependent systematics are crucial. By exposing gradients for arbitrary summary statistics, it enables rapid, high-fidelity inference workflows unachievable with traditional MCMC-based catalog samplers [2602.04977][2211.03852].

## 7. Significance and Context

diffHOD-IA demonstrates the feasibility and utility of making all components of a stochastic forward galaxy population model differentiable at the catalog level. This allows the use of powerful gradient-based inference techniques (e.g., HMC), resulting in 8–20$\times$ speedup in convergence and 8$\times$ greater effective sample size per gradient evaluation compared with Metropolis MCMC, as established in prior studies for HOD-only models [2211.03852]. The inclusion of intrinsic alignments broadens its applicability to weak-lensing cosmology, where mitigation of IA systematics is essential. The framework lays groundwork for embedding physically motivated HOD and IA prescriptions within high-dimensional, fully differentiable simulation pipelines for joint inference over cosmological and galaxy formation parameters.

Source: https://www.emergentmind.com/topics/diffhod-ia