通用 API 快速入门#

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pytensor.tensor as pt
/home/xian/anaconda3/envs/pymc-dev-py39/lib/python3.9/site-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: main is an invalid version and will not be supported in a future release
  warnings.warn(
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

1. 模型创建#

PyMC 中的模型围绕 Model 类展开。它包含所有随机变量 (RVs) 的引用,并计算模型的对数概率及其梯度。通常,您会将其作为 with 上下文的一部分进行实例化:

with pm.Model() as model:
    # Model definition
    pass

我们在下面进一步讨论RVs,但让我们创建一个简单的模型来探索Model类。

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=1, observed=rng.standard_normal(100))
model.basic_RVs
[mu ~ N(0, 1), obs ~ N(mu, 1)]
model.free_RVs
[mu ~ N(0, 1)]
model.observed_RVs
[obs ~ N(mu, 1)]
model.compile_logp()({"mu": 0})
array(-143.03962875)

值得强调的是我们在logp中做出的设计选择。如上所示,logp正在被调用并传递参数,因此它是模型实例的一个方法。更准确地说,它根据模型的当前状态或作为logp参数给出的状态组合一个函数(见下面的示例)。

由于多种原因,我们假设一个Model实例不是静态的。如果你需要在内部循环中使用logp并且它需要是静态的,只需使用类似logp = model.logp的方式。下面是一个示例——注意缓存效果和速度提升:

%timeit model.compile_logp()({"mu": 0.1})
logp = model.compile_logp()
%timeit logp({"mu": 0.1})
83 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
18 µs ± 276 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

2. 概率分布#

每个概率程序都由观测到的和未观测到的随机变量(RVs)组成。观测到的随机变量通过似然分布定义,而未观测到的随机变量通过先验分布定义。在PyMC模块中,概率分布的结构如下所示:

分布

  • pymc:api_distributions_continuous

  • pymc:api_distributions_discrete

  • pymc:api_distributions_multivariate

  • pymc:api_distributions_mixture

  • pymc:api_distributions_timeseries

  • pymc:api_distributions_censored

  • pymc:api_distributions_simulator

未观测的随机变量#

每个未观察到的随机变量都有以下调用签名:名称(字符串),参数关键字参数。因此,可以在模型上下文中这样定义一个正态先验:

with pm.Model():
    x = pm.Normal("x", mu=0, sigma=1)

与模型一样,我们可以评估其对数概率:

pm.logp(x, 0).eval()
array(-0.91893853)

观察到的随机变量#

观测到的RVs与未观测到的RVs定义方式相同,但需要将数据传递到observed关键字参数中:

with pm.Model():
    obs = pm.Normal("x", mu=0, sigma=1, observed=rng.standard_normal(100))

observed 支持列表、numpy.ndarraypytensor 数据结构。

确定性转换#

PyMC 允许你以各种方式自由地对随机变量进行代数运算:

with pm.Model():
    x = pm.Normal("x", mu=0, sigma=1)
    y = pm.Gamma("y", alpha=1, beta=1)
    plus_2 = x + 2
    summed = x + y
    squared = x**2
    sined = pm.math.sin(x)

尽管这些变换可以无缝工作,但它们的结果并不会自动存储。因此,如果你想跟踪一个变换后的变量,你必须使用pm.Deterministic

with pm.Model():
    x = pm.Normal("x", mu=0, sigma=1)
    plus_2 = pm.Deterministic("x plus 2", x + 2)

请注意,plus_2 可以以与上述相同的方式使用,我们只是告诉 PyMC 为我们跟踪这个随机变量。

随机变量列表 / 高维随机变量#

上面我们已经看到了如何创建标量随机变量。在许多模型中,我们希望有多个随机变量。用户有时会尝试创建随机变量列表,如下所示:

with pm.Model():
    # bad:
    x = [pm.Normal(f"x_{i}", mu=0, sigma=1) for i in range(10)]

这可以工作,但它很慢且不推荐使用。相反,我们可以使用坐标

coords = {"cities": ["Santiago", "Mumbai", "Tokyo"]}
with pm.Model(coords=coords) as model:
    # good:
    x = pm.Normal("x", mu=0, sigma=1, dims="cities")

x 现在是一个长度为3的数组,并且这个数组中的每个变量都与一个标签相关联。这将使我们在查看结果时非常容易区分这3个不同的变量。我们可以对这个数组进行索引或进行线性代数运算:

