R/modify_stancode.R

In bentrueman/bgamcar1: Fit Bayesian GAMs With CAR(1) Errors To Censored Data

#' Modify stan code generated by `brms` to fit a CAR(1) model
#'
#' @description N.B., `modify_stancode(method = "car1")` resets the lower bound on the `ar` parameter to zero. It is compatible with priors such as
#' `normal()` (check to make sure it makes the expected modifications before using).
#'
#' @param scode_raw Stan code generated by `brms::make_stancode()`.
#' @param modify Specify which modification to make. The default (and currently the only option) is `"car1"`,
#' which creates the CAR(1) error term.
#' @param var_xcens For multivariate models, specify the names of censored predictor variables as a character vector.
#' @param var_car1 For multivariate models, specify which variable to generate CAR(1) errors for.
#' @param lower_bound An optional lower bound for left-censored parameters. The default is no lower bound.
#' @param lcl A numeric vector of censoring limits for left-censored variables. Alternatively, a list of numeric vectors with censoring limits for
#' each variable and observation; the order should correspond to that of `lcl` and `cens_ind`.
#'
#' @return A Stan program of class `brmsmodel`.
#' @importFrom stringr str_replace str_remove_all
#' @importFrom dplyr %>%
#' @export
#'
#' @examples
#' library("brms")
#' data <- read.csv(paste0(system.file("extdata", package = "bgamcar1"), "/data_car1.csv"))
#' scode <- make_stancode(bf(y ~ ar(time = x)), data)
#' modify_stancode(scode)
modify_stancode <- function(scode_raw, modify = "car1", var_xcens = NULL, var_car1 = NULL, lower_bound = NULL, lcl = NULL) {
  if (!is.null(var_xcens) && is.null(var_car1))
    message("for multivariate models, remember to specify which response gets CAR(1) errors.")
  if (!(is.numeric(lower_bound) | is.null(lower_bound))) stop("'lower_bound' must be numeric")
  if (modify == "car1") scode <- modify_stancode_car1(scode_raw, var_car1) # CAR1 error
  if (modify == "xcens") scode <- modify_stancode_censored(scode_raw, var_xcens, lower_bound, lcl) # left-censored predictors
  scode
}

modify_stancode_car1 <- function(scode_raw, var_car1) {

  ar <- if (is.null(var_car1)) "ar" else paste("ar", var_car1, sep = "_")
  Kar <- if (is.null(var_car1)) "Kar" else paste("Kar", var_car1, sep = "_")
  Err <- if (is.null(var_car1)) "Err" else paste("Err", var_car1, sep = "_")
  mu <- if (is.null(var_car1)) "mu" else paste("mu", var_car1, sep = "_")

  scode <- scode_raw %>%
    # add time difference variable s:
    str_replace(
      "response variable\\\n",
      "response variable\n  vector[N] s;  // CAR(1) exponent\n"
    ) %>%
    # set lower bound of zero on ar param:
    str_replace(
      paste0("vector\\[", Kar, "\\] ", ar, ";"),
      paste0("vector<lower=0, upper=1>[", Kar, "] ", ar, ";")
    ) %>%
    # convert AR process to CAR1:
    str_replace(
      paste0(mu, "\\[n\\] \\+= ", Err, "\\[n, 1:", Kar, "\\] \\* ", ar, ";"),
      paste0(mu, "[n] += ", Err, "[n, 1] * pow(", ar, "[1], s[n]); // CAR(1)")
    )

  class(scode) <- "brmsmodel"

  return(scode)
}

modify_stancode_censored <- function(scode_raw, var_xcens, lower_bound, lcl) {

  stopifnot("must specify censored predictors using the 'var_xcens' argument." = !is.null(var_xcens))

  var_xcens <- make.names(str_remove_all(var_xcens, "[\\._]"))

  n_var_xcens <- length(var_xcens)

  scode <- scode_raw

  if (length(lower_bound) == 1 && is.numeric(lower_bound)) lower_bound <- rep(lower_bound, n_var_xcens)

  for(i in seq_len(n_var_xcens)) {

    # modifications to data block:
    n_cens <- glue("int<lower=0> Ncens_{var_xcens[i]};  // number of left-censored")
    j_cens <- glue("array[Ncens_{var_xcens[i]}] int<lower=1> Jcens_{var_xcens[i]};  // positions of left-censored")
    u <- if (is.list(lcl)) {
      glue("vector[Ncens_{var_xcens[i]}] U_{var_xcens[i]};  // left-censoring limit")
    } else {
      glue("real U_{var_xcens[i]};  // left-censoring limit")
    }
    # modifications to parameters block:
    y_cens <- glue("vector<upper=U_{var_xcens[i]}>[Ncens_{var_xcens[i]}] Ycens_{var_xcens[i]};  // estimated left-censored")
    if (!is.null(lower_bound)) {
      if (!is.na(lower_bound[i])) y_cens <- str_replace(y_cens, "^vector<", glue("vector<lower={lower_bound[i]}, "))
    }
    # modifications to model block:
    yl <- glue("Yl_{var_xcens[i]}[Jcens_{var_xcens[i]}] = Ycens_{var_xcens[i]}; // add imputed left-censored values")

    scode <- scode %>%
      # modifications to data block:
      str_replace(
        "(data \\{\n(.|\n)*?)(?=\n\\})",
        paste(c("\\1", n_cens, j_cens, u), collapse = "\n  ")
      ) %>%
      # modifications to parameters block:
      str_replace(
        "(parameters \\{\n(.|\n)*?)(?=\n\\})",
        paste(c("\\1\n  ", y_cens), collapse = "")
      ) %>%
      # modifications to model block:
      str_replace(
        "(model \\{\n(.|\n)*?)(?=\n    mu_)",
        paste(c("\\1\n    ", yl), collapse = "")
      )

  }

  class(scode) <- "brmsmodel"

  return(scode)

}