R/hmm.R

In bivarhr: Bivariate Hurdle Regression with Bayesian Model Averaging

#' Hidden Markov Model (HMM) for Path Dependence (Counts I and C)
#'
#' Fits a bivariate count Hidden Markov Model (HMM) to the series
#' \code{I} and \code{C} with a self-contained, dependency-free base-R
#' implementation of the Baum-Welch (EM) algorithm. Each latent state
#' \eqn{j} emits \code{I} and \code{C} as \emph{conditionally independent}
#' counts given the state, with either Poisson or Negative Binomial
#' emissions. With Poisson emissions the model coincides with an
#' intercept-only Poisson \code{depmix}; the Negative Binomial family adds
#' a per-state, per-series dispersion parameter and is the appropriate
#' choice for overdispersed counts (the same motivation as the hurdle
#' negative-binomial core of this package).
#'
#' @param DT A \code{data.frame} or \code{data.table} containing at least
#'   the columns \code{I} and \code{C}, interpreted as non-negative
#'   integer count series observed over time.
#' @param nstates Integer; number of latent Markov states (default 3).
#' @param family Character; emission family, either \code{"poisson"}
#'   (default) or \code{"nbinom"} (Negative Binomial, mean--dispersion
#'   parametrization). Negative Binomial is recommended when the counts
#'   are overdispersed.
#' @param seed Integer or \code{NULL}; seed controlling the random EM
#'   restarts. The first restart always uses a deterministic
#'   quantile-based initialization, so results are reproducible even when
#'   \code{seed} is \code{NULL}. The caller's random-number state is
#'   preserved (restored on exit).
#' @param n_starts Integer; number of EM restarts (multi-start) used to
#'   mitigate convergence to a poor local optimum. The fit with the
#'   highest log-likelihood is returned (default 10).
#' @param max_iter Integer; maximum EM iterations per restart (default 500).
#' @param tol Numeric; relative convergence tolerance on the
#'   log-likelihood increase between EM iterations (default 1e-8).
#' @param dir_csv Character scalar or \code{NULL}; directory where the
#'   state-sequence CSV (\code{"hmm_states.csv"}) is written. If
#'   \code{NULL} (default), no CSV is written.
#' @param dir_out Character scalar or \code{NULL}; directory where the
#'   fitted HMM object (\code{"hmm_fit.rds"}) is saved. If \code{NULL}
#'   (default), the object is not saved to disk.
#'
#' @details
#' \strong{Model.} The parameters are an initial-state distribution
#' \eqn{\delta} (length \code{nstates}), a row-stochastic transition
#' matrix \eqn{\Gamma}, per-state means \eqn{\mu^I_j}, \eqn{\mu^C_j} and,
#' for \code{family = "nbinom"}, per-state dispersions \eqn{\phi^I_j},
#' \eqn{\phi^C_j}. The emission density at time \eqn{t} in state \eqn{j}
#' factorizes as
#' \deqn{b_j(t) = f(I_t; \mu^I_j, \phi^I_j)\, f(C_t; \mu^C_j, \phi^C_j),}
#' with \eqn{f} the Poisson or Negative Binomial mass function. As
#' \eqn{\phi \to \infty} the Negative Binomial converges to the Poisson.
#'
#' \strong{Estimation.} The Baum-Welch algorithm is used. A vectorized,
#' \emph{log-space} forward-backward pass computes the smoothed state and
#' transition posteriors (E-step) using the log-sum-exp identity, so the
#' recursion is numerically stable for long series and small emission
#' probabilities. The M-step updates \eqn{\delta} and \eqn{\Gamma} in
#' closed form; the means \eqn{\mu} are the posterior-weighted sample
#' means; for Negative Binomial emissions the dispersions \eqn{\phi} are
#' updated by a bounded one-dimensional maximization of the weighted
#' emission log-likelihood. The EM iteration is monotone in the
#' log-likelihood; convergence is declared on a relative increase below
#' \code{tol}.
#'
#' \strong{Robustness.} Several safeguards are applied: multiple restarts
#' (the first deterministic, the rest random) keep the best fit; empty or
#' near-empty states retain their previous parameters instead of
#' collapsing; the RNG state of the caller is preserved; and the estimated
#' states are relabeled into a canonical order (increasing \eqn{\mu^I}) so
#' that the output is invariant to the arbitrary labeling of hidden states.
#' Model-selection criteria \code{AIC} and \code{BIC} are returned to help
#' choose \code{nstates}.
#'
#' When \code{dir_csv} is supplied, a CSV named \code{"hmm_states.csv"} is
#' written with columns \code{t} and \code{state}; when \code{dir_out} is
#' supplied, the fitted object is saved as \code{"hmm_fit.rds"}.
#'
#' @return An object of class \code{"bivarhr_hmm"}: a list with
#' \itemize{
#'   \item \code{fit}: a list with \code{nstates}, \code{family}, the
#'         estimated parameters (\code{mu_I}, \code{mu_C}, and, for
#'         \code{"nbinom"}, \code{phi_I}, \code{phi_C}), \code{delta},
#'         \code{Gamma}, the maximized \code{logLik}, \code{npar},
#'         \code{AIC}, \code{BIC}, the smoothed state posteriors
#'         (\code{posterior}, a \eqn{T \times} \code{nstates} matrix), and
#'         convergence information (\code{iterations}, \code{converged}).
#'   \item \code{states}: integer vector of Viterbi-decoded latent states.
#' }
#' If estimation fails (e.g., degenerate data), the function returns
#' \code{NULL}.
#'
#' @examples
#' DT <- data.frame(
#'   I = rpois(40, lambda = 4),
#'   C = rpois(40, lambda = 3)
#' )
#'
#' # 'n_starts' is kept small here for speed; the default (10) is more robust.
#' res_hmm <- run_hmm(DT, nstates = 2, seed = 1, n_starts = 2)
#' if (!is.null(res_hmm)) {
#'   print(res_hmm)
#'   table(res_hmm$states)
#' }
#'
#' \donttest{
#' # Negative Binomial emissions for overdispersed counts:
#' res_nb <- run_hmm(DT, nstates = 2, family = "nbinom", seed = 1, n_starts = 2)
#' }
#'
#' @export

