Python（PyStan）で「StanとRでベイズ統計モデリング」の5.1節を実行する

StanのPython バインディングであるPyStanが公開されて久しいですが、検索してもあんまり情報がヒットしません。ちょっと寂しいと思ったので、インストールやtraceplotの出力なども含めて、以下の本の5.1節「重回帰」の一部を実行してみました（ステマです）。

StanとRでベイズ統計モデリング (Wonderful R)

作者:健太郎, 松浦
発売日: 2016/10/25
メディア: 単行本

本自体の紹介は以前の記事を読んでいただければと思います。

「StanとRでベイズ統計モデリング」松浦健太郎という本を書きました - StatModeling Memorandum

インストール

Windows 7 64bit、Python 3系でのインストール手順を説明します。

AnacondaでPythonなどをインストール
コマンドプロンプトからpip install pystanでpystanをインストール
Visual C++ Build Toolsをインストール

以上でPyStanが動くようになります。その他のOSにつきましては公式ページを読んでください。ほぼ同様の手順であり、難しくありません。

なおこの記事で使用したバージョンは次の通りです。

Python 3.5.2 | Anaconda 4.2.0 (64-bit)
PyStan 2.12.0.0

モデルファイル（Stanファイル）の説明

本の5.1節では、「大学の授業の出席率」を応答変数とし、過去の「アルバイトが好きかどうかというアンケート結果」と過去の「テストの点数」を説明変数として重回帰を行います。サポートページであるGitHubリポジトリに各種ファイルを置いてあります。

データやStanファイルの説明は本に任せて、この記事ではPythonからの使い方に注力します。

Stanファイルのコンパイルとサンプリング

pystanを使ったStanファイルのコンパイルとサンプリングは以下のように行います。

import pandas as pd
from pystan import StanModel
import matplotlib.pyplot as plt
import pickle

d = pd.read_csv('input/data-attendance-1.txt')
d.Score /= 200
data = d.to_dict('list')
data.update({'N':len(d)})

stanmodel = StanModel(file='model/model5-3.stan')

# NUTS (No U-Turn Sampler)
fit_nuts = stanmodel.sampling(data=data, n_jobs=1)
mcmc_sample = fit_nuts.extract()
mu_est = mcmc_sample['mu']

# ADVI (Automatic Differentiation Variational Inference)
fit_vb = stanmodel.vb(data=data)
vb_sample = pd.read_csv(fit_vb['args']['sample_file'].decode('utf-8'), comment='#')
vb_sample = vb_sample.drop([0,1])
mu_est = vb_sample.filter(regex='mu\.\d+')

with open('output/model_and_result.pkl', 'wb') as f:
    pickle.dump(stanmodel, f)
    pickle.dump(fit_nuts, f)

6～9行目：　データを読み込んでpandasのDataFrameを作り、そこからStanに渡せるようにdictionaryに変換しています。7行目ではデータのScoreをスケーリングしています。推定したいパラメータのスケールを同じぐらいに揃えるのが目的です（収束が改善する場合があります）。
11行目：　Stanファイルのコンパイルだけを行っています。初期値を変えたりiterを変えたりしてサンプリングをいろいろ試したい場合に便利です。
14行目：　NUTSというアルゴリズムでサンプリングを実行しています。デフォルトだとコア数の分だけ並列で実行されます。しかし、Windowsの場合は無限コンパイル地獄に落ちる不具合があるので、n_jobs=1で並列を止める必要があります。このページを参照。
15～16行目：結果からMCMCサンプルを入手する一例です。
19～22行目：　NUTSではなく、自動変分ベイズ（ADVI）で推定する場合です。結果の形式がまだ洗練されておらず、ファイルから読み込んでからゴミの行を除去する必要があります（簡単ですが）。
24～26行目：　コンパイルされたモデルと結果をpickleで保存しておきます。

実行速度は基本的にC++コンパイラに依存するだけなので、Rの場合（少し古いgcc）と大きな差はなさそうです。

収束の要約とtraceplotの保存

