PyMC: a modern, and comprehensive probabilistic programming framework in Python

Distributions

Because Bayesian statistics involves constructing probabilistic models, PyMC provides robust implementations of all commonly-used probability distributions as classes in its distributions module, as well as a large set of more specialized distributions. The availability of these classes facilitates the construction of probabilistic graphs for forward sampling and inference. Stochastic random variables in PyMC models have a common set of properties as defined by the Distribution base class; these include methods for calculating log-probabilities, log-cumulative probabilities, and sampling random values. Distributions are broadly grouped according to whether they are continuous or discrete and univariate or multivariate, which confers additional properties to them according to this taxonomy (i.e., shape and data type values). In addition to these user-facing features, PyMC’s inference algorithms can manipulate distribution objects to improve their performance; for example, distributions that are bounded in their support are automatically transformed into real-valued spaces in order to improve sampling efficiency. For modeling time series, there are more complex multivariate distributions that account for serial dependence among elements of a sequence of variables, notably the autoregressive process and random walk processes (both Gaussian and Student-T). Also, there are wrapper classes that enforce truncation and censoring bounds for arbitrary base distributions which can be useful, for instance, in the construction of survival models. If a particular distribution of interest is not included in the distributions module, PyMC allows users to specify their own distributions.

Model

As introduced in Section API overview, the Model module provides the Model class, which encapsulates all components of the Bayesian model the user specifies. The user API for a Model instance is through a context manager, which serves to automatically add variables to the model graph, according to the associations between variables by the user; specifically, a random variable passed as an argument to another random variable establishes a parent-child relationship between them. All interaction with a Model instance by the user (e.g., model modification, model fitting) is done through the context manager. While the Model class has a large suite of methods, almost none of them are user-facing but are instead called by other methods and functions in PyMC in order to set up, introspect, or manipulate the model.

As a convenience to users, the Model module includes functions for generating visualizations of the model DAG (see for example Fig. 3), using either GraphViz (Gansner & North, 2000) or NetworkX (Hagberg, Schult & Swart, 2008).

Logprob

This module contains the logic for operating with RandomVariable objects including: Converting RandomVariable graphs into joint log-probability graphs, transforming constrained RandomVariables so their support is on unconstrained spaces, RandomVariable-aware pretty printing, and LaTeX output.

Step methods

PyMC constructs posterior samplers using a Metropolis-within-Gibbs scheme, where blocks of parameters are assigned the same MCMC step algorithm. Distributions define their own default step algorithm, but this may be manually overridden by the user. For instance, in the coal mining disaster example (see Section Coal mining disasters) the Metropolis step method was assigned to the discrete sp variable and NUTS was assigned to the rest of the variables, which are continuous.

The default sampler for all continuous distributions is based on the No-U-Turn sampler (Hoffman & Gelman, 2014)⁴ . In practice, many models use only continuous parameters, and so will be only using this algorithm, which sees particular attention for performance. Some changes from the algorithm as originally described include multinomial sampling of the tree (Betancourt, 2016), and a corrected U-turn check. Other available inference methods include random-walk Metropolis (including specialized versions for binary and categorical parameters), a slice sampler (Neal, 2003), and Differential Evolution Metropolis sampling (Ter Braak, 2006).

Sampling

In addition to the MCMC step methods described in the previous section, which are useful to sample from the posterior distribution, PyMC supports sampling from both the prior predictive and posterior predictive distributions with flexible shapes allowing for in-sample estimates, as well as sample interpolation or extrapolation. This can be done using the pm.sample_prior_predictive or pm.sample_posterior_predictive functions. It is important to remark that PyMC allows using the same model definition to compute posteriors distributions (backward sampling) or predictive distributions (forward sampling), without requiring any intervention from the user.

Additionally, PyMC supports samplers from BlackJAX and NumPyro, as shown in Section JAX-based sampling.

Variational

The variational inference implementation is inspired by Ranganath et al. (2016) and defined as follows:

PyMC implements O—Operator, $p,q$—approximations, $f$—test function $t$—distance function family in a class structure; this is a nice abstraction over variational inference algorithms. One can construct standard Kullback-Leibler (KL) divergence with mean field approximation using this framework. To show that, let’s set operator $O_{KL}^{p,q}$ such that

