# Estimate Hidden Markov Model with both continuous and discrete observations, possible missing data

My code estimates a Hidden Markov Model (HMM) using the expectations-maximization (EM) algorithm. I followed https://www.princeton.edu/~rvan/orf557/hmm080728.pdf with modifications to allow for panel data, time-varying transition matrices for the hidden state, and the possibility of missing observations (missing at random). The code is written in R. My main concerns are correctness and speed. I'm sure the code could be faster in, say, C++, but I wonder whether there's anything I can do to make my R implementation faster. To the best of my knowledge the code does work correctly -- I'm posting here for feedback on speed and any other facets of the code that you'd like to comment on.

Notation: x is the hidden state, y are multivariate normal observations, z are discrete (integer-valued) observations. X follows a Markov process with time-varying transition matrices. I observe a panel of independent realizations of the HMM. The parameters to estimate are the initial distribution for x, the transition matrices for x, the mean vector and covariance matrices of (y | x) and the vector of probabilities (i.e. probability mass function) for (z | x).

File hmm_functions.R:

library(mixtools)  # For rmvnorm
library(parallel)

dmvnorm_with_NA <- function(y, mu, sigma) {
## Density of multivariate normal y (can include NAs) with mean mu and covariance matrix sigma
index_non_NA <- which(!is.na(y))
stopifnot(length(index_non_NA) > 0)  # Can't all be NA
return(dmvnorm(y[index_non_NA], mu[index_non_NA], sigma[index_non_NA, index_non_NA]))
}

rmvnorm_hmm <- function(x, m_matrix, v_array) {
## Returns draws of multivariate y whose distribution depends on hidden x
stopifnot(is.vector(x))
stopifnot(nrow(m_matrix) == dim(v_array)[0])
y <- t(vapply(x, function(x) rmvnorm(1, m_matrix[x, ], v_array[x, , ]), FUN.VALUE=m_matrix[1, ]))
stopifnot(is.matrix(y) && nrow(y) == length(x))
return(y)
}

