Stan: 欠測値のある隠れマルコフモデル

欠測値のある隠れマルコフモデルのコードをStanでかいてみました。

想定したのは、ある動物（N_ind匹）がいて、観察フィールドの内に ある確率(phi)でとどまる一方、ある確率(1-phi)で外へでていく、また、ある確率(1-psi)で外から内に はいってきて、観察フィールド内にいる場合には、ある確率(p)で観察される、というものです。

Stanコードです。Stan Reference Manual 2.17.0の10.6節を参考にしています。欠測の場合には、観測モデルの確率を1としています。

data {
  int<lower = 0> N_ind;
  int<lower = 0> N_t;
  int<lower = 0, upper = 2> Y[N_ind, N_t]; // 0: missing value
}

parameters {
  real<lower = 0, upper = 1> phi;
  real<lower = 0, upper = 1> psi;
  real<lower = 0, upper = 1> p;
}

transformed parameters {
  simplex[2] ps[2];
  simplex[2] po[2];

  ps[1, 1] = phi;
  ps[1, 2] = 1 - phi;
  ps[2, 1] = 1 - psi;
  ps[2, 2] = psi;

  po[1, 1] = p;
  po[1, 2] = 1 - p;
  po[2, 1] = 0;
  po[2, 2] = 1;
}

model {
  real acc[2];
  vector[2] gamma[N_t];

  for (i in 1:N_ind) {
    for (k in 1:2)
      gamma[1, k] = (k == Y[i, 1]);

    for (t in 2:N_t) {
      for (k in 1:2) {
        for (j in 1:2) {
          if (Y[i, t] == 0) {
            acc[j] = gamma[t - 1, j] * ps[j, k];
          } else {
            acc[j] = gamma[t - 1, j] * ps[j, k] * po[k, Y[i, t]];
          }
        }
        gamma[t, k] = sum(acc);
      }
    }
    target += log(sum(gamma[N_t]));
  }
}

Rコードです。

# Hidden Markov model with missing values

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

# State
# 1: inside the field
# 2: outside the field
# Observation
# 1: detected
# 2: not detected

set.seed(20180310)
N_ind <- 20
N_t <- 40

# Transition probability
phi <- 0.6  # Prob. state 1 -> 1
psi <- 0.8  # Prob. state 2 -> 2
pt11 <- phi
pt12 <- 1 - phi
pt21 <- 1 - psi
pt22 <- psi

# Observation probability
p <- 0.9   # Prob. state 1 -> observation 1
po11 <- p
po12 <- 1 - p
po21 <- 0
po22 <- 1

state <- matrix(NA, nrow = N_ind, ncol = N_t)
state[, 1] <- rbinom(N_ind, 1, pt11) + 1
for (t in 2:N_t) {
  for (i in 1:N_ind) {
    state[i, t] <- (2 - state[i, t - 1]) * (rbinom(1, 1, pt12) + 1) +
      (state[i, t - 1] - 1) * (rbinom(1, 1, pt22) + 1)
  }
}

obs <- matrix(NA, nrow = N_ind, ncol = N_t)
for (t in 1:N_t) {
  for (i in 1:N_ind) {
    obs[i, t] <- (2 - state[i, t]) * (rbinom(1, 1, po12) + 1) +
      (state[i, t] - 1) * (rbinom(1, 1, po22) + 1)
  }
}

    
stan_data <- list(N_ind = N_ind,
                  N_t = N_t,
                  Y = obs)
params <- c("phi", "psi", "p")
fit <- stan("hmm_miss.stan", data = stan_data,
            pars = params,
            iter = 2000, warmup = 1000, thin = 1)
print(fit)

## Make missing values
obs[, 15] <- 0
obs[, 25] <- 0

fit2 <- stan("hmm_miss.stan", data = stan_data,
             pars = params,
             iter = 2000, warmup = 1000, thin = 1)
print(fit2)

結果です。だいたいあっているような気がするのですが。

> print(fit)
Inference for Stan model: hmm_miss.
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
phi     0.52    0.00 0.05    0.43    0.49    0.52    0.55    0.64   426 1.00
psi     0.80    0.00 0.02    0.76    0.79    0.81    0.82    0.84   601 1.01
p       0.90    0.00 0.08    0.72    0.86    0.92    0.97    1.00   397 1.00
lp__ -424.64    0.05 1.26 -428.03 -425.26 -424.32 -423.72 -423.14   707 1.00

> print(fit2)
Inference for Stan model: hmm_miss.
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
phi     0.52    0.00 0.05    0.43    0.49    0.52    0.55    0.63  1262    1
psi     0.81    0.00 0.02    0.76    0.79    0.81    0.82    0.84  1226    1
p       0.90    0.00 0.07    0.72    0.86    0.92    0.96    1.00  1038    1
lp__ -424.55    0.03 1.22 -427.58 -425.15 -424.24 -423.65 -423.15  1540    1

タグ：STAn