R/update_hbm.R

In hbsaems: Hierarchical Bayes Small Area Estimation Model using 'Stan'

#' Update a Hierarchical Bayesian Model (hbm) object
#' 
#' @title update_hbm : Update a Hierarchical Bayesian Model (hbm) object
#'
#' @description This function updates an existing hbmfit object generated by `hbm()`, `hbm_beta()`, `hbm_logitnormal()`, and `hbm_lognormal()`.
#' It allows updating the formula, data, and other arguments, following the behavior of `brms::update()`.
#'
#' @param model A `brmsfit`/`hbmfit` object
#' @param newdata (optional) A new dataset with the same structure as the original
#' @param iter (optional) Number of MCMC iterations
#' @param warmup (optional) Number of warmup iterations
#' @param chains (optional) Number of MCMC chains
#' @param cores (optional) Number of cores to use for sampling
#' @param control (optional) A named list of control parameters passed to Stan
#' @param ... Other arguments passed to `update.brmsfit`, except those modifying model structure
#'
#' @return An updated `hbmfit` object
#' 
#' @examples
#' \donttest{
#' library(hbsaems)
#' # Load example data
#' data("data_fhnorm")
#' # Fit initial model
#' model <- hbm(
#' formula = bf(y ~ x1 + x2 + x3),
#' hb_sampling = "gaussian",
#' hb_link = "identity",
#' data = data_fhnorm,
#' chains = 2,
#' iter = 10000,
#' warmup = 2000,
#' cores = 2
#'  )
#' # Update number of  iterations and warmup
#' updated_model <- update_hbm(
#' model,
#' newdata = data_fhnorm,
#' iter = 10000,
#' warmup = 2000,
#'  chains = 2,
#'  cores = 2
#'  )
#' # Check updated model summary
#' summary(updated_model)
#'   }
#' @export

update_hbm <- function(model, 
                       newdata = NULL, 
                       iter = NULL, 
                       warmup = NULL,
                       chains = NULL, 
                       cores = NULL, 
                       control = NULL, 
                       ...) {
  
  model_original <- model
  
  # Ensure the object 
  if (inherits(model, "hbmfit")) {
    model <- model$model  
  }
  if (!inherits(model, "brmsfit")) {
    stop("Input model must be a brmsfit or hbmfit object.")
  }
  
  # Prepare arguments to pass to brms::update
  update_args <- list(object = model)
  
  if (!is.null(newdata)) {
    # Extract all variable names in formula
    if (!is.null(formula(model)$forms)) {
      vars_needed <- all.vars(formula(model)$forms[[1]]$formula)
      formula_text <- deparse(formula(model)$forms[[1]])
    } else {
      vars_needed <- all.vars(formula(model)$formula)
      formula_text <- deparse(formula(model))
    }
    
    # Attempt to detect random effect groupings (e.g. (1 | group))
    re_terms <- unlist(regmatches(formula_text, gregexpr("\\(1 \\| [^)]+\\)", formula_text)))
    
    if (length(re_terms) > 0) {
      grouping_vars <- gsub("\\(1 \\|\\s*|\\)", "", re_terms)
      for (group_var in grouping_vars) {
        if (!group_var %in% colnames(newdata)) {
          message(sprintf("Grouping variable, '%s', not found in newdata, adding default grouping factor.", group_var))
          newdata[[group_var]] <- factor(seq_len(nrow(newdata)))
        }
      }
    }
    
    # Check if there are any variables used in the model formula that are not present in 'newdata'
    absent_vars <- setdiff(vars_needed, colnames(newdata))
    if (length(absent_vars) > 0) {
      stop("The following variables used in the model are not found in 'newdata': ", paste(absent_vars, collapse = ", "))
    }
    
    # Check if any of the required variables in 'newdata' contain missing (NA) values
    has_na <- any(sapply(vars_needed, function(v) any(is.na(newdata[[v]]))))
    if (has_na) {
      stop("Some variables in 'newdata' contain missing (NA) values. Please refit the model using `hbm()` to properly handle missing data.")
    }
    
    update_args$newdata <- newdata
  }
  
  if (!is.null(iter)) update_args$iter <- iter
  if (!is.null(warmup)) update_args$warmup <- warmup
  if (!is.null(chains)) update_args$chains <- chains
  if (!is.null(cores)) update_args$cores <- cores
  if (!is.null(control)) update_args$control <- control
  
  # Call brms update with only the specified arguments
  updated_model <- do.call(update, update_args)
  
  # Wrap back into hbmfit if needed
  if (inherits(model_original, "hbmfit")) {
    model_original$model <- updated_model  
    return(model_original)
  }
  
  return(updated_model)
}