baum_welch <- function(panel_element, params) {
## Baum-Welch for HMM with discrete hidden x
## Observations are multivariate normal y and discrete z, both of which can include NAs
## Written following Ramon van Handel's HMM notes, page 40, algorithm 3.2
## https://www.princeton.edu/~rvan/orf557/hmm080728.pdf
## Careful, his observation index is in {0, 1, ... , n} while I use {1, 2, ... , t_max}
library("mixtools")  # Avoid "no visible binding for global variable 'mixtools'" when compiling
stopifnot(is.list(panel_element))
stopifnot(all(c("y", "z") %in% names(panel_element)))
y <- panel_element$y # Multivariate normal observations z <- panel_element$z  # Discrete observations in {1, 2, ... , |Z|}
stopifnot(is.list(params))
stopifnot(all(c("mu",  # Initial distribution for x
"m_matrix",  # (y | x) ~ Normal(m_matrix[x, ], v_array[x, , ])
"v_array",
"P_list") %in% names(params)))  # Time-varying transition matrix for x
stopifnot(is.list(params$P_list)) stopifnot(length(params$P_list) == nrow(y) - 1)  # Transition count versus observation count
state_space <- seq(1, nrow(params$m_matrix)) # State space for hidden x stopifnot(length(params$mu) == length(state_space))
stopifnot(dim(params$v_array)[1] == length(state_space)) stopifnot(all(dim(params$v_array)[2:3] == ncol(params$m_matrix))) stopifnot(all(vapply(state_space, function(x) all(eigen(params$v_array[x, , ])$values > 0), FUN.VALUE=TRUE))) stopifnot(all(params$mu >= 0))
stopifnot(isTRUE(all.equal(sum(params$mu), 1.0))) # Careful with float comparisons stopifnot(is.matrix(y)) stopifnot(ncol(y) == ncol(params$m_matrix))
t_max <- nrow(y)
stopifnot(t_max > 1)
c <- vector("numeric", t_max)
sigma_y <- lapply(state_space, function(x) { params$v_array[x, , ] }) upsilon <- vapply(state_space, function(x) { dmvnorm_with_NA(y[1, ], mu=as.vector(params$m_matrix[x, ]), sigma=sigma_y[[x]])
}, FUN.VALUE=1)
if(!is.na(z[1])) {
upsilon <- upsilon * params$pr_z[, z[1]] } stopifnot(length(upsilon) == length(state_space)) c[1] <- sum(upsilon * params$mu)
pi_k <- matrix(NA, length(state_space), t_max)
pi_k[, 1] <- upsilon * params$mu / c[1] P_transpose_list <- lapply(params$P_list, t)
upsilon_list <- list()
upsilon_list[[1]] <- upsilon
for(k in seq(2, t_max)) {
## Forward loop
upsilon <- vapply(state_space, function(x) {
dmvnorm_with_NA(y[k, ], mu=params$m_matrix[x, ], sigma=sigma_y[[x]]) }, FUN.VALUE=1) if(!is.na(z[k])) { upsilon <- upsilon * params$pr_z[, z[k]]
}
upsilon_list[[k]] <- upsilon  # Avoid calling dmvnorm again on backward loop
pi_tilde <- upsilon * P_transpose_list[[k-1]] %*% pi_k[, k-1]
c[k] <- sum(pi_tilde)
pi_k[, k] <- pi_tilde  / c[k]
}
beta <- matrix(NA, length(state_space), t_max)
beta[, t_max] <- 1 / c[t_max]
pi_k_n <- matrix(NA, length(state_space), t_max)
pi_k_n[, t_max] <- pi_k[, t_max]
pi_transition <- array(NA, dim=c(rep(length(state_space), 2), t_max - 1))
pi_transition_list <- list()
for(k in seq(1, t_max - 1)) {
## Backward loop
upsilon <- diag(upsilon_list[[t_max - k + 1]], length(state_space), length(state_space))
pi_matrix <- diag(pi_k[, t_max - k], length(state_space), length(state_space))
beta_matrix <- diag(beta[, t_max - k + 1], length(state_space), length(state_space))
P_matrix <- params$P_list[[t_max - k]] beta[, t_max - k] <- P_matrix %*% upsilon %*% beta[, t_max - k + 1] / c[t_max - k] pi_transition[, , t_max - k] <- pi_matrix %*% P_matrix %*% upsilon %*% beta_matrix pi_transition_list[[t_max - k]] <- pi_transition[, , t_max - k] if(!isTRUE(all.equal(sum(pi_transition[, , t_max - k]), 1.0))) { message(sum(pi_transition[, , t_max - k])) } stopifnot(isTRUE(all.equal(sum(pi_transition[, , t_max - k]), 1.0))) pi_k_n[, t_max - k] <- rowSums(pi_transition[, , t_max - k]) } return(list(c=c, beta=beta, pi_k=pi_k, pi_k_n=pi_k_n, pi_transition=pi_transition, pi_transition_list=pi_transition_list)) } baum_welch_compiled <- compiler::cmpfun(baum_welch) em_parameter_estimates <- function(panel, params, max_iter=100, epsilon=10^-3, max_num_cores=4) { ## EM for panel of independent HMM realizations; stop at max_iter or distance < epsilon ## Written following Ramon van Handel's HMM notes page 87, algorithm 6.1, modified for panel data ## https://www.princeton.edu/~rvan/orf557/hmm080728.pdf num_cores <- min(detectCores(), max_num_cores) cluster <- makeCluster(num_cores) # Call stopCluster when done clusterExport(cl=cluster, varlist=c("baum_welch_compiled", "dmvnorm_with_NA"), envir=.GlobalEnv) stopifnot(is.list(panel)) panel_size <- length(panel) stopifnot(all(c("mu", # Initial distribution over hidden x "m_matrix", # Expected values of y | hidden x "v_array", # Variance-covariance of y | hidden x "P_list") %in% names(params))) # Time-varying transition matrix for x stopifnot(nrow(params$pr_z) == nrow(params$x_matrix)) stopifnot(isTRUE(all.equal(rowSums(params$pr_z), rep(1, nrow(params$pr_z))))) stopifnot(all(vapply(panel, function(x) "y" %in% names(x), FUN.VALUE=TRUE))) # Matrix y stopifnot(all(vapply(seq(1, panel_size), function(n) is.matrix(panel[[n]]$y), FUN.VALUE=TRUE)))
stopifnot(all(vapply(panel, function(x) "z" %in% names(x), FUN.VALUE=TRUE)))  # Vector z
stopifnot(all(vapply(seq(1, panel_size), function(n) is.vector(panel[[n]]$z), FUN.VALUE=TRUE))) observation_lengths <- vapply(seq(1, panel_size), function(n) nrow(panel[[n]]$y), FUN.VALUE=1)
stopifnot(all(observation_lengths > 1) && length(unique(observation_lengths)) == 1)
observation_length <- observation_lengths[1]
state_space <- seq(1, length(params$mu)) # State space for hidden x iteration <- 1 while(iteration <= max_iter) { baum <- parLapply(cluster, panel, function(panel_element, params) { baum_welch_compiled(panel_element, params) }, params) P_list <- lapply(seq(1, observation_length - 1), function(time) { numerators <- parLapply(cluster, baum, function(baum_element) { baum_element$pi_transition_list[[time]]
})
denominators <- parLapply(cluster, baum, function(baum_element) {
matrix(baum_element$pi_k_n[, time], length(state_space), length(state_space)) }) numerator <- Reduce("+", numerators) denominator <- Reduce("+", denominators) stopifnot(all(rowSums(denominator) > 0)) # Can fail if a hidden x has zero probability stopifnot(all(denominator > 0)) P <- numerator / denominator stopifnot(isTRUE(all.equal(rowSums(P), rep(1, length(state_space))))) return(P) }) mus <- parLapply(cluster, baum, function(baum_element) { baum_element$pi_k_n[, 1]  # Initial distribution
})
mu <- Reduce("+", mus) / panel_size
stopifnot(isTRUE(all.equal(sum(mu), 1)))
panel_with_baum <- Map(c, panel, baum)  # For updates to m_matrix, v_array and pr_z
stopifnot(length(panel_with_baum) == length(panel) &&
length(panel_with_baum) == length(baum))
m_matrices <- parLapply(cluster, panel_with_baum, function(panel_element) {
## Careful with possible NAs in y
y_without_NA <- ifelse(is.na(panel_element$y), 0, panel_element$y)
stopifnot(identical(dim(y_without_NA), dim(panel_element$y))) return(t(vapply(state_space, function(x) { pi_k_n <- as.vector(panel_element$pi_k_n[x, ])
return(colSums(pi_k_n * y_without_NA))  # Sum over time
}, FUN.VALUE=vector("numeric", ncol(y_without_NA)))))
})
m_matrix_weights <- parLapply(cluster, panel_with_baum, function(panel_element) {
y_non_NA <- 1*(!is.na(panel_element$y)) stopifnot(identical(dim(y_non_NA), dim(panel_element$y)))
return(t(vapply(state_space, function(x) {
pi_k_n <- as.vector(panel_element$pi_k_n[x, ]) return(colSums(pi_k_n * y_non_NA)) # Sum over time }, FUN.VALUE=vector("numeric", ncol(y_non_NA))))) }) m_matrix <- Reduce("+", m_matrices) / Reduce("+", m_matrix_weights) v_arrays <- parLapply(cluster, panel_with_baum, function(panel_element) { v_array <- array(0, c(length(state_space), rep(ncol(m_matrix), 2))) v_list <- lapply(state_space, function(x) { y_residuals <- t(panel_element$y) - as.vector(m_matrix[x, ])
y_residuals_without_NA <- ifelse(is.na(y_residuals), 0, y_residuals)
stopifnot(identical(dim(y_residuals), dim(y_residuals_without_NA)))
weights <- as.vector(panel_element$pi_k_n[x, ]) stopifnot(length(weights) == ncol(y_residuals)) sigma_hat <- y_residuals_without_NA %*% (weights * t(y_residuals_without_NA)) stopifnot(all(!is.na(sigma_hat))) stopifnot(all(dim(sigma_hat) == rep(ncol(m_matrix), 2))) return(sigma_hat) }) for(x in state_space) { v_array[x, , ] <- v_list[[x]] } return(v_array) }) v_array_weights <- parLapply(cluster, panel_with_baum, function(panel_element) { weight <- array(0, c(length(state_space), rep(ncol(m_matrix), 2))) weight_list <- lapply(state_space, function(x) { y_non_NA <- t(1*(!is.na(panel_element$y)))
weights <- as.vector(panel_element$pi_k_n[x, ]) stopifnot(length(weights) == ncol(y_non_NA)) return(y_non_NA %*% (weights * t(y_non_NA))) }) for(x in state_space) { weight[x, , ] <- weight_list[[x]] } return(weight) }) v_array <- Reduce("+", v_arrays) / Reduce("+", v_array_weights) ## Careful, z can contain NAs pr_z_weights <- parLapply(cluster, panel_with_baum, function(panel_element, params) { z_non_NA <- matrix(1 * !is.na(panel_element$z),
nrow(params$pr_z), length(panel_element$z), byrow=T)
rowSums(panel_element$pi_k_n * z_non_NA) # Sum over time }, params) pr_z_matrices <- parLapply(cluster, panel_with_baum, function(panel_element, params) { z_space <- seq(1, ncol(params$pr_z))
column_list <- lapply(z_space, function(curr_z) {
z_indicators <- matrix(!is.na(panel_element$z) & panel_element$z == curr_z,
nrow(params$pr_z), length(panel_element$z), byrow=T)
return(rowSums(panel_element$pi_k_n * z_indicators)) # Sum over time }) return(do.call(cbind, column_list)) }, params) pr_z <- Reduce("+", pr_z_matrices) / Reduce("+", pr_z_weights) stopifnot(isTRUE(all.equal(rowSums(pr_z), rep(1, nrow(pr_z))))) ## Take sup norm between previous parameters and updated values P_list_distances <- vapply(seq_along(params$P_list), function(i) {
max(abs(params$P_list[[i]] - P_list[[i]])) }, FUN.VALUE=1) distance <- max(abs(params$mu - mu),
abs(m_matrix - params$m_matrix), abs(v_array - params$v_array),
abs(pr_z - params$pr_z), P_list_distances) message("distance ", round(distance, 4)) params$mu <- mu
params$m_matrix <- m_matrix params$v_array <- v_array
params$P_list <- P_list params$pr_z <- pr_z
if(distance < epsilon) break
iteration <- iteration + 1
}
stopCluster(cluster)
return(params)
}