with model:
    y = x[0] * x[1]  # indexing is supported
    x.dot(x.T)  # linear algebra is supported

初始化随机变量#

尽管PyMC会自动初始化模型,但有时为随机变量(RVs)定义初始值是有帮助的。这可以通过initval关键字参数来实现:

with pm.Model(coords={"idx": np.arange(5)}) as model:
    x = pm.Normal("x", mu=0, sigma=1, dims="idx")

model.initial_point()
{'x': array([0., 0., 0., 0., 0.])}
with pm.Model(coords={"idx": np.arange(5)}) as model:
    x = pm.Normal("x", mu=0, sigma=1, dims="idx", initval=rng.standard_normal(5))

model.initial_point()
{'x': array([-0.36012097, -0.16168135,  1.07485641, -0.08855632, -0.03857412])}

当尝试识别模型规范或初始化中的问题时,这种技术有时很有用。

3. 推理#

一旦我们定义了模型,我们必须进行推理以近似后验分布。PyMC 支持两种广泛的推理方法:采样和变分推理。

3.1 采样#

MCMC采样算法的主要入口是通过pm.sample()函数。默认情况下,此函数会尝试自动分配正确的采样器。pm.sample()返回一个arviz.InferenceData对象。InferenceData对象可以轻松地从文件中保存/加载,并且可以携带额外的(元)数据,如日期/版本和后验预测样本。查看ArviZ快速入门以了解更多信息。

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=1, observed=rng.standard_normal(100))

    idata = pm.sample(2000)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [mu]
100.00% [6000/6000 00:03<00:00 Sampling 2 chains, 0 divergences]
Sampling 2 chains for 1_000 tune and 2_000 draw iterations (2_000 + 4_000 draws total) took 4 seconds.

如您所见,对于仅包含连续变量的模型,PyMC 会分配 NUTS 采样器,即使在复杂模型中也非常高效。PyMC 还会运行初始调整以找到采样器的好起点参数。这里我们从后验中抽取 2000 个样本,并允许采样器在额外的 1500 次迭代中调整其参数。

如果没有通过 chains 参数设置,链的数量将根据可用 CPU 核心的数量来确定。

idata.posterior.dims
Frozen({'chain': 2, 'draw': 2000})

默认情况下,调优样本会被丢弃。通过设置discard_tuned_samples=False,它们可以被保留,并最终出现在InferenceData对象中的一个单独的组中(即,idata.warmup_posterior)。

您可以使用 chainscores 参数来控制链的并行运行方式:

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=1, observed=rng.standard_normal(100))

    idata = pm.sample(cores=4, chains=6)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (6 chains in 4 jobs)
NUTS: [mu]
100.00% [12000/12000 00:07<00:00 Sampling 6 chains, 0 divergences]
Sampling 6 chains for 1_000 tune and 1_000 draw iterations (6_000 + 6_000 draws total) took 7 seconds.
idata.posterior["mu"].shape
(6, 1000)
# get values of a single chain
idata.posterior["mu"].sel(chain=2).shape
(1000,)

3.2 分析采样结果#

分析采样结果时最常用的图是所谓的轨迹图:

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    sd = pm.HalfNormal("sd", sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=sd, observed=rng.standard_normal(100))

    idata = pm.sample()
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [mu, sd]
100.00% [4000/4000 00:03<00:00 Sampling 2 chains, 0 divergences]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 4 seconds.
az.plot_trace(idata);
../_images/74df7081244e2ea0377564e6ded756f11702b7ab07a138ebbf3fb1a5d14d81ac.png

另一个常见的指标是Gelman-Rubin统计量,或R-hat:

az.summary(idata)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
mu -0.019 0.096 -0.193 0.163 0.002 0.002 1608.0 1343.0 1.0
sd 0.967 0.069 0.835 1.089 0.002 0.001 1836.0 1406.0 1.0

R-hat 也作为 az.plot_forest 的一部分呈现:

az.plot_forest(idata, r_hat=True);
../_images/bd7fcac63f54b2492d34096ec1114c8b06d920e82dfd2816ce0b473d03dafc9b.png

最后,对于受[Kruschke, 2014]启发的后验图,您可以使用:

对于高维模型,查看所有参数的轨迹变得繁琐。当使用NUTS时,我们可以查看能量图来评估收敛问题:

with pm.Model(coords={"idx": np.arange(100)}) as model:
    x = pm.Normal("x", mu=0, sigma=1, dims="idx")
    idata = pm.sample()