In a sense $O_{KL}^{p,q}$ is constant over all $f\in =ℱ$, and yields a constant function without dependency on $f$. Using a short notation we can write

The simplified objective now does not need ${\sup }_{f\in ℱ}$ anymore, and setting the distance function to identity $t_I(x)=x$ we obtain

Setting variational family $q$ to the Mean Field family

completes the derivation.

Under this framework, PyMC also has Stein Variational Gradient (Liu & Wang, 2016) and $\beta -\mathrm{KL}$ (Burgess et al., 2018) objectives. Further extensions are made simpler for researchers and are open for contributions.

ODE

PyMC includes an Ordinary Differential Equations (ODE) module. The API mimics requires a function $f(y,t,\theta )$ which takes as arguments an array of states $\mathbf{y}$, a time argument $t$, and an array of parameters $\theta$. Once a solution to the ODE is found, an array of times at which to evaluate the solution is used.

It is implemented as an PyTensor Op which allows Hamiltonian Monte Carlo to differentiate through the ODE solution. The underlying implementation uses Scipy’s ODE solver which in turn uses the lsoda routine in the ODEPACK library written in FORTRAN.

SMC

Sequential Monte Carlo is a family of Monte Carlo methods. It has wide applications to Bayesian inference for static models and dynamic models, such as sequential time series inference and signal processing (Del Moral, Doucet & Jasra, 2006; Ching & Chen, 2007; Naesseth, Lindsten & Schön, 2019; Chopin & Papaspiliopoulos, 2020). PyMC supports sampling from static models using a Sequential Monte Carlo method, with many kernels, including, Metropolis-Hastings, Independent Metropolis-Hastings, Hamiltonian Monte Carlo. The SMC sampler can be useful for sampling multimodal posteriors, specially when modes are separated by very low probability regions. Additionally, the Independent Metropolis-Hastings kernel can be useful for models where the gradients are not available and for which complexity and/or dimensionality makes Metropolis-Hastings a poor choice. Moreover, SMC computes the marginal likelihood of a model as a by-product of sampling. In this module, we also find the code necessary to run Sequential Monte Carlo for Approximate Bayesian Computation. That is, models where the likelihood is not available explicitly, but we have a method to simulate synthetic data given a set of unknown parameters (Sisson, Fan & Tanaka, 2007; Sunn Sunnåker et al., 2013; Martin, Kumar & Lao, 2021).

Gaussian process

A Gaussian process (GP) can be used as a prior probability distribution over the space of continuous functions for modeling non-linear processes non-parametrically. PyMC includes a specialized API to define and fit GP models, specifically using latent and marginal approximations, and generate predictions for arbitrary inputs. The GP estimation is fully compatible with the MCMC, ADVI, and MAP methods, allowing users to flexibly fit models of various computational complexities while balancing computational estimation time.

Additionally, PyMC includes a number of kernels, including exponentiated quadratic, Matern, and periodic kernels, while allowing for the flexibility of users to define their own kernels if need be.

Coal mining disasters

Between 1851 and 1962, a record number of accidents occurred in coal mines located in the United Kingdom (Jarrett, 1979). It is suspected that the application of certain safety regulations had the effect of reducing the number of accidents. Therefore, we are interested in estimating three quantities: the year in which the rate changed (the switch-point), the rate of accidents prior to regulation, and the rate after the regulation change.

The data is shown in Code Block 1, we have the variable disasters that contains the number of accidents per year and the variable years containing the range of years for which we have data. To encode the data, we are using a Pandas Series with np.nan values for the missing data.

 years = np.arange(1851, 1962)

 disaster_data = pd.Series([4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3,

                 3, 4, 0, 2, 6,

                3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5,

                 5, 3, 4, 2, 5,

                2, 2, 3, 4, 2, 1, 3, np.nan, 2,

                 1, 1, 1, 1, 3,

                0, 0, 1, 0, 1, 1, 0, 0, 3, 1, 0,

                 3, 2, 2, 0, 1,

                1, 1, 0, 1, 0, 1, 0, 0, 0, 2, 1,

                 0, 0, 0, 1, 1,

                 0, 2, 3, 3, 1, np.nan, 2, 1, 1,

                 1, 1, 2, 4, 2,

                0, 0, 1, 4, 0, 0, 0, 1, 0, 0, 0,

                 0, 0, 1, 0, 0,

                 1, 0, 1])