run_hmm <- function(DT, nstates = 3, family = c("poisson", "nbinom"),
                    seed = NULL, n_starts = 10, max_iter = 500, tol = 1e-8,
                    dir_csv = NULL, dir_out = NULL) {

  family <- match.arg(family)

  # ---- Input validation --------------------------------------------------
  y_I <- DT$I
  y_C <- DT$C
  if (is.null(y_I) || is.null(y_C)) {
    stop("run_hmm(): 'DT' must contain columns 'I' and 'C'.", call. = FALSE)
  }
  y_I <- as.numeric(y_I)
  y_C <- as.numeric(y_C)
  if (anyNA(y_I) || anyNA(y_C)) {
    stop("run_hmm(): 'I' and 'C' must not contain missing values.", call. = FALSE)
  }
  if (any(!is.finite(y_I)) || any(!is.finite(y_C))) {
    stop("run_hmm(): 'I' and 'C' must be finite.", call. = FALSE)
  }
  if (any(y_I < 0) || any(y_C < 0)) {
    stop("run_hmm(): 'I' and 'C' must be non-negative counts.", call. = FALSE)
  }
  if (any(abs(y_I - round(y_I)) > 1e-8) || any(abs(y_C - round(y_C)) > 1e-8)) {
    warning("run_hmm(): non-integer values in 'I'/'C' were rounded to counts.",
            call. = FALSE)
  }
  y_I <- round(y_I)
  y_C <- round(y_C)

  Tn <- length(y_I)
  K <- as.integer(nstates)
  if (K < 1L || Tn < 2L) {
    stop("run_hmm(): need nstates >= 1 and at least 2 observations.", call. = FALSE)
  }
  n_distinct <- length(unique(paste(y_I, y_C)))
  if (K > n_distinct) {
    warning(sprintf(
      "run_hmm(): nstates (%d) exceeds the number of distinct (I,C) pairs (%d); states may be unidentified.",
      K, n_distinct), call. = FALSE)
  }

  # ---- Preserve the caller's RNG state -----------------------------------
  if (exists(".Random.seed", envir = .GlobalEnv, inherits = FALSE)) {
    old_seed <- get(".Random.seed", envir = .GlobalEnv, inherits = FALSE)
    on.exit(assign(".Random.seed", old_seed, envir = .GlobalEnv), add = TRUE)
  }
  set.seed(if (is.null(seed)) 1L else seed)

  # ---- Numerically stable log-sum-exp helpers ----------------------------
  # Column-wise logsumexp of a matrix: returns a vector of length ncol.
  col_lse <- function(M) {
    mx <- apply(M, 2L, max)
    mx[!is.finite(mx)] <- 0
    mx + log(colSums(exp(sweep(M, 2L, mx, "-"))))
  }
  # Row-wise logsumexp of a matrix: returns a vector of length nrow.
  row_lse <- function(M) {
    mx <- apply(M, 1L, max)
    mx[!is.finite(mx)] <- 0
    mx + log(rowSums(exp(sweep(M, 1L, mx, "-"))))
  }
  vec_lse <- function(v) {
    m <- max(v)
    if (!is.finite(m)) return(m)
    m + log(sum(exp(v - m)))
  }

  # ---- Emission log-density matrix (T x K) -------------------------------
  log_emission <- function(par) {
    if (family == "poisson") {
      lb <- outer(y_I, par$mu_I, function(y, l) stats::dpois(y, l, log = TRUE)) +
            outer(y_C, par$mu_C, function(y, l) stats::dpois(y, l, log = TRUE))
    } else {
      lb <- matrix(0, Tn, K)
      for (j in seq_len(K)) {
        lb[, j] <- stats::dnbinom(y_I, size = par$phi_I[j], mu = par$mu_I[j], log = TRUE) +
                   stats::dnbinom(y_C, size = par$phi_C[j], mu = par$mu_C[j], log = TRUE)
      }
    }
    lb[!is.finite(lb)] <- -.Machine$double.xmax / 4
    lb
  }

  # Weighted MLE of the Negative Binomial dispension (size) given the mean.
  nb_size <- function(y, w, mu, phi0) {
    sw <- sum(w)
    if (!is.finite(sw) || sw <= .Machine$double.eps) return(phi0)
    neg <- function(logr) -sum(w * stats::dnbinom(y, size = exp(logr), mu = mu, log = TRUE))
    opt <- try(stats::optimize(neg, lower = log(1e-3), upper = log(1e6)), silent = TRUE)
    if (inherits(opt, "try-error") || !is.finite(opt$minimum)) return(phi0)
    exp(opt$minimum)
  }

  # ---- One EM run from a given start -------------------------------------
  em_once <- function(par, delta, Gamma) {
    logGamma <- log(pmax(Gamma, 1e-300))
    logdelta <- log(pmax(delta, 1e-300))
    ll_old <- -Inf
    gamma <- NULL

    for (it in seq_len(max_iter)) {
      logB <- log_emission(par)

      # Forward (log, vectorized over states)
      la <- matrix(-Inf, Tn, K)
      la[1L, ] <- logdelta + logB[1L, ]
      if (Tn >= 2L) for (t in 2:Tn) {
        la[t, ] <- col_lse(la[t - 1L, ] + logGamma) + logB[t, ]
      }
      ll <- vec_lse(la[Tn, ])
      if (!is.finite(ll)) return(NULL)

      # Backward (log, vectorized over states)
      lb <- matrix(-Inf, Tn, K)
      lb[Tn, ] <- 0
      if (Tn >= 2L) for (t in (Tn - 1L):1L) {
        v <- logB[t + 1L, ] + lb[t + 1L, ]
        lb[t, ] <- row_lse(logGamma + matrix(v, K, K, byrow = TRUE))
      }

      # Smoothed posteriors
      lg <- la + lb - ll
      gamma <- exp(lg)
      rs <- rowSums(gamma)
      gamma <- gamma / ifelse(rs > 0, rs, 1)

      # Expected transition counts
      xi_sum <- matrix(0, K, K)
      if (Tn >= 2L) for (t in 1:(Tn - 1L)) {
        m <- matrix(la[t, ], K, K) + logGamma +
             matrix(logB[t + 1L, ] + lb[t + 1L, ], K, K, byrow = TRUE) - ll
        xi_sum <- xi_sum + exp(m)
      }

      # M-step: initial and transition
      delta <- gamma[1L, ]
      row_den <- rowSums(xi_sum)
      Gamma <- xi_sum / ifelse(row_den > 0, row_den, 1)
      zero_rows <- row_den <= 0
      if (any(zero_rows)) Gamma[zero_rows, ] <- 1 / K

      # M-step: emission means (posterior-weighted) with empty-state guard
      gsum <- colSums(gamma)
      new_mu_I <- ifelse(gsum > .Machine$double.eps, colSums(gamma * y_I) / gsum, par$mu_I)
      new_mu_C <- ifelse(gsum > .Machine$double.eps, colSums(gamma * y_C) / gsum, par$mu_C)
      par$mu_I <- pmax(new_mu_I, 1e-8)
      par$mu_C <- pmax(new_mu_C, 1e-8)

      # M-step: Negative Binomial dispersions (1-D weighted maximization)
      if (family == "nbinom") {
        for (j in seq_len(K)) {
          par$phi_I[j] <- nb_size(y_I, gamma[, j], par$mu_I[j], par$phi_I[j])
          par$phi_C[j] <- nb_size(y_C, gamma[, j], par$mu_C[j], par$phi_C[j])
        }
      }

      logGamma <- log(pmax(Gamma, 1e-300))
      logdelta <- log(pmax(delta, 1e-300))

      rel <- (ll - ll_old) / (abs(ll_old) + tol)
      if (it > 1L && is.finite(rel) && rel < tol) {
        return(list(par = par, delta = delta, Gamma = Gamma, logLik = ll,
                    posterior = gamma, iterations = it, converged = TRUE))
      }
      ll_old <- ll
    }
    list(par = par, delta = delta, Gamma = Gamma, logLik = ll_old,
         posterior = gamma, iterations = max_iter, converged = FALSE)
  }

  # ---- Parameter initialization ------------------------------------------
  init_params <- function(restart) {
    if (restart == 1L) {
      probs <- (seq_len(K) - 0.5) / K
      mu_I <- pmax(as.numeric(stats::quantile(y_I, probs, names = FALSE)), 1e-2)
      mu_C <- pmax(as.numeric(stats::quantile(y_C, probs, names = FALSE)), 1e-2)
    } else {
      mu_I <- pmax(stats::runif(K, min(y_I), max(y_I) + 1), 1e-2)
      mu_C <- pmax(stats::runif(K, min(y_C), max(y_C) + 1), 1e-2)
    }
    par <- list(mu_I = mu_I, mu_C = mu_C)
    if (family == "nbinom") {
      par$phi_I <- rep(1, K)
      par$phi_C <- rep(1, K)
    }
    delta <- rep(1 / K, K)
    if (K == 1L) {
      Gamma <- matrix(1, 1, 1)
    } else {
      Gamma <- matrix(0.2 / (K - 1), K, K)
      diag(Gamma) <- 0.8
    }
    list(par = par, delta = delta, Gamma = Gamma)
  }

  # ---- Multi-start EM ----------------------------------------------------
  best <- NULL
  for (s in seq_len(max(1L, n_starts))) {
    st <- init_params(s)
    fit_s <- tryCatch(em_once(st$par, st$delta, st$Gamma),
                      error = function(e) NULL)
    if (!is.null(fit_s) && is.finite(fit_s$logLik) &&
        (is.null(best) || fit_s$logLik > best$logLik)) {
      best <- fit_s
    }
  }
  if (is.null(best)) return(NULL)

  # ---- Canonical relabeling of states (increasing mu_I) ------------------
  ord <- order(best$par$mu_I)
  best$par$mu_I <- best$par$mu_I[ord]
  best$par$mu_C <- best$par$mu_C[ord]
  if (family == "nbinom") {
    best$par$phi_I <- best$par$phi_I[ord]
    best$par$phi_C <- best$par$phi_C[ord]
  }
  best$delta <- best$delta[ord]
  best$Gamma <- best$Gamma[ord, ord, drop = FALSE]
  best$posterior <- best$posterior[, ord, drop = FALSE]

  # ---- Viterbi decoding (log) --------------------------------------------
  logB <- log_emission(best$par)
  logGamma <- log(pmax(best$Gamma, 1e-300))
  logdelta <- log(pmax(best$delta, 1e-300))
  dv <- matrix(-Inf, Tn, K)
  psi <- matrix(0L, Tn, K)
  dv[1L, ] <- logdelta + logB[1L, ]
  if (Tn >= 2L) for (t in 2:Tn) {
    for (j in seq_len(K)) {
      vals <- dv[t - 1L, ] + logGamma[, j]
      psi[t, j] <- which.max(vals)
      dv[t, j] <- max(vals) + logB[t, j]
    }
  }
  states <- integer(Tn)
  states[Tn] <- which.max(dv[Tn, ])
  if (Tn >= 2L) for (t in (Tn - 1L):1L) states[t] <- psi[t + 1L, states[t + 1L]]

  # ---- Model-selection criteria ------------------------------------------
  n_emis <- if (family == "nbinom") 4L else 2L
  npar <- (K - 1L) + K * (K - 1L) + K * n_emis
  aic <- -2 * best$logLik + 2 * npar
  bic <- -2 * best$logLik + log(Tn) * npar

  fit <- list(
    nstates    = K,
    family     = family,
    mu_I       = best$par$mu_I,
    mu_C       = best$par$mu_C,
    phi_I      = if (family == "nbinom") best$par$phi_I else NULL,
    phi_C      = if (family == "nbinom") best$par$phi_C else NULL,
    delta      = best$delta,
    Gamma      = best$Gamma,
    logLik     = best$logLik,
    npar       = npar,
    AIC        = aic,
    BIC        = bic,
    posterior  = best$posterior,
    iterations = best$iterations,
    converged  = best$converged
  )

  if (!is.null(dir_csv)) {
    if (!dir.exists(dir_csv)) dir.create(dir_csv, recursive = TRUE)
    readr::write_csv(data.frame(t = seq_len(Tn), state = states),
                     file.path(dir_csv, "hmm_states.csv"))
  }
  if (!is.null(dir_out)) {
    if (!dir.exists(dir_out)) dir.create(dir_out, recursive = TRUE)
    saveRDS(fit, file.path(dir_out, "hmm_fit.rds"))
  }

  structure(list(fit = fit, states = states), class = "bivarhr_hmm")
}

