R/estimate_bayes.R

In MisRepARMA: Misreported Time Series Analysis

# =============================================================================
# estimate_bayes.R
# =============================================================================

#' Bayesian estimation for misreported ARMA models
#'
#' Internal function that fits a misreported ARMA model using MCMC via JAGS.
#' Called by \code{\link{fitMisRepARMA}} when \code{method = "bayes"}.
#'
#' @param data Numeric vector. The observed (possibly misreported) time series.
#' @param p_AR Non-negative integer. Autoregressive order of the model.
#' @param q_MA Non-negative integer. Moving average order of the model.
#' @param covars Matrix or \code{NULL}. Optional matrix of covariates
#'   (one column per covariate). Default is \code{NULL}.
#' @param misReport Character. Type of misreporting: \code{"U"} for
#'   under-reporting (reporting factor in (0, 1)) or \code{"O"} for
#'   over-reporting (reporting factor > 1). Default is \code{"U"}.
#' @param n_chains Positive integer. Number of parallel MCMC chains.
#'   Default is \code{3}.
#' @param n_iter Positive integer. Total number of MCMC iterations per chain
#'   (including burn-in). Default is \code{2000}.
#' @param n_burnin Non-negative integer. Number of burn-in iterations discarded
#'   at the start of each chain. Default is \code{500}.
#' @param n_thin Positive integer. Thinning interval: only every
#'   \code{n_thin}-th iteration is kept. Default is \code{1} (no thinning).
#' @param q_prior_a Numeric or \code{NULL}. Lower bound of the uniform prior
#'   for the misreporting factor \eqn{q}. If \code{NULL}, determined
#'   automatically from \code{misReport}. Default is \code{NULL}.
#' @param q_prior_b Numeric or \code{NULL}. Upper bound of the uniform prior
#'   for the misreporting factor \eqn{q}. If \code{NULL}, determined
#'   automatically from \code{misReport}. Default is \code{NULL}.
#' @param w_prior_a Positive numeric. Shape parameter \eqn{a} of the
#'   Beta(\eqn{a}, \eqn{b}) prior for the misreporting probability \eqn{w}.
#'   Default is \code{1} (uniform prior together with \code{w_prior_b = 1}).
#' @param w_prior_b Positive numeric. Shape parameter \eqn{b} of the
#'   Beta(\eqn{a}, \eqn{b}) prior for the misreporting probability \eqn{w}.
#'   Default is \code{1} (uniform prior together with \code{w_prior_a = 1}).
#' @param q_init Numeric or \code{NULL}. Starting value for \eqn{q} in the
#'   MCMC chains. If \code{NULL}, set to \code{0.4} for under-reporting or
#'   \code{2.0} for over-reporting. Default is \code{NULL}.
#' @param w_init Numeric or \code{NULL}. Starting value for \eqn{w} in the
#'   MCMC chains. If \code{NULL}, defaults to \code{0.5}. Default is
#'   \code{NULL}.
#'
#' @return A named numeric vector of posterior median estimates with the
#'   following attributes:
#'   \describe{
#'     \item{\code{covars}}{Covariate effect removed before fitting, or
#'       \code{NULL}.}
#'     \item{\code{z}}{Integer vector of estimated misreporting indicators
#'       (0 = correctly reported, 1 = misreported).}
#'     \item{\code{q}}{Posterior median of the misreporting factor \eqn{q}.}
#'     \item{\code{w}}{Posterior median of the misreporting probability
#'       \eqn{w}.}
#'     \item{\code{var}}{Posterior median of the latent process variance
#'       \eqn{\sigma^2}.}
#'     \item{\code{DIC}}{Deviance Information Criterion of the fitted model.}
#'     \item{\code{x_rec}}{Numeric vector with the reconstructed latent
#'       process (posterior medians of \eqn{X_t}).}
#'     \item{\code{jags}}{Full \code{R2jags} output object.}
#'   }
#'
#' @keywords internal
#' @noRd