Code 1. Data for the coal mining disaster example.

How can we build a model for this problem? One approach is to think generatively, that is, create a story of how data may have been generated. Generative stories can be very powerful, informal devices that aid model construction and understanding (Blitzstein & Hwang, 2019).

We can think of our problem as having a moving slider, years to the left of this slider are assigned an average number of accidents while years to the right are assigned a different average number of accidents. A key property of Bayesian modeling is that we consider multiple plausible scenarios that could have generated the data (see Fig. 4).

Next, we need to specify prior probability distributions which quantify the information we have about plausible parameter values before we have observed any data. For the “slider” we will use a discrete uniform distribution that assigns equal probability to all years in a given interval, although other choices are possible. This distribution has two parameters: the lowest possible value, and the highest one. In our problem, those correspond to the year 1851 and 1962, respectively. A range wider than this does not make sense given that we only have data for this particular range, and a narrower range will imply that we have some external information indicating that not all years between 1851 and 1962 are equally plausible candidates.

Our observed data consists of counts, i.e., the number of disasters. Commonly, the Poisson distribution is used for count data. The Poisson distribution is defined using a single parameter that represents the average rate of events (disasters) in our example. As we do not know the rate and want to estimate it, we have to set a prior distribution for it. We do not have much information other than the rate must be positive and “most likely” have a small value, given that coal mining disasters are not very frequent events. One option could be to pick the distribution $Expon(1)$, which says that the expected rate is 1 but is wide enough to allow for lower and higher values.

Using standard statistical notation, we can write the model as:

And using PyMC we can write this model as described in Code Block 2.

 with pm.Model() as mining:

   # Prior for the switch-point

   sp = pm.DiscreteUniform(′sp′,

                 lower=years.min(),

                 upper=years.max())

   # Priors for the rate before (t_0) and after (t_1)

   t_0 = pm.Exponential(′t_0′, 1)

   t_1 = pm.Exponential(′t_1′, 1)

   # We assign the rates according to sp

   rate = pm.Deterministic(″rate″,

                 pm.math.switch(sp < years, t_0, t_1))

   # Likelihood

   acc = pm.Poisson(′acc′, rate, observed=disasters)

   # Backward and forward sampling

   idata = pm.sample()

   pm.compute_log_likelihood(idata)

   idata.extend(pm.sample_prior_predictive())

   idata.extend(pm.sample_posterior_predictive(idata))

Code 2. PyMC model for the coal mining disaster example.

Once defined in PyMC we can get a visual representation of the model as shown in Fig. 3. We can use this visual representation to check that we do not have a semantic error in the model and to communicate the model to others.

One remarkable feature of PyMC is that its syntax is very close to the statistical notation, as we can see by comparing Eq. (6) with Code Block 2. The cases in Eq. (6) are coded using the pm.math.switch (condition, true, false) function, which uses the first argument to select either of the next two arguments.

Missing values are handled concisely by passing a numpy.ndarray, pandas.DataFrame or pandas.Series with NaN values (see Code Block 1) to the observed argument when creating an observed stochastic random variable. This means that we will automatically get a posterior predictive distribution over the missing values. The imputed and observed values are combined into a Deterministic node acc that represents the original vector specified as an observed random variable.

Dirichlet-multinomial distribution

This example demonstrates the use of a Dirichlet compound multinomial distribution to model categorical count data. Models like this one are important in a variety of areas, including natural language processing (Madsen, Kauchak & Elkan, 2005), ecology (Harrison et al., 2019), and Genomics (Holmes, Harris & Quince, 2012; Nowicka & Robinson, 2016).

The Dirichlet-multinomial can be understood as taking draws from a multinomial distribution where each sample has a slightly different probability vector, which is itself drawn from a common Dirichlet distribution. In contrast with the multinomial distribution, which assumes that all observations arise from a single fixed probability vector. This enables the Dirichlet-multinomial to accommodate over-dispersed count data. Here we will discuss a community ecology example. Let’s assume we have observed counts of $k=5$ different tree species in $n=10$ different forests.