simulate_hmm_panel <- function(panel_size, params) {
## Simulate panel_size independent HMMs of length t_max
stopifnot(panel_size > 0)
hmm_list <- lapply(seq(1, panel_size), function(unused) {
return(simulate_hmm(params=params))
})
return(hmm_list)
}

simulate_hmm <- function(params) {
## Simuate HMM where unobserved x follows a Markov chain with time-varying transition matrices
## Conditional on x, observe multivariate normal observations y and discrete observations z
stopifnot(is.list(params))
stopifnot(all(c("mu",  # Initial distribution for x
"m_matrix",  # (y | x) ~ Normal(m_matrix[x, ], v_array[x, , ])
"v_array",
"P_list",  # Time-varying transition matrices for x
"pr_z") %in% names(params)))  # (z | x) ~ Discrete(pr_z[x, ])
t_max <- length(params$P_list) + 1 state_space <- seq(1, length(params$mu))
stopifnot(is.matrix(params$m_matrix)) stopifnot(nrow(params$m_matrix) == length(state_space))
stopifnot(is.array(params$v_array)) stopifnot(dim(params$v_array)[1] == length(state_space))
stopifnot(all(dim(params$v_array[2:3]) == ncol(params$m_matrix)))
stopifnot(all(vapply(state_space,
function(x) all(eigen(params$v_array[x, , ])$values > 0), FUN.VALUE=TRUE)))
x <- simulate_discrete_markov(t_max=t_max, P_list=params$P_list, mu=params$mu)
y <- rmvnorm_hmm(x, params$m_matrix, params$v_array)
z_space <- seq(1, ncol(params$pr_z)) z <- vapply(x, function(x) { sample(z_space, size=1, prob=params$pr_z[x, ])
}, FUN.VALUE=1)
return(list(x=x, y=y, z=z))
}

