PyMCでの階層モデルの実装例

概要
使用環境
階層モデルの実装
実装のポイント
実行結果
参考文献

概要

PyMC5で階層モデルの実装例を紹介します。公式ドキュメントだけではわかりにくかったり、比較的情報の多いPyMC3と仕様が違う部分も多々あるので、実装におけるポイントになった部分を整理して紹介したいと思います。

使用環境

OS	Windows 11
Python	3.10.9
PyMC	5.1.1

環境構築には conda コマンドを使用しています。後述するモデルの図化のため、 python-graphviz もインストールしています。

conda create -n pymc python=3.10
conda activate pymc
conda install conda-forge::pymc
conda install conda-forge::python-graphviz

階層モデルの実装

階層モデルとして『StanとRでベイズ統計モデリング』の「8.4 ロジスティック回帰の階層モデル」の例題を題材にします。この例題では授業での学生の出席の有無を予測するモデルを構築しています。

書籍はRStanで実装されていますが、サポートページにある以下のコードとデータを利用し、同様の階層モデルをPyMCで実装しました。

Rコード
- RStanBook/run-model8-8.R
Stanコード
- RStanBook/model8-8.stan
データ
- RStanBook/data-attendance-4-1.txt
- RStanBook/data-attendance-4-2.txt

実装したコードは以下の通りです。

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm

az.style.use("arviz-docgrid")


class BeyesModel(pm.Model):
    def __init__(self, d1, d2, *args, **kwargs):
        super().__init__(*args, **kwargs)

        N = len(d1)
        C = len(d2["CourseID"].unique())
        I = len(d2)

        self.add_coord(name="person", values=[f"PID{i+1:02}" for i in range(N)])
        self.add_coord(name="course", values=[f"CID{i+1:02}" for i in range(C)])
        self.add_coord(name="record", values=[f"Record{i+1:04}" for i in range(I)])

        score = pm.MutableData(name="Score", value=d1["Score"], dims="person")
        a = pm.MutableData(name="A", value=d1["A"], dims="person")
        w = pm.MutableData(name="Weather", value=d2["Weather"], dims="record")

        b1 = pm.Flat(name="b1")
        b2 = pm.Flat(name="b2")
        b3 = pm.Flat(name="b3")
        b4 = pm.Flat(name="b4")

        s_P = pm.HalfCauchy(name="s_P", beta=10.0)
        b_P = pm.Normal(name="b_P", mu=0.0, sigma=s_P, dims="person")
        x_P = pm.Deterministic(name="x_P", var=b2 * a + b3 * score + b_P, dims="person")

        s_C = pm.HalfCauchy(name="s_C", beta=10.0)
        b_C = pm.Normal(name="b_C", mu=0.0, sigma=s_C, dims="course")
        x_C = pm.Deterministic(name="x_C", var=b_C, dims="course")

        x_J = pm.Deterministic(name="x_J", var=b4 * w, dims="record")
        x = pm.Deterministic(name="x", var=b1 + x_P[d2["PersonID"] - 1] + x_C[d2["CourseID"] - 1] + x_J, dims="record")

        pm.Bernoulli(name="Y", logit_p=x, observed=d2["Y"], dims="record")


def main():
    seed = 0
    rng = np.random.default_rng(seed)

    d1 = pd.read_csv("data-attendance-4-1.txt")
    d2 = pd.read_csv("data-attendance-4-2.txt")

    d1["Score"] = d1["Score"] / 200.0
    d2["Weather"] = d2["Weather"].map({"A": 0.0, "B": 0.2, "C": 1.0})

    model = BeyesModel(d1, d2)
    gv = pm.model_to_graphviz(model)
    gv.render(outfile="model.png")

    idata = pm.sample(
        draws=2000,
        tune=1000,
        random_seed=rng,
        model=model,
    )

    df_summary = az.summary(
        idata,
        var_names=["b1", "b2", "b3", "b4", "s_P", "s_C", "b_C"],
        kind="stats",
        hdi_prob=0.95,
    )
    df_summary.to_csv("summary.csv")

    az.plot_forest(
        idata,
        var_names=["b1", "b2", "b3", "b4", "s_P", "s_C", "b_C"],
        combined=True,
        hdi_prob=0.95,
        figsize=(8, 6),
    )
    plt.savefig("forestplot.png")


if __name__ == "__main__":
    main()

実装のポイント

