The Open Journal of Astrophysics

††thanks<EMAIL_ADDRESS>
#  JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology
Library

J. E. Campagne1,∗ F. Lanusse2 J. Zuntz3
A. Boucaud4 S. Casas5 M. Karamanis6,7 D. Kirkby8 D. Lanzieri9 Y. Li10,11 A.
Peel12 1Université Paris-Saclay, CNRS/IN2P3, IJCLab, 91405 Orsay, France
2Université Paris-Saclay, Université Paris Cité, CEA, CNRS, AIM, 91191, Gif-
sur-Yvette, France  3Institute for Astronomy, University of Edinburgh,
Edinburgh EH9 3HJ, United Kingdom  4Université de Paris, CNRS, Astroparticule
et Cosmologie, F-75013 Paris, France  5Institute for Theoretical Particle
Physics and Cosmology (TTK), RWTH Aachen University, 52056 Aachen, Germany.
6Berkeley Center for Cosmological Physics, University of California, Berkeley,
CA 94720, USA  7Lawrence Berkeley National Laboratory, 1 Cyclotron Road,
Berkeley, CA 94720, USA  8Department of Physics and Astronomy, University of
California, Irvine, CA 92697, USA  9Université Paris Cité, Université Paris-
Saclay, CEA, CNRS, AIM, F-91191, Gif-sur-Yvette, France  10Department of
Mathematics and Theory, Peng Cheng Laboratory, Shenzhen, Guangdong 518066,
China  11Center for Computational Astrophysics & Center for Computational
Mathematics, Flatiron Institute, New York, New York 10010, USA  12Ecole
Polytechnique Fédérale de Lausanne (EPFL), Observatoire de Sauverny, 1290
Versoix, Switzerland

###### Abstract

We present jax-cosmo, a library for automatically differentiable cosmological
theory calculations. jax-cosmo uses the JAX library, which has created a new
coding ecosystem, especially in probabilistic programming. As well as batch
acceleration, just-in-time compilation, and automatic optimization of code for
different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic
differentiation (autodiff) mechanism. Thanks to autodiff, jax-cosmo gives
access to the derivatives of cosmological likelihoods with respect to any of
their parameters, and thus enables a range of powerful Bayesian inference
algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte
Carlo and Variational Inference. In its initial release, jax-cosmo implements
background evolution, linear and non-linear power spectra (using halofit or
the Eisenstein and Hu transfer function), as well as angular power spectra
($C_{\ell}$) with the Limber approximation for galaxy and weak lensing probes,
all differentiable with respect to the cosmological parameters and their other
inputs. We illustrate how automatic differentiation can be a game-changer for
common tasks involving Fisher matrix computations, or full posterior inference
with gradient-based techniques (e.g. Hamiltonian Monte Carlo). In particular,
we show how Fisher matrices are now fast, exact, no longer require any fine
tuning, and are themselves differentiable with respect to parameters of the
likelihood, enabling complex survey optimization by simple gradient descent.
Finally, using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we
demonstrate how jax-cosmo can be combined with Probabilistic Programming
Languages such as NumPyro to perform posterior inference with state-of-the-art
algorithms including a No U-Turn Sampler (NUTS), Automatic Differentiation
Variational Inference (ADVI), and Neural Transport HMC (NeuTra). We show that
the effective sample size per node (1 GPU or 32 CPUs) per hour of wall time is
about 5 times better for a JAX NUTS sampler compared to the well optimized
Cobaya Metropolis-Hasting sampler. We further demonstrate that Normalizing
Flows using Neural Transport are a promising methodology for model validation
in the early stages of analysis.

##  1 Introduction

Bayesian inference has been widely used in cosmology in the form of Monte
Carlo Markov Chains (MCMC) since the work of Knox et al. (2001) and Rubiño-
Martin et al. (2003), and has been the keystone for past and current analysis
thanks partly to packages such as CosmoMC (Lewis & Bridle, 2002), CosmoSIS
(Zuntz et al., 2015), MontePython (Brinckmann & Lesgourgues, 2019), and Cobaya
(Torrado & Lewis, 2019, 2021); see, for instance, the list of citations to
these popular packages for an idea of the wide usage in the community.

Since the development of these MCMC packages, major advances have been made in
automatic differentiation (autodiff) (Baydin et al., 2017; Margossian, 2019),
a set of technologies for transforming pieces of code into their derivatives.

While these tools have especially been applied to neural network optimization
and machine learning (ML) in general, they can also enable classical
statistical methods that require the derivatives of (e.g. likelihood)
functions to operate: we consider such methods in this paper. Autodiff has
been implemented in widely used libraries like Stan (Carpenter et al., 2017),
TensorFlow (Abadi et al., 2015), Julia (Bezanson et al., 2017), and PyTorch
(Paszke et al., 2019).

A recent entrant to this field is the JAX library111https://jax.readthedocs.io
(Bradbury et al., 2018) which has undergone rapid development and can
automatically differentiate native Python and NumPy functions, offering a
speed up to the development process and indeed code runtimes. JAX offers an
easy parallelization mechanism (vmap), just-in-time compilation (jit), and
optimization targeting CPU, GPU, and TPU hardware thanks to the XLA library
(which also powers TensorFlow). These attractive features have driven wide
adoption of JAX in computational research, and motivate us to consider its
usage in cosmology.

JAX contains bespoke reimplementations of packages such as jax.numpy and
jax.scipy, as well as example libraries such as Stax for simple but flexible
neural network development. It has formed the seed for a wider ecosystem of
packages, including, for example: Flax (Heek et al., 2020) a high-performance
neural network library, JAXopt (Blondel et al., 2021) a hardware accelerated,
batchable and differentiable collection of optimizers, Optax (Hessel et al.,
2020) a gradient processing and optimization library, and NumPyro (Phan et
al., 2019; Bingham et al., 2019), a probabilistic programming language (PPL)
that is used in this paper. Other PPL packages such as PyMC (Salvatier et al.,
2016) have switched to a JAX backend in recent versions222A more exhaustive
list and rapidly growing list of packages can be found at https://project-
awesome.org/n2cholas/awesome-jax.

To explore alternative inference methods to the well-known Metropolis-Hasting
likelihood sampler, and in order to use GPU devices in the context of JAX
framework, we have developed the open source jax-cosmo
library333https://github.com/DifferentiableUniverseInitiative/jax_cosmo, which
we present in this paper. The package represents a first step in making the
powerful features described above useful for cosmology; it implements a
selection of theory predictions for key cosmology observables as
differentiable JAX functions.

We give an overview of the code’s design and contents in Section 3. We show
how to use it for rapid and numerically stable Fisher forecasts in Section 4,
in more complete gradient-based cosmological inference with variants of
Hamiltonian Monte Carlo including the No-U-Turn Sampler, and ML-accelerated
Stochastic Variational Inference in Section 5. We discuss and compare these
methods in Section 6 and conclude in Section 7.

##  2 JAX: GPU Accelerated and Automatically Differentiable Python Programming

The aim of this section is to provide a brief technical primer on JAX,
necessary to fully grasp the potential of a cosmology library implemented in
this framework.

##### Automatic Differentiation

Traditionally, two different approaches have been used in cosmology to obtain
derivatives of given numerical expressions. The first approach is to derive
analytically the formula for the derivatives of interest (e.g. Jasche &
Wandelt, 2013; Kostić et al., 2022), with or without the help of tools such as
Mathematica444https://www.wolfram.com/mathematica.. This is typically only
practical, however, for simple analytical models. The second approach is to
compute numerical derivatives by finite differences. This approach can be
applied on black-box models of arbitrary complexity (from typical Boltzmann
codes to cosmological simulations; Villaescusa-Navarro et al., 2020). However
it is notoriously difficult to obtain stable derivatives by finite differences
(e.g. Bhandari et al., 2021; Yahia-Cherif et al., 2021). In addition, their
computational cost does not scale well with the number of parameters (a
minimum of $(2N+1)$ model evaluations is typically required for $N$
parameters), making them impractical whenever derivatives are needed as part
of an outer iterative algorithm.

Automatic differentiation frameworks like JAX take a different approach. They
trace the execution of a given model and decompose this trace into primitive
operations with known derivatives (e.g. multiplication). Then, by applying the
chain rule formula, the computational graph for the derivatives (of any order)
of the model can be built from the known derivatives of every elementary
operation. A new function corresponding to the derivative of the original
function is therefore built automatically for the user. We direct the
interested reader to Baydin et al. (2018) and Margossian (2019) for in-depth
introductions to automatic differentiation.

JAX provides in particular a number of operators (jax.grad, jax.jacobian,
jax.hessian) which can compute derivatives of any function written in JAX:

⬇

# Define a simple function

def f(x):

return y = 5 ∗ x + 2

# Take the derivative

df_dx = jax.grad(f)

# df_dx is a new function that always returns 5

Why is this interesting? autodiff makes it possible to obtain exact gradients
of cosmological likelihoods with respect to all input parameters at the cost
of only two likelihood evaluations.

##### Just In Time Compilation (JIT)

Despite its convenience and wide adoption in astrophysics, Python still
suffers from slow execution times compared to fully compiled languages such as
C/C++. One approach to mitigate these issues and make Python code fast is Just
In Time compilation, which traces the execution of a given Python function the
first time it is called, and compiles it into a fast executable (by-passing
the Python interpreter), which can be transparently used in subsequent calls
to this function. Compared to other strategies for speeding up Python code
such as Cython, JIT allows the user to write plain Python code, and reap the
benefits of compiled code.