estimate_bayes <- function(data, p_AR, q_MA, covars = NULL,
                           misReport = "U",
                           n_chains  = 3,
                           n_iter    = 2000,
                           n_burnin  = 500,
                           n_thin    = 1,
                           q_prior_a = NULL,
                           q_prior_b = NULL,
                           w_prior_a = 1,
                           w_prior_b = 1,
                           q_init    = NULL,
                           w_init    = NULL)
{
  if (!requireNamespace("R2jags", quietly = TRUE))
    stop("Cal instalar el paquet 'R2jags'.")
  
  n <- length(data)
  
  # --- 1. Covariables -------------------------------------------------------
  covars_effect <- NULL
  if (!is.null(covars)) {
    mod_cov       <- lm(data ~ covars)
    covars_effect <- covars %*% coef(mod_cov)[2:length(coef(mod_cov))]
    data          <- data - covars_effect
  }
  
  y_vec   <- as.numeric(data)
  mu_init <- mean(y_vec, na.rm = TRUE)
  
  # --- 2. Noms de parametres -----------------------------------------------
  ar_names <- if (p_AR > 0) paste0("ar", seq_len(p_AR)) else character(0)
  ma_names <- if (q_MA > 0) paste0("ma", seq_len(q_MA)) else character(0)
  
  # --- 3. Limits de q -------------------------------------------------------
  if (misReport == "U") {
    q_lower <- 0.001; q_upper <- 0.999
  } else if (misReport == "O") {
    q_lower <- 1.001; q_upper <- 10.0
  } else stop("misReport ha de ser 'U' o 'O'.")
  
  # --- 4. Matrius de valors passats (dades) ---------------------------------
  # Construim les matrius X_past i eps_past en R i les passem com a dades.
  # Dimensions: n x max(p_AR,1) i n x max(q_MA,1).
  # Quan l'ordre es 0, la matriu es de zeros i el coef fixat a 0, cap efecte.
  p_eff <- max(p_AR, 1L)
  q_eff <- max(q_MA, 1L)
  
  X_past   <- matrix(0.0, nrow = n, ncol = p_eff)
  eps_past <- matrix(0.0, nrow = n, ncol = q_eff)
  for (t in seq_len(n)) {
    for (j in seq_len(p_eff))
      if (t > j) X_past[t, j]   <- y_vec[t - j] - mu_init
    for (k in seq_len(q_eff))
      if (t > k) eps_past[t, k] <- y_vec[t - k] - mu_init
  }
  
  # --- 5. Model JAGS --------------------------------------------------------
  # Clau: definim el model JAGS completament en R com a string,
  # sense cap bucle ni inprod per als coeficients AR/MA.
  # En comptes d'aixo, construim la suma AR i MA terme a terme
  # directament al string del model, amb els indexos literals.
  # Aixo garanteix que JAGS veu cada coeficient com a node escalar
  # independent amb el seu prior, sense cap ambiguitat de longitud.
  
  # Suma AR: "ac1*(X_past[t,1]) + ac2*(X_past[t,2]) + ..."
  if (p_AR > 0) {
    ar_sum_str   <- paste0("ac", seq_len(p_AR), " * X_past[t,", seq_len(p_AR), "]",
                           collapse = " + ")
    ar_prior_str <- paste0("  ac", seq_len(p_AR), " ~ dunif(-0.99, 0.99)\n",
                           collapse = "")
  } else {
    ar_sum_str   <- "0"
    ar_prior_str <- ""
  }
  
  # Suma MA: "tc1*(eps_past[t,1]) + tc2*(eps_past[t,2]) + ..."
  if (q_MA > 0) {
    ma_sum_str   <- paste0("tc", seq_len(q_MA), " * eps_past[t,", seq_len(q_MA), "]",
                           collapse = " + ")
    ma_prior_str <- paste0("  tc", seq_len(q_MA), " ~ dunif(-0.99, 0.99)\n",
                           collapse = "")
  } else {
    ma_sum_str   <- "0"
    ma_prior_str <- ""
  }
  
  # Noms dels coeficients per extreure de sims.matrix
  ac_names_jags <- if (p_AR > 0) paste0("ac", seq_len(p_AR)) else character(0)
  tc_names_jags <- if (q_MA > 0) paste0("tc", seq_len(q_MA)) else character(0)
  
  model_string <- paste0(
    "model {
  # --- Priors ---
  mu      ~ dnorm(mu_data, 0.01)
  tau     ~ dgamma(0.001, 0.001)
  sigma2  <- 1 / tau
  tau_obs ~ dgamma(0.001, 0.001)
  q       ~ dunif(q_lower, q_upper)
  w       ~ dbeta(w_a, w_b)
", ar_prior_str, ma_prior_str,
    "
  # --- Versemblanca ---
  for (t in 1:n) {
    mu_t[t]   <- mu + ", ar_sum_str, " + ", ma_sum_str, "
    X[t]      ~ dnorm(mu_t[t], tau)
    Z[t]      ~ dbern(w)
    mu_obs[t] <- (1 - Z[t]) * X[t] + Z[t] * q * X[t]
    Y[t]      ~ dnorm(mu_obs[t], tau_obs)
  }
}")
  
  # --- 6. Dades per a JAGS --------------------------------------------------
  # Priors per a w: Beta(w_a, w_b).
  # Per defecte Beta(1,1) = uniforme. L'usuari pot passar valors informatius.
  w_a_val <- if (!is.null(w_prior_a)) w_prior_a else 1
  w_b_val <- if (!is.null(w_prior_b)) w_prior_b else 1
  
  jags_data <- list(
    Y        = y_vec,
    n        = as.integer(n),
    mu_data  = mu_init,
    X_past   = X_past,
    eps_past = eps_past,
    q_lower  = q_lower,
    q_upper  = q_upper,
    w_a      = w_a_val,
    w_b      = w_b_val
  )
  
  # --- 7. Parametres a monitoritzar ----------------------------------------
  params_monitor <- c(ac_names_jags, tc_names_jags, "mu", "sigma2", "q", "w", "X")
  
  # --- 8. Valors inicials ---------------------------------------------------
  # Usem estimacio frequentista (q_init, w_init) si disponible,
  # amb soroll petit per diferenciar les cadenes.
  q_start <- if (!is.null(q_init)) q_init else
    if (misReport == "U") 0.4 else 2.0
  w_start <- if (!is.null(w_init)) w_init else 0.5
  
  inits_fun <- function() {
    q_j <- q_start * runif(1, 0.85, 1.15)
    q_j <- if (misReport == "U") min(max(q_j, 0.02), 0.95)
    else                  min(max(q_j, 1.05), 9.5)
    w_j <- min(max(w_start * runif(1, 0.85, 1.15), 0.05), 0.95)
    init <- list(
      mu      = rnorm(1, mu_init, 0.1 * sd(y_vec, na.rm = TRUE)),
      tau     = 1 / var(y_vec, na.rm = TRUE),
      tau_obs = 1 / var(y_vec, na.rm = TRUE),
      w       = w_j,
      Z       = sample(0:1, n, replace = TRUE),
      X       = y_vec,
      q       = q_j
    )
    for (j in seq_len(p_AR)) init[[paste0("ac", j)]] <- runif(1, -0.3, 0.3)
    for (k in seq_len(q_MA)) init[[paste0("tc", k)]] <- runif(1, -0.3, 0.3)
    init
  }
  
  # --- 9. Execucio MCMC ----------------------------------------------------
  fit_jags <- R2jags::jags(
    data               = jags_data,
    inits              = inits_fun,
    parameters.to.save = params_monitor,
    model.file         = textConnection(model_string),
    n.chains           = as.integer(n_chains),
    n.iter             = as.integer(n_iter),
    n.burnin           = as.integer(n_burnin),
    n.thin             = as.integer(n_thin),
    progress.bar       = "text"
  )
  
  # --- 10. Estimacions puntuals (medianes posteriors) ----------------------
  med  <- fit_jags$BUGSoutput$median
  sims <- fit_jags$BUGSoutput$sims.matrix
  
  # Extraiem medianes dels coeficients AR/MA directament de med (escalar)
  ar_medians <- if (p_AR > 0)
    setNames(sapply(ac_names_jags, function(nm) as.numeric(med[[nm]])), ar_names)
  else NULL
  
  ma_medians <- if (q_MA > 0)
    setNames(sapply(tc_names_jags, function(nm) as.numeric(med[[nm]])), ma_names)
  else NULL
  
  est_vec <- c(
    ar_medians,
    ma_medians,
    intercept = as.numeric(med$mu),
    var       = as.numeric(med$sigma2),
    q         = as.numeric(med$q),
    w         = as.numeric(med$w),
    DIC       = fit_jags$BUGSoutput$DIC
  )
  
  # --- 11. Proces latent reconstruit i indicador Z -------------------------
  x_rec <- as.numeric(med$X)
  z_hat <- as.integer(fit_jags$BUGSoutput$mean$Z > 0.5)
  
  # --- 12. Atributs (compatibles amb estimate()) ---------------------------
  attr(est_vec, "covars") <- covars_effect
  attr(est_vec, "z")      <- z_hat
  attr(est_vec, "q")      <- as.numeric(med$q)
  attr(est_vec, "w")      <- as.numeric(med$w)
  attr(est_vec, "var")    <- as.numeric(med$sigma2)
  attr(est_vec, "DIC")    <- fit_jags$BUGSoutput$DIC
  attr(est_vec, "x_rec")  <- x_rec
  attr(est_vec, "jags")   <- fit_jags
  
  return(est_vec)
}