az.plot_energy(idata);
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [x]
100.00% [4000/4000 00:04<00:00 Sampling 2 chains, 0 divergences]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 5 seconds.
../_images/8422788abb1fab86798d0ba9ff90652e9ca219320c97242d6e0fa906139c14a4.png

有关采样器统计信息和能量图的更多信息,请参阅采样器统计信息。有关识别采样问题及其解决方法的更多信息,请参阅使用分歧诊断偏差推理

3.3 变分推断#

PyMC 支持各种变分推断技术。虽然这些方法速度更快,但它们通常也准确性较低,并且可能导致有偏的推断。主要的入口点是 pymc.fit()

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    sd = pm.HalfNormal("sd", sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=sd, observed=rng.standard_normal(100))

    approx = pm.fit()
100.00% [10000/10000 00:00<00:00 Average Loss = 142.01]
Finished [100%]: Average Loss = 142

返回的 Approximation 对象具有多种功能,例如从近似后验中抽取样本,我们可以像常规采样运行一样对其进行分析:

idata = approx.sample(1000)
az.summary(idata)
arviz - WARNING - Shape validation failed: input_shape: (1, 1000), minimum_shape: (chains=2, draws=4)
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
mu -0.023 0.169 -0.338 0.296 0.005 0.004 973.0 880.0 NaN
sd 0.989 0.158 0.694 1.262 0.005 0.004 972.0 1026.0 NaN

The variational 子模块提供了很大的灵活性,可以选择使用哪种变分推断(VI),并遵循面向对象的设计。例如,全秩ADVI估计一个完整的协方差矩阵:

mu = pm.floatX([0.0, 0.0])
cov = pm.floatX([[1, 0.5], [0.5, 1.0]])
with pm.Model(coords={"idx": np.arange(2)}) as model:
    pm.MvNormal("x", mu=mu, cov=cov, dims="idx")
    approx = pm.fit(method="fullrank_advi")
100.00% [10000/10000 00:03<00:00 Average Loss = 0.013113]
Finished [100%]: Average Loss = 0.012772

使用面向对象接口的等效表达式是:

with pm.Model(coords={"idx": np.arange(2)}) as model:
    pm.MvNormal("x", mu=mu, cov=cov, dims="idx")
    approx = pm.FullRankADVI().fit()
100.00% [10000/10000 00:03<00:00 Average Loss = 0.020591]
Finished [100%]: Average Loss = 0.020531
with pm.Model(coords={"idx": np.arange(2)}) as model:
    pm.MvNormal("x", mu=mu, cov=cov, dims="idx")
    approx = pm.FullRankADVI().fit()
100.00% [10000/10000 00:03<00:00 Average Loss = 0.014234]
Finished [100%]: Average Loss = 0.014143
plt.figure()
idata = approx.sample(10000)
az.plot_pair(idata, var_names="x", coords={"idx": [0, 1]});
<Figure size 432x288 with 0 Axes>
../_images/fc6b5bd7b4db1fd17c9a8bfec9fcbaf36039b9ece654beb3fb051b03be3d13c8.png

Stein变分梯度下降(SVGD)使用粒子来估计后验分布:

w = pm.floatX([0.2, 0.8])
mu = pm.floatX([-0.3, 0.5])
sd = pm.floatX([0.1, 0.1])
with pm.Model() as model:
    pm.NormalMixture("x", w=w, mu=mu, sigma=sd)
    approx = pm.fit(method=pm.SVGD(n_particles=200, jitter=1.0))
100.00% [10000/10000 01:20<00:00]
with pm.Model() as model:
    pm.NormalMixture("x", w=[0.2, 0.8], mu=[-0.3, 0.5], sigma=[0.1, 0.1])
plt.figure()
idata = approx.sample(10000)
az.plot_dist(idata.posterior["x"]);
../_images/434538d8660bf2399ebf9df11cbd2b7cec62d8abafc588da625315074b628118.png

有关变分推断的更多信息,请参阅变分推断

4. 后验预测采样#

函数 sample_posterior_predictive() 对保留数据进行预测并执行后验预测检查。

data = rng.standard_normal(100)
with pm.Model() as model:
    mu = pm.Normal("mu", mu=0, sigma=1)
    sd = pm.HalfNormal("sd", sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=sd, observed=data)

    idata = pm.sample()
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [mu, sd]
100.00% [4000/4000 00:03<00:00 Sampling 2 chains, 0 divergences]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 4 seconds.
with model:
    idata.extend(pm.sample_posterior_predictive(idata))