A number of libraries allowing for JIT have already been used in astrophysics,
in particular Numba555https://numba.pydata.org/, or the HOPE library Akeret et
al. (2015) developed specifically for the needs of astrophysics. JAX stands
out compared to these other frameworks in that it relies on the highly
optimized XLA library666https://www.tensorflow.org/xla for executing the
compiled expressions. XLA is continuously developed by Google as part of the
TensorFlow project, for efficient training and inference of large scale deep
learning applications, and as such supports computational backends such as GPU
and Tensor Processing Units (TPU) clusters. The ability to perform
computations directly on GPUs through XLA is one of the major benefits of JAX,
as speed-ups of at least two orders of magnitudes can be expected for typical
parallel linear algebra computations compared to CPU.

In JAX, jitting is achieved by transforming a function with the jax.jit
operation:

⬇

# Redefine our function

def f(x):

return y = 5 ∗ x + 2

# And JIT it

jitted_f = jax.jit(f)

# The first call to jitted_f will be relatively slow

# subsequent calls will be extremely fast

# and run as a compiled code directly on GPU

Why is this interesting? JIT makes it possible to execute entire cosmological
MCMC chains directly on GPUs as compiled code, with orders of magnitude gain
in speedup over Python code.

##### Automatic Vectorization

Another extremely powerful feature of JAX is its ability to automatically
vectorize or paralellize any function. Through the same tracing mechanism used
for automatic differentiation, JAX can decompose a given computation into
primitive operations and add a new batch dimension so that the computation can
be applied to a batch of inputs as opposed to individual ones. In doing so,
the computation will not run sequentially over all entries of the batch, but
truly in parallel making full use of the intrinsic parallel architecture of
modern GPUs.

In JAX automatic vectorization is achieved using the jax.vmap operation:

⬇

# Our function f only applies to scalars

def f(x):

return y = 5 ∗ x + 2

# Applying automatic vectorization

batched_f = jax.vmap(f)

# batched_f now applies to 1D arrays

Again, we stress that in this example, batched_f will not be implemented in
terms of a for loop, but with operations over vectors. The function above is
trivial, but the same approach can be used to parallelize any function, from
Limber integrals, to an entire likelihood. For multi-device use-cases (e.g.,
several GPUs or TPUs), JAX provides pmap which compiles and executes, in
parallel, replicas of the same code on each device. Moreover, recent
experimental developments deal with parallelization of functions over
supercomputer-scale hardware meshes. In the examples detailed in this article,
we have only relied on vmap functionality.

Why is this interesting? Automatic Vectorization makes it possible to
trivially parallelize cosmological likelihood evaluations, to run many
parallel MCMC chains on a single GPU.

##### NumPy API compliance

Finally, the last point to note about JAX, is that it mirrors the NumPy API
(with only a few exceptions). This means in practice that existing NumPy code
can easily be converted to JAX. This is in contrast to other similar
frameworks like TensorFlow, PyTorch, or Julia which all require the user to
learn, and adapt their code to, a new API or even a new language.

Why is this interesting? NumPy compliance implies improved maintainability and
lower barrier to entry for new contributors.

##  3 Capabilities of the jax-cosmo library

In this section, we describe the cosmological modeling provided by jax-cosmo,
and its implementation in JAX. The general design follows that of the Core
Cosmology Library (CCL; Chisari et al., 2019), though in its initial release
jax-cosmo only implements a subset of CCL features and options.

All jax-cosmo data structures are organized as JAX container objects, which
means that two key JAX features are automatically available to them: autodiff
and vmap. The vmap feature enables any operation defined in JAX (including
complicated composite operations) to be applied efficiently as a vector
operation. The autodiff feature further makes it possible to take the
derivative of any operation, automatically transforming a function that takes
$n$ inputs and produces $m$ outputs into a new function that generates an
$m\times n$ matrix of partial derivatives.

jax-cosmo implements, for example, Runge-Kutta solvers for ODEs, as well as
Simpson and Romberg integration routines through this framework, so that we
can automatically compute the derivatives of their solutions with respect to
their inputs. This includes not only the cosmological parameters (the standard
set of $w_{0}w_{a}CDM$ cosmological parameters is exposed, using $\sigma_{8}$
as an amplitude parameter), but also other input quantities such as redshift
distributions as described below.

In the rest of this section we describe the cosmological calculations that are
implemented using these facilities.

###  3.1 Formalism & Implementation

####  3.1.1 Background cosmology

The computation of the evolution of the cosmological background follows a
typical implementation of a Friedmann equation (Percival, 2005, see e.g.):

$E^{2}(a)=\frac{\mathrm{H}^{2}(a)}{\mathrm{H}^{2}_{0}}=\Omega_{m}a^{-3}+\Omega_{k}a^{-2}+\Omega_{de}e^{f(a)}$
(1)

with $a=1/(1+z)$ the scale factor related to the redshift $z$,
$\mathrm{H}(a)=\dot{a}/a$ the Hubble parameter with $\mathrm{H}_{0}$ its
present day value, $\Omega_{m}=\Omega_{cdm}+\Omega_{b}$,
$\Omega_{de}=1-\Omega_{k}-\Omega_{m}$, and

$f(a)=-3(1+w_{0}+w_{a})\ln(a)+3w_{a}(a-1)$ (2)

Notably, the relativistic contributions of massless neutrinos and photon
radiation, as well as the massive neutrino contribution are currently
neglected. From these expressions, in jc.background, are computed the
different cosmological distance functions, such as the radial comoving
distance

$\chi(a)=R_{H}\int_{a}^{1}\frac{\mathrm{d}a^{\prime}}{{a^{\prime}}^{2}E(a^{\prime})}$
(3)

with $R_{H}$ the Hubble radius. $\chi(a)$ evaluation is performed with a
linear interpolation of pre-computed values obtained over a logarithmic grid
of 256 knots regularly spaced in the range [-3,0] using a Runge-Kutta 4th
order solver.

####  3.1.2 Growth of perturbations

Accurate calculations of the growth of cosmic structure perturbations require
solving the Boltzmann equation describing interactions between different
density species. This is a complicated process to achieve at high accuracy,
and is implemented in dedicated codes such as CAMB (Lewis et al., 2000) or
CLASS (Blas et al., 2011). At this stage we have not attempted to implement a
Boltzmann solver in JAX, but instead use the analytic approximation presented
in Eisenstein & Hu (1998). This defines a transfer function $T$ which relates
the primordial matter power spectrum to its late-time non-linear value:

$P(k,z)=P(k,z=\infty)\cdot T^{2}(k,z;\Omega_{m},\Omega_{b},...),$ (4)

through the halofit model by Takahashi et al. (2012) or Smith et al. (2003)
without the neutrino contribution introduced by Bird et al. (2012). No Baryon
feedback modeling is considered yet.

The primordial power spectrum is modelled with the standard form:

$P(k)=Ak^{n_{s}}.$ (5)

The normalisation $A$ is parameterised via $\sigma_{8}$ at $z=0$ as

$A=\sigma_{8}^{2}\times\left(\frac{1}{2\pi^{2}}\int_{0}^{\infty}\frac{\mathrm{d}k}{k}k^{3}P(k)W^{2}(kR_{8})\right)^{-1}$
(6)

with $R_{8}=8\mathrm{Mpc}/h$ and $W(x)$ related to the $j_{1}$ spherical
Bessel function as

$W(x)=\frac{3j_{1}(x)}{x}$ (7)

####  3.1.3 Angular power spectra

The bias and number density functions that enter the kernel terms have been
modelled in many different ways to try to capture possible behavior of these
key systematic errors. Exposing them in JAX-Cosmo allows us easily to deal
with highly complicated models for either of them, opening up many useful
avenues of exploration, as shown in Ruiz-Zapatero et al. (2023).

jax-cosmo is currently focused on predicting projected 2D Fourier-space 2pt
galaxy lensing, clustering, and cross correlations, the $C_{\ell}$ angular
power spectra that are a primary target of upcoming photometric surveys. The
details of the implementation is in jc.angular_cl which deals with the mean
and Gaussian covariance matrix computations.

The angular power spectra $C_{\ell}^{ij}$ for the probes $(i,j)$ and for
redshift bin window selections are computed in the first order Limber
approximation (LoVerde & Afshordi, 2008):

$\displaystyle
C_{\ell}^{i,j}\approx\left(\ell+\frac{1}{2}\right)^{m_{i}+m_{j}}\int_{a_{i,j}}^{1}\frac{\mathrm{d}a}{c^{2}\chi^{2}}\frac{\mathrm{d}\chi}{\mathrm{d}a}K_{i}(\chi)K_{j}(\chi)\,P\left(k=\frac{\ell+1/2}{\chi},z\right),$
(8)

The $m_{i}$ factors are $(0,-2)$ for the galaxy clustering and weak lensing,
respectively, and each $K(z)$ function represents a single tomographic
redshift bin’s number density in $z=(1/a)-1$. The integration lower bound
$a_{ij}$ is defined from the maximum redshift involved and Simpson’s rule is
used with a regular grid of 512 knots. The $K(z)$ tracers are implemented as
two kernel functions:

NumberCounts

$K_{i}(z)=n_{i}(z)\ b(z)\ H(z)$ (9)

where $n_{i}(s)$ is the redshift distribution of the sample (e.g.,
jc.redshift.kde_nz function), and $b(z)$ is the galaxy bias function (see
jc.bias.constant_linear_bias). No redshift space distortions are taken into
account.

WeakLensing

$K_{i}(z)=\left(\frac{3H_{0}^{2}\Omega_{m}}{2c}\right)\left(\frac{(\ell+2)!}{(\ell-2)!}\right)^{1/2}\times\\\
(1+z)\ \chi(z)\int_{z}^{\infty}p_{i}(z^{\prime})\
\frac{\chi(z^{\prime})-\chi(z)}{\chi(z^{\prime})}\
\mathrm{d}z^{\prime}+K_{IA}(z)$ (10)

