分布维度

PyMC 提供了多种方法来指定其分布的维度。本文档提供了概述，并提供了一些用户提示。

术语表

在本文档中，我们将使用术语“维度”来指代维度的概念。以下每个术语在 PyMC 中都有特定的语义和计算定义。虽然我们在此处分享了它们，但在下面的示例中查看时，它们会更有意义。

• 支持维度 → 分布的核心维度

• 批次维度 → 超出分布支持维度的额外维度

• 隐式维度 → 从分布参数的值或形状推导出的维度

• 显式维度 → 由以下参数之一显式定义的维度

  • 形状 → 从分布中抽取的样本数量

  • Dims → 维度名称数组

• Coords → 将维度名称映射到坐标值的字典

from functools import partial

import numpy as np
import pytensor.tensor as pt

import pymc as pm

单变量分布示例

我们可以从最简单的情况开始，即单个正态分布。我们使用 .dist 在 PyMC 模型之外指定一个分布。

normal_dist = pm.Normal.dist()

然后，我们可以使用 draw() 函数从该分布中抽取一个随机样本。

# Just patching the draw function for reproducibility
rng = np.random.default_rng(seed=sum(map(ord, "dimensionality")))
draw = partial(pm.draw, random_seed=rng)

normal_draw = draw(normal_dist)
normal_draw, normal_draw.ndim

(array(0.80189558), 0)

在这种情况下，我们最终得到一个标量值。这意味着正态分布具有标量支持维度，因为您可以提取的最小随机样本是标量，其维度为零。每个分布的支持维度都硬编码为一个属性。

normal_dist.owner.op.ndim_supp

0

显式批次维度

如果需要多个样本，自然而然的倾向是创建相同变量的多个副本并将它们堆叠在一起。

normal_dists = pm.math.stack([pm.Normal.dist() for _ in range(3)])
draw(normal_dists)

array([ 0.83636296, -0.33327414,  0.9434115 ])

更简单地说，可以使用 shape 参数从同一分布族创建一批独立的样本。

normal_dists = pm.Normal.dist(shape=(3,))
draw(normal_dists)

array([ 0.98810294, -0.07003785, -0.37962748])

normal_dists = pm.Normal.dist(shape=(4, 3))
draw(normal_dists)

array([[ 7.99932116e-04, -1.94407945e+00,  3.90451962e-01],
       [ 1.10657367e+00,  6.49042149e-01, -1.09450185e+00],
       [-2.96226305e-01,  1.41884595e+00, -1.31589441e+00],
       [ 1.53109449e+00, -7.73771737e-01,  2.37418367e+00]])

这不仅更简洁，而且产生效率更高的向量化代码。除非我们需要组合来自不同分布族的样本，否则我们很少在 PyMC 中使用堆叠方法。

隐式批次维度

也可以通过传递具有更高维度的参数来创建一批样本，而无需指定形状。

normal_dists = pm.Normal.dist(mu=np.array([0, 0, 0]), sigma=np.array([1, 1, 1]))
draw(normal_dists)

array([ 0.81833093, -0.2891973 ,  1.2399946 ])

这等效于前面的显式形状示例，我们也可以在此处显式传递它。因为我们没有这样做，所以我们将这些批次维度称为隐式维度。

当我们希望参数在批次维度上变化时，这变得非常有用。

draw(pm.Normal.dist(mu=[1, 10, 100], sigma=0.0001))

array([  0.99989975,  10.00009874, 100.00004215])

当参数的形状不同时，它们会像 NumPy 的工作方式一样进行广播。在本例中，sigma 被广播以匹配 mu 的形状。

np.broadcast_arrays([1, 10, 100], 0.0001)

[array([  1,  10, 100]), array([0.0001, 0.0001, 0.0001])]

重要的是要理解 NumPy 广播 的工作原理。当您执行无效操作时，您很容易遇到此类错误

try:
    # shapes of (3,) and (2,) can't be broadcasted together
    draw(pm.Normal.dist(mu=[1, 10, 100], sigma=[0.1, 0.1]))
except ValueError as error:
    print(error)

Could not broadcast dimensions. Incompatible shapes were [(ScalarConstant(ScalarType(int64), data=3),), (ScalarConstant(ScalarType(int64), data=2),)].

组合隐式和显式批次维度

您可以将显式形状维度与隐式批次维度组合使用。如上所述，它们可以提供相同的信息。

normal_dists = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(3,))
draw(normal_dists)

array([ 0.06413633,  1.29893485, -0.48072495])

但是 shape 也可以用于扩展超出任何隐式批次维度。

normal_dists = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(4, 3))
draw(normal_dists)

