PyMC：PyMCでのDistributionとShape

はじめに

モデルを構築する上で基本となる確率分布の使い方を見ていきます。

インポート

import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import arviz as az

import pymc as pm
import pytensor.tensor as pt


# Initialize random number generator
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")

分布の定義

以下のようなコードで平均０、標準偏差0.5の正規分布からサンプリングします。

d = pm.Normal.dist(mu=0, sigma=0.5, shape=(1000,))
samples = pm.draw(d)

sns.displot(samples)

また、pm.logp(distribution, value)により任意の値に対して生成確率の対数値を計算することができます。

np.exp(pm.logp(d, 0).eval())

0.7978845608028654

distributionと次元

分布の次元としてはSupport dimensionsとBatch dimensionsがあります。

Support dimensions

確率分布ごとに固有の次元で、分布が出力するサンプルの種類を意味しています。例えば、サンプリングされる値がスカラーの場合は0になりますし、ベクトルの場合は1になります。ここで「1」というのはベクトルの次元ではなく、1次元の配列という意味です。同様に、行列の場合は2で、３次元のテンソルの場合は3になります。
以下例を見ていきます。

d = pm.Normal.dist()
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

d = pm.MvNormal.dist(mu=np.zeros(2), cov=np.eye(2))
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

M = 3
N = 5
d = pm.MatrixNormal.dist(mu=np.zeros((3, 5)), rowcov=np.eye(3), colchol=np.eye(5))
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

-1.1292054972116046
support dimensions : 0
[-0.71230056 -1.54551037]
support dimensions : 1
[[ 0.65398482 -0.41187789 -0.93498174 -0.10377283 -1.45674856]
 [ 0.6190954   1.31758234  0.29038815 -1.39005479 -0.67170904]
 [ 0.64175279 -0.00501192 -1.03533078 -0.21913851 -0.26720093]]
support dimensions : 2

Batch dimensions

PyMCでは同時に複数の分布を定義し使用することができます。これらは同時に定義しているだけでお互いに独立にサンプリングされたものになります。この時の形状のことをBatch dimensionsと呼びます。この時単純に１次元の正規分布から独立に複数回サンプリングしていることと同じであるため、Support dimensionsは変わらず０のままになります。

# 暗黙的にbatch dimensionsを指定
d = pm.Normal.dist(mu=np.zeros(3), sigma=np.ones(3))
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

# 明示的にbatch dimensionsを指定
d = pm.Normal.dist(shape=(3,))
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

d = pm.Normal.dist(shape=(3,3))
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

[ 2.34605032 -0.28283419 -0.38987575]
support dimensions : 0
[-1.82800188  1.74844484 -0.64056904]
support dimensions : 0
[[-0.73550669 -1.85438216  0.75037397]
 [ 0.47106989  0.80295611 -2.0991726 ]
 [ 0.33004138  0.11099869 -0.62621057]]
support dimensions : 0

上記のコードは以下のようにpm.math.stack()を使用した時と同様の操作になります。

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

# pm.math.stack自体は別の分布をつなげたい時に使用できる
d = pm.math.stack([pm.Normal.dist(), pm.StudentT.dist(nu=1)])
samples = pm.draw(d)
print(samples)

[-1.12677307  0.56926807 -0.61839486]
[ 0.11892697 -0.72002848]

Shapeをより詳細に

shapeで指定する形状はshape = batchの形状 + supportの形状のように、supportの形状の左側にバッチの形状を指定する必要があります。

# ２次元正規分布を独立して２回サンプリングする
batch_shape = (2,)
support_shape = (2,)
shape = batch_shape + support_shape
print(shape)

d = pm.MvNormal.dist(mu=np.zeros(2), cov=np.eye(2), shape=shape)
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")


# support shapeは無視されるため、タイプヒントの意味しかない
batch_shape = (2,)
support_shape = (4,)
shape = batch_shape + support_shape
print(shape)

d = pm.MvNormal.dist(mu=np.zeros(2), cov=np.eye(2), shape=shape)
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")


# support shapeの左側がbatch shapeであるため、batch_shapeは(2,3)回2次元正規分布からサンプリングする
batch_shape = (2,3)
support_shape = (2,)
shape = batch_shape + support_shape
print(shape)

d = pm.MvNormal.dist(mu=np.zeros(2), cov=np.eye(2), shape=shape)
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

# shapeにおけるsupport shapeは無視されるため、冗長である場合はbatch shapeを指定できるsizeを使用する
batch_shape = (2,3)
d = pm.MvNormal.dist(mu=np.zeros(2), cov=np.eye(2), size=batch_shape)
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

(2, 2)
[[-0.3218661   0.29324166]
 [ 2.36724481 -0.53960187]]
support dimensions : 1
(2, 4)
[[-0.33417838  0.96077246]
 [-0.42918312  0.12768932]]
support dimensions : 1
(2, 3, 2)

# MatrixNormの場合はsupport shapeが(3, 5)になるため、2回(3,5)次元多次元正規分布からサンプリングする
batch_shape = (2,)
support_shape = (3, 5)
shape = batch_shape + support_shape
print(shape)

M = 3
N = 5
d = pm.MatrixNormal.dist(mu=np.zeros((3, 5)), rowcov=np.eye(3), colchol=np.eye(5), shape=shape)
samples = pm.draw(d)

print(samples)
print(f"support dimensions : {d.owner.op.ndim_supp}")

(2, 3, 5)
[[[-3.08260754 -1.42353719 -0.38391965 -1.34518582 -1.74453836]
  [-0.18962124  0.47155496 -0.28134255  1.14528603 -0.04782457]
  [ 1.3174886   1.92490691 -0.57543177 -0.3242891  -0.87854472]]

 [[ 1.78671704 -0.19819367 -0.97942314  0.12953045 -0.39287415]
  [ 1.08665596 -1.47828118 -0.96593289  0.99432209 -1.08570123]
  [ 0.00671384  1.53035117 -2.53245528  0.95806438 -0.15679972]]]
support dimensions : 2

最後に

今回は分布とその形状に関して試してみました。Supportとバッチを両方指定する場合はややこしいですが、shape = batchの形状 + supportの形状であることを分かっていれば問題ないと思います。