Emulation and Assessment of Gradient-Based Samplers in Cosmology

We assess gradient-based samplers like the No-U-Turn Sampler ($\tt{NUTS}$) compared to traditional Metropolis-Hastings algorithms in tomographic $3\times 2$ point analyses using DES Year 1 data and a simulated LSST-like survey. These studies involve 20 and 32 nuisance parameters, respectively. We implement a differentiable forward model using $\tt{JAX-COSMO}$ and derive parameter constraints using $\tt{NUTS}$ and Metropolis-Hastings algorithms. $\tt{NUTS}$ shows a relative efficiency gain of $\mathcal{O}(10)$ in terms of effective samples per likelihood evaluation but only a factor of $\sim 2$ in terms of computational time due to the higher gradient computation cost. We validate these results with analytical multivariate distributions, concluding that $\tt{NUTS}$ can be beneficial for sampling high-dimensional parameter spaces in Cosmology, though the efficiency gain is modest for moderate dimensions ($\mathcal{O}(50)$).