simulate_discrete_markov <- function(t_max, P_list, mu) {
## Simulate distrete markov chain with transitions P_list, initial distribution mu
stopifnot(t_max > 1)
stopifnot(is.list(P_list) && length(P_list) == t_max- 1)
stopifnot(all(vapply(seq_along(P_list), function(i) {
isTRUE(all.equal(rowSums(P_list[[i]]), rep(1, nrow(P_list[[i]]))))
}, FUN.VALUE=TRUE)))
P_list_dims <- vapply(seq_along(P_list), function(i) {
dim(P_list[[i]])
}, FUN.VALUE=c(0, 1))
stopifnot(length(unique(as.vector(P_list_dims))) == 1)  # Same nrow, ncol in all P
state_space <- seq(1, nrow(P_list[[1]]))
stopifnot(is.vector(mu) && all(mu >= 0) && length(mu) == length(state_space))
stopifnot(isTRUE(all.equal(sum(mu), 1.0)))
x <- vector("numeric", t_max)
x[1] <- sample(state_space, 1, prob=mu)
for (t in seq(2, t_max)) {
x[t] <- sample(state_space, 1, prob=P_list[[t - 1]][x[t - 1], ])
}
stopifnot(all(x %in% state_space))
return(x)
}


File hmm_estimation.R:

