PyMCのモデル宣言におけるクラス継承の実装例

概要

PyMCのモデルの宣言において、公式チュートリアル等では、下記のような with ステートメントを使った記法がよく使用されています。

with Model() as model:  # model specifications in PyMC are wrapped in a with-statement
    # Define priors
    sigma = HalfCauchy("sigma", beta=10)
    intercept = Normal("Intercept", 0, sigma=20)
    slope = Normal("slope", 0, sigma=20)

    # Define likelihood
    likelihood = Normal("y", mu=intercept + slope * x, sigma=sigma, observed=y)

    # Inference!
    # draw 3000 posterior samples using NUTS sampling
    idata = sample(3000)

上記の例では、 with ステートメント内で宣言された変数は自動的に model 内の変数として扱われています。

しかし、pymc.Model のドキュメントでも解説されている通り、PyMCでは pymc.Model クラスを継承したモデルの宣言もサポートされています。モデルを個別のクラスとして宣言することで、名前空間を切り分けて記述することができ、保守性の高いコードになります。

この記事では、クラス継承によるPyMCのモデルの宣言の使用例を紹介したいと思います。

使用環境

ここでは以下の環境で検証しています。

OS	Windows 11
Python	3.10.9
PyMC	5.1.1

環境構築には conda コマンドを使用しています。PyMC 5.1.1はconda-forgeリポジトリで配布されているので、下記のようにチャンネルを指定することでインストール可能です。

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

※チャンネル指定を行わないと、anacondaリポジトリのPyMC 2.3.8がインストールされます。

with ステートメントを使った実装

まず、基本の実装として、モデル宣言に with ステートメントを使用したコードを実装してみます。公式のGLMチュートリアルを参考に作成しました。

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

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-darkgrid")


def main():
    RANDOM_SEED = 8927
    rng = np.random.default_rng(RANDOM_SEED)

    size = 200
    true_intercept = 1
    true_slope = 2
    true_sigma = 0.5

    x = np.linspace(0, 1, size)
    true_regression_line = true_intercept + true_slope * x
    y = true_regression_line + rng.normal(scale=true_sigma, size=size)
    df = pd.DataFrame(dict(x=x, y=y))
    print(df)

    with pm.Model() as model:
        sigma = pm.HalfCauchy("sigma", beta=10)
        intercept = pm.Normal("Intercept", 0, sigma=20)
        slope = pm.Normal("slope", 0, sigma=20)

        pm.Normal("y", mu=intercept + slope * df.x, sigma=sigma, observed=df.y)

        idata = pm.sample(draws=3000, random_seed=rng)

    az.plot_trace(idata)
    plt.savefig("traceplot.png")


if __name__ == "__main__":
    main()

出力として以下のようなトレースプロットが得られます。

pymc.Model の継承を利用したコード

上記の with ステートメントを使用したコードを、pymc.Model を継承したクラスを宣言して書き換えてみます。実装したコードは以下の通りです。

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-darkgrid")


class GeneralizedLinearModel(pm.Model):
    def __init__(self, df, *args, **kwargs):
        super().__init__(*args, **kwargs)

        sigma = pm.HalfCauchy("sigma", beta=10)
        intercept = pm.Normal("Intercept", 0, sigma=20)
        slope = pm.Normal("slope", 0, sigma=20)

        pm.Normal("y", mu=intercept + slope * df.x, sigma=sigma, observed=df.y)


def main():
    RANDOM_SEED = 8927
    rng = np.random.default_rng(RANDOM_SEED)

    size = 200
    true_intercept = 1
    true_slope = 2
    true_sigma = 0.5

    x = np.linspace(0, 1, size)
    true_regression_line = true_intercept + true_slope * x
    y = true_regression_line + rng.normal(scale=true_sigma, size=size)
    df = pd.DataFrame(dict(x=x, y=y))
    print(df)

    model = GeneralizedLinearModel(df=df)
    idata = pm.sample(draws=3000, random_seed=rng, model=model)

    az.plot_trace(idata)
    plt.savefig("traceplot.png")


if __name__ == "__main__":
    main()

出力のトレースプロットは全く同じものが得られましたので省略します。

実装のポイントとしては、 pm.Model を継承したサブクラス GeneralizedLinearModel を宣言しているところが大きな違いです。ここでは、一般的なクラス継承の処理と同様にコンストラクタでスーパークラスのコンストラクタを呼び出しています。その後の各変数の宣言については with ステートメントで行っていたものと全く同じです。

class GeneralizedLinearModel(pm.Model):
    def __init__(self, df, *args, **kwargs):
        super().__init__(*args, **kwargs)

        sigma = pm.HalfCauchy("sigma", beta=10)
        intercept = pm.Normal("Intercept", 0, sigma=20)
        slope = pm.Normal("slope", 0, sigma=20)

        pm.Normal("y", mu=intercept + slope * df.x, sigma=sigma, observed=df.y)

そしてメイン処理では宣言したサブクラスをインスタンス化し、 model に代入しています。ここで観測データの df を引数としてモデルに渡しています。

model = GeneralizedLinearModel(df=df)

最後のサンプリングですが、with ステートメントを使用しない場合には、以下のように model 引数でインスタンス化したモデルを指定する必要があります。（with ステートメントを使用していた場合ではモデルが自動的に選択されるため指定不要でした。）

idata = pm.sample(draws=3000, random_seed=rng, model=model)

以上のように実装することで、モデルの宣言を個別のクラスで行うことができます。

このシンプルな例では with ステートメントを使用した方がわかりやすいのかもしれませんが、以下の記事で紹介しているような階層モデルなどでは扱う確率変数が増え、変数の管理が煩雑になってきます。その際、モデルを個別のクラスとして宣言しておくことで変数の名前空間が分けられるため、保守しやすいコードが実装できます。