習作;混合正規分布モデル(Gaussian Mixture Model) by rstan

Stanのマニュアルに載っている簡単なモデルなんだけど,最後の「各要素の所属確率」を算出するのになぜか手間取った。softmax関数を使うときの型の問題でした。ちぇ。

Stanコードはこの通り。generated quantitiesのところでsoftmax関数を使います。

data{
  int<lower =2> K; #number of clusters
  int<lower =1> N; #number of observations
  real X[N]; #observed data
}

parameters{
  vector[K] mu;
  vector<lower =0,upper=10>[K] sig2;
  simplex[K] theta;
}

transformed parameters{
  vector[K] ps[N];
  for(n in 1:N){
    for(k in 1:K){
      ps[n,k] = log(theta[k])+normal_lpdf(X[n]|mu[k],sig2[k]);
    }
  }
  
}

model{
  sig2 ~ cauchy(0,2.5);
  mu ~ normal(0,100);
  for(n in 1:N){
    target += log_sum_exp(ps[n]);
  }
}

generated quantities{
  simplex[K] u[N]; #class membership probability
  for(n in 1:N){
    u[n] = softmax(ps[n]);
  }
}

キックするコードは次の通り。

library(rstan)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

# 1 dimensional Mixed Gauss
#  K個の正規分布を混ぜる

K < - 3
mean1 <- 5
mean2 <- -5
mean3 <- 0
sd1 <- 1
sd2 <- 2
sd3 <- 3

N1 <- 500
N2 <- 300
N3 <- 200
N <- N1+N2+N3
X <-c(rnorm(N1,mean1,sd1),rnorm(N2,mean2,sd2),rnorm(N3,mean3,sd3))

### MCMC!
stanmodel <- stan_model("develop/gmm_single.stan",model_name="GMM")
standata <- list(K=K,N=N,X=X)
max.iter <- 3000
burn.in <- 1000
itr <- max.iter - burn.in
C <- 4


# vbは速いしラベルスイッチングとかないから嬉しい
fit_vb <- vb(stanmodel,data=standata)
print(fit_vb,pars=c("mu","sig2","theta"))



# 一つならラベルスイッチングは起きない
fit_sp <- sampling(stanmodel,data=standata,chain=1,iter=max.iter) 
# Rhatも優秀
print(fit_sp,pars=c("mu","sig2","theta"))
print(fit_sp,pars=c("u"))


# チェインを複数にするとラベルスイッチングが起きる 
fit_sp <- sampling(stanmodel,data=standata,chain=C,iter=max.iter,warmup=burn.in)
# Rhatが悪くなってもくじけちゃいけない
print(fit_sp,pars=c("mu","sig2","theta"))
# 目で見ると綺麗なラベルスイッチングが確認できるよ。
traceplot(fit_sp,pars="mu")