where $K_{IA}(z)$ is an optional kernel function to deal with the Intrinsic
Alignment. The implementation of this term currently follows Joachimi et al.
(2011), and is given by:

$K_{IA}(z)=\left(\frac{(\ell+2)!}{(\ell-2)!}\right)^{1/2}\ p_{i}(z)\ b(z)\
H(z)\ \frac{C\ \Omega_{m}}{D(z)}$ (11)

with $C\approx 0.0134$ being a dimensionless constant and $D(z)$ the growth
factor of linear perturbations. The integration is performed using Simpson’s
rule with an upper bound given by the maximum redshift involved.

Because, like the other ingredients, the NumberCounts and WeakLensing kernels
are implemented as JAX objects, all the integrals involved in these
computations can be differentiated with respect to the cosmological parameters
and the number densities, using autodiff and jax-cosmo’s implementation of
integration quadrature. An example is given in the context of DES Y1 3x2pts
analysis (Sec 5.1).

###  3.2 Validation against the Core Cosmology Library (CCL)

To illustrate the different features available with the present version of the
library (jax-cosmo 0.1, which is available in the Python Package Index
PyPI777https://pypi.org/), we have written a companion notebook  to compare it
to the well-validated Core Cosmology Library (Chisari et al.,
2019)888https://ccl.readthedocs.io, version 2.5.1.. As examples, Figures 1, 2
and 3 show the radial comoving distance (Eq. 3), the non-linear matter power
spectrum computation, and the angular power spectrum for galaxy-galaxy lensing
(Eq. 8) using the NumberCounts and WeakLensing kernel functions. jax-cosmo
features a suite of validation tests against CCL, automatically validating the
precision of all computations to within the desired numerical accuracy; the
relative differences between the two libraries are at the level of few
$10^{-3}$ or better.

These numerical differences are mostly due to different choices of integration
methods and accuracy parameters (e.g. number of quadrature points). Increasing
these parameters leads to performance degradation for jax-cosmo but increases
the XLA compilation memory requirements significantly, especially for the
angular power spectra computation. Since these differences are likely to be
within the tolerance of the current generation of cosmological surveys, this
trade-off is an acceptable one.

Despite this high precision with respect to the CCL library, the current jax-
cosmo implementation will have to be reviewed to tackle Stage IV accuracy
requirements. For instance the non-limber $C_{\ell}$ computation (Leonard et
al., 2023; Campagne et al., 2017) would be a limitation needing revision in a
future release.

Figure 1: Comparison of the radial comoving distance between CCL and jax-
cosmo. More plots are available in the companion notebook . Figure 2:
Comparison of the non-linear matter power spectrum (halofit function) between
CCL and jax-cosmo. Also shown is the linear power spectrum. Figure 3:
Comparison of the shape-position angular power spectrum between CCL and jax-
cosmo.

##  4 Fisher Information Matrices Made Easy

As a first illustration of the value of autodiff in cosmological computations,
we present in this section a few direct applications involving the computation
of the Fisher information matrix, widely used in cosmology (Tegmark et al.,
1997; Stuart & Ord, 1991).
Not only does the computation of the Fisher matrix become trivial, but the
Fisher matrix itself becomes differentiable, allowing in particular for
powerful survey optimization applications.

###  4.1 Instantaneous Fisher Forecasts

Fisher matrices are a key tool in cosmology forecasting and experimental
planning. By computing the Hessian matrix of a likelihood with respect to its
parameters, we can find a Gaussian approximation to a posterior, which is
usually sufficient for comparing the constraining power of different
experimental configurations. As noted above, computing Fisher matrices is
notoriously error-prone, since finite difference approximations to likelihoods
must be carefully tuned for convergence, and observable calculations can be
numerically unstable. Autodiff can help avoid this challenge.

We first illustrate, with an artificial case study, the computation of a
Fisher matrix using two methods with the autodiff ability of JAX. For the
detailed implementation, the reader is invited to look at the following
companion notebook . In this example we use four tracer redshift
distributions: two to define WeakLensing kernels and two for NumberCounts
kernels. Then, the $10$ angular power spectra $C_{\ell}^{p,q}$
($p,q:1,\dots,4$) are computed for $50$ logarithmically-spaced angular moments
between $\ell=10$ and $\ell=1000$ using Equation 8. The Gaussian covariance
matrix is computed simultaneously. A dataset is obtained from the computation
of the $C_{\ell}^{p,q}$ with a fiducial cosmology. Then, the following snippet
shows the log likelihood function $\mathcal{L}(\theta)$ implementation
considering a constant covariance matrix ($\theta$ stands for the set of
cosmological parameters).

⬇

@jax.jit

def likelihood(p):

# Create a new cosmology at these parameters

cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])

# Compute mean and covariance of angular Cls

mu, C = jc.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, tracers,
sparse=True)

# Return likelihood value assuming constant covariance, so we stop the
gradient

# at the level of the precision matrix, and we will not include the logdet
term

# in the likelihood

P = jc.sparse.inv(jax.lax.stop_gradient(C))

r = data - mu

return -0.5 ∗ r.T @ jc.sparse.sparse_dot_vec(P, r)

The jc.sparse functions are implementations of block matrix computations: a
sparse matrix is represented as a 3D array of shape $(n_{y},n_{x},n_{diag})$
composed of $n_{y}\times n_{x}$ square blocks of size $n_{diag}\times
n_{diag}$. The jax.jit decorator builds a compiled version of the function on
first use.

The first approach to obtaining approximate 1-sigma contours of the two
parameters ($\Omega_{c},\sigma_{8}$) with a Fisher matrix uses the Hessian of
the log-likelihood as follows:

$F_{i,j}=-\frac{\partial^{2}\mathcal{L}(\theta)}{\partial\theta_{i}\partial\theta_{j}}\qquad(\theta_{1}=\Omega_{c},\theta_{2}=\sigma_{8})$
(12)

which is accomplished in two lines of JAX code:

⬇

hessian_loglike = jax.jit(jax.hessian(likelihood))

F = - hessian_loglike(params)

The second approach to computing the Fisher matrix, which is restricted to
Gaussian likelihoods but is more commonly used in the field because of its
numerical stability, is to define a function that computes the summary
statistic mean $\mu(\ell;\theta)$; the Fisher matrix elements are then:

$F_{i,j}=\sum_{\ell}\frac{\partial\mu(\ell)}{\partial\theta_{i}}^{T}C^{-1}(\ell)\frac{\partial\mu(\ell)}{\partial\theta_{j}}$
(13)

where $C^{-1}(\ell)$ is the covariance matrix computed with the fiducial
cosmology. This can be computed in JAX as:

⬇

# We define a parameter dependent mean function

@jax.jit

def jc_mean(p):

cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])

# Compute signal vector

mu = jc.angular_cl.angular_cl(cosmo, ell, tracers)

# We want mu in 1d to match the covariance matrix

return mu.flatten()

# We compute its jacobian with JAX, and we JIT it for efficiency

jac_mean = jax.jit(jax.jacfwd(jc_mean))

# We can now evaluate the jacobian at the fiducial cosmology

dmu = jac_mean(params)

# Now we can compose the Fisher matrix

F = jc.sparse.dot(dmu.T, jc.sparse.inv(cov), dmu)

JAX implementations of the two methods agree to near perfect accuracy, as
shown in Figure 4.

Figure 4: Comparison of the two methods to compute the Fisher matrix: Eq. 12
(method 1) and Eq. 13 (method 2).

It is worth noting that in the two methods for Fisher matrix computations
described above, the user does not need to vary individual parameter values to
compute the 1st or 2nd order derivatives; this is in contrast to the usual
finite difference methods. As an illustration, we have used CCL to compute the
Fisher matrix via Equation 13. To do so, the Jacobian
($\partial\mu/\partial\theta_{\alpha}$) is computed with centered order 4
finite differences available in the Numdifftools Python
package999https://numdifftools.readthedocs.io. Using the parameter values
spaced by ($10^{-6}$, $10^{-2}$, $10^{-1}$) one obtains three different
approximation of the 1-sigma contours as shown on Figure 5. The contour that
agrees best with the jax-cosmo method is obtained with the intermediate
spacing parameter $10^{-2}$, implying that the user must tune this parameter
carefully. Although very simple, this case study demonstrates the significant
challenge of using finite difference methods for computing the Fisher matrix,
as has been shown for instance in a more advanced case study in Bhandari et
al. (2021).

Figure 5: Fisher matrices (Eq.12) estimated using the CCL for the angular
power spectra estimation and finite difference method to get the Jacobian with
different spacing of parameters ($10^{-6}$, $10^{-2}$, $10^{-1}$). For
comparison, the contour obtained with jax-cosmo is reproduced from Figure 4 in
black.

###  4.2 Survey Optimization by FoM Maximization

Fisher forecasts are also commonly used in survey and analysis strategy, where
running a full posterior analysis for each possible choice would be
unfeasible. The inverse area of a projection of a Fisher matrix in parameters
of interest can be used as a metric for survey constraining power, such as in
the Dark Energy Task Force report (Albrecht et al., 2006a).

jax-cosmo was used in a recent example of such a process, for the the LSST-
DESC 3x2pt tomography optimization challenge (Zuntz et al., 2021), where the
best methodology for assignment of galaxies to tomographic bins was assessed
using several such figures of merit and related calculations. The jax-cosmo
metric proved to be stable and fast.

Because JAX functions are differentiable with respect to all their inputs,
including survey configuration parameters (e.g. depth, area, etc), we can even
compute the derivative of an FoM with respect to these inputs, allowing for
rapid and complete survey optimization.

###  4.3 Massive Optimal Compression in 3 Lines

