EM アルゴリズムとベイズ

MikuHatsune2017-03-02

EM アルゴリズムベイズという話が出てきたので、やってみる。
題材はこちら
状況としては、ABOの血液型で、どんな血液型を持っているかは観測できるが、その血液型population を生み出したアレル頻度は一体どのようなものだろうか、これを推定したい、ということ。
EM アルゴリズムベイズも、観測できない潜在パラメータ\theta を、手持ちのデータからなんとかして推定しようという試み。
 
EM では適当にABO のアレル頻度P_a, P_b, P_o として、実際にABO 血液型を観測した人数に当てはめて、アレルの期待値を出して、それを更新して人数に当てはめてアレルの期待値を出して、…を繰り返す。
(ここで、ABO 血液型がどのように生じるかという遺伝学、分子生物学的な知識は既にあるものとし、ボンベイ型やシスAB 型みたいなレアなものは考えないとする。)
すると、リンクのような、A:B:AB:O=160:80:40:120のときに、アレル頻度をEM で求めると

dat <- c(A=160, B=80, AB=40, O=120)
N <- sum(dat)*2

# 初期値
pa <- 0.2
pb <- 0.2
po <- 1-pa-pb

ABO <- c(pa, pb, po)

niter <- 10
for(i in 1:niter){
  # E-step
  pA <- c(pa*(pa+po) / (pa*(pa+2*po)), 0, pa*pb/(2*pa*pb), 0)
  nA <- sum(pA * dat)

  pB <- c(0, pb*(pb+po)/(pb*(pb+2*po)), pa*pb/(2*pa*pb), 0)
  nB <- sum(pB * dat)

  nO <- sum(dat) - nA - nB
  # M-step
  pabo <- c(nA, nB, nO)/sum(dat)
  
  pa <- pabo[1]
  pb <- pabo[2]
  po <- pabo[3]

  ABO <- rbind(ABO, pabo)
}

matplot(ABO, type="l", lwd=3, lty=1, ylim=c(0, 1), xlab="iteration")
legend("topright", legend=c("A", "B", "O"), lty=1, col=1:3, lwd=3)

アルゴリズムの停止は、更新値がどれだけ更新されなくなったかという収束限界値をいれてもいいし、更新回数でもよい。
この場合ではほんの数回でほぼ収束する。

 
ベイズ的にやってみる。rstan でやるならば、アレルABO は足して1になるのでsimplex からサンプリングされ(dirichlet)、実際のABO の血液型は、そのアレル頻度から理論的(HWE)に推定されるA/B/AB/O 血液型の確率から非負の整数値がサンプリングされる(multinomial)とする。このとき、A/B/AB/O 血液型は足して1になる。というのは、普通に考えてそうだし、A/B/AB/O を生み出すABO アレルがそもそも足して1 なので計算上も勝手にそうなる。

library(rstan)
stanmodel <- '
data{
  int<lower=0> N[4];
}
parameters{
  simplex[3] pABO;
}
transformed parameters{
  simplex[4] p;
  p[1] <- pABO[1]*pABO[1] + 2*pABO[1]*pABO[3]; 
  p[2] <- pABO[2]^2 + 2*pABO[2]*pABO[3];
  p[3] <- 2*pABO[1]*pABO[2];
  p[4] <- pABO[3]^2;
}
model{
  N ~ multinomial(p);
}
'
m <- stan_model(model_code=stanmodel)
fit <- sampling(m, list(N=dat))
4 chains, each with iter=2000; warmup=1000; thin=1; 
post-warmup draws per chain=1000, total post-warmup draws=4000.

           mean se_mean   sd    2.5%     25%     50%     75%   97.5% n_eff Rhat
pABO[1]    0.29    0.00 0.02    0.26    0.28    0.29    0.30    0.33  3005    1
pABO[2]    0.16    0.00 0.01    0.14    0.15    0.16    0.17    0.19  2609    1
pABO[3]    0.54    0.00 0.02    0.50    0.53    0.54    0.56    0.58  2525    1
p[1]       0.40    0.00 0.02    0.36    0.39    0.40    0.42    0.45  3058    1
p[2]       0.20    0.00 0.02    0.17    0.19    0.20    0.22    0.24  2830    1
p[3]       0.10    0.00 0.01    0.08    0.09    0.10    0.10    0.11  2462    1
p[4]       0.30    0.00 0.02    0.25    0.28    0.30    0.31    0.34  2522    1
lp__    -516.70    0.03 1.06 -519.64 -517.09 -516.35 -515.96 -515.70  1139    1

実行結果は、EM の場合とほぼ一緒である。
実際にサンプリングされたABO アレル頻度の散布図である。上の推定結果からもわかるように、非常に狭い範囲で推定されている。図でもx+y+z=1 の平面上の極狭い部分にしか点が存在しない。
モデルとしてはかなりシンプルなので推定が楽だし、ベイズを使うとEM では初期値を(A,B,O)=(0.2,0.2,0.6) と選んだのがなぜそれを選んだの?(先行研究から選ぶもんだが) という疑問が拭えないので、ベイズで無情報分布でうまくいったのでいいことにする。
もちろん、EM でも初期値をディリクレ分布から適当に選びまくって推定して、どんな初期値でも同じような結果に収束する(今回はするだろうけど)ことを確認してもいいし、逆にベイズのほうでも、事前に先行研究から範囲が分かっていれば、ABO アレルのsimplex サンプリングに偏ったdirichlet 分布から発生されてもよいはず。