array([[-0.49526775, -0.94608062,  1.66397913],
       [ 0.703617  ,  0.66713031,  0.80725231],
       [ 0.19219926,  1.62987906,  2.30590873],
       [ 1.83763939, -0.19878079,  1.46751553]])

请注意，由于广播规则，显式批次维度必须始终“位于”任何隐式维度的“左侧”。因此，在前面的示例中，shape=(4, 3) 是有效的，但 shape=(3, 4) 无效，因为 mu 参数可以广播到第一个形状，但不能广播到第二个形状。

try:
    draw(pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(3, 4)))
except ValueError as error:
    print(error)

shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (3, 4) and arg 1 with shape (1, 3).
Apply node that caused the error: normal_rv{"(),()->()"}(RNG(<Generator(PCG64) at 0x7FB323BFA0A0>), [3 4], [[0 1 2]], [[1]])
Toposort index: 0
Inputs types: [RandomGeneratorType, TensorType(int64, shape=(2,)), TensorType(int64, shape=(1, 3)), TensorType(int8, shape=(1, 1))]
Inputs shapes: ['No shapes', (2,), (1, 3), (1, 1)]
Inputs strides: ['No strides', (8,), (24, 8), (1, 1)]
Inputs values: [Generator(PCG64) at 0x7FB323BFA0A0, array([3, 4]), array([[0, 1, 2]]), array([[1]], dtype=int8)]
Outputs clients: [['output'], ['output']]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

如果您需要正态变量具有 shape=(3, 4)，您可以在定义后转置它。

transposed_normals = pm.Normal.dist(mu=np.array([0, 1, 2]), sigma=1, shape=(4, 3)).T
draw(transposed_normals)

array([[ 1.36252056,  0.90337366, -1.83306938, -1.04031058],
       [ 0.09757005, -0.03093604,  3.29729122, -0.86869013],
       [ 3.51136436, -0.33437459,  1.93223367,  3.71535763]])

提示

重要的是不要混淆在分布定义中设置的维度与在下游操作（如转置、索引或广播）中设置的维度。当使用 PyMC 进行采样时（无论是通过前向采样还是 MCMC），随机样本将始终来自分布形状。请注意，在以下示例中，尽管两个变量具有相同的最终形状，但实际采样的“随机”样本数量却不同。

vector_normals = pm.Normal.dist(shape=(3,))
broadcasted_normal = pt.broadcast_to(pm.Normal.dist(), (3,))

draw(vector_normals), draw(broadcasted_normal)

(array([-0.73397401,  2.54543846, -1.14705529]),
 array([-0.45755879, -0.45755879, -0.45755879]))

多变量分布示例

某些分布根据定义在评估时返回多个值。这可能是值向量、矩阵或任意多维张量。一个例子是多元正态分布，它始终返回一个向量（一维数组）。

mvnormal_dist = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3))
mvnormal_draw = draw(mvnormal_dist)
mvnormal_draw, mvnormal_draw.ndim

(array([1.29866199, 1.01091254, 0.08414986]), 1)

与任何分布一样，支持维度被指定为固定属性

mvnormal_dist.owner.op.ndim_supp

1

即使您指定具有单个维度的 MvNormal，您也会得到一个向量！

smallest_mvnormal_dist = pm.MvNormal.dist(mu=[1], cov=[[1]])
smallest_mvnormal_draw = draw(smallest_mvnormal_dist)
smallest_mvnormal_draw, smallest_mvnormal_draw.ndim

(array([0.55390975]), 1)

隐式支持维度

在我们刚刚看到的 MvNormal 示例中，支持维度实际上是隐式的。我们没有在任何地方指定我们想要 3 个或 1 个样本的向量。这是从 mu 和 cov 的形状推断出来的。因此，我们将其称为隐式支持维度。我们可以通过使用形状来更明确地说明。

explicit_mvnormal = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), shape=(3,))
draw(explicit_mvnormal)

array([-0.68893796,  1.10911095, -0.30443374])

警告

但是，请注意，在撰写本文时，形状对于支持维度而言被简单地忽略了。它仅作为标记预期维度的“类型提示”。

ignored_shape_mvnormal = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), shape=(4,))
draw(ignored_shape_mvnormal)

array([0.57262853, 0.34230354, 1.96818163])

显式批次维度

与单变量分布一样，我们可以添加显式批次维度。我们将使用另一个向量分布来说明这一点：多项分布。以下代码片段定义了一个由五个独立多项分布组成的矩阵，每个分布都是大小为 3 的向量。

draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 3)))

array([[0, 2, 3],
       [0, 2, 3],
       [1, 0, 4],
       [0, 1, 4],
       [0, 1, 4]])

警告