Once the Fisher matrix has been accurately estimated, the MOPED101010Massively
Optimised Parameter Estimation and Data compression algorithm can be used to
compress data sets with minimal information loss (Heavens et al., 2000;
Zablocki & Dodelson, 2016; Heavens et al., 2017). In the case of the constant
covariance matrix the algorithm compresses data in a way that is lossless at
the Fisher matrix level (i.e. Fisher matrices estimated using the compressed
and full data are identical, by construction) which reduces a possibly large
data vector $\mu$ of size $N$ to $M$ numbers, where $M$ is the number of
parameters $\theta_{i}$ considered. For instance, in the previous section,
$N=500$ as $\mu=(C_{\ell}^{p,q})$ and $M=2$ for $(\Omega_{c},\sigma_{8})$.

The algorithm computes by iteration $M$ vectors of size $N$ such that (taking
the notation of the previous section)

$b_{i}=\frac{C^{-1}\mu_{,i}-\sum_{j=1}^{i-1}(\mu_{,i}^{T}b_{j})b_{j}}{\sqrt{F_{i,i}-\sum_{j=1}^{i-1}(\mu_{,i}^{T}b_{j})^{2}}}$
(14)

where $\mu_{,i}=\partial\mu/\partial\theta_{i}$ ($i=1,\dots,M$). The vectors
$(b_{i})_{i\leq M}$ satisfy the following orthogonality constraint

$b_{i}^{T}Cb_{j}=\delta_{i,j}$ (15)

The algorithm is similar to the Gram-Schmidt process, using the constant
covariance matrix $C$ to define the scalar product $\langle
b_{i},b_{j}\rangle$. Then, the original data set $x=C_{\ell}^{p,q}$ is
compressed in a data set composed of $M$ numbers $y_{i}$ according to

$y_{i}=b_{i}^{T}x$ (16)

These numbers are uncorrelated and of unit variance and this construction
ensures that the log-likelihood of $y_{i}$ given $\theta_{i}$ is identical to
that of $x$ up to second order, meaning that the Fisher matrices derived from
the two parameters should be identical, and in general the $y$ values should
lose very little information compared to the full likelihood.

In problems where a (constant) covariance matrix is estimated from
simulations, the number of such simulations required for a given accuracy
typically scales with the number of data points used. MOPED therefore greatly
reduces this number, often by a factor of hundreds. Since the uncertainty in
the covariance due to a finite number of simulations must be accounted for
(Sellentin & Heavens, 2018; Hartlap et al., 2007), this reduction can also
offset any loss of information from the compression. Inaccuracies in the full
covariance matrix used in the data compression result only in information
loss, not bias, as does mis-specification of the fiducial parameter set
$\theta_{i}$ in equation 14.

Another key advantage of the MOPED algorithm is to eliminate the need for
large covariance matrix inversion of size $N\times N$ requiring $O(N^{3})$
operations. This inversion takes place not only for the Fisher matrix
computation (Eq. 13), but more importantly in the log-likelihood computation
(see the snippet in the previous section). The MOPED algorithm reduces the
complexity to $O(M^{3})$ operations.

To give an illustration, the following snippet uses the mock data set and
results on the Fisher matrix computation from the previous section, to obtain
the MOPED compressed data set composed of two parameters $(y_{0},y_{1})$ with
maximal information on $(\Omega_{c},\sigma_{8})$:

⬇

# Orthogonal vectors

b0 = jc.sparse.dot(C_inv,dmu[:,0])/jax.numpy.sqrt(F[0,0])

a = dmu[:,1].T @ b0

b1 = (jc.sparse.dot(C_inv,dmu[:,1]) - a ∗ b0)/jax.numpy.sqrt(F[1,1]-a∗a)

# MOPED vectors

y0 = b0.T @ data

y1 = b1.T @ data

Then, the log-likelihood can be implemented as:

⬇

@jax.jit

def compressed_likelihood(p):

# Create a new cosmology at these parameters

cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])

# Compute mean Cl

mu = jc.angular_cl.angular_cl(cosmo, ell, tracers).flatten()

# likelihood using the MOPED vector

return -0.5 ∗ ((y0 - b0.T @ mu)∗∗2 + (y1 - b1.T @ mu)∗∗2)

The comparison between contour lines obtained with the original likelihood
(uncompressed data set) and the MOPED version are shown in Figure 6 for the
case study of the previous section. Close to the negative likelihood minimum,
the lines agree very well.

Figure 6: Illustration of the log-likelihood contours obtained with an
uncompressed data set (plain lines) and a MOPED version (dashed lines).

The case where the covariance matrix depends upon the parameters is discussed
by Heavens et al. (2017) and leads to similar advantages of the MOPED
compression algorithm.

##  5 Posterior Inference made fast by Gradient-Based Inference Methods

In the following sections we review more statistical methods which directly
benefit from the automatic differentiablity of jax-cosmo likelihoods. We
demonstrate a range of gradient-based methods, from Hamiltonian Monte Carlo
(HMC), and its No-U-Turn Sampler (NUTS) variant, to Stochastic Variational
Inference (SVI). We further report their respective computational costs in a
DES-Y1 like analysis. All the methods have been implemented using the NumPyro
probabilistic programming language (PPL).

###  5.1 Description of the DES-Y1 exercise

From the DES Year 1 lensing and clustering data release111111http://desdr-
server.ncsa.illinois.edu/despublic/y1a1_files/chains/2pt_NG_mcal_1110.fit we
have extracted the $N(z)$ distributions of the four source and five lens
samples. We normalize the sources to $[1.47,1.46,1.50,0.73]$ effective number
of sources per $\mathrm{arcmin}^{2}$. These distributions are modelled in jax-
cosmo using a kernel density estimation in the jc.redshift.kde_nz function and
are presented in Figure 7.

Figure 7: Distributions of the sources and lenses for the different redshift bins considered. Table 1: Priors on the 21 variables of the DES-Y1 of the 3x2pt likelihood (number counts and shear). parameter | prior
---|---
Cosmology
$\Omega_{c}$ | $\mathcal{U}[0.10,0.9]$
$\sigma_{8}$ | $\mathcal{U}[0.40,1.0]$
$\Omega_{b}$ | $\mathcal{U}[0.03,0.07]$
$w_{0}$ | $\mathcal{U}[-2.00,-0.33]$
$h$ | $\mathcal{U}[0.55,0.91]$
$n_{s}$ | $\mathcal{U}[0.87,1.07]$
Intrinsic Alignment
$A$ | $\mathcal{U}[-5,5]$
$\eta$ | $\mathcal{U}[-5,5]$
Lens Galaxy Bias
$(b_{i})_{i:1,\dots,5}$ | $\mathcal{U}[0.8,3.0]$
Shear Calibration Systematics
$(m_{i})_{i:1,\dots,4}$ | $\mathcal{N}[0.012,0.023]$
Source photo-$z$ shift
$dz_{1}$ | $\mathcal{N}[0.001,0.016]$
$dz_{2}$ | $\mathcal{N}[-0.019,0.013]$
$dz_{3}$ | $\mathcal{N}[0.009,0.011]$
$dz_{4}$ | $\mathcal{N}[-0.018,0.022]$

Using the NumPyro PPL, we then set up a forward model following the DES Y1
Pipeline (Abbott et al., 2018). Here we show some key elements of the
implementation; the details and inference examples may be found in the
following notebook . The model parameters and their prior distributions are
shown in Table 1. For instance the $\Omega_{c}$ parameter is treated as random
variable as follows

⬇

Omega_c = numpyro.sample("Omega_c", Uniform(0.1, 0.9))

Then, we generate mock angular power spectra $C_{\ell}$ from the auto & cross
correlations of the Number Counts and Weak Lensing probes computed with the
gaussian_cl_covariance_and_mean function in the jc.angular_cl module. The code
reads:

⬇

# Define the lensing and number counts probe

probes = [jc.probes.WeakLensing(nzs_s_sys,

ia_bias=b_ia,

multiplicative_bias=m),

jc.probes.NumberCounts(nzs_l, b)]

cl, C = gaussian_cl_covariance_and_mean(cosmo,

ell, probes,

f_sky=0.25, sparse=True)

with cosmo an instance of the jc.Cosmology setting the cosmological parameters
generated with the priors, and ell (i.e. $\ell$) a series of 50 angular modes.
After encapsulating the code above with input sampled parameters (using
NumPyro distribution classes) in a function model, we can generate our mock
data (cl_obs). This comes from this model function evaluated at a fiducial
cosmology, with random noise generated by NumPyro:

⬇

fiducial_model = numpyro.condition(model,

{"Omega_c":0.2545, "sigma8":0.801,

"h":0.682, "Omega_b":0.0485, "w0":-1.,"n_s":0.971,

"A":0.5,"eta":0.,

"m1":0.0,"m2":0.0,"m3":0.0,"m4":0.0,

"dz1":0.0,"dz2":0.0,"dz3":0.0,"dz4":0.0,

"b1":1.2,"b2":1.4,"b3":1.6,"b4":1.8,"b5":2.0

})

with seed(rng_seed=42):

cl_obs, P, C = fiducial_model()

These theoretical $C_{\ell}$ are in turn conditioned on the mock $C_{\ell}$:

⬇

observed_model = numpyro.condition(model, {"cl": cl_obs})

How to use this model to perform inference is described in the sections 5.3
and 5.4.

###  5.2 Vanilla Hamiltonian Monte Carlo

Hamiltonian Monte Carlo (HMC) is an MCMC-type method particularly suited to
drawing samples from high dimensional parameter spaces. It was introduced in
Duane et al. (1987) and developed extensively since. See Betancourt (2017) for
a full review; we describe very basic features here.

