1
\$\begingroup\$

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)
\$\endgroup\$
  • \$\begingroup\$ 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. \$\endgroup\$ – Curt F. May 28 '16 at 21:03
  • \$\begingroup\$ @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. \$\endgroup\$ – Adrian May 28 '16 at 21:10

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.