#' Print method for bivarhr_hmm objects
#'
#' @param x An object of class \code{"bivarhr_hmm"} returned by
#'   \code{\link{run_hmm}()}.
#' @param ... Ignored; present for S3 compatibility.
#'
#' @return Invisibly returns \code{x}.
#'
#' @export
print.bivarhr_hmm <- function(x, ...) {
  f <- x$fit
  cat(sprintf("Bivariate %s HMM with %d states\n",
              ifelse(f$family == "nbinom", "Negative Binomial", "Poisson"),
              f$nstates))
  cat(sprintf("logLik = %.2f | AIC = %.2f | BIC = %.2f | npar = %d\n",
              f$logLik, f$AIC, f$BIC, f$npar))
  cat(sprintf("EM %s in %d iterations\n",
              ifelse(f$converged, "converged", "stopped (max_iter)"),
              f$iterations))
  rates <- data.frame(state = seq_len(f$nstates),
                      mu_I = round(f$mu_I, 3), mu_C = round(f$mu_C, 3))
  if (f$family == "nbinom") {
    rates$phi_I <- round(f$phi_I, 3)
    rates$phi_C <- round(f$phi_C, 3)
  }
  cat("State-specific emission parameters:\n")
  print(rates, row.names = FALSE)
  invisible(x)
}