pyMCで線形回帰分析

pyMCの記事4記事目です。これまでの記事では観測値だけ与えてそれを生成する確率分布を考えてきましたが、今回は観測値だけでなく何か特徴量を持つデータを考えます。
その最も単純な例として1変数の線形回帰をやってみましょう。特徴量が増えて重回帰分析になってもほとんど同じように対応できるので汎用性は高いと思います。

データの準備

何のデータを使ってもいいのですが、今回はscikit-learnのiris使います。3種類のアヤメのうち、virginicaに絞って、petal length (cm)からsepal length (cm)を予測するモデルを考えてみましょう。(相関係数が0.86くらいあって予測が簡単なのです。)

次のようにしてデータを取得します。

%%pycodestyle
import pandas as pd
from sklearn.datasets import load_iris


# irisデータ取得
df = pd.DataFrame(iris.data, columns=iris.feature_names)
df["label"] = iris.target
# virginica に絞る
df = df[df.label == 2].reset_index(drop=True)

# 回帰分析に使う特徴量xと目的変数yを取得
x = df["petal length (cm)"].values
y = df["sepal length (cm)"].values

モデルの実装

データが揃ったらpyMCでモデルを作っていきます。

回帰係数を a, 定数項をb、誤差をε とするとモデルの式はこのようになりますね。
$$y= ax + c + \epsilon$$

ベイズで行いますので、それぞれに事前分布が必要です。
a, c の事前分布は期待値が0、標準偏差が10の正規分布としましょう。
そして、誤差項εは、期待値が0で、標準偏差がσの正規分布に従うとし、このσは標準偏差が10の半正規分布に従うとします。

これを実装してきますが、今回新たに使うのは、特徴料等の定数を格納する pm.ConstantData と 数式を定義できるpm.Deterministic です。

参考:
pymc.ConstantData — PyMC dev documentation
pymc.Deterministic — PyMC v5.6.0 documentation

正確には、 ConstantData の方は使わなくてもいいのですが、明示的に書いておくとモデルを可視化した時に定数部分も表示されるので便利です。

実際にコードを見ていただくと使い方がわかると思うのでやっていきましょう。例によって、Graphvizで可視化してJpyterで表示しています。

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


model = pm.Model()

with model:
    x_data = pm.ConstantData("x_data", x)
    y_data = pm.ConstantData("y_data", y)

    # 回帰係数と定数項
    a = pm.Normal("a", mu=0, sigma=10)
    c = pm.Normal("c", mu=0, sigma=10)

    # yの期待値
    mu = pm.Deterministic("mu", a*x_data + c)
    # 誤差
    sigma = pm.HalfNormal("sigma", sigma=10)

    # 観測値
    obs = pm.Normal("obs", mu=mu, sigma=sigma, observed=y_data)

# モデルの可視化
g = pm.model_to_graphviz(model)
display(g)

参考ですが、ConstantDataを使わない場合はこうなります。

model2 = pm.Model()

with model2:

    # 回帰係数と定数項
    a = pm.Normal("a", mu=0, sigma=10)
    c = pm.Normal("c", mu=0, sigma=10)

    # yの期待値
    mu = pm.Deterministic("mu", a*x + c)
    # 誤差
    sigma = pm.HalfNormal("sigma", sigma=10)

    # 観測値
    obs = pm.Normal("obs", mu=mu, sigma=sigma, observed=y)

g = pm.model_to_graphviz(model2)
display(g)

モデルができたのでサンプリングして結果を見ていきましょう。

with model:
    trace = pm.sample(random_seed=42, chains=2)


display(az.summary(trace, var_names=["a", "c", "sigma"]))
"""
	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
a	0.995	0.084	0.828	1.143	0.003	0.002	641.0	636.0	1.0
c	1.066	0.470	0.216	1.988	0.019	0.013	648.0	679.0	1.0
sigma	0.331	0.035	0.259	0.391	0.001	0.001	554.0	349.0	1.0
"""

az.plot_trace(trace, var_names=["a", "c", "sigma"], compact=False)
plt.tight_layout()

いい感じに推定できていますね。

回帰直線の可視化

せっかく単回帰したので、回帰直線を可視化してみたいと思います。上記のsummaryのa,cで可視化してもいいのですがせっかくなのでサンプリングの各ステップの値で可視化してみましょう。

# a, cの各ステップの値を取得
a_list = trace.posterior.a.values.ravel().reshape(-1, 1)
c_list = trace.posterior.a.values.ravel().reshape(-1, 1)

# xの範囲
x_values = np.array([4.4, 7.0])

# a, c の各値からyの値を計算
y_preds = x_values * a_list + c_list

# 回帰直線と元のデータを可視化
for y_pred in y_preds:
    plt.plot(x_values, y_pred, lw=1, alpha=0.01, c="c")
plt.scatter(x, y)
plt.show()

なかなか妥当な結果が得られましたね。

まとめ

今回は線形回帰を題材として取り上げましたが、線形回帰に限らず特徴量を使うモデリングは同じようにして実装していくことができます。pm.Deterministicを使うと一気に実装の幅が広がりますのでぜひ試してみてください。