## Simulate from HMM with known parameters and recover using EM

source("./hmm_functions.R")

set.seed(9237)

P1 <- rbind(c(0.7, 0.2, 0.1),
c(0.2, 0.7, 0.1),
c(0.1, 0.1, 0.8))  # Transitions for hidden x between periods 1 and 2
P2 <- rbind(c(0.6, 0.1, 0.3),
c(0.2, 0.3, 0.5),
c(0.5, 0.4, 0.1))  # Transitions for hidden x between periods 2 and 3
P3 <- rbind(c(0.1, 0.1, 0.8),
c(0.2, 0.7, 0.1),
c(0.6, 0.1, 0.3))  # Transitions for hidden x between periods 3 and 4
P_list <- list(P1, P2, P3)
m_matrix <- rbind(c(0, 5, 10, 5),
c(2, 5, 12, 8),
c(0, 0, 5, 5))  # Observe (y | x) ~ Normal(m_matrix[x, ], v_array[x, , ])
v_list <- list(rbind(c(2, 1, 1, 1),
c(1, 3, 1, 1),
c(1, 1, 4, 1),
c(1, 1, 1, 5)),
rbind(c(5, 2, 1, 1),
c(2, 3, 1, 1),
c(1, 1, 4, 1),
c(1, 1, 1, 1)),
3*diag(ncol(m_matrix)))
v_array <- array(0, c(nrow(m_matrix), rep(ncol(m_matrix), 2)))
for(x in seq_len(nrow(m_matrix))) {
v_array[x, , ] <- v_list[[x]]
}
mu <- rep(1/nrow(m_matrix), nrow(m_matrix))
pr_z <- rbind(c(0.60, 0.25, 0.05, 0.10),
c(0.10, 0.75, 0.05, 0.10),
c(0.05, 0.45, 0.45, 0.05))  # (z | x) ~ Discrete(pr_z[x, ])
params0 <- list(P_list=P_list, m_matrix=m_matrix, v_array=v_array, pr_z=pr_z, mu=mu)

