* (62) Y. Nomura, _Black Hole Interior in Unitary Gauge Construction_ , _Phys. Rev. D_ 103 (2021) 066011 [2010.15827].
  * (63) P.-S. Hsin, L. V. Iliesiu and Z. Yang, _A violation of global symmetries from replica wormholes and the fate of black hole remnants_ , _Class. Quant. Grav._ 38 (2021) 194004 [2011.09444].
  * (64) N. Gaddam and N. Groenenboom, _Soft graviton exchange and the information paradox_ , 2012.02355.
  * (65) S. Pasterski and H. Verlinde, _HPS meets AMPS: how soft hair dissolves the firewall_ , _JHEP_ 09 (2021) 099 [2012.03850].
  * (66) S. Raju, _Lessons from the information paradox_ , _Phys. Rept._ 943 (2022) 2187 [2012.05770].
  * (67) R. Bousso and E. Wildenhain, _Gravity/ensemble duality_ , _Phys. Rev. D_ 102 (2020) 066005 [2006.16289].
  * (68) P. Cheng and Y. An, _Soft black hole information paradox: Page curve from Maxwell soft hair of a black hole_ , _Phys. Rev. D_ 103 (2021) 126020 [2012.14864].
  * (69) K. Su, P. Zhang and H. Zhai, _Page curve from non-Markovianity_ , _JHEP_ 21 (2020) 156 [2101.11238].
  * (70) E. E. Flanagan, _Order-Unity Correction to Hawking Radiation_ , _Phys. Rev. Lett._ 127 (2021) 041301 [2102.04930].
  * (71) H. Geng, S. Lüst, R. K. Mishra and D. Wakeham, _Holographic BCFTs and Communicating Black Holes_ , _jhep_ 08 (2021) 003 [2104.07039].
  * (72) A. Almheiri, R. Mahajan, J. Maldacena and Y. Zhao, _The Page curve of Hawking radiation from semiclassical geometry_ , _JHEP_ 03 (2020) 149 [1908.10996].
  * (73) D. N. Page, _Information in black hole radiation_ , _Phys. Rev. Lett._ 71 (1993) 3743 [hep-th/9306083].
  * (74) D. N. Page, _Time Dependence of Hawking Radiation Entropy_ , _JCAP_ 09 (2013) 028 [1301.4995].
  * (75) A. Almheiri, T. Hartman, J. Maldacena, E. Shaghoulian and A. Tajdini, _The entropy of Hawking radiation_ , _Rev. Mod. Phys._ 93 (2021) 035002 [2006.06872].
  * (76) H. Geng, A. Karch, C. Perez-Pardavila, S. Raju, L. Randall, M. Riojas et al., _Information Transfer with a Gravitating Bath_ , _SciPost Phys._ 10 (2021) 103 [2012.04671].
  * (77) L. Randall and R. Sundrum, _A Large mass hierarchy from a small extra dimension_ , _Phys. Rev. Lett._ 83 (1999) 3370 [hep-ph/9905221].
  * (78) L. Randall and R. Sundrum, _An Alternative to compactification_ , _Phys. Rev. Lett._ 83 (1999) 4690 [hep-th/9906064].
  * (79) A. Karch and L. Randall, _Locally localized gravity_ , _JHEP_ 05 (2001) 008 [hep-th/0011156].
  * (80) A. Karch, E. Katz and L. Randall, _Absence of a VVDZ discontinuity in AdS(AdS)_ , _JHEP_ 12 (2001) 016 [hep-th/0106261].
  * (81) S. Ryu and T. Takayanagi, _Holographic derivation of entanglement entropy from AdS/CFT_ , _Phys. Rev. Lett._ 96 (2006) 181602 [hep-th/0603001].
  * (82) S. Ryu and T. Takayanagi, _Aspects of Holographic Entanglement Entropy_ , _JHEP_ 08 (2006) 045 [hep-th/0605073].
  * (83) V. E. Hubeny, M. Rangamani and T. Takayanagi, _A Covariant holographic entanglement entropy proposal_ , _JHEP_ 07 (2007) 062 [0705.0016].
  * (84) M. Rangamani and T. Takayanagi, _Holographic Entanglement Entropy_ , vol. 931. Springer, 2017, 10.1007/978-3-319-52573-0, [1609.01287].
  * (85) E. Verheijden and E. Verlinde, _From the BTZ black hole to JT gravity: geometrizing the island_ , _JHEP_ 11 (2021) 092 [2102.00922].
  * (86) H. Geng, A. Karch, C. Perez-Pardavila, S. Raju, L. Randall, M. Riojas et al., _Inconsistency of Islands in Theories with Long-Range Gravity_ , 2107.03390.
  * (87) W. Israel, _Singular hypersurfaces and thin shells in general relativity_ , _Nuovo Cim. B_ 44S10 (1966) 1.
  * (88) M. Banados, C. Teitelboim and J. Zanelli, _The Black hole in three-dimensional space-time_ , _Phys. Rev. Lett._ 69 (1992) 1849 [hep-th/9204099].
  * (89) M. Banados, M. Henneaux, C. Teitelboim and J. Zanelli, _Geometry of the (2+1) black hole_ , _Phys. Rev. D_ 48 (1993) 1506 [gr-qc/9302012].
  * (90) S. Carlip, _The (2+1)-Dimensional black hole_ , _Class. Quant. Grav._ 12 (1995) 2853 [gr-qc/9506079].
  * (91) J. M. Maldacena, _Eternal black holes in anti-de Sitter_ , _JHEP_ 04 (2003) 021 [hep-th/0106112].
  * (92) J. D. Brown and M. Henneaux, _Central Charges in the Canonical Realization of Asymptotic Symmetries: An Example from Three-Dimensional Gravity_ , _Commun. Math. Phys._ 104 (1986) 207.
  * (93) I. Affleck and A. W. W. Ludwig, _Universal noninteger ’ground state degeneracy’ in critical quantum systems_ , _Phys. Rev. Lett._ 67 (1991) 161.
  * (94) T. Takayanagi, _Holographic Dual of BCFT_ , _Phys. Rev. Lett._ 107 (2011) 101602 [1105.5165].
  * (95) R. Jackiw, _Lower Dimensional Gravity_ , _Nucl. Phys. B_ 252 (1985) 343.
  * (96) C. Teitelboim, _Gravitation and Hamiltonian Structure in Two Space-Time Dimensions_ , _Phys. Lett. B_ 126 (1983) 41.
  * (97) J. Maldacena, D. Stanford and Z. Yang, _Conformal symmetry and its breaking in two dimensional Nearly Anti-de-Sitter space_ , _PTEP_ 2016 (2016) 12C104 [1606.01857].
  * (98) T. Hartman and J. Maldacena, _Time Evolution of Entanglement Entropy from Black Hole Interiors_ , _JHEP_ 05 (2013) 014 [1303.1080].
  * (99) S. H. Shenker and D. Stanford, _Black holes and the butterfly effect_ , _JHEP_ 03 (2014) 067 [1306.0622].
  * (100) P. Calabrese and J. L. Cardy, _Entanglement entropy and quantum field theory_ , _J. Stat. Mech._ 0406 (2004) P06002 [hep-th/0405152].
  * (101) P. Calabrese and J. L. Cardy, _Entanglement entropy and quantum field theory: A Non-technical introduction_ , _Int. J. Quant. Inf._ 4 (2006) 429 [quant-ph/0505193].
  * (102) M. Fujita, T. Takayanagi and E. Tonni, _Aspects of AdS/BCFT_ , _JHEP_ 11 (2011) 043 [1108.5152].
  * (103) S. Cooper, M. Rozali, B. Swingle, M. Van Raamsdonk, C. Waddell and D. Wakeham, _Black hole microstate cosmology_ , _JHEP_ 07 (2019) 065 [1810.10601].
  * (104) M. Van Raamsdonk, _Building up spacetime with quantum entanglement_ , _Gen. Rel. Grav._ 42 (2010) 2323 [1005.3035].
  * (105) A. R. Brown, H. Gharibyan, G. Penington and L. Susskind, _The Python’s Lunch: geometric obstructions to decoding Hawking radiation_ , _JHEP_ 08 (2020) 121 [1912.00228].
  * (106) N. Engelhardt, G. Penington and A. Shahbazi-Moghaddam, _A world without pythons would be so simple_ , _Class. Quant. Grav._ 38 (2021) 234001 [2102.07774].
  * (107) N. Engelhardt, G. Penington and A. Shahbazi-Moghaddam, _Finding Pythons in Unexpected Places_ , 2105.09316.
  * (108) N. Engelhardt and A. C. Wall, _Decoding the Apparent Horizon: Coarse-Grained Holographic Entropy_ , _Phys. Rev. Lett._ 121 (2018) 211301 [1706.02038].
  * (109) N. Engelhardt and A. C. Wall, _Coarse Graining Holographic Black Holes_ , _JHEP_ 05 (2019) 160 [1806.01281].
  * (110) P. Calabrese and J. L. Cardy, _Time-dependence of correlation functions following a quantum quench_ , _Phys. Rev. Lett._ 96 (2006) 136801 [cond-mat/0601225].
  * (111) P. Calabrese and J. Cardy, _Quantum Quenches in Extended Systems_ , _J. Stat. Mech._ 0706 (2007) P06008 [0704.1880].
  * (112) R. Emparan, A. M. Frassino and B. Way, _Quantum BTZ black hole_ , _JHEP_ 11 (2020) 137 [2007.15999].
  * (113) R. Emparan, A. M. Frassino, M. Sasieta and M. Tomašević, _Holographic Complexity of Quantum Black Holes_ , 2112.04860.
  * (114) J. Hernandez and R. C. Myers, _Island complexity in doubly holographic models_ , in preparation.
  * (115) S. D. Mathur, _The Fuzzball proposal for black holes: An Elementary review_ , _Fortsch. Phys._ 53 (2005) 793 [hep-th/0502050].
  * (116) B. Guo, M. R. R. Hughes, S. D. Mathur and M. Mehta, _Contrasting the fuzzball and wormhole paradigms for black holes_ , 2111.05295.
  * (117) R. M. Wald, _Black hole entropy is the Noether charge_ , _Phys. Rev. D_ 48 (1993) R3427 [gr-qc/9307038].
  * (118) A. M. Polyakov, _Quantum Geometry of Bosonic Strings_ , _Phys. Lett. B_ 103 (1981) 207.
  * (119) K. Skenderis and S. N. Solodukhin, _Quantum effective action from the AdS / CFT correspondence_ , _Phys. Lett. B_ 472 (2000) 316 [hep-th/9910023].
  * (120) R. C. Myers, _Black hole entropy in two-dimensions_ , _Phys. Rev. D_ 50 (1994) 6412 [hep-th/9405162].
  * (121) A. Almheiri and J. Polchinski, _Models of AdS 2 backreaction and holography_, _JHEP_ 11 (2015) 014 [1402.6334].
  * (122) J. Engelsöy, T. G. Mertens and H. Verlinde, _An investigation of AdS 2 backreaction and holography_, _JHEP_ 07 (2016) 139 [1606.03438].
  * (123) A. Almheiri and B. Kang, _Conformal Symmetry Breaking and Thermodynamics of Near-Extremal Black Holes_ , _JHEP_ 10 (2016) 052 [1606.04108].
  * (124) E. Witten, _Anti-de Sitter space, thermal phase transition, and confinement in gauge theories_ , _Adv. Theor. Math. Phys._ 2 (1998) 505 [hep-th/9803131].
  * (125) G. W. Gibbons and S. W. Hawking, _Action Integrals and Partition Functions in Quantum Gravity_ , _Phys. Rev. D_ 15 (1977) 2752.
  * (126) S. W. Hawking, _The path-integral approach to quantum gravity_ , in _General Relativity: An Einstein Centenary Survey_ , S. W. Hawking and W. Israel, eds., (Cambridge, UK), Univ. Pr., 1979.
  * (127) S. W. Hawking and D. N. Page, _Thermodynamics of Black Holes in anti-De Sitter Space_ , _Commun. Math. Phys._ 87 (1983) 577.
  * (128) C. Germani and G. P. Procopio, _Two-dimensional Quantum Black Holes, Branes in BTZ and Holography_ , _Phys. Rev. D_ 74 (2006) 044012 [hep-th/0605068].
  * (129) R. Emparan, C. V. Johnson and R. C. Myers, _Surface terms as counterterms in the AdS / CFT correspondence_ , _Phys. Rev. D_ 60 (1999) 104001 [hep-th/9903238].
  * (130) M. Eune, W. Kim and S.-H. Yi, _Hawking-Page phase transition in BTZ black hole revisited_ , _JHEP_ 03 (2013) 020 [1301.0395].
  * (131) Y. S. Myung, _Phase transition between the BTZ black hole and AdS space_ , _Phys. Lett. B_ 638 (2006) 515 [gr-qc/0603051].
  * (132) A. Maloney and E. Witten, _Quantum Gravity Partition Functions in Three Dimensions_ , _JHEP_ 02 (2010) 029 [0712.0155].
  * (133) P. Nayak, A. Shukla, R. M. Soni, S. P. Trivedi and V. Vishal, _On the Dynamics of Near-Extremal Black Holes_ , _JHEP_ 09 (2018) 048 [1802.09547].
  * (134) D. Harlow and D. Jafferis, _The Factorization Problem in Jackiw-Teitelboim Gravity_ , _JHEP_ 02 (2020) 177 [1804.01081].
  * (135) E. Witten, _Deformations of JT Gravity and Phase Transitions_ , 2006.03494.