100.00% [2000/2000 00:00<00:00]
fig, ax = plt.subplots()
az.plot_ppc(idata, ax=ax)
ax.axvline(data.mean(), ls="--", color="r", label="True mean")
ax.legend(fontsize=10);
/home/xian/anaconda3/envs/pymc-dev-py39/lib/python3.9/site-packages/IPython/core/pylabtools.py:151: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
  fig.canvas.print_figure(bytes_io, **kw)
../_images/39efd09b12953741f958e0a4e30c960050d4f5134413a47538ee821e2f1d7766.png

4.1 对保留数据的预测#

在许多情况下,您希望对未见过的/保留的数据进行预测。这在概率机器学习和贝叶斯深度学习中尤为重要。PyMC 包含一个 pm.MutableData 容器,以帮助实现这些用途。它是一个围绕 pytensor.shared 变量的包装器,并允许稍后更改数据的值。否则,pm.MutableData 对象可以像任何其他 numpy 数组或张量一样使用。

这一区别非常重要,因为PyMC中的所有模型在内部都是巨大的符号表达式。当你直接将原始数据传递到模型中时,你实际上是允许PyTensor将这些数据视为常量,并在这样做有意义的情况下对其进行优化。如果你以后需要更改这些数据,你可能无法在更大的符号表达式中指向它。使用pm.MutableData提供了一种在符号表达式中指向特定位置并更改其内容的方法。

x = rng.standard_normal(100)
y = x > 0

coords = {"idx": np.arange(100)}
with pm.Model() as model:
    # create shared variables that can be changed later on
    x_obs = pm.MutableData("x_obs", x, dims="idx")
    y_obs = pm.MutableData("y_obs", y, dims="idx")

    coeff = pm.Normal("x", mu=0, sigma=1)
    logistic = pm.math.sigmoid(coeff * x_obs)
    pm.Bernoulli("obs", p=logistic, observed=y_obs, dims="idx")
    idata = pm.sample()
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [x]
100.00% [4000/4000 00:03<00:00 Sampling 2 chains, 0 divergences]
Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 3 seconds.

现在假设我们要对未见过的数据进行预测。为此,我们需要更改 x_obsy_obs 的值。理论上我们不需要设置 y_obs,因为我们想要预测它,但它必须与 x_obs 的形状相匹配。

with model:
    # change the value and shape of the data
    pm.set_data(
        {
            "x_obs": [-1, 0, 1.0],
            # use dummy values with the same shape:
            "y_obs": [0, 0, 0],
        },
        coords={"idx": [1001, 1002, 1003]},
    )

    idata.extend(pm.sample_posterior_predictive(idata))
100.00% [2000/2000 00:00<00:00]
idata.posterior_predictive["obs"].mean(dim=["draw", "chain"])
<xarray.DataArray 'obs' (idx: 3)>
array([0.0215, 0.488 , 0.982 ])
Coordinates:
  * idx      (idx) int64 1001 1002 1003

参考资料#

[1]

约翰·克鲁施克。贝叶斯数据分析:R、JAGS和Stan的教程。学术出版社,2014年。

%load_ext watermark
%watermark -n -u -v -iv -w -p pytensor,aeppl,xarray
Last updated: Fri Jun 03 2022

Python implementation: CPython
Python version       : 3.9.13
IPython version      : 8.4.0

pytensor: 2.6.2
aeppl : 0.0.31
xarray: 2022.3.0

arviz     : 0.12.1
numpy     : 1.22.4
pymc      : 4.0.0b6
pytensor    : 2.6.2
matplotlib: 3.5.2

Watermark: 2.3.1

许可证声明#

本示例库中的所有笔记本均在MIT许可证下提供,该许可证允许修改和重新分发,前提是保留版权和许可证声明。

引用 PyMC 示例#

要引用此笔记本,请使用Zenodo为pymc-examples仓库提供的DOI。

重要

许多笔记本是从其他来源改编的:博客、书籍……在这种情况下,您应该引用原始来源。

同时记得引用代码中使用的相关库。

这是一个BibTeX的引用模板:

@incollection{citekey,
  author    = "<notebook authors, see above>",
  title     = "<notebook title>",
  editor    = "PyMC Team",
  booktitle = "PyMC examples",
  doi       = "10.5281/zenodo.5654871"
}

渲染后可能看起来像: