Poisson Hidden Markov Model

PoisHMM.R

# 須山『ベイズ推論による機械学習入門』（講談社）の実装例です

softmax <- function(x){
  maxx <- max(x)
  exp(x-maxx)/sum(exp(x-maxx))
}
logp_x <-function(x,lambda,loglambda){
  x*loglambda-lambda
}
logsumexp <- function(x){
  maxx <- max(x)
  maxx + log(sum(exp(x-maxx)))
}

#完全分解変分推論
VBHMMP <- function(x,alpha = c(1,1),beta = c(1,1),a = 1,b = 1){
  shat <- matrix(1/2,N,2)
  Elambda <- rgamma(2,a,1/mean(x))
  Eloglambda <- log(Elambda)
  Elogpi <- log(c(0.5,0.5))
  ElogA <- log(matrix(c(0.5,0.5,0.5,0.5),2,2))
  for (j in 1:1000) {
    lp <- logp_x(x[1],Elambda,Eloglambda)+Elogpi+
      as.vector(ElogA%*%shat[2,])
    shat[1,] <- softmax(lp)
    for (i in 2:(N-1)) {
      lp <- logp_x(x[i],Elambda,Eloglambda)+
        as.vector(shat[i-1,]%*%ElogA)+
        as.vector(shat[i,]%*%ElogA)
      shat[i,] <- softmax(lp)
    }
    lp <- logp_x(x[N],Elambda,Eloglambda)+
      as.vector(shat[N-1,]%*%ElogA)
    shat[N,] <- softmax(lp)
    ahat <- colSums(shat*x)+a
    bhat <- colSums(shat)+b
    Elambda <-ahat/bhat
    Eloglamba <- digamma(ahat)-log(bhat)
    alphahat <- shat[1,]+alpha
    Elogpi <- digamma(alphahat)-digamma(sum(alphahat))
    betahat <-t(shat[1:(N-1),])%*%shat[2:N,]+beta
    ElogA <- sweep(digamma(betahat),1,digamma(rowSums(betahat)),"-")
  }
  Epi <- alphahat/sum(alphahat)
  EA <- sweep(betahat,1,rowSums(betahat),"/")
  list(lambda=Elambda,pi=Epi,A=EA,s=shat)
}

fb <- function(x, lambda, loglambda, ln_pi, ln_A){
  smallvalue <- 0
  El_A <- exp(ln_A)
  N <- length(x)
  lp <- matrix(0, ncol = 2, nrow = N)
  for(i in 1:N){
    lp[i,] <- logp_x(x[i], lambda, loglambda)
  }
  ln_ZZ = array(0,dim=c(2, 2, N))
  ln_alpha = matrix(0,N,2)
  ln_beta = matrix(0,N,2)
  ln_st = numeric(N)
  ln_alpha[1,] = ln_pi + lp[1,]
  ln_st[1] = logsumexp(ln_alpha[1,])
  ln_alpha[1,] = ln_alpha[1,] - ln_st[1]
  for (i in 2:N) {
    ln_alpha[i,] = as.vector(log(El_A%*%exp(ln_alpha[i-1,]))) + lp[i,]
    ln_st[i] = logsumexp(ln_alpha[i,])
    ln_alpha[i,] = ln_alpha[i,] - ln_st[i]
  }
  for(i in (N-1):1){
    ln_beta[i,] = as.vector(log(t(El_A)%*%exp(ln_beta[i+1,] + lp[i+1,]+smallvalue)))
    ln_beta[i,] = ln_beta[i,] - ln_st[i+1] 
  }
  ln_Z = ln_alpha + ln_beta
  for (i in 1:(N-1)){
    ln_ZZ[,,i] = matrix(ln_alpha[i,],2,2,byrow = TRUE) + ln_A + matrix(ln_beta[i+1,]+ lp[i+1,],2,2)
    ln_ZZ[,,i] = ln_ZZ[,,i] - ln_st[i+1]
  }
  return(list(s=exp(ln_Z),ss=exp(ln_ZZ)))
}

#構造化変分推論
SVBHMMP <- function(x,alpha = c(1,1),beta = c(1,1),a = 1,b = 1){
  lp <- matrix(0,N,2)
  shat <- matrix(1/2,N,2)
  Elambda = rgamma(2,a,1/mean(x))
  Eloglambda <- log(Elambda)
  Elogpi <- log(c(0.5,0.5))
  ElogA <- log(matrix(c(0.5,0.5,0.5,0.5),2,2))
  for (j in 1:1000) {
    tmps <- fb(x,Elambda,Eloglambda,Elogpi,ElogA)
    shat <- tmps$s
    ahat <- colSums(tmps$s*x)+a
    bhat <- colSums(tmps$s)+b
    Elambda <-ahat/bhat
    Eloglambda <- digamma(ahat)-log(bhat)
    alphahat <- tmps$s[1,]+alpha
    Elogpi <- digamma(alphahat)-digamma(sum(alphahat))
    betahat <-apply(tmps$ss,1:2,sum)+beta
    ElogA <- sweep(digamma(betahat),1,digamma(rowSums(betahat)),"-")
  }
  Epi <- alphahat/sum(alphahat)
  EA <- sweep(betahat,1,rowSums(betahat),"/")
  list(lambda=Elambda,pi=Epi,A=EA,s=tmps$s)
}
A=matrix(c(0.9,0.1,0.1,0.9),2,2,byrow = TRUE)
lambda <- c(10,20)
N <- 100
s <- integer(N)
s[1] <- sample.int(2,1,prob=A[1,])
for(i in 2:N){
  s[i] <- sample.int(2,1,prob=A[s[i-1],])
}
x <- rpois(N,lambda[s])
plot(x,col=s)
hist(x)
out1 <- VBHMMP(x)
print(out1$A)
print(out1$lambda)
out2 <- SVBHMMP(x)
print(out2$A)
print(out2$lambda)