Our data is arranged in a two-dimensional matrix of integers where each row, indexed by $i\in (\mathrm{0...}n-1)$, is an observation (different forest), and each column $j\in (\mathrm{0...}k-1)$ is a category (tree species).

We could write this model as shown in Code Block 4, where we have a multinomial likelihood counts with a Dirichlet prior p. Furthermore, the prior is parameterized in terms of two hyper-priors, a Dirichlet distribution for the expected fraction of each category, frac and a log-normal conc controlling the concentration of the Dirichlet prior or in terms of the data the level of over-dispersion.

 with pm.Model() as model_dm_explicit:

   frac = pm.Dirichlet(″frac″, a=np.ones(k))

   conc = pm.LogNormal(″conc″, mu=1, sigma=1)

   p = pm.Dirichlet(″p″, a=frac ∗ conc)

   counts = pm.Multinomial(″counts″,

                 n=total_count, p=p,

                 observed=observed_counts))

Code 4. An explicit Dirichlet-multinomial distribution model.

The model in Code Block 4 is semantically and syntactically correct. Nevertheless, we can rewrite the model in an equivalent but slightly different form as shown in Code Block 5. The first change is that we are using pm.DirichletMultinomial distribution which takes into account that the Dirichlet distribution is conjugate to the multinomial and therefore there’s a computationally more efficient way to sample from, i.e. using a marginalized closed-form. The second modification with respect to Code Block 4 is that we are going to use labeled coordinates (coords) and dimensions (dims) which provides a tighter integration of PyMC with ArviZ.

 coords = {″tree″: trees, ″forest″: forests}

 with pm.Model(coords=coords) as model_dmm:

   frac = pm.Dirichlet(″frac″, a=np.ones(k), dims=″tree″)

   conc = pm.LogNormal(″conc″, mu=1, sigma=1)

   counts = pm.DirichletMultinomial(″counts″,

         n=total_count, a=frac ∗ conc,

         observed=observed_counts, dims=(″forest″, ″tree″))

   idata_dmm = pm.sample()

Code 5. A marginalized Dirichlet-multinomial model augmented with coordinates and dimensions.

Working with labeled arrays reduces the cognitive load of working with multidimensional arrays, reducing the chance of making errors and reducing frustration. In Code Block 5 by defining the frac Dirichlet parameters using dims="tree" we are guaranteed to have the dimensions of the prior matching the number of trees in our dataset. The advantage of using labels also extends to the post-inference processing stage. For example, idata_dmm.posterior.sel(tree="pine") will return the subset of the posterior related to pine and idata_dmm.posterior_predictive.counts.sel(tree="pine") will do the same for the predictive counts of pine.

Additionally, automatic labeling becomes possible. For instance, after sampling from the model in Code Block 5, calling az.plot_posterior(idata_dmm) generates Fig. 8. Notice how the frac parameter is meaningfully labeled with the names of the trees. The alternative would be integer labels with no intrinsic meaning.

JAX-based sampling

The most recent major version of PyMC is built on top of the PyTensor Python package, which allows the definition, optimization, and efficient evaluation of mathematical expressions involving multi-dimensional arrays. PyMC models, through PyTensor, can be compiled to C, Numba and JAX, and in principle, other computational backends could be added with relatively little effort. This allows for the efficient and fast evaluation of the log-probability density (see Section PyTensor for details). Still, the samplers accessible by calling pm.sample() are coded in Python and NumPy. A good approach to improve the performance of PyMC’s samplers is to write them in PyTensor. This will reduce the overhead of calling Python code and—more importantly—enable a series of optimizations due to PyTensor’s ability to manipulate graphs, including the ability to customize the sampler based on patterns in the model structure. Details of such optimizations are out of the scope of this article and will be discussed in a future manuscript.