The Open Journal of Astrophysics submitted XXX; accepted YYY

††thanks: ⋆ E-mail<EMAIL_ADDRESS>† E-mail:
<EMAIL_ADDRESS>
#  CosmoPower-JAX: high-dimensional Bayesian inference
with differentiable cosmological emulators

D. Piras⋆1 A. Spurio Mancini†2 1 Département de Physique Théorique, Université
de Genève, 24 quai Ernest Ansermet, 1211 Genève 4, Switzerland  2 Mullard
Space Science Laboratory, University College London, Holmbury St. Mary,
Dorking, Surrey, RH5 6NT, UK

(28th August 2024)

###### Abstract

We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower
framework, which accelerates cosmological inference by building neural
emulators of cosmological power spectra. We show how, using the automatic
differentiation, batch evaluation and just-in-time compilation features of
JAX, and running the inference pipeline on graphics processing units (GPUs),
parameter estimation can be accelerated by orders of magnitude with advanced
gradient-based sampling techniques. These can be used to efficiently explore
high-dimensional parameter spaces, such as those needed for the analysis of
next-generation cosmological surveys. We showcase the accuracy and
computational efficiency of CosmoPower-JAX on two simulated Stage IV
configurations. We first consider a single survey performing a cosmic shear
analysis totalling 37 model parameters. We validate the contours derived with
CosmoPower-JAX and a Hamiltonian Monte Carlo sampler against those derived
with a nested sampler and without emulators, obtaining a speed-up factor of
$\mathcal{O}(10^{3})$. We then consider a combination of three Stage IV
surveys, each performing a joint cosmic shear and galaxy clustering (3x2pt)
analysis, for a total of 157 model parameters. Even with such a high-
dimensional parameter space, CosmoPower-JAX provides converged posterior
contours in 3 days, as opposed to the estimated 6 years required by standard
methods. CosmoPower-JAX is fully written in Python, and we make it publicly
available to help the cosmological community meet the accuracy requirements
set by next-generation surveys. github

##  1 Introduction

Bayesian inference of cosmological parameters from next-generation large-scale
structure (LSS) and cosmic microwave background (CMB) surveys such as Euclid
(Laureijs et al., 2011)111https://www.euclid-ec.org/, the Vera Rubin
Observatory (Ivezić et al., 2019)222https://www.lsst.org/, the Nancy Grace
Roman Space Telescope (Spergel et al., 2015)333https://roman.gsfc.nasa.gov/,
the Simons Observatory (Ade et al., 2019)444https://simonsobservatory.org/,
CMB-S4 (Abazajian et al., 2016)555https://cmb-s4.org/, and CMB-HD (Sehgal et
al., 2019)666https://cmb-hd.org/, will require the exploration of high-
dimensional parameter spaces – $\mathcal{O}(100)$ parameters and higher –
necessary to accurately model the physical signals and their several
systematic contaminants. Sampling the posterior distribution in these high-
dimensional spaces represents a significant computational challenge for Markov
Chain Monte Carlo (MCMC) algorithms (Roberts et al., 1997; Katafygiotis &
Zuev, 2008; Liu, 2009), which are traditionally used in cosmological analyses
(Lewis & Bridle, 2002; Audren et al., 2013; Brinckmann & Lesgourgues, 2019;
Torrado & Lewis, 2021). Gradient-based inference methods, such as Hamiltonian
Monte Carlo (HMC, Duane et al., 1987; Neal, 1996) and variational inference
(VI, Hoffman et al., 2013; Blei et al., 2017), manage to concentrate the
sampling in regions of high posterior mass, even in large parameter spaces,
provided one has efficient access to accurate derivatives of the likelihood
function with respect to the model parameters (Brooks et al., 2011; Neal,
2011; Zhang et al., 2017; Betancourt, 2017). Combining differentiable and
computationally inexpensive likelihood functions with gradient-based inference
techniques is thus crucial to efficiently obtain unbiased posterior
distributions in high-dimensional parameter spaces.

Recently, there has been growing interest in the use of machine learning
emulators to accelerate cosmological parameter estimation (e.g. Mootoovaloo et
al., 2020; Aricò et al., 2022; Zennaro et al., 2021; Nygaard et al., 2022;
Günther et al., 2022; Bonici et al., 2022; Eggemeier et al., 2022). Spurio
Mancini et al. (2022) (SM22 hereafter), in particular, developed CosmoPower, a
suite of neural network emulators of cosmological power spectra that replaces
the computation of these quantities traditionally performed with Einstein-
Boltzmann solvers such as the Code for Anisotropies in the Microwave
Background (CAMB, Lewis & Challinor 2011) or the Cosmic Linear Anisotropy
Solving System (CLASS, Blas et al. 2011). In SM22 the authors show how
Bayesian inference of cosmological parameters can be accelerated by several
orders of magnitude using CosmoPower; the speed-up becomes particularly
relevant when the emulators are employed within an inference pipeline that can
be run on graphics processing units (GPUs).