HMC samples a space by generating particle trajectories through it, using the
log-posterior as the negative potential energy of a particle at each point $q$
in the space. Associated with $q$, we introduce an auxiliary $p$ variable as
Hamiltonian momentum such that

$-\log{\cal P}(q)=V(q)\quad H(q,p)=V(q)+U(p)$ (17)

where $U(p)$ is a kinetic energy-like term defined by

$U(p)=p^{T}M^{-1}p$ (18)

where $M$ is a mass matrix which should be set to approximate the covariance
of the posterior. At each sample, a trajectory is initialized with a random
momentum $p$, and then Hamilton’s equations are integrated:

$\displaystyle\frac{\mathrm{d}p}{\mathrm{d}t}$ $\displaystyle=-\frac{\partial
V}{\mathrm{d}q}=\frac{\partial\log{\cal P}}{\mathrm{d}q}$ (19)
$\displaystyle\frac{\mathrm{d}q}{\mathrm{d}t}$ $\displaystyle=+\frac{\partial
U}{\mathrm{d}p}=M^{-1}p$ (20)

This is also used to set the scale of the random initial velocities. These
differential equations may be integrated numerically, taking $L$ small steps
of the leapfrog algorithm:

$\displaystyle p_{n+\frac{1}{2}}$
$\displaystyle=p_{n}-\frac{\varepsilon}{2}\frac{\partial
V}{\mathrm{d}q}(q_{n})$ (21) $\displaystyle q_{n+1}$
$\displaystyle=q_{n}+\varepsilon M^{-1}p_{n+\frac{1}{2}}$ (22) $\displaystyle
p_{n+1}$ $\displaystyle=p_{n+\frac{1}{2}}-\frac{\varepsilon}{2}\frac{\partial
V}{\mathrm{d}q}(q_{n+1})$ (23)

where $\varepsilon$ is a step size parameter.

Formally, the set of $n_{\mathrm{dim}}$ momenta are treated as new parameters,
and after some number of integration steps a final point in the trajectory is
compared to the initial one, and a Metropolis-Hastings acceptance criterion on
the total energy $H(q,p)$ is applied. If the trajectory is perfectly simulated
then this acceptance is unity, since energy is conserved; applying it allows a
relaxation of the integration accuracy.

The gradients $\partial\log{\cal P}/\mathrm{d}q$ can be estimated using finite
differences, but this requires at least $2n_{\mathrm{dim}}+1$ posterior
evaluations per point, greatly slowing it in high dimension, and as with the
Fisher forecasting is highly prone to numerical error. Automatically
calculating the derivative, as in jax-cosmo, makes it feasible and efficient.

Metropolis-Hastings, and related methods like emcee (Goodman & Weare, 2010;
Foreman-Mackey et al., 2013), suffer as dimensionality increases, as the
region of high probability mass (the typical set) becomes a very small
fraction of the total parameter space volume. At high enough dimension they
become a slow random walk around the space and cannot remain in typical set
regions. The dynamics of HMC allows it to make large jumps that nonetheless
stay within the region of high posterior.

The tricky part of HMC is that the leapfrog algorithm needs tuning to set the
number of steps as well as the step size of integration. The next section
examines a solution to this problem: the No-U-Turn HMC version.

###  5.3 NUTS

The No-U-Turn Sampler (NUTS) variant of the traditional HMC sampler was
introduced in Hoffman & Gelman (2014). It aims to help finding a new point
$x_{i+1}$ from the current $x_{i}$ by finding good and dynamic choices for the
leapfrog integration parameters in the root HMC algorithms, the step size
$\varepsilon$ and the number of steps $L$.

NUTS iterates the leapfrog algorithm not for a fixed $L$, but until the
trajectory starts to “double back” and return to previously visited region, at
the cost of increasing the number of model evaluations. The user has to set a
new parameter (max_tree_depth) which gives as a power of $2$ the maximum
number of model calls at each generation.

Both sampler HMC and NUTS are available in the NumPyro library. After the
forward model creation for the DES-Y1 3x2pt exercise described in section 5.1:

  * •

we apply a transformation to the cosmological, intrinsic alignment and bias
parameters to use a consistent uniform prior $\mathcal{U}[-5,5]$ (Table 1);

  * •

we use a structured mass matrix $M$ in a block diagonal form with the blocks
as the following sets of parameters
$(\Omega_{b},\Omega_{c},\sigma_{8},w_{0},h)$ and $(b_{i})_{i:1\dots 5}$. The
remaining parameters have uncorrelated masses. This matrix structure is
motivated by the expected degree of parameter correlation as shown for
instance in Figure 10.

We ran the NUTS sampler using numpyro.infer.NUTS on the DES Y1 likelihood,
with 16 chains of 1,000 samples each after a warm-up phase consisting of 200
samples, with the max_tree_depth set to seven (i.e. 128 steps for each
iteration). Using the “vectorized” numpyro option we ran all 16 chains
simultaneously on a single GPU, made possible by the JAX vmap mechanism. If
one has several GPU devices available, then the using the JAX paralellization
mechanism (pmap), it is further possible to launch the vectorized sampling
across the devices, and get back all the MCMC chains. However, these
experiments have all been undertaken on single GPUs, either an NVidia Titan Xp
(12GB RAM) on a desktop or an NVidia V100 (32GB RAM) at the IN2P3 Computing
Centre121212https://cc.in2p3.fr/en/. The elapsed time for these experiments
was 20 hours.

The results in terms of relative effective sample sizes (ESS) are detailed in
Table 2 while the confidence level (CL) contours are presented in Figure 8. We
compare to a reference sample from the highly-optimized Cobaya Metropolis-
Hastings implementation (Torrado & Lewis, 2019, 2021), which is widely used in
cosmology and which we ran for around 40 hours on CPU to obtain the set of
contours shown.

There is a dramatic improvement of the ESS by about a factor of 10 using the
NUTS sampler compared to Cobaya, with very good agreement between the CL
contours. It is worth mentioning that the mass matrix structure described
above increases the sampling efficiency by about a factor of two.

The speed of the sampling could be further improved: we have tested using the
parameter max_tree_depth=5 and found convergence in five hours, showing a
linear scaling in this parameter while keeping the sampling efficiencies at a
high level; the user is highly encouraged to tune this critical parameter.

Figure 8: Constraints (30%, 68%, 90%) on 4 of the 21 parameters of a simulated
DES-Y1 3x2pt likelihood, using the Cobaya Metropolis-Hastings sampler (full
curve in blue, 140,000 samples), the NUTS sampling described in section 5.3
(dashed curve in red; 16,000 samples), and a SVI MVN optimisation followed by
a Neural Transport parametrisation to run a NUTS sampling (dashed-dotted curve
in green, 200 samples).

###  5.4 Stochastic Variational Inference

We now explore Stochastic Variational Inference (Hoffman et al., 2013; Zhang
et al., 2019), another inference algorithm enabled by auto-differentiation. If
we write $p(z)$ the prior, $p(\mathcal{D}|z)$ the likelihood and
$p(\mathcal{D})$ the marginal likelihood, then thanks to Bayes theorem we have
$p(z|\mathcal{D})=p(z)p(\mathcal{D}|z)/p(\mathcal{D})$ as the posterior
distribution of a model with latent variables $z$ and a set of observations
$\mathcal{D}$. Variational Inference (VI) aims to find an approximation to
this distribution, i.e. $p(z|\mathcal{D})\approx q(z;\lambda)$, by determining
the variational parameters $\lambda$ of a predefined distribution. To do so,
one uses the Kullback-Leibler divergence of the two distributions
$KL(q(z;\lambda)||p(z|\mathcal{D}))$ leading to the following relation

$\displaystyle\log p(\mathcal{D})$
$\displaystyle=\mathtt{ELBO}+KL(q(z;\lambda)||p(z|\mathcal{D}))$ (24)
$\displaystyle\mathrm{with}\ \mathtt{ELBO}$
$\displaystyle\equiv-\mathbb{E}_{q(z;\lambda)}\left[\log
q(z;\lambda)\right]+\mathbb{E}_{q(z;\lambda)}\left[\log
p(z,\mathcal{D})\right]$ (25)

which defines the evidence lower bound (ELBO) that one aims to maximize to get
the $\lambda$ values. So, the optimal variational distribution satisfies

$q(z;\lambda^{\ast})=\underset{q(z;\lambda)}{\mathrm{argmax}}\
\mathtt{ELBO}=\underset{\lambda}{\mathrm{argmin}}\ \mathcal{L}(\lambda)$ (26)

The function $\mathcal{L}(\lambda)$ is the cost function used in practice. It
is composed of two parts:

$\mathcal{L}(\lambda)=\underbrace{\mathbb{E}_{q(z;\lambda)}\left[\log
q(z;\lambda)\right]}_{guide}-\underbrace{\mathbb{E}_{q(z;\lambda)}\left[\log
p(z,\mathcal{D})\right]}_{model}$ (27)

where the guide in the NumPyro library (i.e. the parameterised function $q$)
may be a multi-variate Gaussian distribution (MVN) for instance.

Using the auto-differentiation tool, one can use “black-box” guides (aka
automatic differentiation variational inference). As stated by the authors of
(Kucukelbir et al., 2017) ADVI specifies a variational family appropriate to
the model, computes the corresponding objective function, takes derivatives,
and runs a gradient-based or coordinate-ascent optimization. First we define a
invertible differentiable transformation $T$ of the original latent variables
$z$ into new variables $\xi$, such $\xi=T(z)$ and $z=T^{-1}(\xi)$, where the
new $\xi$ parameters are unbounded, $\xi_{i}\in(-\infty,\infty)$ and so the
subsequent minimization step can be performed with no bound constraints. The
cost function then reads

$\mathcal{L}(\lambda)=\underbrace{\mathbb{E}_{q(\xi;\lambda)}\left[\log
q(\xi;\lambda)\right]}_{guide}-\underbrace{\mathbb{E}_{q(\xi;\lambda)}\left[\log
p(\xi,\mathcal{D})\right]}_{model}$ (28)

with

$p(\xi,\mathcal{D}):=p(T^{-1}(\xi),\mathcal{D})|J_{T^{-1}}(\xi)|$ (29)

which includes the Jacobian of the $T^{-1}$ transformation. The evaluation of
the expectations during gradient descent of the loss Eq. 28 can be done using
what is called elliptical standardisation or re-parametrization trick or
coordinate transformation (see references in Kucukelbir et al. 2015). Let us
illustrate the method using an invertible transformation $S_{\lambda}$ such
that $S_{\lambda}(\xi)=\zeta$, where $\zeta\sim\mathcal{N}(0,I)$131313Notice
that $z=T^{-1}(\xi)=(T^{-1}\circ S_{\lambda}^{-1})(\zeta)=(S_{\lambda}\circ
T)^{-1}(\zeta)=F_{\lambda}(\zeta)$.. The Jacobian of this distribution is 1 by
definition (volume conservation), so the loss function reads

$-\mathcal{L}(\lambda)=\underbrace{\mathbb{E}_{\zeta\sim\mathcal{N}(0,I)}\left[\log
p(T^{-1}(S_{\lambda}^{-1}(\zeta)),\mathcal{D})+\log|J_{T^{-1}}(S_{\lambda}^{-1}(\zeta))|\right]}_{model}\\\
+\underbrace{\mathbb{H}[q(\xi;\lambda)]}_{guide}$ (30)

where $\mathbb{H}(q)\equiv\mathbb{E}_{q(\xi;\lambda)}\left[\log
q(\xi;\lambda)\right]$ is the Shannon entropy of the $q$ distribution; its
gradient can be computed once for all for a given $q$ distribution family and
reused in any user model. Then, to get $\nabla_{\lambda}\mathcal{L}$, the
$\nabla$ operator can be put inside the expectation which leads to141414To
simplify the notation, $T^{-1}(S_{\lambda}^{-1}(\zeta))$ has been replaced by
$z$.

$-\nabla_{\lambda}\mathcal{L}(\lambda)=\mathbb{E}_{\zeta\sim\mathcal{N}(0,I)}\left\\{\left[\nabla_{z}\log
p(z,\mathcal{D})\times\nabla_{\xi}[T^{-1}(\xi)]\right.\right.\\\
+\left.\left.\nabla_{\xi}\log|J_{T^{-1}}(\xi)|\right]\times\nabla_{\lambda}S_{\lambda}^{-1}(\zeta)\right\\}+\nabla_{\lambda}\mathbb{H}[q(\xi;\lambda)]$
(31)

An implementation example using the NumPyro library may be found in this
companion notebook: . Once the optimisation is done one can obtain i.i.d. $z$
samples from the $q(z,\lambda^{\ast})$ distribution applying the inverse of
the $S_{\lambda^{\ast}}$ and $T$ transformations. Using the same DES-Y1
simulation as in previous section, we use both a Multivariate Normal
distribution (MVN) and a Block Neural Autoregressive Flow (B-NAF) (De Cao et
al., 2020) as guides to approximate the true posterior (model). The B-NAF
architecture is composed of a single flow using a block autoregressive
structure with 2 hidden layers of 8 units each. The SVI optimization has been
performed with the Adam optimizer (Kingma & Ba, 2015) and a learning rate set
to $10^{-3}$. We have stopped the optimization after 20,000 (30,000) steps to
ensure a stable ELBO loss convergence when using the B-NAF (MVN) guides, and
no tuning of the optimizer learning rate scheduling or other parameters was
performed. This takes about 2 or 3 hours on the NVidia V100 GPU scaling,
depending on the number of steps.

In figure 9, we compare the contours obtained with Cobaya (as in figure 8) and
those obtained with the MVN and B-NAF guided SVI. As noted in Dhaka et al.
(2020) one challenge with variational inference is assessing how close the
variational approximation $q(z,\lambda^{\ast})$ is to the true posterior
distribution. It is not in the scope of this article to elaborate a
statistical diagnosis; rather we show that both guided SVI exhibit rather
similar contours and both estimate are close to the Cobaya posterior sampling.
The difference is that these SVI approximate posteriors have been obtained in
a much shorter time, and can serve as starting point for a NUTS sampler as
described in the next section.

Figure 9: Same configuration as in Figure 8 but using sampling of approximated
posterior using SVI with MVN and B-NAF guides (see text) compared to Cobaya
sampling.

####  5.4.1 Neural Transport

If the SVI method can be used as is to get $z$ i.i.d. samples from the
$q(z,\lambda^{\ast})$ distribution as shown on the previous section, the
Neural Transport MCMC method (Parno & Marzouk, 2018; Hoffman et al., 2019) is
an efficient way to boost HMC efficiency, especially in target distribution
with unfavourable geometry where for instance the leapfrog integration
algorithm has to face squeezed joint distributions for a subset of variables.
From SVI, one obtains a first approximation of the target distribution, and
this approximation is used to choose a better transform $T$ to map the
parameter space to a more convenient one $z=F_{\lambda}(\zeta)$ (e.g.
$F_{\lambda}=T^{-1}\circ S^{-1}_{\lambda}$) is such that

$q(z;\lambda)\rightarrow
q(\zeta;\lambda):=q(F_{\lambda}(\zeta))|J_{F_{\lambda}}(\zeta)|$ (32)

where $F_{\lambda}$ with the optimal $\lambda^{\ast}$ maps the best-fitting
$q(z;\lambda^{\ast})$ to a geometrically simple function like a unit
multivariate normal distribution. So, one can use a HMC sampler (e.g. NUTS)
based on $p(\zeta;\mathcal{D})$ distribution, initialized with $\zeta$ samples
from $q(\zeta;\lambda^{\ast})$, to get a Markov Chain of $N$ samples
$(\zeta_{i})_{i<N}$. Then, from the transformation
$z_{i}=F_{\lambda^{\ast}}(\zeta_{i})$ one finally obtain a Markov Chain with
$(z_{i})_{i<N}$ samples.

Figure 10: Same configuration as in Figure 8 but for the highly correlated
five lens galaxy bias.

We have used the MVN guided SVI described in the previous section on the same
DES-Y1 analysis described above. NUTS was run using 1 chain of 200 samples (a
fast configuration), one chain with 1000 samples and a set of ten chains with
1000 samples each combined into a single run of 10,000 samples (each setup
began all chains with 200 samples for initialisation). All NUTS sampling was
performed with dense mass matrix optimisation without the special block
structuring and with max_tree_depth=5 which is different than the default NUTS
setting described in section 5.3. The elapsed time for each of the 3 setups
was 50 minutes, 150 minutes and 5 hours, respectively. Naturally, more samples
lead to better contour precision. But what is illustrative is the fast
configuration results as shown in the Figures 8 and 10, compared to Cobaya and
NUTS results presented in section 5.3. The results are good even for the
highly correlated lens bias parameters. It is noticeable that running NUTS
with the same settings but without the SVI Neural Transport phase has
demonstrated a rather poor behaviour with only 200 samples.

Results in terms of sampling efficiency are shown in Table 2. SVI followed by
Neutral Transport gives high efficiency at low number of samples which may be
particularly useful during early phases of model development.

###  5.5 Sampling efficiency

Table 2: The relative effective sample size (ESS) in percent computed by the Arviz library (Kumar et al., 2019) from: (a): Cobaya 140,000 samples; (b): NUTS sampler with 16 chains of 1,000 samples each and 200 warm-up and max_tree_depth=7 (Section 5.3); (c) SVI Multivariate Normal followed by NUTS and Neural Transform with one chain of 200 samples and 200 warm-up. The ESS can be larger than 100% in some cases. (Section 5.4.1). | $\Omega_{b}$ | $\Omega_{c}$ | $\sigma_{8}$ | $w_{0}$ | $h$ | $n_{s}$ | $A$ | $\eta$
---|---|---|---|---|---|---|---|---
(a) | $3.1$ | $2.5$ | $2.9$ | $2.9$ | $2.6$ | $3.1$ | $3.1$ | $2.8$
(b) | $48.1$ | $45.6$ | $36.2$ | $33.4$ | $52.8$ | $50.1$ | $68.8$ | $48.8$
(c) | $84.0$ | $28.0$ | $26.0$ | $20.5$ | $80.0$ | $110.5$ | $58.5$ | $29.5$

Looking at the results of Table 2, a key question is: what is the HMC/NUTS
gain compared to the highly optimized Cobaya sampler? One useful metric is the
number of effective samples per model evaluation:

$\eta=\frac{n_{\mathrm{eff}}}{n_{\mathrm{eval}}}=\frac{N_{s}\times\varepsilon}{N_{s}\times
n_{\mathrm{step}}}=\frac{\varepsilon}{n_{\mathrm{step}}}$ (33)

with $N_{s}$ the total number of samples, $\varepsilon$ the effective sampler
efficiency and $n_{\mathrm{step}}$ the number of steps (calls) per generated
sample. For Cobaya we find $\varepsilon\approx 3\%$ with $n_{\mathrm{step}}=1$
while for the NUTS sampler $\varepsilon\approx 50\%$ but at the expense of
$n_{\mathrm{step}}=2^{5}$ or more. With better tuning of the sampling
parameters we would expect the $\eta$ values for both methods to become more
comparable, but at this intermediate dimensionality the gain from HMC/NUTS
compared to a standard MCMC sampler is small. The power of these approaches
will become most evident at higher dimensionality still, such as when
marginalizing over increasingly complex systematic models. NUTS also makes
post-processing simpler, since samples are nearly uncorrelated, removing the
need for a thinning step, which is a rather delicate procedure, needing know-
how to be conducted correctly (South et al., 2022; Owen, 2017).

As the dimensionality of cosmological models increases, methods like HMC/NUTS
that by construction are more efficient will become increasingly important.
Moreover, using SVI with neural reparametrisation offers an effective way to
undertake a progressive validation of a model with rather modest sample set
(e.g. starting with 200 samples) producing good enough marginal contours in
few hours. In practice, this validation phase can save time before producing
sizeable batch for final analysis. The authors have not investigated higher
dimensional ($O(10^{2})$ parameters) or multi-modal problems, but the key
argument in favour of HMC/NUTS sampling is that it exploits the geometry of
the typical set of the posterior distribution automatically, unlike the
standard random walk of Metropolis-Hasting sampling. Furthermore, using
reparametrisation one can adapt to poor geometry cases (e.g. Hoffman et al.,
2019).

##  6 General discussion

We have demonstrated that the derivatives made available in JAX-Cosmo can be
applied in a range of different methodologies to accelerate and improve the
stability of cosmological parameter estimation. We expect these improvements
to be increasingly important over the coming years, as higher statistical
precision pushes us towards more complicated models. Even though the standard
cosmological model has now remained static for several decades, models for
systematics such as redshift distribution uncertainty and galaxy bias are
becoming more and more complex. Using derivative-based methods to constrain
them can ensure that we propagate their uncertainties in full.

With these successes in mind, we now turn to important limitations of jax-
cosmo, and to questions about its place in the cosmology software ecosystem.

The first essential barrier to a fully-fledged automatically differentiable
cosmology library is the need for a differentiable Boltzmann solver to compute
the CMB or matter power spectra. At this stage, jax-cosmo relies on the
analytic Eisenstein & Hu fitting formula for the latter, which is not accurate
enough for Stage IV (Albrecht et al., 2006b) requirements, and it does not
include models beyond $\Lambda$CDM. Existing solvers such as CLASS (Blas et
al., 2011) or CAMB (Lewis et al., 2000) are large and complex codes which are
not easily reimplemented in an autodiff framework and therefore cannot be
directly integrated in jax-cosmo .

A first option to resolve this issue would be to implement from scratch a new
Boltzmann code in a framework that supports automatic differentiation. This is
the approach behind works such as
Bolt.jl151515https://github.com/xzackli/Bolt.jl which provides a simplified
Boltzmann solver in Julia, or PyCosmo (Refregier et al., 2018) which is based
on the SymPy symbolic mathematics library and could be relatively compatible
with JAX. However, even if very promising, both of these options thus far
remain limited. While we do believe an automatically differentiable Boltzmann
code is the best option, it seems that the cost of developing such a code
remains very high at this time.

A second approach would be to develop emulators of a fully-fledged Boltzmann
code. Emulators based on neural networks or Gaussian processes are themselves
automatically differentiable with respect to cosmological parameters. In fact,
the literature is now rich in examples of such emulators (e.g. Angulo et al.,
2021; Aricò et al., 2022; Günther et al., 2022; Nygaard et al., 2022; Spurio
Mancini et al., 2022; Albers et al., 2019; Schneider et al., 2011, and
references therein). After validating their accuracy against a reference CAMB
or CLASS implementation, they could be directly integrated as a plug-and-play
replacement for the computation of the matter power spectrum. At this time, it
seems that using emulators will be the most straightforward approach to bring
more accurate models to jax-cosmo. We believe, though, that one of the reason
for the wide diversity in this is a lack of standardization - a unified
interface and validation suite for such methods would provide a much simpler
comparison between them and enable wider usage.

Connected to this discussion about emulators, a point that could be discussed
is whether one even needs a library like jax-cosmo if one can build emulators
of a CCL likelihood for use inside gradient-based inference algorithms. While
this could indeed be feasible, and of similar cost as just making an emulator
of the matter power spectrum, the drawback of this approach is that many
analysis parameters and choices become hard coded in the emulator. Since the
model for the linear matter power spectrum is typically kept fixed in
practical analyses, all the of the choices related in particular to
systematics modeling (e.g. photometric redshift errors, galaxy bias, or
intrinsic galaxy alignments) will vary significantly in the process of
developing the analysis. Building an emulator for the likelihood would require
the emulator to be retrained every time the likelihood is changed.

Another aspect worth discussing are the prospects for scaling and speeding up
cosmological inference in practice, given tools such as jax-cosmo. As we
illustrated in the previous section, gradient-based inference techniques yield
significantly less correlated MCMC chains, scale better than any other known
sampling method as dimension increases, and can provide very fast approximate
posteriors if needed. jax-cosmo is also well suited to aid in the
parallelisation of likelihoods and algorithms, especially on multiple GPUs,
which will become increasingly important as the high-performance computing
landscape evolves.

Finally, how does jax-cosmo position itself against classical codes such as
CCL? While we are convinced of the benefits of a JAX implementation, we expect
CCL and other key codes to remain critical as standard cosmology
implementations. Ultimately, a natural transition may occur towards
differentiable frameworks like jax-cosmo when they reach the ability to run
fully-fledged Stage IV likelihoods.

##  7 Conclusions & Prospects

In this paper, we have presented jax-cosmo, a cosmology library implemented in
the JAX framework that enables the automatic computation of the derivatives of
cosmological likelihoods with respect to their parameters, and greatly speeds
up likelihood evaluations thanks to automatic parallelisation and just-in-time
compilation on GPUs. Currently, jax-cosmo contains a small set of features
corresponding to a DES-Y1 3x2pt analysis. Being an open source project,
contributions of new features for additional scientific areas such as CMB or
spectroscopic galaxy clustering are warmly welcome.

To demonstrate the value of an automatically differentiable library, we have
illustrated with concrete examples how Fisher matrices, which are notoriously
unstable and require extensive and careful fine tuning, can now be computed
robustly and at much lower cost. In addition, Fisher matrices becomes
themselves differentiable, which allows for Figure of Merit optimization by
gradient descent, making survey optimization extremely straightforward.

Going beyond Fisher forecasts, we have also compared simple Metropolis-
Hastings to several gradient-based inference techniques (Hamiltonian Monte-
Carlo, No-U-Turn-Sampler, and Stochastic Variational Inference). We have shown
that the posterior samples with gradient-based methods can reproduce classical
methods very efficiently, and can provide approximate posteriors very rapidly.
These inference techniques can scale to hundreds of dimensions, and may become
necessary in Stage IV analysis, as the number of nuisance parameters is likely
to become large.

The next extensions to this framework will be the inclusion of additional
cosmological probes, as well as the integration of emulators for the matter
power spectrum trained on CAMB or CLASS, as a means to go beyond the current
analytic Eisenstein & Hu model.

In the spirit of reproducible research, all results presented in this paper
can be reproduced with code contained in the following GitHub repository:

https://github.com/DifferentiableUniverseInitiative/jax-cosmo-paper/

## Credit authorship contribution statement

A. Boucaud: Software and comments. J.E Campagne: Conceptualization,
Methodology, Software, Validation, Writing, Visualization. S. Casas: Software
contribution for growth rate and power spectra. M. Karamanis: Software for
Gaussian likelihood computation. D. Kirkby: Software and validation for sparse
linear algebra. F. Lanusse: Conceptualization, Methodology, Software,
Validation, Writing, Project administration. D. Lanzieri: Software and
validation for redshift distribution. Y. Li: Software contributions. A. Peel:
Software and validation for spline interpolations in JAX. J. Zuntz:
Investigation, Writing.

## Acknowledgements

Some of the numerical experiments have been conducted at the IN2P3 Computing
Center (CC-IN2P3 - Lyon/Villeurbanne - France) funded by the Centre National
de la Recherche Scientifique.