An alternative to PyMC’s PyTensor-based samplers is samplers written in JAX. Using these samplers, all the operations needed to compute a posterior can be performed under JAX, reducing the Python overhead during sampling and leveraging all JAX performance improvements and features like the ability to sample on GPUs or TPUs. Currently, PyMC offers NUTS JAX samplers via NumPyro (Phan, Pradhan & Jankowiak, 2019) or BlackJAX (BlackJax devs, 2022) with the functions pm.sample_numpyro_nuts and pm.sample_blackjax_nuts, respectively. Significantly, BlackJAX and NumPyro can both be used because in PyMC the modeling language is decoupled from the inference methods; BlackJAX and NumPyro only require a log-probability density function written in JAX. This demonstrates that samplers can be developed independently of PyMC and then be made available to users of the library.

In the following example, we compare PyMC with its default Python/NumPy NUTS sampler, PyMC running the BlackJAX NUTS sampler, and PyMC running the NumPyro sampler. We also include cmdstanpy (Lee et al., 2017), a command line interface to Stan. The motivation for these comparisons is not to provide an exhaustive benchmark over a wide range of models and datasets but instead to give an example of the attainable speed-ups when using PyMC with a JAX-based sampler.

Suppose we are interested in ranking tennis players from 1968 until now⁵ . To do so, we can use the Bradley-Terry model. The central idea of this model is that each player has a latent skill $\theta$. When players $i$ and $j$ play each other, player $i$ wins with probability $p(i_{\mathrm{w}\mathrm{i}\mathrm{n}\mathrm{s}}\mid {\theta }_i,{\theta }_j)=\mathrm{l}\mathrm{o}\mathrm{g}\mathrm{i}\mathrm{s}\mathrm{t}\mathrm{i}\mathrm{c}({\theta }_i-{\theta }_j)$. For example, if player $i$ has a skill value of 1 and player $j$ has a skill value of −1, then the Bradley-Terry model implies that the player $i$ beats the player $j$ with probability $\mathrm{l}\mathrm{o}\mathrm{g}\mathrm{i}\mathrm{s}\mathrm{t}\mathrm{i}\mathrm{c}(2)\approx$88.1%.

For the priors we set:

when $\sigma =1$, 95% of the players’ skills distribution will fall between −2 and 2, which we consider quite plausible under this model and data, by setting a prior over $\sigma$ we are also not ruling out smaller or larger distributions.

 with pm.Model() as tennis_model:

   sd = pm.HalfNormal(″sd″, sigma=1.0)

   skill_raw = pm.Normal(″skill_raw″, 0.0, sigma=1.0,

                shape=(n_players,)

   )

   skill = pm.Deterministic(″skill″, skill_raw * sd)

    logit_skill = skill[winner_ids] - skill[loser_ids]

   win = pm.Bernoulli(″win″, logit_p=logit_skill,

              observed=np.ones(winner_ids.shape[0])

   )

Code 6. A hierarchical, non-centered Bradley-Terry model for ranking tennis players.

We run the tennis_model in Code Block 6 using the NUTS sampler under six conditions:

• pymc: PyMC with the default sampler

• cmdstanpy: CmdStanPy

• pymc_numpyro_gpu_vectorized: PyMC with NumPyro NUTS sampler on the GPU, running chains using vmap.

• pymc_numpyro_cpu_parallel: PyMC with NumPyro sampler on the CPU, running chains using pmap.

• pymc_blackax_gpu_vectorized: PyMC with BlackJAX NUTS sampler on the GPU, running chains using vmap.

• pymc_blackjax_cpu_parallel: PyMC with BlackJAX NUTS sampler on the CPU, running chains using pmap.

To see how the runtime changes with the size of the dataset, we choose different start years for the fits: 2020, 2019, 2015, 2010, 2000, 1990, 1980, and finally 1968. This means datasets ranging from 3,620 observations to 160,420.

For all conditions, we run 1,000 warm-up steps and 1,000 draws per chain, for a total of four chains.

Figure 9 shows the effective sample size per second for all the samplers previously mentioned. The values are an average of four separate runs. CmdStanPy performs better than PyMC on smaller datasets, but is slower on larger ones. PyMC with either the BlackJAX or NumPyro backends performs the best on the CPU, shown in yellow and magenta respectively. These JAX-based samplers have similar performance. On the other hand, when running on the GPU, the samplers are more efficient on larger datasets. Among the largest dataset, PyMC with NumPyro and BlackJax on the vectorized GPU performs the best, while PyMC with its default sampler and CmdStanPy (both on the CPU) show the worst results.

Tensors

