Stanの隠れマルコフモデルでrow_stochastic_matrixをつかう

CmdStan 2.36がリリースされました。このバージョンでは、row_stochastic_matrixという制限付きの型がくわわりました。これは各行が確率単体(simplex)となっている行列です。（同時に、column_stochastic_matrixとsum_to_zero_vectorもくわわりました。それぞれ、列が確率単体の行列、合計すると0になるベクトルです。）そこで、以前「Stanのhmm_marginalによる隠れマルコフモデルのあてはめ」で使用したモデルをcolumn_stochastic_matrixをつかって、かきなおしてみました。

模擬データ

前回と同じ隠れマルコフモデルで模擬データを生成します。

モデル

潜在状態、観測値とも3値をとります。

潜在状態の推移確率は、1→1: 0.5, 1→2: 0.4, 1→3: 0.1, 2→1: 0.3, 2→2: 0.6, 2→3: 0.1, 3→1: 0.1, 3→2: 0.1, 3→3: 0.8とします。

\[ \left(\begin{array}{lll}0.5 & 0.4 & 0.1 \\ 0.3 & 0.6 & 0.1 \\ 0.1 & 0.1 & 0.8 \end{array}\right) \]

潜在状態から観測値への出力確率は、1→1: 0.9, 1→2: 0.09, 1→3: 0.01, 2→1: 0.1, 2→2: 0.8, 2→3: 0.1, 3→1: 0.05, 3→2: 0.05, 3→3: 0.9とします。

\[ \left(\begin{array}{lll}0.9 & 0.09 & 0.01 \\ 0.1 & 0.8 & 0.1 \\ 0.05 & 0.05 & 0.9 \end{array}\right) \]

模擬データの生成

上のモデルで模擬データを生成します。時点数Tを200、潜在状態をx、観測値をyとします。乱数のseedは前回とはかえました。

library(extraDistr)
set.seed(231)

transition <- matrix(c(0.5,  0.4,  0.1,
                       0.3,  0.6,  0.1,
                       0.1,  0.1,  0.8), ncol = 3, byrow = TRUE)
emission   <- matrix(c(0.9,  0.09, 0.01,
                       0.1,  0.8,  0.1,
                       0.05, 0.05, 0.9), ncol = 3, byrow = TRUE)
T <- 200

## prepare data
x <- rep(0, T)
y <- rep(0, T)

### latent state
x[1] <- 1
for (t in 2:T) {
  x[t] <- rcat(1, transition[x[t - 1], ])
}

### observation
for (t in 1:T) {
  y[t] <- rcat(1, emission[x[t], ])
}

Stanのモデル

なにしろリリースされたばかりのCmdStan 2.36.0をつかいたいので、CmdStanRをつかいます。

library(cmdstanr)
## This is cmdstanr version 0.8.1
## - CmdStanR documentation and vignettes: mc-stan.org/cmdstanr
## - CmdStan path: /usr/local/cmdstan
## - CmdStan version: 2.36.0

モデルコード

モデルコードです。前回同様、出力確率は既知として、遷移確率を推定します。

7行目と18行目でrow_stochastic_matrixをつかっています。前回のコードでは、遷移確率行列のGammaを定義するのに、いったん各行をsimplexとして定義して、そのあとで行列にまとめていましたが、その部分がすっきりとしました。

// hidden Morkov model using hmm_marginal function
// and row_stochastic_matrix

data {
  int<lower=0> T;                   // number of observations
  array[T] int<lower=1, upper=3> Y; // observations
  row_stochastic_matrix[3, 3] p;    // emission probability
}

transformed data {
  matrix[3, T] log_omega; // log density for each output

  for (t in 1:T)
    log_omega[1:3, t] = log(p[, Y[t]]);
}

parameters {
  row_stochastic_matrix[3, 3] Gamma; // transition matrix
  simplex[3] rho;                    // initial state probability
}

model {
  target += hmm_marginal(log_omega, Gamma, rho);
}

実行

コンパイルとあてはめ

コンパイルして、あてはめます。

model <- cmdstan_model("models/hmm-with-row_stochastic_matrix.stan")
fit <- model$sample(data = list(T = T, Y = y, p = emission),
                    iter_warmup = 1000, iter_sampling = 1000,
                    chains = 4, parallel_chains = 4,
                    refresh = 0)
## Running MCMC with 4 parallel chains...
## 
## Chain 1 finished in 1.0 seconds.
## Chain 2 finished in 1.1 seconds.
## Chain 3 finished in 1.1 seconds.
## Chain 4 finished in 1.0 seconds.
## 
## All 4 chains finished successfully.
## Mean chain execution time: 1.1 seconds.
## Total execution time: 1.1 seconds.

結果

推定された遷移確率の事後分布の要約です。

fit$print("Gamma")
##    variable mean median   sd  mad   q5  q95 rhat ess_bulk ess_tail
##  Gamma[1,1] 0.56   0.56 0.08 0.08 0.43 0.68 1.00     4958     3090
##  Gamma[2,1] 0.38   0.37 0.10 0.10 0.22 0.55 1.00     4915     3242
##  Gamma[3,1] 0.13   0.13 0.06 0.06 0.04 0.24 1.00     3968     1921
##  Gamma[1,2] 0.24   0.23 0.08 0.08 0.11 0.37 1.00     4495     3110
##  Gamma[2,2] 0.48   0.48 0.10 0.11 0.31 0.65 1.00     4836     3185
##  Gamma[3,2] 0.14   0.14 0.06 0.06 0.05 0.25 1.00     4444     2364
##  Gamma[1,3] 0.21   0.20 0.06 0.06 0.11 0.32 1.01     6198     2685
##  Gamma[2,3] 0.14   0.13 0.07 0.07 0.04 0.27 1.00     4456     2037
##  Gamma[3,3] 0.72   0.72 0.07 0.07 0.61 0.83 1.00     6043     3482