本の4.4.8項「収束診断をファイルへ出力する」に相当します。pystanのデフォルトのプロット関数を使ってfit_nuts.plot()のようにするとtraceplotを描くことができますが、4つのchainを線種を変えて描くのではなくて1つのchainにまとめて描いてしまったり、表示がつぶれたりする不満点があります。代わりの関数としては、PyMCの開発チームが作成中のmcmcplotlibに期待していますが、例も使い方も載ってないので躊躇してしまいます。ここでは自作でRの{coda}パッケージ風のtraceplotを、複数ページのpdfで作成することにしました。

import pickle
import sys
import numpy as np
import math
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.backends.backend_pdf import PdfPages

with open('output/model_and_result.pkl', 'rb') as f:
    stanmodel = pickle.load(f)
    fit_nuts = pickle.load(f)

## summary
with open('output/fit-summary.txt', 'w') as f:
    f.write(str(fit_nuts))

## traceplot
palette = sns.color_palette()
ms = fit_nuts.extract(permuted=False, inc_warmup=True)
iter_from = fit_nuts.sim['warmup']
iter_range = np.arange(iter_from, ms.shape[0])
paraname = fit_nuts.sim['fnames_oi']
num_pages = math.ceil(len(paraname)/4)

with PdfPages('output/fit-traceplot.pdf') as pdf:
    for pg in range(num_pages):
        plt.figure()
        for pos in range(4):
            pi = pg*4 + pos
            if pi >= len(paraname): break
            plt.subplot(4, 2, 2*pos+1)
            plt.tight_layout()
            [plt.plot(iter_range + 1, ms[iter_range,ci,pi], color=palette[ci]) for ci in range(ms.shape[1])]
            plt.title(paraname[pi])
            plt.subplot(4, 2, 2*(pos+1))
            plt.tight_layout()
            [sns.kdeplot(ms[iter_range,ci,pi], color=palette[ci]) for ci in range(ms.shape[1])]
            plt.title(paraname[pi])
        pdf.savefig()
        plt.close()

9～11行目：　pickleで保存した結果を読み込んでいます。
14～15行目：　MCMCサンプルの要約（Rhatなどの収束診断を含む）をファイルに出力しています。
18～40行目：　traceplotを描いています。
- 18行目： chainごとの色をseabornパッケージのデフォルトカラーパレットで決めています。このままですとchainの数が6を超えたらエラーが出ます（色が巡回するようにすればOKと思います）。
- 19行目：　extract関数はMCMCサンプルを取り出す関数ですが、デフォルトだと各chainを混ぜて、iterationの順番も混ぜて出力します。開発チームの想定と異なる使われ方を防ぐためだそうです（ここ参照）。traceplotをプロットするにはchainごとにiterationの順番を保持したままのMCMCサンプルが欲しいので、このように引数を指定しています。
- 20～21行目：　ここではRの{rstan}パッケージと同様に、warmup期間を除いた部分のtraceplotと密度関数を描くようにしました。もしwarmup期間も含めてプロットしたい場合にはiter_from = 0としてください。
- 22～23行目：　複数ページのpdfを作るためにパラメータ数を取得しています。1ページに4つのパラメータを描きます。
- 25～40行目：　残りはmatplotlibの話なのでここでは省略します。複数ページのpdfの作成方法は公式ページを参考にしました。

結果の図は以下になります（pdfの1ページ目です）。

f:id:StatModeling:20201106182522p:plain

実測値と予測値のプロット

本の5.1.8項「図によるモデルのチェック」に相当します。図5.3を描きます。横軸に実測値（データ）・縦軸に予測分布の中央値と予測区間をプロットしたグラフです。

このようなグラフを描くにはMCMCサンプルを集計したり加工したりする必要があります。岩波DS1のGitHub リポジトリ内に、sinhrksさん（pandasのコミッタ）に書いてもらったコードがありますので、よくそれらを参考にしています（example1,example2,example3）。

Pythonコードは以下になりました。

import pickle
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

with open('output/model_and_result.pkl', 'rb') as f:
    stanmodel = pickle.load(f)
    fit_nuts = pickle.load(f)

d_ori = pd.read_csv('input/data-attendance-1.txt')
ms = fit_nuts.extract()

quantile = [10, 50, 90]
colname = ['p' + str(x) for x in quantile]
d_qua = pd.DataFrame(np.percentile(ms['y_pred'], q=quantile, axis=0).T, columns=colname)
d = pd.concat([d_ori, d_qua], axis=1)
d0 = d[d.A == 0]
d1 = d[d.A == 1]

palette = sns.color_palette()
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot([0,0.5], [0,0.5], 'k--', alpha=0.7)
ax.errorbar(d0.Y, d0.p50, yerr=[d0.p50 - d0.p10, d0.p90 - d0.p50],
    fmt='o', ecolor='gray', ms=10, mfc=palette[0], alpha=0.8, marker='o')
ax.errorbar(d1.Y, d1.p50, yerr=[d1.p50 - d1.p10, d1.p90 - d1.p50],
    fmt='o', ecolor='gray', ms=10, mfc=palette[1], alpha=0.8, marker='^')
ax.set_aspect('equal')
ax.set_xlim(0, 0.5)
ax.set_ylim(0, 0.5)
ax.set_xlabel('Observed')
ax.set_ylabel('Predicted')
fig.savefig('output/fig5-3.png')

14～19行目：　MCMCサンプルの分位点（ここでは10%点, 50%点, 90%点）を計算してDataFrameにしています。行や列をしぼるのがラクで、描くときに便利だからです。

残りはmatplotlibの話なのでここでは省略します。結果の図は以下になります。

f:id:StatModeling:20201106182514p:plain

MCMCサンプルの散布図行列

同じく本の5.1.8項「図によるモデルのチェック」の、図5.5に相当する図を描きます。Pythonで散布図行列を描くにはseabornパッケージのPairGridを使うのが便利です。ただし、上三角行列の部分はカスタマイズする必要があります。ここを参考にしました。

import pickle
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import math
from matplotlib.patches import Ellipse
from scipy import stats

with open('output/model_and_result.pkl', 'rb') as f:
    stanmodel = pickle.load(f)
    fit_nuts = pickle.load(f)

d_ori = pd.read_csv('input/data-attendance-1.txt')
ms = fit_nuts.extract()

def corrfunc(x, y, **kws):
    r, _ = stats.spearmanr(x, y)
    ax = plt.gca()
    ax.axis('off')
    ellcolor = plt.cm.RdBu(0.5*(r+1))
    txtcolor = 'black' if math.fabs(r) < 0.5 else 'white'
    ax.add_artist(Ellipse(xy=[.5, .5], width=math.sqrt(1+r), height=math.sqrt(1-r), angle=45,
        facecolor=ellcolor, edgecolor='none', transform=ax.transAxes))
    ax.text(.5, .5, '{:.0f}'.format(r*100), color=txtcolor, fontsize=28,
        horizontalalignment='center', verticalalignment='center', transform=ax.transAxes)

d = pd.DataFrame({'b1':ms['b1'], 'b2':ms['b2'], 'b3':ms['b3'], 'sigma':ms['sigma'],
                  'mu1':ms['mu'][:,0], 'mu50':ms['mu'][:,49], 'lp__':ms['lp__']},
                  columns=['b1', 'b2', 'b3', 'sigma', 'mu1', 'mu50', 'lp__'])
sns.set(font_scale=2)
g = sns.PairGrid(d)
g = g.map_lower(sns.kdeplot, cmap='Blues_d')
g = g.map_diag(sns.distplot, kde=False)
g = g.map_upper(corrfunc)
g.fig.subplots_adjust(wspace=0.05, hspace=0.05)
for ax in g.axes.flatten():
    for t in ax.get_xticklabels():
        _ = t.set(rotation=40)
g.savefig('output/fig5-5.png')

17～26行目：上三角行列を作成しています。Spearmanの順位相関係数に応じて、楕円の太さ細さと色を変えています。数値は相関係数×100を表示しています。
36行目： subplotの間の空白を狭くしています。空いていると少し見づらかったので。
37～39行目： X軸ラベルを少し傾けています。PairGridなので少しややこしいです。

結果の図は以下になります。

f:id:StatModeling:20201106182518p:plain

まとめ

pystanを使ってPythonからStanを実行してみました。思ったより簡単です。
numpyやpandasはMCMCサンプルと相性がよくて楽しいです。
作図の部分はRを使う場合と同様で、自分である程度頑張る必要があります。matplotlib力が試されます。
「StanとRでベイズ統計モデリング」をPythonで実行するのは十分可能と思いました。