tests/testthat/test.baum_welch.R

In ricky-kotecha/rkHMM:

test_that("test_em 2 variables", {
  means <- list(
    c(1, 1),
    c(1, 1)
  )
  st_devs = list(
    c(3, 5),
    c(3, 5)
  )
  corrs = list(
    c(0.9),
    c(0.1)
  )
  test_data <- get_multivariable_test_data(
    means = means,
    st_devs = st_devs,
    corrs = corrs
  )

  n_states <- 2
  init_p <- c(0.8, 0.2)
  init_trn_mtx <- matrix(c(0.9, 0.1, 0.1, 0.9), n_states, n_states)

  mv_fit <- fit_multivariable(
    time_series = test_data,
    n_states = n_states,
    init_p = init_p,
    init_trn_mtx = init_trn_mtx,
    verbose = FALSE,
    random.start = FALSE,
    use_rk = FALSE
  )

  bw_fit <- baum_welch(
    data = test_data,
    n_states = n_states,
    init_p = init_p,
    init_trn_mtx = init_trn_mtx,
    init_trn_model_data = data.frame(1),
    init_trn_model_formula = ~ 1,
    verbose = FALSE,
    random.start = FALSE,
    use_rk = TRUE
  )

  expect_true(
    all(abs(as.numeric(mv_fit$means[[1]]) - as.numeric(bw_fit$response[[2]]$mean)) < 1e-8) |
      all(abs(as.numeric(mv_fit$means[[1]]) - as.numeric(bw_fit$response[[1]]$mean)) < 1e-8)
  )

  expect_true(
    all(abs(as.numeric(mv_fit$means[[2]]) - as.numeric(bw_fit$response[[2]]$mean)) < 1e-8) |
      all(abs(as.numeric(mv_fit$means[[2]]) - as.numeric(bw_fit$response[[1]]$mean)) < 1e-8)
  )

  expect_true(
    all(abs(as.numeric(mv_fit$cov_matrices[[1]]) - as.numeric(bw_fit$response[[2]]$cov_mtx)) < 1e-8) |
      all(abs(as.numeric(mv_fit$cov_matrices[[1]]) - as.numeric(bw_fit$response[[1]]$cov_mtx)) < 1e-8)
  )

  expect_true(
    all(abs(as.numeric(mv_fit$cov_matrices[[2]]) - as.numeric(bw_fit$response[[2]]$cov_mtx)) < 1e-8) |
      all(abs(as.numeric(mv_fit$cov_matrices[[2]]) - as.numeric(bw_fit$response[[1]]$cov_mtx)) < 1e-8)
  )

  trn_mtx_1 <- bw_fit$trn_mtx
  trn_mtx_2 <- bw_fit$trn_mtx
  diag(trn_mtx_2) <- rev(diag(trn_mtx_2))
  expect_true(
    all(abs(trn_mtx_1 - mv_fit$transition_matrix) < 1e-8) |
      all(abs(trn_mtx_2 - mv_fit$transition_matrix) < 1e-8)
  )
})