同样，请注意，形状对支持维度没有影响

draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 4)))

array([[2, 0, 3],
       [1, 1, 3],
       [0, 2, 3],
       [0, 1, 4],
       [1, 0, 4]])

出于同样的原因，您必须始终将显式批次维度定义在支持维度的“左侧”。以下代码的行为将不如预期。

draw(pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(3, 5)))

array([[0, 1, 4],
       [0, 0, 5],
       [3, 1, 1]])

如果您需要多项变量具有 shape=(3, 5)，您可以在定义后转置它。

transposed_multinomials = pm.Multinomial.dist(n=5, p=[0.1, 0.3, 0.6], shape=(5, 3)).T
draw(transposed_multinomials)

array([[2, 1, 1, 0, 2],
       [0, 3, 1, 0, 1],
       [3, 1, 3, 5, 2]])

隐式批次维度

与单变量分布一样，我们可以为每个批次维度使用不同的参数

multinomial_dist = pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6])
draw(multinomial_dist)

array([[0, 2, 3],
       [1, 4, 5]])

这等效于更冗长的

draw(pm.Multinomial.dist(n=[5, 10], p=[[0.1, 0.3, 0.6], [0.1, 0.3, 0.6]]))

array([[1, 2, 2],
       [0, 3, 7]])

如果您熟悉 NumPy 广播规则，您可能会好奇 PyMC 是如何实现这一点的。朴素广播在这里不起作用

try:
    np.broadcast_arrays([5, 10], [0.1, 0.3, 0.6])
except ValueError as exc:
    print(exc)

shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (2,) and arg 1 with shape (3,).

为了理解发生了什么，我们需要引入参数核心维度的概念。分布参数的核心维度是参数定义分布所需的最小维度数。在多项分布中，n 必须至少是标量整数，但 p 必须至少是表示每个类别结果概率的向量。因此，对于多项分布，n 具有 0 个核心维度，而 p 具有 1 个核心维度。

因此，如果我们有一个由两个 n 组成的向量，我们实际上应该将 p 向量广播到具有两个此类向量的矩阵中，并将每个 n 与 p 的每个广播行配对。这与 np.vectorize 的工作方式完全相同。

def core_multinomial(n, p):
    print(">>", n, p)
    return draw(pm.Multinomial.dist(n, p))


vectorized_multinomial = np.vectorize(core_multinomial, signature="(),(p)->(p)")
vectorized_multinomial([5, 10], [0.1, 0.3, 0.6])

>> 5 [0.1 0.3 0.6]
>> 10 [0.1 0.3 0.6]

array([[2, 2, 1],
       [1, 1, 8]])

每个分布参数的核心维度也硬编码为每个分布的属性

multinomial_dist.owner.op.ndims_params

[0, 1]

ndim_supp 和 ndims_params 实际上都是从类似 numpy 的签名中提取的

multinomial_dist.owner.op.signature

'(),(p)->(p)'

隐式批次维度仍然必须遵守广播规则。以下示例无效，因为 n 的批次维度为 shape=(2,)，而 p 的批次维度为 shape=(3,)，它们无法一起广播。

try:
    draw(pm.Multinomial.dist(n=[5, 10], p=[[0.1, 0.3, 0.6], [0.1, 0.3, 0.6], [0.1, 0.3, 0.6]]))
except ValueError as error:
    print(error)

Could not broadcast dimensions. Incompatible shapes were [(ScalarConstant(ScalarType(int64), data=2),), (ScalarConstant(ScalarType(int64), data=3),)].

组合隐式和显式批次维度

您可以并且应该将来自多维参数的隐式维度与显式形状信息组合使用，这更容易理解。

draw(pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6], shape=(2, 3)))

array([[1, 1, 3],
       [2, 1, 7]])

显式批次维度仍然可以扩展超出任何隐式批次维度。同样，由于广播的工作方式，显式批次维度必须始终“位于左侧”。以下情况无效，因为 n 的批次维度为 shape=(2,)，无法广播到 shape=(2, 4) 的显式批次维度。

try:
    draw(pm.Multinomial.dist(n=[5, 10], p=[0.1, 0.3, 0.6], shape=(2, 4, 3)))
except ValueError as error:
    print(error)

