\section{Introduction}

Bayesian parameter estimation is a problem usually encountered when testing physical models against some experimental data. In these cases, inference from the posterior of the model is often done using a sampling algorithm. The posterior's likelihood in realistic scenarios is usually a complex hierarchical model, involving costly computations of intermediate quantities and possibly multiple likelihoods from different data sets. Each part of the calculation may have its own parameter dependencies, from the underlying parameters of main physical interest to nuisance parameters describing the experimental model or uncertainties in different parts of the theoretical calculation. Different subsets of parameters may have very different computational costs, depending on which parts of the model have to be recalculated when each parameter is varied. If subsets of the parameters can be updated with low computational cost, so there is a significant speed hierarchy, this can be exploited to greatly improve sampling efficiency. However, the simplest applications of standard sampling algorithms usually fail to exploit this property.

\texttt{Cobaya}\xspace (COde for BAYesian Analysis; meaning ``guinea pig'' in Spanish) is a general framework for defining a pipeline of interdependent calculations, automatically exploiting the hierarchical nature of the model when sampling from its parameters. The main advantage of Cobaya resides in its ability to automatically analyse the dependency structure of the different model components, and to block parameters according to which parts of the computation (directly or indirectly) depend on them, sorting these parameters blocks in an optimal way so that looping over them takes the least amount of time and no part of the calculation is unnecessarily repeated (intermediate results are automatically cached). This also allows for oversampling of fast blocks according to their speed and use of sampling algorithms (such as fast dragging~\cite{Neal04}) that can efficiently exploit the presence of very fast parameter subspaces. \texttt{Cobaya}\xspace includes its own MCMC algorithms (adapted from \texttt{CosmoMC}\xspace \cite{Lewis:2013hha}) and an interface to the nested sampler \texttt{PolyChord}\xspace \cite{2015MNRAS.450L..61H,2015MNRAS.453.4384H}. In addition to Monte Carlo samples, \texttt{Cobaya}\xspace includes a importance-weighting tool, and interfaces with popular minimizers for posterior and likelihood maximization (from \texttt{scipy}\xspace \cite{scipy} and \texttt{Py-BOBYQA}\xspace implementation \cite{2018arXiv180400154C} of the \texttt{BOBYQA}\xspace algorithm \cite{Powell2009})\footnote{The structure is however general for samplers where posterior derivatives are not available: additional samplers could be used by providing a compatible sampler class implementation.} \texttt{Cobaya}\xspace is also ready for high-performance computation in clusters, featuring MPI parallelization and batch job generation.

Calculations can be separated into physically distinct but interdependent theoretical calculations and likelihood evaluations, where any part of the calculation can be provided by separately-maintained external Python packages each with its own parameter dependencies. This allows theory modellers and analysers of experimental data to provide and maintain different cleanly separated parts of the calculation rather than relying on a single monolithic package that does everything.

Though \texttt{Cobaya}\xspace is a completely generic framework, it includes out-of-the-box support for cosmological parameter estimation via a series of interfaces to a number of cosmological theory codes (\texttt{CAMB}\xspace \cite{Lewis:1999bs,Howlett:2012mh} and \texttt{CLASS}\xspace \cite{Lesgourgues:2011re,Blas:2011rf})and data likelihoods. In that regard, it compares to \texttt{CosmoMC}\xspace \cite{Lewis:2002ah,Lewis:2013hha}, and later \texttt{MontePython}\xspace \cite{2013JCAP...02..001A,Brinckmann:2018cvx} and \texttt{CosmoSIS}\xspace \cite{Zuntz:2014csq}, amongst others. \texttt{Cobaya}\xspace does not re-distribute most of the original likelihoods, data sets or theory codes, but provide installers for stock distributions. \texttt{Cobaya}\xspace can also be used as a simple cosmological likelihood wrapper, providing stand-alone cosmological posteriors that can be incorporated into the user's own statistical or machine-learning pipeline.

\texttt{Cobaya}\xspace is distributed as a standard \texttt{Python}\xspace package\footnote{\url{https://pypi.org/project/cobaya/}}, which allows for easy management of dependencies and simple roll-out of updates. Extensive documentation is provided at \url{https://cobaya.readthedocs.io}. Users are also welcome to contribute through the central git repository at \url{https://github.com/CobayaSampler/cobaya}.

In the rest of the paper, we describe the probabilistic model underlying \texttt{Cobaya}\xspace's framework (Sec.~\ref{sec:probmodel}), comment on our main goals and consequent design choices (Sec.~\ref{sec:goals}), and describe the main structure of the code (Sec.~\ref{sec:structure}). We then briefly discuss use on high-performance computers (Sec.~\ref{sec:hpc}) and describe the specific realistic case of using \texttt{Cobaya}\xspace\ for cosmological data analysis (Sec.~\ref{sec:cosmo}). We explain the technical details of our new optimized parameter blocking in Appendix~\ref{app:blocking}, and display some of the source code necessary to reproduce the used case presented in Sec.\ \ref{sec:cosmo} in Appendix \ref{sec:cosmo_demo_source}.


\section{Probabilistic model}\label{sec:probmodel}

Bayesian inference on model parameters is usually done by performing Monte Carlo sampling on the posterior $\prob{\theta\,|\,\mathcal{D},\mathcal{M}}$, where $\mathcal{M}$ is some model parameterized by $\theta$ following some \emph{prior} $\prior{\theta\,|\,\mathcal{M}}$, and $\mathcal{D}$ is some data set we aim to model using $\mathcal{M}$. The probability for observing the data $\mathcal{D}$ given some particular set of model parameters is the \emph{likelihood}  $\like{\mathcal{D}\,|\,\mathcal{M}(\theta)}$. These quantities are related through the \emph{Bayes theorem}:
\begin{equation}
  \label{eq:bayestheorem}
  \prob{\theta\,|\,\mathcal{D},\mathcal{M}} =
  \frac{\like{\mathcal{D}\,|\,\mathcal{M}(\theta)}\prior{\theta\,|\,\mathcal{M}}}
       {\prob{\mathcal{D}\,|\,\mathcal{M}}}
  \eqpunctsp ,
\end{equation}
where $\prob{\mathcal{D}\,|\,\mathcal{M}}$ is the marginal likelihood or \emph{evidence}, which equals the integral of the numerator in the right-hand side.

In simple physical applications, the data that appears in the likelihood is expressed in terms of some summary \emph{observable} $\vector{O}_\mathcal{D}$. In those cases, the likelihood consists of a comparison between the observed $\vector{O}_\mathcal{D}$ and the one computed from a deterministic \emph{theoretical model} $\mathcal{T}$, $\vector{O}_\mathcal{T}$, filtered through the simulated pipeline $\mathcal{F}_\mathcal{E}$ of the experiment following some \emph{experimental model} $\mathcal{E}$. For a simple Gaussian likelihood, this would look like
\begin{equation}
  \label{eq:like}
  \vector{O}_\mathcal{D} \sim \normal{\mathcal{F}_\mathcal{E}(\vector{O}_\mathcal{T})}{\Sigma_{\mathcal{E},\mathcal{T}}}
  \eqpunctsp ,
\end{equation}
where the covariance matrix $\Sigma_{\mathcal{E},\mathcal{T}}$, dependent on both the theoretical and the experimental model, parameterizes those deviations.

This simple example illustrates how a natural \emph{speed hierarchy} appears in the parameter space: whenever we vary parameters of the theoretical model, we would have to recompute both the theoretical prediction for the observable $\vector{O}_\mathcal{T}$ \emph{and} the likelihood, whereas a variation in the experimental model parameters only requires recomputation of the likelihood (as long as we cache and reuse the computed $\vector{O}_\mathcal{T}$). The likelihood calculation for fixed theoretical model is also often much faster than computing the theoretical prediction $\vector{O}_\mathcal{T}$. Exploiting this speed hierarchy efficiently can save a considerable amount of time when sampling from the posterior (see appendix \ref{app:blocking}).

In general, physical inference pipelines may have several components and dependencies, without a sharp distinction between a theoretical and an experimental model. In \texttt{Cobaya}\xspace we allow the experimental likelihoods to communicate with one or more theory codes (and possibly other likelihoods) to calculate the required observables and any necessary derived parameters. Theory codes can depend on sampled parameters and on the outputs of other theory codes, so we allow for general (non-circular) dependencies between the different components so that the calculation can be cleanly modularized. This also allows any resulting hierarchy in the speed of the likelihood calculation under changes in different parameter variations to be exploited efficiently.

When various theoretical codes are available that can produce the same quantities, interfaces can be choice-agnostic so that they can be exchanged easily for testing and comparison.


\section{Goals and design choices}\label{sec:goals}

We aimed at the following development and usability goals, and made design choices accordingly:

\paragraph{Modularity:} The different actors in the Bayesian model (prior, likelihood, theory model, Monte Carlo sampler), are individual objects which share as little information as possible: prior and sampler know which parameters are sampled, but they do not know about any other fixed parameter of the model; the likelihood, in turn, deals only with input/output parameters and observables, not caring about their prior; the sampler does now know which individual likelihood (or theory code) understands which parameters, and lets a \emph{model} wrapper manage that. These design choices impose some compromises, but are a fair price for making the code more easily extendable and maintainable.

\paragraph{Rapid prototyping:} We have attempted to lower as much as possible the barrier to adding new priors, likelihoods, theory codes or samplers, or modify existing ones, all without touching \texttt{Cobaya}\xspace's source code or having to write a lot of \textit{wrapping} code. In order to achieve this, we have developed an API (i.e.\ \emph{application programming interface}) that offloads the need by the user to write specific code to handle communication between different parts of the pipeline, to cache intermediate results, or to exploit the parameter speed hierarchy. The user just needs to describe what a pipeline component needs and provides, and the actual calculation of the provided quantities. Creating such an API has been made much easier by writing our code in \texttt{Python}\xspace, which allows us to discover parameter dependencies of generic functions (used as priors, likelihoods or parameter redefinitions) via its \textit{introspection} capabilities. Python is able to interface with libraries written in other languages, allowing use of external compiled theory codes and likelihoods, and also has useful computational capabilities (e.g.\ handling infinities directly). Though Python, as in interpreted language, is in general slower than compiled languages such as C or Fortran, this will not necessarily have a significant impact on computational speed, since most heavy computations are usually performed by calling functions in optimized packages based on a compiled back-end library (such as \texttt{numpy}\xspace). Not-easy-to-offload computations (e.g.\ those containing long explicit loops) can also be implemented using popular just-in-time compilers or language extensions such as NUMBA \cite{Lam:2015:NLP:2833157.2833162} or Cython \cite{behnel2010cython}, or coded in C/Fortran and interfaced to Python.

\paragraph{Supporting external components:}
Using Python makes it simple to load theory or likelihood codes from other Python packages. \texttt{Cobaya}\xspace supports direct use of theory and likelihood classes defined in other packages simply by referencing their fully-qualified name in the package, making it straightforward for modellers or experimental collaborations to release their own code packages that can be used directly. External likelihood codes can either inherit directly from \texttt{Cobaya}\xspace\ classes (which also support direct instantiation and use outside of Cobaya), or include a separate compatible wrapping class that can be used with \texttt{Cobaya}\xspace.


More information on design choices can be found in the documentation at \url{https://cobaya.readthedocs.io}.

\section{Structure of the code}\label{sec:structure}

\texttt{Cobaya}\xspace's main structure is shown in Fig.~\ref{fig:structure}. The main two classes are the Bayesian \texttt{Model}, and the Monte Carlo \texttt{Sampler} (or, more broadly, any analysis tool that operates on the model). This section describes the elements shown there.

\begin{figure}[ht]
  \includegraphics[width=0.8\columnwidth]{img_diagram.pdf}
  \caption{Simplified structure of \texttt{Cobaya}\xspace's source, showing classes (squares) and parameters (ellipses). See section \ref{sec:structure} for a description of each class and parameter role. The arrows between \texttt{TheoryCollection} and \texttt{LikelihoodCollection} represent computed quantities and parameters that can be exchanged arbitrarily between theories and likelihoods.
  }
  \label{fig:structure}
\end{figure}

\subsection{Input}

For each particular run, the input must describe the model and its parameter space in enough detail, and also specify the analysis tool that will be used. Ideally, the syntax of this input would reproduce the structure of the code as closely as possible (it makes it easier to swap different input chunks), and should be as literate (self-descriptive), human-readable and easy to remember as possible.

\texttt{Cobaya}\xspace's input takes the form of a \texttt{Python}\xspace dictionary, and can be serialized in plain text in the YAML format \cite{YAML} (as long as its elements can be serialized in plain text). We show an example in Fig.~\ref{fig:input}, including specification of parameter roles and dynamic parameter definition, as well as definition of the likelihood and prior for the simple case where everything is just a basic Python function.

\begin{figure}[ht]
  \centering
  \lstinputlisting[language=yaml, frame=single]{src_input_example.yaml}
  \caption{Example input in plain text (YAML). It defines a Gaussian-ring likelihood with radius $1$ and standard deviation $0.02$, over the combination of a uniform prior $(x,y)\in(0,2)^2$ (notice the two possible different specifications used for \texttt{x} and \texttt{y}) and a Gaussian prior of standard deviation $0.3$ along the $x=y$ direction (n.b.: simple 1D priors are defined in \texttt{params}, while multidimensional ones are defined in \texttt{prior}). The likelihood, the multidimensional prior and the derived parameters \texttt{r} and \texttt{theta} are given as \texttt{Python}\xspace functions (here source strings, but can be assigned \texttt{Python}\xspace functions directly when working in a \texttt{Python}\xspace file or shell -- for source strings \texttt{scipy.stats} and \texttt{numpy} are pre-imported as \texttt{stats} and \texttt{np} resp.). The results of this MCMC sample will be written in a folder called \texttt{chains} with file name prefix \texttt{ring}, as per the \texttt{output} option. The resulting densities can be seen in Fig.~\ref{fig:results}.}
  \label{fig:input}
\end{figure}

Describing complex models and pipelines can take a large amount of input. To facilitate these cases, we have implemented an inheritance system by which defaults for every class are automatically loaded. Options only need to be specified when their values are different from the default ones. In the example of Fig.\ \ref{fig:input}, for the MCMC sampler only the stopping criterion is specified, while the rest of the options are either inherited from the defaults of this sampler (which can be retrieved via a shell script ``\texttt{cobaya-doc mcmc}''), or automated from other options (e.g.\ a diagonal proposal covariance matrix constructed from the parameters' \texttt{proposal} property).


\subsection{Bayesian \texttt{Model}}

The Bayesian \texttt{Model} consists of the Bayesian prior and likelihood (this last one including theory and experimental likelihood components), as described in section \ref{sec:probmodel}. It also contains a \texttt{Parameterization} layer that manages the flow of parameters to and from the likelihood and prior, and computes the dynamically-defined ones, and a \texttt{Provider} that handles the exchange of parameters and computed quantities between different theory and likelihood components.

A \texttt{Model} instance can be passed as an argument to a sampler, or it can be integrated by the user into an external pipeline (using its API to access the member classes shown in Fig.\ \ref{fig:structure} and described below).


\subsubsection{Parameterization}

The \texttt{Parameterization} class controls the flow of parameters into and out of the \texttt{Model}, taking into account their roles with respect to different parts of the code.

On their way out of the model (i.e.\ from the point of view of the \texttt{Sampler}), parameters can play three different roles:
\begin{itemize}
\item \emph{Sampled} parameters are the ones whose value is to be varied and explored by the \texttt{Sampler} or the user-defined pipeline. They are identified in the input by having a defined prior.
\item \emph{Fixed} parameters are those whose value is not going to change, and are needed as input by some piece of the \texttt{Model}.
\item \emph{Derived} parameters are arbitrary functions of the rest of the parameters at every step, and are tracked and stored for the user's convenience. Functions defining them can be provided \emph{on the fly} (like \texttt{r} and \texttt{theta} in the example in Fig.~\ref{fig:input}) or can be implicitly defined inside the code of a theory or likelihood.
\end{itemize}

The \texttt{Parameterization} class processes the parameters to turn them into \emph{input} parameters for the likelihoods and theories, and requests from them the \emph{output} parameters that are needed to compute \emph{derived} parameters that cannot be computed directly from inputs.

The parameterization layer also manages other properties of the parameters, such as their labels (used for plots).


\subsubsection{Prior}

Priors for the sampled parameters can be specified in two different ways:
\begin{itemize}
\item The \texttt{prior} keyword for each parameter defines a separable product of 1D priors. They are interfaced directly from the 1D continuous distributions of \texttt{scipy.stats}\footnote{\url{https://docs.scipy.org/doc/scipy/reference/stats.html\#continuous-distributions}}. In the example in Fig.\ \ref{fig:input}, these are the \emph{uniform} priors $(x,y)\in(0,2)^2$.
\item Additional, possibly-multi-dimensional priors can be defined under the global \texttt{prior} block. In the example in Fig.\ \ref{fig:input}, this is the Gaussian prior along $x=y$.
\end{itemize}

\subsubsection{Likelihood and theory}
\label{sec:liketheo}

\begin{figure*}[ht]
  \centering
  \begin{subfigure}{0.5\textwidth}
\begin{python}
import numpy as np
import scipy.stats as stats

from cobaya.theory import Theory
from cobaya.likelihood import Likelihood


class ToPolarCoords(Theory):
    params = {'x': None, 'y': None}

    def get_can_provide(self):
        return ['r', 'theta']

    def calculate(self, state, want_derived=True,
                  **params_values):
        x = params_values['x']
        y = params_values['y']
        state['r'] = np.sqrt(x**2 + y**2)
        state['theta'] = np.arctan(y / x)


class GaussianRingLikelihood(Likelihood):
    params = {'mean_radius': None,
              'std_radius': None}

    def get_requirements(self):
        return ['r']

    def logp(self, **params_values):
        r = self.provider.get_result('r')
        mean_radius = \
            params_values['mean_radius']
        std_radius = params_values['std_radius']
        return stats.norm.logpdf(
            r, loc=mean_radius, scale=std_radius)
\end{python}

  \end{subfigure}\hfill
  \begin{subfigure}{0.48\textwidth}
    \lstinputlisting[language=yaml, frame=single, linewidth=0.99\linewidth]{src_input_external.yaml}
  \end{subfigure}
  \caption{
    Example similar to the one in Fig.\ \ref{fig:input}, now using \texttt{Cobaya}\xspace classes to split the computation in two: the transformation between orthogonal and polar coordinates, and the likelihood in terms of polar coordinates. Here we let the mean radius of the gaussian ring vary over a narrow prior, to illustrate \texttt{Cobaya}\xspace's automated blocking: when sampling using MCMC or \texttt{PolyChord}\xspace, jumps in the $(x, y)$ directions will be alternated with jumps in \texttt{mean\_radius}. After every jump on $(x, y)$, the resulting intermediate product $r$ is cached, so that $r$ does not need to be recomputed when only \texttt{mean\_radius} is varied. In this trivial example, the intermediate quantity $r$ exchanged between the theory and the likelihood is just a real number, but it could be any arbitrarily complicated and many-dimensional numerical quantity, as well as a general Python object. \texttt{Cobaya}\xspace knows about the interdependency between the likelihood, that needs $r$, and the theory, that computes $r$, via the respective declarations in methods \texttt{Likelihood.get\_requirements} and \texttt{Theory.get\_can\_provide}.}
  \label{fig:external}
\end{figure*}

In general there can be multiple theory and likelihood components in a single model, with each theory component calculating some quantity required as input to the likelihoods or to another theory component. Likelihoods are just a special subclass of a general Theory class that directly return likelihood values. Both inherit a caching layer that increases efficiency when steps in the parameter space are blocked in particular ways, such that only the piece that depends on the varied parameters (directly or via a dependency) needs to be recomputed.

A model class instance holds a list of the required likelihood and theory component instances, and uses their interface methods or introspection to work out their dependencies and hence required execution order in order to calculate the final likelihood. Components may also calculate derived parameters that are useful to include in the output parameter chain (or may be used by other components).

In  \texttt{Cobaya}\xspace's input, likelihoods and theory components can be specified as \texttt{Python}\xspace functions, mentions to \emph{internal} classes (i.e.\ those distributed with \texttt{Cobaya}\xspace, either as stand-alone versions or as wrappers to external code/data), inherited classes, or mentions to classes distributed in external packages.

A simple example illustrating the use of \texttt{Cobaya}\xspace Theory and Likelihood classes can be seen in Fig.\ \ref{fig:external}. This example highlights one of the biggest advantages of \texttt{Cobaya}\xspace's Bayesian model API with respect to most generalist statistical languages: the specification of what is needed and what is provided by each stage of the pipeline is enough to produce effective posterior exploration. Automatically, \texttt{Cobaya}\xspace will separate the sampled parameters in two blocks according to their dependencies, [\texttt{x}, \texttt{y}] and [\texttt{mean\_radius}], it will sort them optimally (in this simple case [\texttt{mean\_radius}] comes obviously last), it will prompt the sampler, either MCMC of \texttt{PolyChord}\xspace, to vary the parameters block-wise, and it will make sure that the intermediate quantity \texttt{r} is cached so that variations of \texttt{mean\_radius} only do not require its recomputation.


\subsection{Sampler}

Monte Carlo samplers in \texttt{Cobaya}\xspace take models and explore their \emph{sampled} parameters.

\texttt{Cobaya}\xspace implements adaptive fast-slow-optimized MCMC samplers, translated from \texttt{CosmoMC}\xspace \cite{Lewis:2002ah,Lewis:2013hha}. This includes a Metropolis-Hastings MCMC \cite{Metropolis53,Hastings70} sampler, with optional oversampling of fast parameters, and a variation of it that performs \emph{dragging} \cite{Neal04} along the fastest directions. In many physical applications, complex interdependencies between multiple theory codes and experimental likelihoods produce a slow-fast parameter hierarchy. Efficient exploitation of this hierarchy, while keeping the ability to propose steps in random directions of the parameter space, is crucial to efficient sampling \cite{Lewis:2002ah, Lewis:2013hha}.

Previous implementations have required a manual specification of the parameter blocking and oversampling configuration. For \texttt{Cobaya}\xspace\ we have implemented a new automated optimization scheme that blocks parameters according to their role in the different stages of the likelihood, measures the computational cost of varying parameters in each of these blocks, and determines the optimal parameter block sorting so that varying all parameters in a row takes the smallest amount of time possible, taking into account both parameter mixing (necessary to explore possible degeneracies) and the amount of oversampling requested by the user. For details on this algorithm see Appendix \ref{app:blocking}.

Chain sampling convergence is quantified using a modified version of the Gelman-Rubin $R-1$ statistic \cite{Gelman92,Brooks98}, either over several chains run in parallel, or over the latest chunks of a single chain. The $R-1$ statistic quantifies the variance in the means of parameters estimated from different chains (or chain subsets), and termination of the chain can be specified by giving a target small value for $R-1$. Where tail exploration is important, an equivalent stopping target can also be set for the dispersion of confidence limits.

In addition, \texttt{Cobaya}\xspace contains a wrapper for the nested sampler \texttt{PolyChord}\xspace \cite{2015MNRAS.450L..61H,2015MNRAS.453.4384H}, which can also estimate model evidences and explore complicated multi-modal likelihood surfaces. \texttt{PolyChord}\xspace\ can also exploit the parameter speed hierarchy determined by \texttt{Cobaya}\xspace's parameter blocking. \texttt{Cobaya}\xspace provides an installer for \texttt{PolyChord}\xspace, as well as for all supported external dependencies. Wrappers for additional samplers may be implemented in the future.

Under the \textit{sampler} category, \texttt{Cobaya}\xspace also includes interfaces to minimizers (those in \texttt{scipy}\xspace \cite{scipy} and the \texttt{Py-BOBYQA}\xspace implementation \cite{2018arXiv180400154C} of the \texttt{BOBYQA}\xspace algorithm \cite{Powell2009}), a simple test function to evaluate different quantities of the model at fixed points, and an importance-reweighting tool.

\subsection{Analysis -- interface to \texttt{GetDist}\xspace}

\texttt{Cobaya}\xspace manages Monte Carlo samples as wrapped \texttt{DataFrame} objects from \texttt{Pandas}\xspace~\cite{mckinney-proc-scipy-2010}. When written to the hard drive, they can be stored in plain text as parameter tables, including the corresponding probabilities and sample weights. Both the sample objects and output files can be easily loaded and analysed with the user's tool of choice.

The suggested analysis package is \texttt{GetDist}\xspace\footnote{\url{https://github.com/cmbant/getdist/}} \cite{Lewis:2019xzd}, which can load \texttt{Cobaya}\xspace results transparently. \texttt{GetDist}\xspace provides summary statistics including confidence intervals, density estimates (via optimized kernel density estimation), and convergence diagnostics, plotting tools, and a graphical user interface. Examples of \texttt{GetDist}\xspace outputs can be seen in Fig.~\ref{fig:results}, obtained from the model described by the input in Fig.~\ref{fig:input}.

\begin{figure}[ht]
  \centering
  \includegraphics[width=0.9\columnwidth]{img_triangle.pdf}
  \includegraphics[width=0.9\columnwidth]{img_rtheta.pdf}
  \caption{Results from sampling using the inputs shown in Fig.\ \ref{fig:input}, analysed by \texttt{GetDist}\xspace: The upper figure shows a \emph{triangle plot} combining 2D posterior contours (enclosing 68\% and 95\% of the probability) and marginalized 1D posteriors for the sampled parameters $(x,y)$. The lower plot shows the 1D posteriors for the derived parameters $(r,\theta)$. All the posteriors are shown normalized to the same maximum.}
  \label{fig:results}
\end{figure}


\section{High performance computing with \texttt{Cobaya}\xspace}
\label{sec:hpc}

\texttt{Cobaya}\xspace is tailored for high-performance computing on clusters. The overhead of \texttt{Cobaya}\xspace per posterior evaluation is $\alt 0.2\,\mathrm{ms}$, with no strong dependence on the dimensionality for $d<100$; this means that for theory or likelihood computations that take more than a few milliseconds to compute, any number of evaluations while sampling their posterior in \texttt{Cobaya}\xspace would not take significantly longer than evaluating the same code in a simple loop the same number of times. For higher dimensions, the components of the overhead that have a stronger dependence on dimensionality start taking over, but usually so does computation time of likelihoods that depend on an increasingly high number of parameters.


\subsection{Parallelization}

\texttt{Cobaya}\xspace takes advantage of hybrid MPI+threading parallelization. All samplers are MPI-aware and use communication between parallel processes to improve sampling efficiency (e.g.\ MCMC uses all chains to learn a better covariance matrix and to assess convergence). Each single MPI process carries an independent \texttt{Model}, which takes advantage of threading to accelerate its likelihood and theory computations. Some of the sampler-intrinsic operations can also take advantage of threading (e.g.\ recomputing the proposal covariance matrix). For most sampling algorithms MPI usage is a very small fraction of the computing time, so interconnect speed is largely irrelevant and the code will run well on cheap commodity (or virtualized) clusters.

Threading is usually leveraged via \texttt{numpy}\xspace's \cite{numpy} interface to a threading-aware linear algebra library such as \texttt{OpenBLAS}\xspace\footnote{\url{https://www.openblas.net/}}, or via the use of externally linked compiled codes where necessary.

\subsection{Batch runs}

In physical applications, it is usual to test a number of theoretical models against different combinations of data sets. Running the sampler for every combination in this \texttt{grid} of theory--data requires a large amount of very redundant input: it differs little between cases in the grid.

In order to facilitate that, we have incorporated a piece of code first implemented for \texttt{CosmoMC}\xspace, that can define such a \texttt{grid} starting from common definitions plus two lists of variations: theoretical models and data sets. This code prepares input files containing all possible (or allowed) combinations, creates a set of nested \texttt{model/data\_set} folder structures to store the results, and generates scripts to be submitted to a cluster queue.

\subsection{Cloud computing}

Users without access to a cluster may be able to pay for on-demand online computing resources, or cloud computing. We have explored and documented\footnote{\url{https://cobaya.readthedocs.io/en/latest/cluster_amazon.html}} \texttt{Cobaya}\xspace's usage in one popular solution, Amazon's EC2.

\section{Cosmology}
\label{sec:cosmo}

Although \texttt{Cobaya}\xspace\ is a general analysis tool, its development was originally motivated for its use in cosmology, using various standard codes for calculating cosmological observables and a variety of publicly available likelihoods from different data sets. We have included tools specific to this use case and have used it for testing on real-world problems. The use of any sampler in a cosmological context presents a series of challenges, some already presented in section \ref{sec:goals}, and others intrinsic to Cosmology:
\\

\textbf{Input complexity:} To deal with the large amount of information needed to describe cosmological models and data sets, we have created a graphical application to generate input based on the user's choice for each piece of the cosmological model (including usual priors in the literature), and the data sets used. Also, this tool automatically selects an optimal sampler configuration for each particular problem (including picking an appropriate proposal covariance matrix for MCMC from a database of standard cosmological runs).
\\

\textbf{External dependencies:}
In general, we do not modify and re-distribute existing publicly available theory and likelihood codes, but instead interface stock versions through light wrappers\footnote{We do re-implement a few dependencies (e.g.\ the $H_0$, BAO and DES cosmological likelihoods), specially when the cost of interfacing very simple external code outweighs that of maintaining our own version, or because of other potential benefits, such as extensibility, re-usability for other experimental pipelines, etc.}. This allows us to keep \texttt{Cobaya}\xspace's source light and more easily maintainable, letting us focus on the statistical and user experience aspects. We mitigate the added complication of having to download external code by providing a one-line installer for all those external packages, that works across multiple systems (GNU/Linux, macOS and Windows, subject to OS compatibility of each of those external codes). The installer script can take an input file and install everything necessary to run it, including theory and likelihood codes, and data sets.
\\

\textbf{User-modified external code:} By not repackaging and re-distributing our own version of external codes, we impose no entry barrier for user-modified versions: users just need to provide paths to the installation folders of their modified versions.
\\

\textbf{Alternative external codes with the same role:} We have made the interface between experimental likelihoods and cosmological theory codes agnostic to the theory code used (as long as the quantities requested by other parts of the code can be computed by all of the alternatives).
\\

\textbf{Computational cost:} For cosmological data probing non-background cosmology, a likelihood evaluation for given cosmological parameters typically requires running at least a Boltzmann code to calculate the perturbation transfer functions (taking a second or more depending on number of cores and the specific model), plus possibly additional steps such as modelling non-linear evolution, calculating angular power spectra, correlation functions, etc.\ from the transfer functions. The second step can be fast in linear theory, or seconds when non-linear modelling is involved or complex observables. The original \texttt{CosmoMC}\xspace\ code provided a split between cosmological parameters affecting the transfer functions and parameters governing in the initial power spectrum, allowing the latter to be sampled very quickly for fixed values of the transfer function parameters~\cite{Lewis:2002ah}, but assumed linear theory. In \texttt{Cobaya}\xspace\ the parameter dependencies and possible speed hierarchy can also be exploited by splitting the theory calculation up into a transfer function calculation, and subsequent calculations involving initial power spectrum parameters and non-linear evolution parameters. \texttt{Cobaya}\xspace's \texttt{CAMB}\xspace\ wrapper explicitly supports this with non-linear modelling based on variations of \textsc{halofit}~\cite{Smith:2002dz,Takahashi:2012em,Mead:2015yca}, generalizing the approach used in \texttt{CosmoMC}\xspace. Although the speed hierarchy is modest once non-linear corrections to the CMB lensing\footnote{A Limber approximation for the non-linear corrections could likely speed this up with very small loss of accuracy, but is not currently implemented.} and other spectra are included, the saving can be significant since at the very least the cost of computing the transfer functions does not have to be paid every time power spectrum and non-linear evolution parameters are varied.

Despite the small additional Python overhead, \texttt{Cobaya}\xspace cosmology MCMC chains runs using the $\Lambda$CDM model with \texttt{CAMB}\xspace and {\it Planck} data~\cite{PCP2018}  have comparable computation cost to using \texttt{CosmoMC}\xspace. In particular, assuming a good starting guess of the proposal covariance matrix,\footnote{This is done to reduce random noise in the test, by removing the very variable initial fraction of samples during which most of the proposal covariance matrix is learned. Also notice that in \texttt{CosmoMC}\xspace (and also in \texttt{MontePython}\xspace and \texttt{CosmoSIS}\xspace at the moment of writing) the proposal covariance matrix guess needs to be given by hand, while in Cobaya it is automatically selected from a database for a range of cosmological configurations.} both \texttt{Cobaya}\xspace and \texttt{CosmoMC}\xspace consistently reach Gelman-Rubin convergence of $R-1\alt 0.02$ in 8 hours on four cores per chain and 4 chains. Equivalence with \texttt{CosmoMC}\xspace as the \textit{de facto} standard in the field should be enough, specially when considering the additional advantages of \texttt{Cobaya}\xspace. In particular, \texttt{Cobaya}\xspace's model API (see section \ref{sec:liketheo}) and its optimal-sorting algorithm (see Appendix \ref{app:blocking}) guarantee that as pipelines become more complex and are further modularised, \texttt{Cobaya}\xspace will consistently outperform simpler approaches.

For example, \texttt{Cobaya}\xspace\ would be more efficient when varying a larger number of primordial power spectrum parameters with non-linear modelling, e.g.\ when performing primordial power spectrum reconstruction (which is explicitly supported in \texttt{Cobaya}\xspace\ by using a separate primordial power spectrum Theory class; currently compatible with \texttt{CAMB}\xspace only). For application to future data the calculation could be further modularized. For example, the non-linear modelling could be separated into a separate Theory class taking inputs from a Boltzmann code and calculating non-linear power spectra. This could be implemented in different ways and used consistently by likelihoods using different class implementations for cross-comparison. The calculation of actual observables from the underlying power spectra, for example by codes such as CCL\cite{Chisari:2018vrw}, could be implemented as a separate Theory class taking non-linear power spectra and  outputs from a Boltzmann Theory code, and using them to calculate observables such as tomographic lensing correlation functions required as the theoretical model input used by likelihoods.

\subsection{An example use case in Cosmology}
\label{sec:cosmo_use_case}

As a demonstration of some of \texttt{Cobaya}\xspace's capabilities when applied to a cosmological scenario, we present here an inference problem in which we search for a primordial power spectrum feature of the form
\begin{equation}
  \label{eq:feature}
 
  \frac{\Delta P(k)}{P_0(k)} = A \exp^{-0.5 \left[\frac{\log_{10}\left(\frac{k}{k_c}\right)}{w}\right]^2} \sin\left(2 \pi \frac{k}{l}\right)
  \,,
\end{equation}
where $P_0(k)$ is a power-law nearly-scale-invariant scalar power spectrum, that has been injected into a fiducial cosmological model. This ansatz has four parameters: a relative \emph{amplitude} $A$ of the feature, a \emph{wavelength} $l$, and an \emph{envelope centre} $k_c$ and \emph{envelope log-width} $w$. We forecast how well the feature can be characterized using highly idealized versions of the Planck and forthcoming Simons Observatory (SO) data~\cite{Ade:2018sbj,PL2018}. To use the maximal effective sky area at each multipole we separate the likelihood into low- and high-$\ell$ ranges, with different sensitivities, sky fractions, and polarization combinations (we will see that a single likelihood wrapper can incorporate all of these).

We try to constrain/recover the injected feature parameters in two scenarios: one using standard lensed CMB power spectra, and in which we assume that a certain amount of delensing~\cite{Sherwin:2015baa,Larsen:2016wpa,Green:2016cjr} is possible. Large-scale lensing modes shift the scales of small-scale oscillation modes leading to them being smoothed out in the lensed power spectrum. If lensing multipoles $L\alt 200$ can recovered using an external tracer like the Cosmic Infrared Background (CIB) or internally reconstructed, this smoothing effect can be substantially reduced by delensing. As a simplified model we simply assume that the lensing power can be reduced by a constant factor of $0.7$ using a combination of internal delensing, CIB, and large-scale structure\footnote{The detectability of small-scale features is not very sensitive to whether or not the $L>200$ lensing modes are also delensed, since they largely contribute to adding power to both spectra.}. We do not attempt to model the statistics of the delensed power spectrum in detail, and take the log likelihood of both delensed and lensed spectra to be given by an $f_{\rm sky}$-scaled ideal full-sky mean log likelihood.

We demonstrate here how \texttt{Cobaya}\xspace can be used to build a pipeline for this use case. To do this, we take advantage of \texttt{Cobaya}\xspace's \texttt{CAMB}\xspace interface for computing the CMB power spectrum given a primordial spectrum, that we will define in a separate Theory class; we will also create a generic Likelihood class using data that we have generated based on a fiducial CMB power spectrum containing a feature. The step-by-step building of this pipeline, would proceed in the following way (the source code of a possible implementation is included in Appendix \ref{sec:cosmo_demo_source}):
\begin{itemize}

\item We define a Theory class that computes the primordial power spectrum including the feature (see \texttt{FeaturePrimordialPk} in Fig.\ \ref{fig:src_demo_theory}). This class, inheriting from the \texttt{Cobaya}\xspace generic Theory prototype class, only needs to define two methods: one returning what this class can compute and starting with \texttt{get\_}, in this case \texttt{get\_primordial\_scalar\_pk}; and another called \texttt{calculate} taking the parameters defining the power spectrum and computing it, in the example of Fig.\ \ref{fig:src_demo_theory}, simply a wrapper around a Python function that defines a linearly-oscillating feature with a Gaussian envelope on top of a nearly-scale-invariant power spectrum. The \emph{class attribute} \texttt{params} defines what parameters among all the ones declared in the input are to be passed to this class (and also which derived parameters are made available by this likelihood). The rest of the class attributes (in our example \texttt{n\_samples\_wavelength} and \texttt{k\_pivot}), define parameters that can be set in the declaration of the class in the \texttt{Cobaya}\xspace input (see Fig.\ \ref{fig:src_demo_input_delensed}), and their default values if they are omitted.

\item Regardless of the definition of the parameters passed to the new Theory class, we may want to redefine them e.g.\ to sample them over a log-uniform prior. We show how to do that for the feature \texttt{amplitude}, its \texttt{wavelength} and the \texttt{centre} of its envelope in Fig.\ \ref{fig:src_demo_input_delensed}: we define the respective dummy log-parameters in the \texttt{params} block, indicate that they shall not be passed to the pipeline with the option ``drop: True'', and define the original ones as a function of them. These redefinitions are specially useful when we cannot easily modify the source code of the original Theory code or likelihood (which is a common occurrence for complex codes such as official experimental likelihoods).

\item We may want to impose additional priors that depend on more than one parameter. In our example, we would like to ban combinations of amplitude and position and width of the feature envelope that produce no significant trace in the observable CMB multipole range. To do that, we can define the corresponding log-prior function e.g.\ in the same file as the Theory code (see \texttt{logprior\_high\_k} in Fig.\ \ref{fig:src_demo_theory}), and reference it in the input as shown in Fig.\ \ref{fig:src_demo_input_delensed}.

\item The likelihood segment of the pipeline is usually the most complicated one for research cases that do not use public data as-is, or public likelihood codes. Users can take two alternative approaches: on the one hand, they can write generic (i.e.\ non-\texttt{Cobaya}\xspace-dependent) code to generate or load some data and compute their likelihood (possibly assuming a parameterized experimental model), and then wrap it in a \texttt{Cobaya}\xspace Likelihood class; alternatively they can inherit from one of \texttt{Cobaya}\xspace's cosmological likelihoods and override the necessary methods. In both cases, the methods needing defining/overriding are \texttt{get\_requirements}, declaring which cosmological observables/quantities will be requested from the Theory code(s), and \texttt{logp}, that retrieves said quantities and passes them to the user-defined probability function (or does this computation directly). These two methods for our use case can be seen in Fig.\ \ref{fig:src_demo_likelihood}. Here, we inherit from one of \texttt{Cobaya}\xspace's likelihoods, \texttt{CMBlikes}, that takes care of loading the necessary data and computing a mean log-likelihood given the polarized CMB power spectra. We override \texttt{get\_requirements}  to add the \emph{results} object of \texttt{CAMB}\xspace (\texttt{CAMBdata}, from which we will extract the partially-lensed CMB spectra), to the list of CMB polarizations and multipole ranges defined by the data with which the likelihood was initialized (and set here via a \texttt{super} call). This shows the flexibility of \texttt{Cobaya}\xspace's wrappers: even if \texttt{Cobaya}\xspace does not know about a particular observable, it can still be retrieved from the theory code by hand. We also need to override the \texttt{logp} method of this class to retrieve the partially-lensed power spectra instead of the fully-lensed ones.

Notice that likelihoods inheriting from an existing cosmological \texttt{Cobaya}\xspace likelihood can work outside the \texttt{Cobaya}\xspace pipeline: they can e.g.\ be used as standalone objects that can evaluate data probabilities given an arbitrary realization of the required observable. As with the Theory codes, we make use of \emph{class attributes} to specify which options can be defined in the input at run time, which allows us to write a single likelihood wrapper for our four likelihoods, with the respective data specified via \texttt{dataset\_file}, and the lensing scenario defined by the value of \texttt{Alens\_delens} ($0.3$ for the delensed case, and $1.0$ for the non-delensed one).

\item On the sampler side, here MCMC, thanks to our use of \texttt{Cobaya}\xspace's class wrappers, the parameter speed hierarchy is taken advantage of automatically: parameters are split into a \emph{primordial power spectrum} block (those sent to \texttt{FeaturePrimordialPk}) and a \emph{CMB transfer functions} block (those sent to \texttt{CAMB}\xspace), and these last ones, as they are the slowest, are placed first. Since we are exploring both feature and baseline $\Lambda$CDM parameters, and the latter will have approximately similar posteriors in delensed and non-delensed cases (since we are injecting a feature for which we don't expect strong correlations with background cosmology), we may want to perform an initial run with $\Lambda$CDM cosmology and a liberal convergence criterion, that we will use to generate a covariance matrix (a standard output file of \texttt{Cobaya}\xspace's MCMC sampler). We can then set this file as the proposal \texttt{covmat} for the \texttt{mcmc} sampler (see Fig.\ \ref{fig:src_demo_input_delensed}). The \texttt{mcmc} sampler will then create the total covariance matrix, using the provided one (including correlations) for the $\Lambda$CMD parameters, and the \texttt{proposal} property of each parameter for the feature parameters.
\end{itemize}

Notice that when defining this pipeline, we did not need to modify \texttt{Cobaya}\xspace or \texttt{CAMB}\xspace, but simply write the code carrying out the calculations in a non-\texttt{Cobaya}\xspace-dependent way, and code \texttt{Cobaya}\xspace wrappers for these calculations, what they require, and what they make available. The advantages of this approach are on the one hand that it is more accessible for researchers that do not know the inner workings of \texttt{Cobaya}\xspace (or \texttt{CAMB}\xspace), and on the other that distributing a small set of files (for reproducibility and archiving purposes\footnote{Reproducibility is ensured despite possible backwards-incompatible changes in the codes used, since \texttt{Cobaya}\xspace stores both its version and those of the external codes used (here \texttt{CAMB}\xspace) in the output files.}) is more practical and scalable than distributing whole modified versions of codes.\footnote{The source code to reproduce this example can be downloaded from \url{https://github.com/CobayaSampler/paper_cosmo_demo}}

In order to test our pipeline, we have produced both a delensed and a non-delensed MCMC run based on a fiducial cosmological model containing a feature as defined in Eq.\ \eqref{eq:feature} with fiducial parameter values $A=0.1$, $l=0.008\,\mathrm{Mpc}^{-1}$, $k_c=0.2\,\mathrm{Mpc}^{-1}$, $w=0.1\,\mathrm{Mpc}^{-1}$. For the rest of the cosmological model, we assume Planck-like $\Lambda$CDM \cite{Aghanim:2018eyx} with parameter values given by the best fit to the low-$\ell$ TT+EE, high-$\ell$ TT+TE+EE and lensing power spectrum Planck 2018 likelihoods. We vary both feature and $\Lambda$CDM parameters, with a prior for the feature parameters centered around the fiducial values (the aim for this demonstration is to characterize the injected feature, not to detect its presence). We show the result for the delensed scenario in Fig.\ \ref{fig:cosmo_demo} (for the sake of brevity, the marginalized 1- and 2-d posteriors for the rest of the cosmological parameters are not shown).\footnote{The sampling process took 72h (wall-clock time) to converge (Gelman-Rubin $R-1<0.01$) running 4 parallel chains with 4 threads each, on an Intel Xeon CPU E5-2640v3. The running time is significantly higher than in the test case discussed in Sec.\ \ref{sec:cosmo} due to the increased precision necessary to run this pipeline (at the levels of transfer function computation, CMB multipole integration, and lensing potential computation), which makes the full pipeline in this setting take $\approx 7s$.} Not surprisingly, we find a degeneracy between the feature, at the same time, (1) being centered towards higher $k$, (2) having a wider envelope, and (3) having a bigger amplitude. We cannot usually map the end of this degeneracy because it would involve considering amplitudes $\Delta P(k) / P(k)\sim 1$, which would contradict the assumption implicit in Eq.\ \ref{eq:feature} that the feature is a small perturbation on top of vanilla slow-roll dynamics producing nearly-scale-invariant power-law primordial spectrum.

\begin{figure*}[ht]
  \centering
  \includegraphics[width=.8\textwidth]{img_cosmo_demo.pdf}
  \caption{1- and 2-d marginalized posteriors for the feature parameters in the delensed scenario ($\Lambda$CDM cosmological parameters were sampled but are not shown here). There exists a degeneracy towards simultaneous large amplitude, wide envelope and large-$k$ center reaching the prior boundaries (see discussion in main text).}
  \label{fig:cosmo_demo}
\end{figure*}

In the non-delensed case, not shown in Fig.\ \ref{fig:cosmo_demo}, the MCMC chains mostly reproduce the prior region in which the feature does not leave a significant imprint on top of baseline $\Lambda$CDM (i.e.\ small feature amplitudes, envelopes centered outside the CMB window, or envelopes being too narrow). This is due to the small imprint the fiducial feature leaves on top of $\Lambda$CDM when full lensing is considered.

This example use case is presented as a demonstration of an array of \texttt{Cobaya}\xspace's features, and it should not be taken as proof that delensing significantly enhances primordial feature recovery. While this is probably true for a small region of the parameter space, one would need to perform a more detailed study, using more realistically-generated data and more realistic delensing, and characterizing the $\chi^2$ distribution over a number of simulations to assess detectability. The resulting posterior would also be more irregular than the smooth one shown in Fig.\ \ref{fig:cosmo_demo}, and possibly multimodal, in which case it would make sense to use \texttt{PolyChord}\xspace instead as a sampler.


\section{Conclusion}

\texttt{Cobaya}\xspace is a general modularized framework for Bayesian analysis of models with complex pipelines, with simple and structured input specification, and a minimal API for interfacing external theories and likelihoods. Model posteriors can be sampled, maximized and importance-reweighted. \texttt{Cobaya}\xspace's main advantage is its ability to automatically account for interdependencies between the different components of the pipeline, and to use these dependencies and the computational costs of varying each parameter to optimally sort the parameter blocks of the model so that the posterior is efficiently explored, all without the need for specific code by the user.

\texttt{Cobaya}\xspace is well suited for high-performance computing thanks to its low overhead, MPI parallelization, and batch-running of grids of jobs. In a future release, we plan to implement HPC-enabled containerization capabilities (\texttt{Docker}+\texttt{Shifter} and \texttt{Singularity}), as well as to improve the batch-running tools.

\section{Acknowledgements}

Thanks to J.\ Lesgourgues and W.\ Handley for support on interfacing \texttt{CLASS}\xspace and \texttt{PolyChord}\xspace, respectively. Thanks to G.\ Ca\~nas Herrera, A.\ Finke, X. Garrido, S.\ Heimersheim, L.\ Hergt, M.S.\ Madhavacheril, V.\ Miranda, T.\ Morton, J.\ Zunz and many other for help debugging and patching Cobaya.

We acknowledge support from the European Research Council under the European Union's Seventh Framework Programme (FP/2007-2013) / ERC Grant Agreement No. [616170].
AL is also supported by the UK STFC grants ST/P000525/1 and ST/T000473/1.