継承クラスによるモデルの宣言

pymc.Model を継承したクラスによりモデルを宣言しています。階層モデルでは扱う変数が多く、個別のクラスとすることで名前空間の管理が楽になります。

継承クラスによるモデルの宣言についてはこちらの記事で紹介しています。

モデルの図化

モデルは pymc.model_to_graphviz により Digraph オブジェクトを生成することで、ファイル出力を行う render() メソッドが利用できます。

gv = pm.model_to_graphviz(model)
gv.render(outfile="model.png")

出力したモデルの図は以下の通りです。

モデルの図化は宣言したモデルが意図通りに構成されているか確認するのに便利です。

`dims` による次元の設定

モデルが有する次元やインデックスに命名できる機能として coords や dims があります。

pymc.Model のコンストラクタの coords 引数や、 add_coord() メソッドにより次元やインデックスを定義することができます。実装したモデルでは、学生、授業、出席データの次元があるので、以下のように定義しています。

self.add_coord(name="person", values=[f"PID{i+1:02}" for i in range(N)])
self.add_coord(name="course", values=[f"CID{i+1:02}" for i in range(C)])
self.add_coord(name="record", values=[f"Record{i+1:04}" for i in range(I)])

変数の宣言の際に、ここで定義した次元名を dims で指定することで、各次元のサイズやインデックスを指定することができます。

s_P = pm.HalfCauchy(name="s_P", beta=10.0)
b_P = pm.Normal(name="b_P", mu=0.0, sigma=s_P, dims="person")
x_P = pm.Deterministic(name="x_P", var=b2 * a + b3 * score + b_P, dims="person")

このように次元を指定しておくと、前述のモデルの図化において次元の情報が記載されたり、指定した次元と計算結果が一致しない時にエラーが送出されてデバッグしやすくなるなどのメリットがあります。

モデルへのデータ変数の登録

尤度を算出する確率変数に対して observed で観測データを渡すなど、モデル内で確率変数ではないデータも扱う必要があります。numpyの配列やpandasのデータフレームを渡すこともできますが、そういったデータをモデルに登録するためのクラスとして pymc.Data があります。

pymc.Data は mutable という引数によりメモリの共有の可否を指定できますが、このオプションを使用する代わりに pymc.ConstantData や pymc.MutableData を利用するのが推奨されています。そのため、この実装でも pymc.MutableData を利用してデータフレームのデータをモデルに登録しています。また、このデータに対しても dims 引数で次元を明示するようにしています。

score = pm.MutableData(name="Score", value=d1["Score"] / 200.0, dims="person")
a = pm.MutableData(name="A", value=d1["A"], dims="person")
w = pm.MutableData(name="Weather", value=d2["Weather"].map({"A": 0.0, "B": 0.2, "C": 1.0}), dims="record")

こうしてデータをモデルに登録することで、モデルの図化において入力データの情報が記載されたり、pymc.set_data を用いてモデル宣言後に動的にデータを変更することができるようになります。

実行結果

正しい計算ができているかを確認するため、出力結果を確認します。

フォレストプロット forestplot.png は以下のようになりました。点で示されるのが事後平均、太い線が50%信用区間、細い線が95%信用区間になります。

summary.csv に出力した数値としては以下の通りです。

	mean	sd	hdi_2.5%	hdi_97.5%
b1	0.817	0.597	-0.363	2.001
b2	-0.802	0.153	-1.105	-0.503
b3	1.69	0.586	0.565	2.847
b4	-0.779	0.143	-1.051	-0.495
s_P	0.353	0.089	0.186	0.536
s_C	1.411	0.417	0.769	2.228
b_C[CID01]	0.423	0.529	-0.61	1.484
b_C[CID02]	-1.371	0.494	-2.324	-0.392
b_C[CID03]	-1.005	0.5	-2.019	-0.035
b_C[CID04]	1.339	0.524	0.288	2.361
b_C[CID05]	-0.283	0.494	-1.234	0.712
b_C[CID06]	-1.936	0.505	-2.955	-0.969
b_C[CID07]	0.865	0.511	-0.122	1.894
b_C[CID08]	-0.556	0.494	-1.556	0.402
b_C[CID09]	1.184	0.526	0.168	2.242
b_C[CID10]	1.265	0.528	0.242	2.342

書籍の図8-9と概ね同等の結果が得られており、意図通り計算できているようです。

概要