## References

  * Abadi et al. (2015) Abadi, M., Agarwal, A., Barham, P., et al. 2015, TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems, software available from tensorflow.org
  * Abbott et al. (2018) Abbott, T. M. C., Abdalla, F. B., Alarcon, A., et al. 2018, Phys. Rev. D, 98, 043526
  * Akeret et al. (2015) Akeret, J., Gamper, L., Amara, A., & Refregier, A. 2015, Astronomy and Computing, 10, 1
  * Albers et al. (2019) Albers, J., Fidler, C., Lesgourgues, J., Schöneberg, N., & Torrado, J. 2019, J. Cosmology Astropart. Phys, 2019, 028
  * Albrecht et al. (2006a) Albrecht, A., Bernstein, G., Cahn, R., et al. 2006a, arXiv e-prints, astro
  * Albrecht et al. (2006b) Albrecht, A., Bernstein, G., Cahn, R., et al. 2006b, arXiv e-prints, astro
  * Angulo et al. (2021) Angulo, R. E., Zennaro, M., Contreras, S., et al. 2021, Monthly Notices of the Royal Astronomical Society, 507, 5869
  * Aricò et al. (2022) Aricò, G., Angulo, R., & Zennaro, M. 2022, Open Research Europe, 1
  * Baydin et al. (2017) Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. 2017, J. Mach. Learn. Res., 18, 5595–5637
  * Baydin et al. (2018) Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. 2018, Journal of Marchine Learning Research, 18, 1
  * Betancourt (2017) Betancourt, M. 2017, arXiv e-prints [[arXiv]1701.02434]
  * Bezanson et al. (2017) Bezanson, J., Edelman, A., Karpinski, S., & Shah, V. B. 2017, SIAM review, 59, 65
  * Bhandari et al. (2021) Bhandari, N., Leonard, C. D., Rau, M. M., & Mandelbaum, R. 2021, arXiv e-prints, arXiv:2101.00298
  * Bingham et al. (2019) Bingham, E., Chen, J. P., Jankowiak, M., et al. 2019, J. Mach. Learn. Res., 20, 28:1
  * Bird et al. (2012) Bird, S., Viel, M., & Haehnelt, M. G. 2012, Monthly Notices of the Royal Astronomical Society, 420, 2551
  * Blas et al. (2011) Blas, D., Lesgourgues, J., & Tram, T. 2011, J. Cosmology Astropart. Phys, 2011, 034
  * Blondel et al. (2021) Blondel, M., Berthet, Q., Cuturi, M., et al. 2021, arXiv preprint arXiv:2105.15183
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., et al. 2018, JAX: composable transformations of Python+NumPy programs
  * Brinckmann & Lesgourgues (2019) Brinckmann, T. & Lesgourgues, J. 2019, Physics of the Dark Universe, 24, 100260
  * Campagne et al. (2017) Campagne, J. E., Neveu, J., & Plaszczynski, S. 2017, A&A, 602, A72
  * Carpenter et al. (2017) Carpenter, B., Gelman, A., Hoffman, M. D., et al. 2017, Journal of Statistical Software, 76, 1–32
  * Chisari et al. (2019) Chisari, N. E., Alonso, D., Krause, E., et al. 2019, ApJS, 242, 2
  * De Cao et al. (2020) De Cao, N., Aziz, W., & Titov, I. 2020, in Proceedings of Machine Learning Research, Vol. 115, Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, ed. R. P. Adams & V. Gogate (PMLR), 1263–1273
  * Dhaka et al. (2020) Dhaka, A. K., Catalina, A., Andersen, M. R., et al. 2020, in Advances in Neural Information Processing Systems, ed. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, & H. Lin, Vol. 33 (Curran Associates, Inc.), 10961–10973
  * Duane et al. (1987) Duane, S., Kennedy, A. D., Pendleton, B. J., & Roweth, D. 1987, Physics Letters B, 195, 216
  * Eisenstein & Hu (1998) Eisenstein, D. J. & Hu, W. 1998, The Astrophysical Journal, 496, 605
  * Foreman-Mackey et al. (2013) Foreman-Mackey, D., Hogg, D. W., Lang, D., & Goodman, J. 2013, PASP, 125, 306
  * Goodman & Weare (2010) Goodman, J. & Weare, J. 2010, Communications in Applied Mathematics and Computational Science, 5, 65
  * Günther et al. (2022) Günther, S., Lesgourgues, J., Samaras, G., et al. 2022, Journal of Cosmology and Astroparticle Physics, 2022, 035
  * Hartlap et al. (2007) Hartlap, J., Simon, P., & Schneider, P. 2007, A&A, 464, 399
  * Heavens et al. (2000) Heavens, A. F., Jimenez, R., & Lahav, O. 2000, MNRAS, 317, 965
  * Heavens et al. (2017) Heavens, A. F., Sellentin, E., de Mijolla, D., & Vianello, A. 2017, MNRAS, 472, 4244
  * Heek et al. (2020) Heek, J., Levskaya, A., Oliver, A., et al. 2020, Flax: A neural network library and ecosystem for JAX
  * Hessel et al. (2020) Hessel, M., Budden, D., Viola, F., et al. 2020, Optax: composable gradient transformation and optimisation, in JAX!
  * Hoffman et al. (2019) Hoffman, M., Sountsov, P., Dillon, J. V., et al. 2019, arXiv e-prints [[arXiv]1903.03704]
  * Hoffman et al. (2013) Hoffman, M. D., Blei, D. M., Wang, C., & Paisley, J. 2013, J. Mach. Learn. Res., 14, 1303–1347
  * Hoffman & Gelman (2014) Hoffman, M. D. & Gelman, A. 2014, J. Mach. Learn. Res., 15, 1593–1623
  * Jasche & Wandelt (2013) Jasche, J. & Wandelt, B. D. 2013, MNRAS, 432, 894
  * Joachimi et al. (2011) Joachimi, B., Mandelbaum, R., Abdalla, F. B., & Bridle, S. L. 2011, A&A, 527, A26
  * Kingma & Ba (2015) Kingma, D. P. & Ba, J. 2015, in 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, ed. Y. Bengio & Y. LeCun
  * Knox et al. (2001) Knox, L., Christensen, N., & Skordis, C. 2001, ApJ, 563, L95
  * Kostić et al. (2022) Kostić, A., Nguyen, N.-M., Schmidt, F., & Reinecke, M. 2022, arXiv e-prints, arXiv:2212.07875
  * Kucukelbir et al. (2015) Kucukelbir, A., Ranganath, R., Gelman, A., & Blei, D. M. 2015, in Proceedings of the 28th International Conference on Neural Information Processing Systems \- Volume 1, NIPS’15 (Cambridge, MA, USA: MIT Press), 568–576
  * Kucukelbir et al. (2017) Kucukelbir, A., Tran, D., Ranganath, R., Gelman, A., & Blei, D. M. 2017, J. Mach. Learn. Res., 18, 430–474
  * Kumar et al. (2019) Kumar, R., Carroll, C., Hartikainen, A., & Martin, O. 2019, Journal of Open Source Software, 4, 1143
  * Leonard et al. (2023) Leonard, C. D., Ferreira, T., Fang, X., et al. 2023, The Open Journal of Astrophysics, 6, 8
  * Lewis & Bridle (2002) Lewis, A. & Bridle, S. 2002, Phys. Rev. D, 66, 103511
  * Lewis et al. (2000) Lewis, A., Challinor, A., & Lasenby, A. 2000, ApJ, 538, 473
  * LoVerde & Afshordi (2008) LoVerde, M. & Afshordi, N. 2008, Phys. Rev. D, 78, 123506
  * Margossian (2019) Margossian, C. C. 2019, WIREs Data Mining and Knowledge Discovery, 9, e1305
  * Nygaard et al. (2022) Nygaard, A., Brinch Holm, E., Hannestad, S., & Tram, T. 2022, arXiv e-prints, arXiv:2205.15726
  * Owen (2017) Owen, A. B. 2017, Journal of Computational and Graphical Statistics, 26, 738
  * Parno & Marzouk (2018) Parno, M. D. & Marzouk, Y. M. 2018, SIAM/ASA Journal on Uncertainty Quantification, 6, 645
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., et al. 2019, in Advances in Neural Information Processing Systems 32 (Curran Associates, Inc.), 8024–8035
  * Percival (2005) Percival, W. J. 2005, A&A, 443, 819
  * Phan et al. (2019) Phan, D., Pradhan, N., & Jankowiak, M. 2019, arXiv [[arXiv]1912.11554]
  * Refregier et al. (2018) Refregier, A., Gamper, L., Amara, A., & Heisenberg, L. 2018, Astronomy and Computing, 25, 38
  * Rubiño-Martin et al. (2003) Rubiño-Martin, J. A., Rebolo, R., Carreira, P., et al. 2003, MNRAS, 341, 1084
  * Ruiz-Zapatero et al. (2023) Ruiz-Zapatero, J., Hadzhiyska, B., Alonso, D., et al. 2023, arXiv e-prints, arXiv:2301.11978
  * Salvatier et al. (2016) Salvatier, J., Wiecki, T. V., & Fonnesbeck, C. 2016, PeerJ Computer Science, 2, e55
  * Schneider et al. (2011) Schneider, M. D., Holm, Ó., & Knox, L. 2011, ApJ, 728, 137
  * Sellentin & Heavens (2018) Sellentin, E. & Heavens, A. F. 2018, MNRAS, 473, 2355
  * Smith et al. (2003) Smith, R. E., Peacock, J. A., Jenkins, A., et al. 2003, MNRAS, 341, 1311
  * South et al. (2022) South, L. F., Riabiz, M., Teymur, O., & Oates, C. J. 2022, Annual Review of Statistics and Its Application, 9, null
  * Spurio Mancini et al. (2022) Spurio Mancini, A., Piras, D., Alsing, J., Joachimi, B., & Hobson, M. P. 2022, MNRAS, 511, 1771
  * Stuart & Ord (1991) Stuart, A. & Ord, K. 1991, Kendall’s advanced theory of statistics, 5th edn., ed. E. Arnold, Vol. 2, Classical Inference and Relationship
  * Takahashi et al. (2012) Takahashi, R., Sato, M., Nishimichi, T., Taruya, A., & Oguri, M. 2012, ApJ, 761, 152
  * Tegmark et al. (1997) Tegmark, M., Taylor, A. N., & Heavens, A. F. 1997, ApJ, 480, 22
  * Torrado & Lewis (2019) Torrado, J. & Lewis, A. 2019, Cobaya: Bayesian analysis in cosmology
  * Torrado & Lewis (2021) Torrado, J. & Lewis, A. 2021, J. Cosmology Astropart. Phys, 2021, 057
  * Villaescusa-Navarro et al. (2020) Villaescusa-Navarro, F., Hahn, C., Massara, E., et al. 2020, ApJS, 250, 2
  * Yahia-Cherif et al. (2021) Yahia-Cherif, S., Blanchard, A., Camera, S., et al. 2021, A&A, 649, A52
  * Zablocki & Dodelson (2016) Zablocki, A. & Dodelson, S. 2016, Phys. Rev. D, 93, 083525
  * Zhang et al. (2019) Zhang, C., Bütepage, J., Kjellström, H., & Mandt, S. 2019, IEEE Transactions on Pattern Analysis and Machine Intelligence, 41, 2008
  * Zuntz et al. (2021) Zuntz, J., Lanusse, F., Malz, A. I., et al. 2021, The Open Journal of Astrophysics, 4, 13
  * Zuntz et al. (2015) Zuntz, J., Paterno, M., Jennings, E., et al. 2015, Astronomy and Computing, 12, 45