An additional advantage in using machine learning emulators is that they
efficiently provide accurate derivatives with respect to their input
parameters. This is made possible by the automatic differentiation features
implemented in the libraries routinely used to build these emulators, such as
TensorFlow (Abadi et al., 2015), JAX (Frostig et al., 2018) or PyTorch (Paszke
et al., 2019). Automatic differentiation (autodiff, Bartholomew-Biggs et al.,
2000; Neidinger, 2010; Baydin et al., 2018; Paszke et al., 2017; Margossian,
2019) can be used in gradient-based sampling algorithms to compute the
derivatives of the likelihood function, provided the emulators are inserted
within likelihood functions that can also be automatically differentiated. In
practice, this means that the entire software implementation of the
likelihoods should be written using primitives that belong to the
aforementioned libraries. This is, for example, the main idea behind jax-cosmo
(Campagne et al., 2023), a JAX-based library recently developed to compute
cosmological observables, validated against the widely used Core Cosmology
Library (CCL, Chisari et al., 2019). For other examples of gradient-based
inference in cosmology, see e.g. Hajian (2007); Taylor et al. (2008); Jasche
et al. (2010); Jasche & Wandelt (2013); Lavaux & Jasche (2015); Nguyen et al.
(2021); Valade et al. (2022); Kostić et al. (2022); Loureiro et al. (2023);
Porqueres et al. (2023); Ruiz-Zapatero et al. (2023).

In this paper, we first develop an implementation of CosmoPower built using
purely JAX primitives. We train the neural network emulators on the same data
used in SM22, achieving an accuracy similar to that obtained by the
TensorFlow-based implementation of SM22. We additionally compare the
derivatives of the power spectra with respect to the input cosmological
parameters computed using autodiff against those obtained with numerical
differences and CAMB, showing a good agreement overall. We then couple
CosmoPower-JAX with jax-cosmo to build likelihood functions completely written
using the JAX library; this allows us to run parameter estimation pipelines
using gradient-based algorithms on GPUs over very large parameter spaces. In
particular, we consider two examples of future Stage IV survey configurations.
In the first one, we study a single cosmic shear survey with typical Stage IV
specifications; we compare the posterior contours obtained with our JAX-based
pipeline against those obtained with a standard pipeline based on CCL, CAMB
and the nested sampler PolyChord (Handley et al., 2015a, b). In the second
case, we consider a joint analysis of three Stage IV surveys, each carrying
out a joint shear - galaxy clustering (3x2pt) analysis. We show how
CosmoPower-JAX enables efficient exploration of the 157-dimensional parameter
space required for the modelling of the observables at a fraction of the time
that would be required with traditional methods. In both experiments,
inference is performed using the No U-Turn Sampler (NUTS) HMC variation
(Hoffman & Gelman, 2014), as implemented in the NumPyro library (Bingham et
al., 2019; Phan et al., 2019).

The structure of this paper is as follows. In Sec. 2 we first describe our new
JAX emulators and how they are trained on the same data used in SM22. We then
review gradient-based methods for posterior inference and describe our JAX-
based cosmological likelihoods. In Sec. 3 we validate the accuracy of our
emulators and report our inference results in the two simulated Stage IV
configurations described above. We conclude in Sec. 4. We make CosmoPower-JAX,
which is fully written in Python, publicly
available.777https://github.com/dpiras/cosmopower-jax

##  2 Methods

###  2.1 CosmoPower-JAX

CosmoPower is a suite of neural network emulators of cosmological power
spectra, namely matter and cosmic microwave background power spectra (albeit
the emulation framework is highly flexible and can be used to emulate any
other cosmological quantity; Burger et al., 2023; Gong et al., 2023, see
e.g.). Power spectra represent the key component of 2-point statistics
analysis of cosmological fields, and are typically computed by Einstein-
Boltzmann solvers such as CAMB or CLASS. These computations usually represent
the main bottleneck in parameter estimation pipelines, particularly when high
accuracy is required or when the cosmological scenario considered is not
standard (see e.g. Spurio Mancini & Pourtsidou, 2022; Balkenhol et al., 2022;
Bolliet et al., 2023).

CosmoPower-JAX is a JAX-based implementation of CosmoPower; as such, it
implements the same two emulation methods present in the original CosmoPower
software. These are either a direct neural network mapping between
cosmological parameters and logarithmic spectra, or a neural network mapping
between cosmological parameters and the coefficients of a principal component
analysis (PCA) of the spectra; we refer to SM22 for a detailed discussion. In
our release version of CosmoPower-JAX these differences are dealt with
internally; therefore, the user can simply obtain fast predictions of
cosmological power spectra in three lines of code.

The neural networks we employ are made of four layers of 512 nodes each, with
the same activation function described in SM22. In the case of the CMB
temperature-polarisation and lensing power spectra, we preprocess the spectra
using PCA, keeping 512 and 64 components, respectively. We optimise the mean-
squared-error loss function using Adam (Kingma & Ba, 2015), and keep a
constant batch size of 512. The starting learning rate is 10-2, which we
decrease by a factor of 10 (until 10-6) every time the validation loss does
not improve for more than 40 consecutive epochs.

We choose to re-implement and re-train the CosmoPower emulators in JAX to
demonstrate that it is possible to build neural networks with this library;
nonetheless, in our public implementation of CosmoPower-JAX it is also
possible to upload weights and biases from previously-trained models and
obtain JAX predictions using only the forward pass of the network. Building
the emulators using the JAX library unlocks its automatic-differentiation
features, as well as the efficient batch evaluation and just-in-time
compilation. Importantly, JAX is built around the popular NumPy library
(Harris et al., 2020), which facilitates its use and portability. We refer the
reader to Frostig et al. (2018) and Campagne et al. (2023) for a complete
description of the features of JAX, especially in the context of cosmological
analyses.

###  2.2 Data

We consider the same datasets as in SM22 to train and validate our emulators.
Using CAMB, we produce a total of $\sim 2\times 10^{5}$ matter power spectra
at 420 wavenumbers in the interval $k\in\left[10^{-5}-10\right]$ Mpc-1 and
with varying redshift $z\in\left[0,5\right]$. We consider both the linear
matter power spectrum $P_{\delta\delta}^{\rm{L}}$ and the non-linear
correction $P_{\delta\delta}^{\rm{NL-CORR}}$, such that the non-linear matter
power spectrum $P_{\delta\delta}$ can be written as:

$\displaystyle
P_{\delta\delta}(k,z)=P_{\delta\delta}^{\textrm{L}}(k,z)\,P_{\delta\delta}^{\textrm{NL-
CORR}}(k,z)\ ,$ (1)

where $P_{\delta\delta}^{\textrm{NL-CORR}}$ is computed with HMcode (Mead et
al., 2015, 2016). This non-linear correction introduces two extra baryonic
parameters: the minimum halo concentration $c_{\textrm{min}}$, and the halo
bloating $\eta_{0}$, which we vary when performing Bayesian inference.

Using CAMB we also produce $\sim 5\times 10^{5}$ CMB power spectra, including
the temperature $C^{\textrm{TT}}_{\ell}$, polarisation
$C^{\textrm{EE}}_{\ell}$, temperature-polarisation $C^{\textrm{TE}}_{\ell}$
and lensing potential $C^{\phi\phi}_{\ell}$power spectra in the interval
$\ell\in\left[2,2508\right]$. All emulators are trained on the same parameter
range indicated in table 1 of SM22. In both cases, we use $\sim 90$% of the
data to train the neural networks (of which we leave 10% for validation), and
evaluate the trained models on the remaining $\mathcal{O}(10^{4})$ test
spectra.

###  2.3 Gradient-based Bayesian inference

The goal of Bayesian inference is to sample the posterior distribution
$p(\boldsymbol{\theta}|\boldsymbol{d})$ of some parameters
$\boldsymbol{\theta}$ given observed data $\boldsymbol{d}$. Bayes’ theorem
relates the posterior distribution to the likelihood function
$p(\boldsymbol{d}|\boldsymbol{\theta})$:

$\displaystyle
p(\boldsymbol{\theta}|\boldsymbol{d})=\frac{p(\boldsymbol{d}|\boldsymbol{\theta})p(\boldsymbol{\theta})}{p(\boldsymbol{d})},$
(2)

where $p(\boldsymbol{\theta})$ expresses prior knowledge on the parameters
$\boldsymbol{\theta}$, and the evidence $p(\boldsymbol{d})$ is a normalisation
factor commonly ignored in parameter estimation tasks.

Drawing from $p(\boldsymbol{\theta}|\boldsymbol{d})$ requires the use of
stochastic samplers to generate a list of samples drawn from the posterior
distribution. MCMC and nested sampling algorithms are the two main classes of
samplers typically used in cosmological applications. Widely used examples of
MCMC methods in cosmology include the Metropolis-Hastings (with its
variations) and affine ensemble algorithms (Lewis & Bridle, 2002; Lewis, 2013;
Karamanis et al., 2022; Goodman & Weare, 2010; Foreman-Mackey et al., 2013),
whereas the most widely adopted nested sampling algorithms are multimodal
nested sampling and slice-based nested sampling (Feroz & Hobson, 2008; Feroz
et al., 2009, 2019; Handley et al., 2015a, b).

Metropolis-Hastings and nested sampling methods can struggle to explore large
parameter spaces, since, as the dimensionality increases, sampling from the
typical set of the posterior distribution becomes exponentially harder
(Betancourt, 2017). In these challenging inference scenarios, it is useful to
resort to gradient-based algorithms, which manage to concentrate the sampling
in regions of high posterior mass despite the large number of model
parameters. A popular gradient-based method is Hamiltonian Monte Carlo (HMC,
Duane et al., 1987; Neal, 1996), which formulates the problem of sampling the
posterior distribution as the dynamical evolution of a particle with position
$\boldsymbol{\theta}$ and momentum $\boldsymbol{p}$. The Hamiltonian of the
system is written as:

$\mathcal{H}(\boldsymbol{\theta},\boldsymbol{p})=\frac{1}{2}\boldsymbol{p}^{T}\mathcal{M}^{-1}\boldsymbol{p}+\textbf{U}(\boldsymbol{\theta})\
,$ (3)

