Testing a hidden Markov model function incorporated into CmdStan 2.24

hmm_test3.R

#
# HMM Test for CmdStan 2.24
#

library(rstan)
options(mc.cores = parallel::detectCores())
library(cmdstanr)
set_cmdstan_path("/usr/local/cmdstan")
library(extraDistr)

# transition probability
gam <- matrix(c(0.9, 0.25, 0.1, 0.75), 2, 2)

# emission probability
phi <- matrix(c(0.7, 0.05, 0.3, 0.95), 2, 2)

# length of the time series
N <- 1000

# generate data
# x:  latent state {1, 2}
# x0: initial latent state {1, 2}
# y:  output {1, 2}
# log_omega: log probability(y | x=k, phi)
set.seed(123)
y <- x <- vector("numeric", N)
x0 <- 1
x[1] <- rcat(1, gam[x0, ])
for (n in 2:N)
  x[n] <- rcat(1, gam[x[n - 1], ])
y <- rcat(N, phi[x, ])

# fitting
data <- list(N = N, Y = y)
model <- cmdstanr::cmdstan_model("hmm_test3.stan")
init_fun <- function() {
  phi1 <- runif(1, 0.8, 1)
  phi2 <- runif(1, 0.8, 1)
  rho1 <- runif(1, 0, 1)
  list(phi = matrix(c(phi1, 1 - phi1, 1 - phi2, phi2), 2, 2, byrow = TRUE),
       rho = c(rho1, 1 - rho1))
}
fit <- model$sample(data, init = init_fun, chains = 4,
                    iter_warmup = 1000, iter_sampling = 1000)
fit$summary(variables = c("Gamma", "phi", "rho"))
stanfit <- rstan::read_stan_csv(fit$output_files())
traceplot(stanfit, pars = c("p", "phi", "rho"))

hmm_test3.stan

data {
  int N;
//  int K = 2;
  int<lower = 1, upper = 2> Y[N];
}

parameters {
  simplex[2] p[2];    // transition probabilities
  simplex[2] phi[2];  // emission probabilities
  simplex[2] rho;     // initial state probabilities
}

transformed parameters {
  matrix<lower = 0, upper = 1>[2, 2] Gamma
    = append_row(p[1]', p[2]');
}

model {
  matrix[2, N] log_omega;

  for (n in 1:N)
    for (x in 1:2)
      log_omega[x, n] = log(phi[x][Y[n]]);
  target += hmm_marginal(log_omega, Gamma, rho);
}