markov_sanity_check <- simulate_discrete_markov(t_max=length(P_list) + 1, P_list=P_list, mu=mu)
hmm_sanity_check <- simulate_hmm(params=params0)

panel <- simulate_hmm_panel(panel_size=10000, params=params0)  # Simulation speed isn't crucial
baum_sanity_check <- baum_welch(panel_element=panel[[1]], params=params0)
round(baum_sanity_check$pi_k_n, 3) # Posterior over x given y and z, compare to panel[[1]]$x

panel_with_NA <- panel  # Data missing at random with no dependence on x, y, z
for(idx in seq_along(panel)) {
time_indices <- seq_len(nrow(panel[[idx]]$y)) for (time in time_indices) { random_cols <- sample(seq_len(ncol(panel[[idx]]$y)), size=floor(ncol(panel[[idx]]$y)/2)) panel_with_NA[[idx]]$y[time, random_cols] <- NA
}
random_time <- sample(time_indices, size=1)
panel_with_NA[[idx]]$z[random_time] <- NA } stopifnot(any(is.na(panel_with_NA[[1]]$y)))
stopifnot(!all(is.na(panel_with_NA[[1]]$y))) ## Call these lines manually using lineprof params_hat0 <- em_parameter_estimates(panel_with_NA, params0) # I'd like this to run faster ## Estimation above started at true parameters, try again starting at incorrect parameters uniform_P <- matrix(1 / nrow(m_matrix), nrow(m_matrix), nrow(m_matrix)) m_matrix_incorrect <- m_matrix + rbind(c(0, 1, 1, -2), c(2, 1, 2, 2), c(0, 0, -1, -1)) pr_z_incorrect <- 0.6*pr_z + 0.4*matrix(1/ncol(pr_z), nrow(pr_z), ncol(pr_z)) params1 <- list(P_list=list(uniform_P, uniform_P, uniform_P), m_matrix=m_matrix_incorrect, pr_z=pr_z_incorrect, v_array=v_array, mu=mu) params_hat1 <- em_parameter_estimates(panel_with_NA, params1) ## Examine distances between params_hat1, params_hat0 and true parameters, should be small ## Careful with these comparisons, I'm assuming ordering of hidden x did not change round(abs(params_hat0$P_list[[1]] - params0$P_list[[1]]), 3) round(abs(params_hat0$m_matrix - params0$m_matrix), 3) round(abs(params_hat0$v_array[1, , ] - params0$v_array[1, , ]), 3) round(abs(params_hat0$pr_z - params0$pr_z), 3) ## Same comparisons as above, now params_hat0 versus params_hat1, expect smaller distances round(abs(params_hat0$P_list[[1]] - params_hat1$P_list[[1]]), 3) round(abs(params_hat0$m_matrix - params_hat1$m_matrix), 3) round(abs(params_hat0$v_array[1, , ] - params_hat1$v_array[1, , ]), 3) round(abs(params_hat0$pr_z - params_hat1\$pr_z), 3)

• I think you're likely to get more feedback if you were to break your question into chunks, with maybe one or several related functions from your code per chunk. The downside of doing this is that you'll prevent anyone from being able to offer feedback on how the whole thing fits together, but I'm fairly confident you'll get a lot more traction with more "bite-sized" questions. – Curt F. May 28 '16 at 21:03
• @CurtF. thank you, point well taken -- I agree that my question in its current form is quite long, and that answering it would require a fairly large chunk of someone's time. – Adrian May 28 '16 at 21:10