operands could not be broadcast together with remapped shapes [original->remapped]: (1,2)  and requested shape (2,4)
Apply node that caused the error: multinomial_rv{"(),(p)->(p)"}(RNG(<Generator(PCG64) at 0x7FB323BF9460>), [2 4], [[ 5 10]], [[[0.1 0.3 0.6]]])
Toposort index: 0
Inputs types: [RandomGeneratorType, TensorType(int64, shape=(2,)), TensorType(int64, shape=(1, 2)), TensorType(float64, shape=(1, 1, 3))]
Inputs shapes: ['No shapes', (2,), (1, 2), (1, 1, 3)]
Inputs strides: ['No strides', (8,), (16, 8), (24, 24, 8)]
Inputs values: [Generator(PCG64) at 0x7FB323BF9460, array([2, 4]), array([[ 5, 10]]), array([[[0.1, 0.3, 0.6]]])]
Outputs clients: [['output'], ['output']]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

使用模型图检查维度

分布通常在 PyMC 模型内部使用，因此有一些工具可以帮助推理该上下文中的分布形状。

with pm.Model() as pmodel:
    mu = pm.Normal("x", mu=0, shape=(3))
    sigma = pm.HalfNormal("sigma")
    y = pm.Normal("y", mu=mu, sigma=sigma)

for rv, shape in pmodel.eval_rv_shapes().items():
    print(f"{rv:>11}: shape={shape}")

          x: shape=(3,)
sigma_log__: shape=()
      sigma: shape=()
          y: shape=(3,)

理解和调试 PyMC 中维度的更强大的工具是 model_to_graphviz() 函数。我们可以读取 Graphviz 输出而不是检查数组输出，以了解变量的维度。

pm.model_to_graphviz(pmodel)

在上面的示例中，每个框（或板）左下角的数字表示其中分布的维度。如果分布在任何带有数字的框之外，则它具有标量形状。

让我们使用此工具来回顾隐式和显式维度

with pm.Model() as pmodel:
    pm.Normal("scalar (support)")
    pm.Normal("vector (implicit)", mu=[1, 2, 3])
    pm.Normal("vector (explicit)", shape=(4,))

pm.model_to_graphviz(pmodel)

Dims

PyMC 支持 dims 的概念。对于许多随机变量，哪个维度对应于哪个“真实世界”概念可能会变得令人困惑，例如，观察次数、处理单元数等。dims 参数是一个额外的人类可读标签，可以传达此含义。当单独使用时，dims 必须与显式 shape 信息结合使用。

with pm.Model() as model:
    pm.Normal("profit", shape=(3,), dims="year")

pm.model_to_graphviz(model)

当在模型本身中指定 coords 时，dims 可以变得越来越强大。这为每个 dim 条目提供了唯一的标签，使其更有意义。

coords = {"year": [2020, 2021, 2022]}
with pm.Model(coords=coords) as model:
    pm.Normal("profit", dims="year")

pm.model_to_graphviz(model)

在这种情况下，分布的维度实际上可以由使用的 dims 定义。我们不必传递 shape 或定义隐式批次维度。

让我们回顾一下多元正态分布示例的不同维度风格。

coords = {
    "batch": [0, 1, 2, 3],
    "support": ["A", "B", "C"],
}
with pm.Model(coords=coords) as model:
    pm.MvNormal("vector", mu=[0, 0, 0], cov=np.eye(3), dims=("support",))
    pm.MvNormal("matrix (implicit)", mu=np.zeros((4, 3)), cov=np.eye(3), dims=("batch", "support"))
    pm.MvNormal(
        "matrix (explicit)", mu=[0, 0, 0], cov=np.eye(3), shape=(4, 3), dims=("batch", "support")
    )

pm.model_to_graphviz(model)

提示

对于最终模型发布，我们建议使用 dims 和 coords，因为标签将传递给 arviz.InferenceData。这既是最佳实践透明度，也是其他人的可读性。它在单个开发人员工作流程中也很有用，例如，在存在 3 维或更高维度分布的情况下，它将有助于指示哪个维度对应于哪个模型概念。

调试形状问题的提示

虽然我们提供了所有这些工具以方便使用，并且 PyMC 尽最大努力理解用户意图，但混合维度工具的结果可能并不总是导致预期的最终维度。有时，模型可能在采样之前不会指示错误，或者根本不指示问题。当处理维度时，特别是更复杂的维度时，我们建议

• 在采样前使用 model_to_graphviz 可视化您的模型

• 使用 draw 或 sample_prior predictive 尽早捕获错误

• 检查返回的 az.InferenceData 对象以确保所有数组大小均符合预期

• 在跟踪错误时使用质数定义形状。

%load_ext watermark

%watermark -n -u -v -iv -w -p pytensor

Last updated: Tue Jun 25 2024

Python implementation: CPython
Python version       : 3.11.8
IPython version      : 8.22.2

pytensor: 2.20.0+3.g66439d283.dirty

pymc    : 5.15.0+1.g58927d608
numpy   : 1.26.4
pytensor: 2.20.0+3.g66439d283.dirty

Watermark: 2.4.3