where $\mathcal{M}$ is a “mass matrix”, which is assumed to be the covariance
matrix of a zero-mean multivariate Gaussian distribution used to sample
$\boldsymbol{p}$, and the potential energy $\textbf{U}(\boldsymbol{\theta})$
is defined as:

$\textbf{U}(\boldsymbol{\theta})\equiv-\ln
p(\boldsymbol{\theta}|\boldsymbol{d})\ .$ (4)

Starting from an initial point in the phase space
$(\boldsymbol{\theta},\boldsymbol{p})$, it is possible to collect samples of
the posterior distribution by numerically solving the Hamilton’s equations
using a leapfrog algorithm. A new proposed state
$(\boldsymbol{\theta}^{*},\boldsymbol{p}^{*})$ is then either accepted or
rejected based on the new energy value
$\mathcal{H}(\boldsymbol{\theta}^{*},\boldsymbol{p}^{*})$. HMC can thus be
seen as a modification to the original Metropolis-Hastings algorithm, which
leverages the analogy with a Hamiltonian system to find a better proposal
distribution and results in a higher acceptance rate.

A practical difficulty of HMC is the need for tuning of the hyperparameters
governing the leapfrog numerical integration of the Hamiltonian dynamics.
These are the number and the size of steps to be taken by the integrator
before the sampler changes direction to a new random one. These numbers can be
particularly hard to tune and lead to inefficient sampling of the posterior.
The No U-Turn Sampler (NUTS, Hoffman & Gelman, 2014) tackles this issue by
forcing the sampler to avoid U-turns in parameter space. In particular, rather
than fixing the number of integration steps, NUTS adjusts the integration
length by running the leapfrog algorithm until the trajectory starts to return
to previously-visited regions of the parameter space; while the number of
model evaluations increases, this leads to a higher acceptance rate, and
therefore to faster and more reliable convergence.

In our experiments, we use the NUTS implementation provided by the NumPyro
library (Bingham et al., 2019; Phan et al., 2019). We set the integration step
size to $10^{-3}$, and find no difference in the final results by increasing
or decreasing this value by one order of magnitude. Additionally, we set the
maximum number of model evaluations before sampling a new momentum to $2^{7}$,
and specify a block mass matrix according to the expected correlations among
parameters. Finally, to further improve the geometry of the problem we perform
a reparameterisation of the inference variables: all of the parameters with a
Gaussian prior distribution are decentered (Gorinova et al., 2020), while all
of those with a uniform prior are rescaled to $\mathcal{U}\left[-5,5,\right]$.
All original prior distributions are reported in Table 1.

We run the NUTS algorithm on three NVIDIA A100 GPUs with 80 GB of memory.
Using multiple GPUs allows us to showcase the pmap feature of JAX to
distribute the NUTS chains over multiple GPU devices. We report below an
example snippet to run inference over multiple GPUs using NumPyro, where the
parallelisation happens simultaneously over a single device and across
multiple devices.

⬇

def do_mcmc(rng_key, n_vectorized=8):

# All posterior details are defined in this kernel

nuts_kernel = NUTS(log_posterior)

mcmc = MCMC(

nuts_kernel,

num_chains=n_vectorized,

chain_method="vectorized"

)

mcmc.run(rng_key)

return {∗∗mcmc.get_samples()}

# Number of devices to pmap over

n_parallel = jax.local_device_count()

rng_ks = jax.random.split(PRNGKey(rng_seed), n_parallel)

samples = pmap(do_mcmc)(rng_ks)

In this example, 8 chains are run on each GPU, and the total number of GPUs
available is determined at run time.

###  2.4 Likelihoods

| Parameter | Prior range | Fiducial value
---|---|---|---
Cosmology | $\omega_{\mathrm{b}}=\Omega_{\mathrm{b}}h^{2}$ |  $\mathcal{U}$[0.01875, 0.02625] | 0.02242
$\omega_{\mathrm{cdm}}=\Omega_{\mathrm{cdm}}h^{2}$ |  $\mathcal{U}$[0.05, 0.255] | 0.11933
$h$ |  $\mathcal{U}$[0.64, 0.82] | 0.6766
$n_{\textrm{s}}$ |  $\mathcal{U}$[0.84, 1.1] | 0.9665
$\mathrm{ln}10^{10}A_{\textrm{s}}$ |  $\mathcal{U}$[1.61, 3.91] | 3.047
Baryons | $c_{\mathrm{min}}$ |  $\mathcal{U}$[2, 4] | 2.6
$\eta_{0}$ |  $\mathcal{U}$[0.5, 1] | 0.7
Nuisance S1 |  $A^{\mathrm{S1}}_{\mathrm{IA},i}$ $i=1,\dots,10$ |  $\mathcal{U}$[$-6$, 6] | $1-0.1i$
$D^{\mathrm{S1}}_{z_{i},\textrm{source}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-4})$ | 0
$m^{\mathrm{S1}}_{i}$ $i=1,\dots,10$ | $\mathcal{N}(0.01,0.02)$ | 0.01
$D^{\mathrm{S1}}_{z_{i},\textrm{lens}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-4})$ | 0
$b^{\mathrm{S1}}_{i}$ $i=1,\dots,10$ |  $\mathcal{U}$[0.1, 5] | 1
Nuisance S2 |  $A^{\mathrm{S2}}_{\mathrm{IA},i}$ $i=1,\dots,10$ |  $\mathcal{U}$[0, 10] | $6.45-0.5i$
$D^{\mathrm{S2}}_{z_{i},\textrm{source}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-3})$ | 0
$m^{\mathrm{S2}}_{i}$ $i=1,\dots,10$ | $\mathcal{N}(0,0.002)$ | 0
$D^{\mathrm{S2}}_{z_{i},\textrm{lens}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-3})$ | 0
$b^{\mathrm{S2}}_{i}$ $i=1,\dots,10$ |  $\mathcal{U}$[0.8, 3] | 1.2+0.1$i$
Nuisance S3 |  $A^{\mathrm{S3}}_{\mathrm{IA},i}$ $i=1,\dots,10$ |  $\mathcal{U}$[0, 6] | $1.1-0.1i$
$D^{\mathrm{S3}}_{z_{i},\textrm{source}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-3})$ | 0
$m^{\mathrm{S3}}_{i}$ $i=1,\dots,10$ | $\mathcal{N}(0,0.003)$ | 0
$D^{\mathrm{S3}}_{z_{i},\textrm{lens}}\hfill i=1,\dots,10$ | $\mathcal{N}(0,10^{-3})$ | 0
$b^{\mathrm{S3}}_{i}$ $i=1,\dots,10$ |  $\mathcal{U}$[0.8, 3] | 1.25+0.05$i$
Table 1: Prior distributions and fiducial values of the cosmological and
nuisance parameters for the simulated analyses performed in this work. Uniform
(Gaussian) prior distributions are indicated with $\mathcal{U}$
($\mathcal{N}$). For the cosmological parameters and the baryonic feedback
parameters ($c_{\mathrm{min}},\eta_{0}$) the prior range corresponds to the
range of validity of our emulators. The superscripts S1, S2 and S3 refer to
three simulated Stage IV surveys, whose details we discuss in Sect. 2.4.

To test the computational efficiency and accuracy of our CosmoPower-JAX
emulators, we run examples of cosmological inference on large parameter
spaces, such as those that will characterise the analysis of future surveys.
We consider a power spectrum analysis of both cosmic shear alone, as well as
in cross-correlation with galaxy clustering (3x2pt), which in recent years has
become the standard method to combine information on the cosmic shear and
galaxy clustering field (e.g. Joachimi & Bridle, 2010; Sanchez et al., 2021;
Heymans et al., 2021; Abbott et al., 2022). Redshift information is included
in the analysis through tomography, assigning galaxies to one of $N_{\rm
bins}$ bins based on their photometric redshift. The angular power spectra of
the probes are jointly modelled computing all of the unique correlations
between tomographic bins $i,j=1,\dots,N_{\rm bins}$. In the following we
briefly review the modelling of the power spectra for cosmic shear, galaxy
clustering and their cross-correlation (“galaxy-galaxy lensing”), as well as
describe the Stage IV surveys we take into consideration. We follow the
notation of SM22.

Angular power spectra for the three probes can be expressed as integrals of
the matter power spectrum $P_{\delta\delta}(k,z)$ (a function of wavevector
$k$ and redshift $z$), weighted by pairs of window functions $W$ for the shear
($\gamma$), clustering ($\mathrm{n}$) and intrinsic alignment ($\mathrm{I}$)
fields:

$\displaystyle
C_{ij}^{AB}(\ell)=\int_{0}^{\chi_{\mathrm{H}}}\mathrm{d}\chi\,\frac{W_{i}^{A}\,W_{j}^{B}}{\chi^{2}}\,P_{\delta\delta}\left(k=\frac{\ell+1/2}{\chi},z\right),$
(5)

where $\\{A,B\\}=\\{\gamma,\mathrm{n},\mathrm{I}\\}$, $\chi$ is the comoving
distance (itself a function of redshift), and the upper limit of integration
is the Hubble radius $\chi_{\rm H}=c/H_{0}$, with $c$ the speed of light and
$H_{0}$ the Hubble constant. In Eq. (5) we assumed the extended Limber
approximation (LoVerde & Afshordi, 2008) which connects Fourier scales $k$
with angular multipoles $\ell$.

The angular power spectrum $C_{ij}^{\epsilon\epsilon}(\ell)$ of the cosmic
shear signal for tomographic bins $i$ and $j$ is a combination of a pure shear
contribution and an intrinsic alignment contamination:

$\displaystyle
C_{ij}^{\epsilon\epsilon}(\ell)=C_{ij}^{\gamma\gamma}(\ell)+C_{ij}^{\gamma\mathrm{I}}(\ell)+C_{ij}^{\mathrm{I}\gamma}(\ell)+C_{ij}^{\mathrm{I}\mathrm{I}}(\ell)\
.$ (6)

For a tomographic redshift bin distribution
$n_{i,,\textrm{source}}(z),i=1,\dots,N_{\rm bins}$, the weighting function
$W_{i}^{\gamma}(\chi)$ for the cosmic shear field $\gamma$ is given by:

$\displaystyle W_{i}^{\gamma}(\chi)=\frac{3\,H_{0}^{2}\,\Omega_{\rm
m}}{2\,c^{2}}\frac{\chi}{a}\int_{\chi}^{\chi_{\rm
H}}\mathrm{d}\chi^{\prime}\,n_{i,\textrm{source}}(\chi^{\prime})\,\frac{\chi^{\prime}-\chi}{\chi^{\prime}}\
,$ (7)

where $\Omega_{\rm m}$ is the matter density parameter and $a$ is the scale
factor. To model the window function for the intrinsic alignment field
$\mathrm{I}$, we start from the commonly-used non-linear alignment model
(Hirata & Seljak, 2004; Joachimi et al., 2011):

$\displaystyle W_{i}^{\rm I}(\chi)=-A_{\rm IA}\left(\frac{1+z}{1+z_{\rm
p}}\right)^{\eta_{\rm IA}}\frac{C_{1}\,\rho_{\rm cr}\,\Omega_{\rm
m}}{D(\chi)}\,n_{i,\textrm{source}}(\chi)\ ,$ (8)

with two free parameters $A_{\rm IA}$ and $\eta_{\rm IA}$; the linear growth
factor $D(\chi)$, the critical matter density $\rho_{\rm cr}$, a constant
$C_{1}$ and a pivot redshift $z_{\rm p}$ also enter Eq. (8). We modify the
redshift dependence in Eq. (8) by setting $\eta_{\rm IA}=0$ and using one
amplitude parameter $A_{\mathrm{IA},i}$ for each redshift bin $i$. This gives
more flexibility to the modelling, at the cost of increasing the number of
free parameters. Thus, the window function for the intrinsic alignment field
reads:

$\displaystyle W_{i}^{\rm I}(\chi)=-A_{\mathrm{IA},i}\,\frac{C_{1}\,\rho_{\rm
cr}\,\Omega_{\rm m}}{D(\chi)}\,n_{i,\textrm{source}}(\chi)\ .$ (9)

| S1 | S2 | S3
---|---|---|---
$f_{\textrm{sky}}$ | 0.35 | 0.05 | 0.45
$\bar{n}_{\textrm{source}}$ [gal/arcmin2] | 30 | 50 | 50
$\bar{n}_{\textrm{lens}}$ [gal/arcmin2] | 40 | 55 | 60
$\sigma_{\epsilon}$ | 0.30 | 0.25 | 0.35
Table 2: Specifications for the three Stage IV surveys S1, S2 and S3 that we
consider in this work. All details can be found in Sect. 2.4.

We include a multiplicative bias in our modelling of the cosmic shear signal.
This keeps into account that the ellipticities measured in a survey are
typically biased with respect to an ideal measurement, due to e.g. point
spread function modelling errors, selection biases or detector effects
(Huterer et al., 2006; Amara & Réfrégier, 2008; Kitching et al., 2015;
Mandelbaum et al., 2018; Taylor & Kitching, 2018; Mahony et al., 2022; Cragg
et al., 2022). We consider a multiplicative bias coefficient $m_{i}$ for each
redshift bin $i$; therefore, the multiplicative bias rescales the cosmic shear
power spectrum $C_{ij}^{\epsilon\epsilon}(\ell)$ by a factor
$(1+m_{i})(1+m_{j})$. Another important source of uncertainty is the imperfect
knowledge of the redshift distributions $n_{i,\textrm{source}}(z)$. Errors in
these distributions have traditionally been modelled through the addition of a
shift parameter $D_{z_{i},\textrm{source}}$ for each redshift bin $i$, which
changes the mean of the bin redshift distribution by shifting the distribution
$n_{i,\textrm{source}}(z)$ to
$n_{i,\textrm{source}}^{\prime}(z)=n_{i,\textrm{source}}(z-D_{z_{i},\textrm{source}})$
(see e.g. Eifler et al., 2021).

Figure 1: Derivatives of the matter power spectrum components as defined in
Eq. (1) with respect to cosmological parameters for random samples in the test
set. Solid red lines represent finite differences calculated using CAMB and a
five-point stencil as in Eq. (13), while dashed blue lines are obtained with
the CosmoPower-JAX emulator and automatic differentiation (autodiff). A single
derivative computed with finite differences requires $\mathcal{O}(10)$ s, as
opposed to the $\mathcal{O}(0.1)$ s required for the derivatives with respect
to all cosmological parameters obtained using CosmoPower-JAX.

For the modelling of the galaxy clustering power spectrum
$C_{ij}^{\mathrm{n}\mathrm{n}}(\ell)$ we assume a linear bias model between
galaxy and dark matter density, with a free parameter $b_{i}$ for each
redshift bin $i$. The window function $W_{i}^{\rm n}$ for the galaxy
clustering field $\mathrm{n}$ is then given by:

$\displaystyle W_{i}^{\rm n}(\chi)=b_{i}\,n_{i,\textrm{lens}}(\chi)\ ,$ (10)

where for the galaxy clustering we use a different sample of redshift
distributions $n_{i,\textrm{lens}}(\chi)$, which is also not perfectly known
and therefore includes a set of shifts $D_{z_{i},\textrm{lens}}$. We do not
take into account redshift space distortions. Finally, the galaxy-galaxy
lensing power spectrum is given by

$\displaystyle
C_{ij}^{\mathrm{n}\epsilon}(\ell)=C_{ij}^{\mathrm{n}\gamma}(\ell)+C_{ij}^{\mathrm{n}\mathrm{I}}(\ell)\
.$ (11)

The first experiment we run is a simulated cosmic shear analysis of a Stage IV
survey (dubbed “S1”). We assume 10 tomographic bins for
$n_{\textrm{source}}(z)$, with galaxies following the distribution:

$\displaystyle n(z)\propto z^{2}\exp\left(-\left(z/z_{0}\right)^{1.5}\right).$
(12)

with $z_{0}=0.64$ (Smail et al., 1994; Joachimi & Bridle, 2010). We consider a
Gaussian covariance matrix (e.g. Tutusaus et al., 2020) with a sky fraction
$f_{\textrm{sky}}=0.35$, a surface density of galaxies
$\bar{n}_{\textrm{source}}=30$ galaxies per arcmin2, and an observed
ellipticity dispersion $\sigma_{\epsilon}=0.3$. The $C(\ell)$ spectra are
computed for 30 bin values between $\ell_{\textrm{min}}=30=$ and
$\ell_{\textrm{max}}=3000$.

Figure 2: Same as in Fig. 1 for the CMB temperature power spectrum
$C_{\ell}^{\textrm{TT}}$. Figure 3: Inference results for the cosmic shear
analysis of a simulated Stage IV survey with a total of 37 parameters. The
contours show the 68% and 95% credibility contours and 1-D marginals for all
of the parameters. The inset table reports the plot range and corresponding
parameter name of every axis symbol; all parameters are described in Sect.
2.4, while in Table 1 we report the prior distributions and fiducial values.
The contours in blue are obtained with the No U-Turn Sampler (NUTS) combined
with our CosmoPower-JAX emulators. The contours in red are produced using the
PolyChord sampler together with the Boltzmann code CAMB. The CosmoPower-JAX
results are obtained over 3 graphics processing units (GPUs) in about 1 day,
while the PolyChord+CAMB chain ran on 48 CPU cores for about 4 months, so that
the total speed-up factor in terms of CPU/GPU hours is $\mathcal{O}(10^{3})$.
When replacing CAMB with the original CosmoPower emulator (Spurio Mancini et
al., 2022), we obtain the contours in grey, which required about 8 days on 48
CPU cores, and show an excellent agreement overall. All of the contours shown
in this paper are plotted using GetDist (Lewis, 2019).

We use the nested sampler PolyChord to sample the posterior distribution and
run the inference pipeline within Cobaya (Lewis, 2013). To compute the
theoretical predictions for the cosmological observables we use the Core
Cosmology Library (CCL, Chisari et al., 2019). The three upper blocks of Table
1 summarise the prior distributions assumed for the parameters used in this
analysis. These include:

  * •

five standard $\Lambda$CDM cosmological parameters, namely the baryon density
$\omega_{\mathrm{b}}=\Omega_{\mathrm{b}}h^{2}$, the cold dark matter density
$\omega_{\mathrm{cdm}}=\Omega_{\mathrm{cdm}}h^{2}$, the Hubble parameter $h$,
the scalar spectral index $n_{\textrm{s}}$, and the primordial power spectrum
amplitude $\mathrm{ln}10^{10}A_{\textrm{s}}$;

  * •

the HMcode parameters $c_{\textrm{min}}$, $\eta_{0}$;

  * •

one intrinsic alignment amplitude $A_{\textrm{IA}}$ for each of the ten
redshift bins;

  * •

one shift parameter $D_{z,\textrm{source}}$ for each of the ten source
redshift distributions;

  * •

one multiplicative bias coefficient $m$ for each of the ten redshift bins,

for a total of 37 model parameters.

We run a second experiment in which we consider the joint analysis of three
surveys, dubbed “S1”, “S2” and “S3”, each one performing a 3x2pt analysis over
10 redshift bins over different $z$ ranges. $n_{\textrm{source}}(z)$ and
$n_{\textrm{lens}}(z)$ are sampled from Eq. (12), with different galaxy
densities $\bar{n}_{\textrm{source}}$ and $\bar{n}_{\textrm{lens}}$ reported
in Table 2. In this analysis, the cosmological and baryon parameters are
unique, while each survey has its own set of nuisance parameters to model
redshift distribution shifts, multiplicative biases and galaxy bias. With
respect to the cosmic shear analysis, for each survey we add:

  * •

one shift parameter $D_{z,\textrm{lens}}$ for each of the ten lens redshift
distributions;

  * •

one linear galaxy bias coefficient $b$ for each of the ten redshift bins,

for a total of 157 parameters. The prior distributions and fiducial values of
the nuisance parameters for all surveys are reported in Table 1, while we
report all survey specifications in Table 2.

##  3 Results

###  3.1 Emulator accuracy

Having used the same dataset and training settings, the accuracy of our JAX
emulators is in excellent agreement with the results presented in SM22, i.e.
the difference with respect to the CAMB test spectra is always sub-percent for
both the matter power spectrum and the CMB probes. We do not show the accuracy
plots here for brevity, and refer the reader to SM22 for an extended
comparison.

In Fig. 1 we then compare the derivatives of the linear matter power spectrum
and the non-linear boost with respect to cosmological parameters, using CAMB
and finite differences with a five-point stencil as a benchmark. These
numerical derivatives rely on a simple Taylor expansion of the function $f(x)$
being considered, and are calculated as:

$f^{\prime}(x)\simeq\frac{-f(x+2s)+8f(x+s)-8f(x-s)+f(x-2s)}{12s}\ ,$ (13)

where $s$ is a small step size. We choose the step size for each derivative by
ensuring that the numerical derivative would only change at the sub-percent
level, or within machine precision, when increasing or decreasing $s$ by one
order of magnitude.

The derivatives computed by our JAX emulators are efficient to obtain and show
excellent agreement with the numerical ones, despite the latter being prone to
numerical instabilities due to the choice of the step size. A single
derivative of the linear matter power spectrum with finite differences takes
$\mathcal{O}(10)$ s on a CPU core, since it requires four independent calls to
CAMB. On the other hand, on the same CPU core we can obtain derivatives with
respect to all cosmological parameters and for batches of tens of spectra in
$\mathcal{O}(0.1)$ s; we expect the speed-up to be even more substantial when
running on GPU. In Fig. 2 we repeat this comparison for the
$C_{\ell}^{\textrm{TT}}$ power spectrum, showing again good agreement and high
computational efficiency for all cosmological parameters.

###  3.2 Bayesian inference

Here we present and comment on the results of running our gradient-based
inference pipelines. We emphasise that, while the survey configurations that
we consider are representative of the ones expected from upcoming surveys, our
goal is not to obtain official forecast predictions, but rather to focus on
the challenge represented by the large parameter space required by these
analyses, and how to implement viable solutions to efficiently explore these
spaces.

In Fig. 3 we show the posterior contours for the cosmological and nuisance
parameters in the case of a cosmic shear analysis of a single Stage IV-like
survey. The results obtained with NUTS and the CosmoPower-JAX emulators on
three GPUs in one day are in excellent agreement with those yielded by the
PolyChord sampler using CAMB and CCL, which however require about 4 months on
48 2.60 GHz Intel Xeon Platinum 8358 CPU cores. The total speed-up factor is
then $\mathcal{O}(10^{3})$ in terms of CPU/GPU hours, and
$\mathcal{O}(10^{2})$ in terms of actual elapsed time. We also run PolyChord
with the original CosmoPower emulator (written using TensorFlow) on the same
48 CPU cores, obtaining overlapping contours in about 8 days. This suggests
that a $\mathcal{O}(10)$ acceleration is provided by the neural emulator, with
the NUTS sampler and the GPU hardware providing the remaining
$\mathcal{O}(10)$ speed-up (in terms of the actual elapsed time).

Figure 4: Inference results for the joint analysis of three simulated Stage IV surveys, each performing a 3x2pt analysis, for a total of 157 model parameters. Here we show the posterior contours for the cosmological parameters, while we report the marginal distribution of the nuisance parameters in Fig. 5. The blue contours represent the 68% and 95% credibility contours and 1-D marginals. These are obtained with the No U-Turn Sampler (NUTS) combined with our CosmoPower-JAX emulators in about 3 days on 3 graphics processing units (GPUs). Dashed grey lines represent the fiducial values, reported in Table 1 together with the prior distributions. Figure 5: Marginal posterior distributions obtained in about 3 days on 3 graphics processing units (GPUs) with our CosmoPower-JAX emulators and the No U-Turn Sampler (NUTS) for the nuisance parameters of the the joint analysis of three simulated Stage IV surveys. Each survey performs a 3x2pt analysis, for a total of 157 parameters. The corresponding cosmological constraints can be found in Fig. 4. Dashed grey lines represent the fiducial values, reported in Table 1 together with the prior distributions. | $\omega_{\mathrm{b}}$ | $\omega_{\mathrm{cdm}}$ | $h$ | $n_{\textrm{s}}$ | $\mathrm{ln}10^{10}A_{\textrm{s}}$
---|---|---|---|---|---
(a) | 0.02 | $2\cdot 10^{-5}$ | $1\cdot 10^{-5}$ | $2\cdot 10^{-5}$ | $4\cdot 10^{-5}$
(b) | 0.1 | $2\cdot 10^{-5}$ | $1\cdot 10^{-5}$ | $2\cdot 10^{-5}$ | $3\cdot 10^{-5}$
(c) | 0.02 | 0.02 | 0.02 | 0.03 | 0.03
(d) | 0.03 | 0.03 | 0.03 | 0.03 | 0.03
Table 3: Effective sample size per likelihood call (in percentage) $\eta$, as
defined in Eq. (14), for:
(a) PolyChord+CAMB, cosmic-shear only;
(b) PolyChord+CosmoPower, cosmic-shear only;
(c) NUTS+CosmoPower-JAX, cosmic-shear only;
(d) NUTS+CosmoPower-JAX, joint analysis.
Except for $\omega_{\mathrm{b}}$, which is unconstrained in the case of the
cosmic shear analysis, the combination of NUTS and CosmoPower-JAX yields
values of $\eta$ higher by up to three orders of magnitude, confirming the
higher efficiency of the sampling approach.

In Fig. 4 and Fig. 5 we then show the posterior contours for the 7
cosmological parameters and the 150 nuisance parameters, respectively, of the
joint analysis of the three Stage IV-like surveys described in Sect. 2.4.
Obtaining the corresponding contours using PolyChord and CAMB would be
extremely computationally expensive: optimistically, assuming a linear (rather
than exponential, see e.g. Bishop 1995) scaling of the run time with the
number of inference parameters, we would expect to be able to obtain such
contours in at least 6 years on 48 CPU cores. Therefore, we only show the
results with NUTS and CosmoPower-JAX, which we obtain in about 3 days over 3
GPUs, for a total speed-up factor of $\mathcal{O}(10^{4})$ in terms of CPU/GPU
hours. These results are unbiased and in good agreement with the fiducial
values of our analysis.

We further quantify the efficiency of the different sampling approaches by
computing the effective number of samples per likelihood call, namely:

$\eta=\frac{n_{\textrm{eff}}}{n_{\textrm{eval}}}\ ,$ (14)

where $n_{\textrm{eff}}$ is the effective sample size (ESS), computed with the
ArviZ library (Kumar et al., 2019), and $n_{\textrm{eval}}$ is the total
number of likelihood calls. As we report and discuss in Table 3, $\eta$ is
almost always $\mathcal{O}(10^{3})$ times bigger when using the NUTS sampler
with CosmoPower-JAX with respect to the PolyChord sampler, and does not
degrade even when considering 157 inference parameters.

##  4 Conclusions

We presented CosmoPower-JAX, a JAX-based implementation of the cosmological
neural emulator CosmoPower (Spurio Mancini et al., 2022). CosmoPower-JAX
adopts the same emulation methods of the original TensorFlow-based version of
CosmoPower, but implements them using the JAX library. The key feature of JAX
is the ability to automatically differentiate functions written in common
Python libraries, such as NumPy. Additionally, JAX allows NumPy software to be
dynamically compiled, evaluated in batch, and run on graphics processing units
(GPUs) and tensor processing units (TPUs). In this paper, we made use of these
JAX features to show how they can accelerate Bayesian cosmological inference
by orders of magnitude.

We started by training a set of CosmoPower-JAX emulators of matter and cosmic
microwave background (CMB) power spectra, using the same datasets of Spurio
Mancini et al. (2022). We compared the accuracy of our JAX-based emulators
with the Boltzmann code CAMB, showing excellent agreement both at the power
spectra level as well as at the level of their derivatives. Crucially, rather
than being obtained with finite differences, which are prone to numerical
instabilities and require fine-tuning of the step size, derivatives of the
power spectra with CosmoPower-JAX are obtained straightforwardly and
efficiently with the built-in automatic differentiation (autodiff) features
provided by JAX. We note that the usefulness of our differentiable power
spectra emulators is not limited to 2-point-statistics analyses of
cosmological fields: for example, field-level inference pipelines (e.g.
Makinen et al., 2021; Makinen et al., 2022; Loureiro et al., 2023; Porqueres
et al., 2023) greatly benefit from efficient access to power spectra and their
gradients.

We combined our emulators with the JAX-based cosmological library jax-cosmo
(Campagne et al., 2023) to write likelihood functions that are fully
differentiable by virtue of their pure-JAX implementation. Running these
likelihoods within inference pipelines using gradient-based Monte Carlo
algorithms like the No U-Turn Sampler (NUTS), we showed how to use CosmoPower-
JAX to perform cosmological Bayesian inference over parameter spaces with
$\mathcal{O}(100)$ parameters, using multiple graphics processing units (GPUs)
in parallel. We performed a cosmic shear analysis for a Stage IV-like survey
configuration with a total of 37 parameters, showing good agreement with
standard nested sampling inference performed with CAMB, while being
$\mathcal{O}(10^{3})$ times faster. We finally showed that with CosmoPower-JAX
and NUTS it is possible to perform a combined Bayesian inference on three
different Stage IV-like surveys, each performing a joint cosmic shear and
galaxy clustering analysis (3x2pt), for a total of 157 parameters, of which
150 represent nuisance parameters used to model systematic effects for each
survey. We obtained unbiased posterior contours in about 3 days on 3 GPUs, as
opposed to an optimistic estimate of at least 6 years needed to obtain the
same results with traditional sampling methods and standard Boltzmann codes,
for a total speed-up factor of $\mathcal{O}(10^{4})$ in terms of CPU/GPU
hours.

Developing fully differentiable, GPU-accelerated implementations of key
analysis pipelines is of the utmost importance to tackle the massive
computational requirements imposed by upcoming Stage IV surveys. To this
purpose, we envision the use of differentiable libraries such as TensorFlow,
JAX and PyTorch to become standard practice in cosmological software
implementations, not least because they can all be used within the Python
programming language, which is becoming the de facto standard programming
language in both the cosmological and machine learning communities.

Libraries such as CosmoPower-JAX and jax-cosmo are therefore of paramount
importance for the success of next generation surveys. To this purpose, we
plan to integrate the public release of CosmoPower-JAX with the jax-cosmo
library. In parallel, we will explore improved gradient-based posterior
samplers, such as those described in Ver Steeg & Galstyan (2021); Park (2021);
Robnik et al. (2022); Wong et al. (2023), to further improve the sampling
efficiency, and thus the usefulness of our differentiable emulators.

## Acknowledgements

We are grateful to Erminia Calabrese, Benjamin Joachimi and Jason McEwen for
useful discussions and feedback on this work. DP was supported by a Swiss
National Science Foundation (SNSF) Professorship grant (No. 202671). ASM
acknowledges support from the MSSL STFC Consolidated Grant ST/W001136/1. Part
of the computations were performed on the Baobab cluster at the University of
Geneva. This work has been partially enabled by funding from the UCL
Cosmoparticle Initiative. The authors are pleased to acknowledge that part of
the work reported on in this paper was performed using the Princeton Research
Computing resources at Princeton University which is a consortium of groups
led by the Princeton Institute for Computational Science and Engineering
(PICSciE) and Office of Information Technology’s Research Computing.

## Data Availability

We make the CosmoPower-JAX emulators available in this GitHub repository
(https://github.com/dpiras/cosmopower-jax, also accessible by clicking the
icon github).

## References

  * Abadi et al. (2015) Abadi M., et al., 2015, TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems, https://www.tensorflow.org/
  * Abazajian et al. (2016) Abazajian K. N., et al., 2016, arXiv e-prints, p. arXiv:1610.02743
  * Abbott et al. (2022) Abbott T. M. C., et al., 2022, Phys. Rev. D, 105, 023520
  * Ade et al. (2019) Ade P., et al., 2019, Journal of Cosmology and Astroparticle Physics, 2019, 056–056
  * Amara & Réfrégier (2008) Amara A., Réfrégier A., 2008, Monthly Notices of the Royal Astronomical Society, 391, 228
  * Aricò et al. (2022) Aricò G., Angulo R., Zennaro M., 2022, Open Research Europe, 1
  * Audren et al. (2013) Audren B., Lesgourgues J., Benabed K., Prunet S., 2013, JCAP, 1302, 001
  * Balkenhol et al. (2022) Balkenhol L., et al., 2022, arXiv e-prints, p. arXiv:2212.05642
  * Bartholomew-Biggs et al. (2000) Bartholomew-Biggs M., Brown S., Christianson B., Dixon L., 2000, Journal of Computational and Applied Mathematics, 124, 171
  * Baydin et al. (2018) Baydin A. G., Pearlmutter B. A., Radul A. A., Siskind J. M., 2018, Journal of Machine Learning Research, 18, 1
  * Betancourt (2017) Betancourt M., 2017, arXiv e-prints, p. arXiv:1701.02434
  * Bingham et al. (2019) Bingham E., et al., 2019, J. Mach. Learn. Res., 20, 28:1
  * Bishop (1995) Bishop C., 1995, Neural networks for pattern recognition. Oxford University Press, USA
  * Blas et al. (2011) Blas D., Lesgourgues J., Tram T., 2011, J. Cosmology Astropart. Phys, 2011, 034
  * Blei et al. (2017) Blei D. M., Kucukelbir A., McAuliffe J. D., 2017, Journal of the American Statistical Association, 112, 859
  * Bolliet et al. (2023) Bolliet B., Spurio Mancini A., Hill J. C., Madhavacheril M., Jense H. T., Calabrese E., Dunkley J., 2023, arXiv e-prints, p. arXiv:2303.01591
  * Bonici et al. (2022) Bonici M., Biggio L., Carbone C., Guzzo L., 2022, arXiv e-prints, p. arXiv:2206.14208
  * Brinckmann & Lesgourgues (2019) Brinckmann T., Lesgourgues J., 2019, Physics of the Dark Universe, 24, 100260
  * Brooks et al. (2011) Brooks S., Gelman A., Jones G., Meng X.-L., eds, 2011, Handbook of Markov Chain Monte Carlo. Chapman and Hall/CRC, doi:10.1201/b10905, https://doi.org/10.1201%2Fb10905
  * Burger et al. (2023) Burger P. A., et al., 2023, A&A, 669, A69
  * Campagne et al. (2023) Campagne J.-E., et al., 2023, Open J. Astrophys., 6, 1
  * Chisari et al. (2019) Chisari N. E., et al., 2019, The Astrophysical Journal Supplement Series, 242, 2
  * Cragg et al. (2022) Cragg C., Duncan C. A. J., Miller L., Alonso D., 2022, Monthly Notices of the Royal Astronomical Society, 518, 4909
  * Duane et al. (1987) Duane S., Kennedy A., Pendleton B. J., Roweth D., 1987, Physics Letters B, 195, 216
  * Eggemeier et al. (2022) Eggemeier A., Camacho-Quevedo B., Pezzotta A., Crocce M., Scoccimarro R., Sánchez A. G., 2022, Monthly Notices of the Royal Astronomical Society, 519, 2962
  * Eifler et al. (2021) Eifler T., et al., 2021, Monthly Notices of the Royal Astronomical Society, 507, 1514
  * Feroz & Hobson (2008) Feroz F., Hobson M. P., 2008, Monthly Notices of the Royal Astronomical Society, 384, 449
  * Feroz et al. (2009) Feroz F., Hobson M. P., Bridges M., 2009, Monthly Notices of the Royal Astronomical Society, 398, 1601
  * Feroz et al. (2019) Feroz F., Hobson M. P., Cameron E., Pettitt A. N., 2019, The Open Journal of Astrophysics, 2
  * Foreman-Mackey et al. (2013) Foreman-Mackey D., Hogg D. W., Lang D., Goodman J., 2013, Publications of the Astronomical Society of the Pacific, 125, 306
  * Frostig et al. (2018) Frostig R., Johnson M., Leary C., 2018, in SysML (2018). https://mlsys.org/Conferences/doc/2018/146.pdf
  * Gong et al. (2023) Gong Z., Halder A., Barreira A., Seitz S., Friedrich O., 2023, arXiv e-prints, p. arXiv:2304.01187
  * Goodman & Weare (2010) Goodman J., Weare J., 2010, Communications in Applied Mathematics and Computational Science, 5, 65
  * Gorinova et al. (2020) Gorinova M., Moore D., Hoffman M., 2020, in III H. D., Singh A., eds, Proceedings of Machine Learning Research Vol. 119, Proceedings of the 37th International Conference on Machine Learning. PMLR, pp 3648–3657, https://proceedings.mlr.press/v119/gorinova20a.html
  * Günther et al. (2022) Günther S., Lesgourgues J., Samaras G., Schöneberg N., Stadtmann F., Fidler C., Torrado J., 2022, Journal of Cosmology and Astroparticle Physics, 2022, 035
  * Hajian (2007) Hajian A., 2007, Phys. Rev. D, 75, 083525
  * Handley et al. (2015a) Handley W. J., Hobson M. P., Lasenby A. N., 2015a, Monthly Notices of the Royal Astronomical Society: Letters, 450, L61
  * Handley et al. (2015b) Handley W. J., Hobson M. P., Lasenby A. N., 2015b, Monthly Notices of the Royal Astronomical Society, 453, 4384
  * Harris et al. (2020) Harris C. R., et al., 2020, Nature, 585, 357
  * Heymans et al. (2021) Heymans C., et al., 2021, A&A, 646, A140
  * Hirata & Seljak (2004) Hirata C. M., Seljak U., 2004, Phys. Rev. D, 70, 063526
  * Hoffman & Gelman (2014) Hoffman M. D., Gelman A., 2014, J. Mach. Learn. Res., 15, 1593–1623
  * Hoffman et al. (2013) Hoffman M. D., Blei D. M., Wang C., Paisley J., 2013, Journal of Machine Learning Research, 14, 1303
  * Huterer et al. (2006) Huterer D., Takada M., Bernstein G., Jain B., 2006, Monthly Notices of the Royal Astronomical Society, 366, 101
  * Ivezić et al. (2019) Ivezić Ž., et al., 2019, ApJ, 873, 111
  * Jasche & Wandelt (2013) Jasche J., Wandelt B. D., 2013, Monthly Notices of the Royal Astronomical Society, 432, 894
  * Jasche et al. (2010) Jasche J., Kitaura F. S., Li C., Enßlin T. A., 2010, Monthly Notices of the Royal Astronomical Society, 409, 355
  * Joachimi & Bridle (2010) Joachimi B., Bridle S. L., 2010, A&A, 523, A1
  * Joachimi et al. (2011) Joachimi B., Mandelbaum R., Abdalla F. B., Bridle S. L., 2011, A&A, 527, A26
  * Karamanis et al. (2022) Karamanis M., Beutler F., Peacock J. A., Nabergoj D., Seljak U., 2022, Monthly Notices of the Royal Astronomical Society, 516, 1644
  * Katafygiotis & Zuev (2008) Katafygiotis L., Zuev K., 2008, Probabilistic Engineering Mechanics, 23, 208
  * Kingma & Ba (2015) Kingma D. P., Ba J., 2015, in Bengio Y., LeCun Y., eds, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. http://arxiv.org/abs/1412.6980
  * Kitching et al. (2015) Kitching T. D., Taylor A. N., Cropper M., Hoekstra H., Hood R. K. E., Massey R., Niemi S., 2015, Monthly Notices of the Royal Astronomical Society, 455, 3319
  * Kostić et al. (2022) Kostić A., Nguyen N.-M., Schmidt F., Reinecke M., 2022, arXiv e-prints, p. arXiv:2212.07875
  * Kumar et al. (2019) Kumar R., Carroll C., Hartikainen A., Martin O., 2019, Journal of Open Source Software, 4, 1143
  * Laureijs et al. (2011) Laureijs R., et al., 2011, arXiv e-prints, p. arXiv:1110.3193
  * Lavaux & Jasche (2015) Lavaux G., Jasche J., 2015, Monthly Notices of the Royal Astronomical Society, 455, 3169
  * Lewis (2013) Lewis A., 2013, Phys. Rev. D, 87, 103529
  * Lewis (2019) Lewis A., 2019, arXiv e-prints, p. arXiv:1910.13970
  * Lewis & Bridle (2002) Lewis A., Bridle S., 2002, Phys. Rev. D, 66, 103511
  * Lewis & Challinor (2011) Lewis A., Challinor A., 2011, CAMB: Code for Anisotropies in the Microwave Background, Astrophysics Source Code Library, record ascl:1102.026 (ascl:1102.026)
  * Liu (2009) Liu J., 2009, Monte Carlo Strategies in Scientific Computing. Springer, doi:10.1007/978-0-387-76371-2
  * LoVerde & Afshordi (2008) LoVerde M., Afshordi N., 2008, Physical Review D, 78
  * Loureiro et al. (2023) Loureiro A., Whiteway L., Sellentin E., Silva Lafaurie J., Jaffe A. H., Heavens A. F., 2023, The Open Journal of Astrophysics, 6, 6
  * Mahony et al. (2022) Mahony C., Fortuna M. C., Joachimi B., Korn A., Hoekstra H., Schmidt S. J., Collaboration L. D. E. S., 2022, Monthly Notices of the Royal Astronomical Society, 513, 1210
  * Makinen et al. (2021) Makinen T. L., Charnock T., Alsing J., Wandelt B. D., 2021, Journal of Cosmology and Astroparticle Physics, 2021, 049
  * Makinen et al. (2022) Makinen T. L., Charnock T., Lemos P., Porqueres N., Heavens A. F., Wandelt B. D., 2022, The Open Journal of Astrophysics, 5
  * Mandelbaum et al. (2018) Mandelbaum R., et al., 2018, Monthly Notices of the Royal Astronomical Society, 481, 3170
  * Margossian (2019) Margossian C. C., 2019, WIREs Data Mining and Knowledge Discovery, 9, e1305
  * Mead et al. (2015) Mead A. J., Peacock J. A., Heymans C., Joudaki S., Heavens A. F., 2015, Monthly Notices of the Royal Astronomical Society, 454, 1958
  * Mead et al. (2016) Mead A. J., Heymans C., Lombriser L., Peacock J. A., Steele O. I., Winther H. A., 2016, Monthly Notices of the Royal Astronomical Society, 459, 1468
  * Mootoovaloo et al. (2020) Mootoovaloo A., Heavens A. F., Jaffe A. H., Leclercq F., 2020, Monthly Notices of the Royal Astronomical Society, 497, 2213
  * Neal (1996) Neal R. M., 1996, Bayesian Learning for Neural Networks, Vol. 118 of Lecture Notes in Statistics. Springer-Verlag
  * Neal (2011) Neal R., 2011, in , Handbook of Markov Chain Monte Carlo. Chapman and Hall/CRC, pp 113–162, doi:10.1201/b10905
  * Neidinger (2010) Neidinger R. D., 2010, SIAM Review, 52, 545
  * Nguyen et al. (2021) Nguyen N.-M., Schmidt F., Lavaux G., Jasche J., 2021, Journal of Cosmology and Astroparticle Physics, 2021, 058
  * Nygaard et al. (2022) Nygaard A., Brinch Holm E., Hannestad S., Tram T., 2022, arXiv e-prints, p. arXiv:2205.15726
  * Park (2021) Park J., 2021, arXiv e-prints, p. arXiv:2111.06871
  * Paszke et al. (2017) Paszke A., et al., 2017, in NIPS 2017 Workshop on Autodiff. https://openreview.net/forum?id=BJJsrmfCZ
  * Paszke et al. (2019) Paszke A., et al., 2019, in , Advances in Neural Information Processing Systems 32. Curran Associates, Inc., pp 8024–8035, http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * Phan et al. (2019) Phan D., Pradhan N., Jankowiak M., 2019, arXiv preprint arXiv:1912.11554
  * Porqueres et al. (2023) Porqueres N., Heavens A., Mortlock D., Lavaux G., Makinen T. L., 2023, arXiv e-prints, p. arXiv:2304.04785
  * Roberts et al. (1997) Roberts G. O., Gelman A., Gilks W. R., 1997, The Annals of Applied Probability, 7, 110
  * Robnik et al. (2022) Robnik J., De Luca G. B., Silverstein E., Seljak U., 2022, arXiv e-prints, p. arXiv:2212.08549
  * Ruiz-Zapatero et al. (2023) Ruiz-Zapatero J., Hadzhiyska B., Alonso D., Ferreira P. G., García-García C., Mootoovaloo A., 2023, arXiv e-prints, p. arXiv:2301.11978
  * Sanchez et al. (2021) Sanchez J., Nicola A., Alonso D., Chang C., DESC L., 2021, Bulletin of the AAS, 53
  * Sehgal et al. (2019) Sehgal N., et al., 2019, in Bulletin of the American Astronomical Society. p. 6 (arXiv:1906.10134), doi:10.48550/arXiv.1906.10134
  * Smail et al. (1994) Smail I., Ellis R. S., Fitchett M. J., 1994, Monthly Notices of the Royal Astronomical Society, 270, 245
  * Spergel et al. (2015) Spergel D., et al., 2015, arXiv e-prints, p. arXiv:1503.03757
  * Spurio Mancini et al. (2022) Spurio Mancini A., Piras D., Alsing J., Joachimi B., Hobson M. P., 2022, Monthly Notices of the Royal Astronomical Society, 511, 1771
  * Spurio Mancini & Pourtsidou (2022) Spurio Mancini A., Pourtsidou A., 2022, Monthly Notices of the Royal Astronomical Society: Letters, 512, L44
  * Taylor & Kitching (2018) Taylor A. N., Kitching T. D., 2018, Monthly Notices of the Royal Astronomical Society, 477, 3397
  * Taylor et al. (2008) Taylor J. F., Ashdown M. A. J., Hobson M. P., 2008, Monthly Notices of the Royal Astronomical Society, 389, 1284
  * Torrado & Lewis (2021) Torrado J., Lewis A., 2021, Journal of Cosmology and Astroparticle Physics, 2021, 057
  * Tutusaus et al. (2020) Tutusaus I., et al., 2020, A&A, 643, A70
  * Valade et al. (2022) Valade A., Hoffman Y., Libeskind N. I., Graziani R., 2022, Monthly Notices of the Royal Astronomical Society, 513, 5148
  * Ver Steeg & Galstyan (2021) Ver Steeg G., Galstyan A., 2021, in Advances in Neural Information Processing Systems.
  * Wong et al. (2023) Wong K., Gabrié M., Foreman-Mackey D., 2023, Journal of Open Source Software, 8, 5021
  * Zennaro et al. (2021) Zennaro M., Angulo R. E., Pellejero-Ibáñez M., Stücker J., Contreras S., Aricò G., 2021, arXiv e-prints, p. arXiv:2101.12187
  * Zhang et al. (2017) Zhang C., Bütepage J., Kjellström H., Mandt S., 2017, IEEE Transactions on Pattern Analysis and Machine Intelligence, PP