pystanでglassoを実装してみる

はじめに

変数間のinteractionを発見するのに使われるglassoの拡張モデルを作る準備としてstanでglassoを実装してみます．ただglassoを使うだけなら

sklearn.covariance.graphical_lasso — scikit-learn 0.23.1 documentation

を使うのが一番手っ取り早いです．

stanファイルの実装

stanでパラメタ推定を行う場合はモデルを記述するstanファイルを用意する必要があります． 詳細は省きますが，glassoに置けるL1正則化項は正規分布の精度行列パラメタの事前分布と見なせます． この事前分布は精度行列の各要素を対角成分と非対角成分に分けて，それぞれ独立同分布とみなすものです． ただし，この事前分布は精度行列としての条件(逆行列＝分散共分散行列が半正定値行列)を満たすことが保証されないのでstanでの実装においてはcov_matrix型を定義することで対処します． このcov_matrix型で宣言されたパラメタは分散共分散行列としての条件を満たすことが保証されます． そして実装したのが以下になります．

// glasso.stan
data {
    int<lower=0> n;
    int<lower=0> p;
    vector[p] x[n];
    real<lower=0> lambda;
    // real<lower=0> prior_parameter_for_lambda;
}

transformed data {
}

parameters {
    cov_matrix[p] invsigma;
    vector[p] mu;
    // real<lower=0> lambda;
}

transformed parameters {
    real<lower=0> invlambda;
    real<lower=0> lambda_half;
    invlambda = 1/lambda;
    lambda_half = lambda/2;
}

model {
    for (i in 1:n) {
        x[i] ~ multi_normal_prec(mu, invsigma);
    }
    for (j1 in 1:p) {
        for (j2 in 1:j1) {
            if (j1==j2) {
                invsigma[j1,j2] ~ exponential(lambda_half);
            } else {
                invsigma[j1,j2] ~ double_exponential(0,invlambda);
            }
        }
    }
    //Prior for lambda
    // lambda ~ exponential(prior_parameter_for_lambda);
}

generated quantities {
}

こちらの論文&コードを参考にしました(なんでバイオインフォやねんというのはたまたま思い浮かんだのがこの論文でして...)

人工データによる実験

生成過程から人工データを作成して真のパラメタを当てられるか検証します． まず，importとデータの次元を決めます．

import pandas as pd
import numpy as np
import sys
import seaborn as sns
sns.set()
import matplotlib.pyplot as plt
import pystan

N = 1000
p = 5

精度行列を作ります．ただし，精度行列をglassoの生成過程通りに作るとスパースにならないので，categorical分布で0か否かを生成してから値を生成します．

prec_matrix = np.zeros([p,p])

for i in range(p):
    for j in range(i+1):
        if i == j:
            prec_matrix[i, j] = np.random.gamma(100,1)
        else: 
            prec_matrix[i, j] = np.random.choice(2, p=[0.8, 0.2])
            prec_matrix[i, j] = prec_matrix[i, j] * np.random.normal(0,100)
            prec_matrix[j, i] = prec_matrix[i, j]

固有値の最小値を出力し，負でないことを確認します(負であれば再生成します)．

np.linalg.eigvals(np.linalg.inv(prec_matrix)).min()
# 0.005406751950232075

精度行列を可視化します．

sns.heatmap(prec_matrix, cmap='coolwarm', center=0)

データを生成します．平均は0にします．

x = []
for i in range(N):
    x.append(np.random.multivariate_normal(np.zeros(5), np.linalg.inv(prec_matrix)))
x = pd.DataFrame(x)

罰則項の係数は1にしてみます．

stan_data = {'n': N, 'p': p, 'x': x, 'lambda': 1}
fit = pystan.stan(file='./glasso.stan', data=stan_data, iter=1000, chains=1)
print('Sampling finished.')

一応MAP推定をしたいので，サンプル列からposteriorが最大のindexを取ってきて，該当するサンプルを可視化します．

la  = fit.extract(permuted=True)
names = fit.model_pars
map_index = pd.Series(fit.get_logposterior()[0]).idxmax()
la  = fit.extract(permuted=True)
invsigma_sample = pd.DataFrame(la[names[0]].reshape(-1, la[names[0]][0].size))
# burn_in期間の500を引いたindexを使用
estimated_prec_matrix = pd.DataFrame(np.array(invsigma_sample.iloc[map_index - 500]).reshape(p, p))
sns.heatmap(estimated_prec_matrix, cmap='coolwarm', center=0, annot=True)

うまく推定できていますね． ただ，posteriorが最大になるサンプルを取ってきているだけなので0に見える部分は0ではなかったりします．

おわりに

glassoはposteriorを最大化した場合にスパースになりやすいという手法なのでMCMCで行うのは邪道かもしれませんね．精度行列の非対角成分にラプラス分布の代わりに正規分布をおいてもたいして変わらないと思います．しかし，MCMCで近似された事後分布の形はおおよそ正しい訳ですから，考案したモデルをスクラッチから実装する前の確認としては十分有用なのでは無いかと思います．