Now we will discuss some core concepts of PyTensor and how they relate to PyMC using a few simple code examples. For a more detailed description of PyTensor please refer to the official documentation (https://PyTensor.readthedocs.io/en/latest/).

To begin, in lines 1 and 2 of Code Block 7 we define two PyTensor tensors, the first one being a scalar (0 dimension) and the second being a vector (1 dimension). On line 3, a new tensor z is created by adding x and y, and then taking the natural logarithm. Then, on line 4 we add a name to the newly created tensor, as this can be useful to easily reference tensors when working with many of them.

It is important to note that contrary to other objects like Python integers, floats or NumPy arrays, the operations described in Code Block 7 do not immediately lead to numerical values. Instead, they are just a symbolic specification of the operations we want to compute. This example literally just defines the abstract expression, take the logarithm of the sum of x and y. To actually perform computations, we first need to define a function, and then call it with the proper input as shown in lines 6 and 7.

 x = pt.scalar(name=″x″)

 y = pt.vector(name=″y″)

 z = pt.log(x + y)

 z.name = ″log(x + y)″

 f = pytensor.function(inputs=[x, y],  outputs=z)

 f(x=0, y=[1,  np.e])

 >>> array([0., 1.])

Code 7. Definition and call of a PyTensor function. Notice that the tensors x, y, and z have been previously defined. When debugging it may be useful to avoid defining a function and instead perform a direct evaluation of the tensor, like z.eval(x: 0, y:[1, np.e]).

This separation of the abstract definitions of mathematical expressions and the actual computation of those expressions is central to PyTensor and hence PyMC. When defining a PyMC model, we are just defining a PyTensor computational graph that we will later use to obtain quantities like prior predictive samples, posterior samples, log-probabilities, etc. This separation is useful as PyTensor can automatically optimize the mathematical operations inside a graph. For example, if we define w = pt.exp(pt.log(x + y)), PyTensor will first simplify the graph to w = x + y and then perform the computation. Other optimizations include constant propagation, replacing numerically unstable operations with numerically stable versions, avoiding computing the same quantity more than once, and efficient sparse matrix multiplication.

Random variables

We now show how to manually generate samples from PyMC distribution and evaluate their log-probability. On line 1 of Code Block 8 we define a Normal distribution with mean 0 and standard deviation 1. On line 2 we take 1,000 draws from that distribution, using the pm.draw(.) function. Finally, on line 3 we use the pm.logp(.) function to compute the log-probability of the samples generated in the previous step. We use the eval() method to obtain the actual values, instead of a symbolic representation. Figure 10 shows the results.

 x = pm.Normal.dist(mu=0, sigma=1)

 x_draws = pm.draw(x, draws=1_000)

 x_logp = pm.logp(rv=x, value=[x_draws]).eval()

Code 8. In line 1 a PyMC distribution is used to specify a symbolic random variable from a corresponding to a Normal distribution. The pm.draw(.) call in line 2 invokes PyTensor graph compilation while accounting for automatic updating of random number generators, and returns an array with 1,000 draws of variable x. Using pm.logp(.), the log-probability densities of x for each element in x_draws are derived symbolically, and then evaluated by the call to the .eval() method.

While the example from Code Block 8 may seem trivial, it is important to note that the variable x which is passed to pm.draw(.) or pm.logp(.) can be the result of a symbolic computation involving multiple random variables and other tensors. A simple example is given in Code Block 9 where the random variable b depends on another random variable a, and the variable x is a tensor variable that merely depends on other variables, some of which represent random variables. To a much larger extent, this is how computational graphs of PyMC models are handled behind the scenes, and how users can access properties such as the log-probability densities of probabilistic graphs built with PyTensor.

 a = pm.Uniform.dist()

 b = pm.Normal.dist(mu=a, sigma=1)

 x = b + pt.as_tensor([2, 3, 4])

 x_draws = pm.draw(x, draws=1_000)

 x_logp = pm.logp(rv=x, value=[x_draws]).eval()

Code 9. Random variables such as the ones defined in lines 1 and 2 behave like tensor variables and can be used as such in standard operations such as addition (line 3). The sampling and derivation of log-probabilities (lines 4 and 5) can operate on any graph that involves random variables.