R/nimble_inputs.R

In vlandau/tempo: Temporally-Explicit Models for Phenological Observations

#' @importFrom utils combn
nimble_inputs <- function(y, covariates, beta_prior_mean, beta_prior_sd,
                          beta_start, beta_random_indices, group_ids,
                          sd_eps_start, correlated, pp_checks) {
  y <- as.matrix(y[, 2:3])
  ### prep NIMBLE inputs for basic case ###
  pheno_data <- list(
    y = y,
    X = covariates,
    censored = rep(1, dim(covariates)[1]),
    z = rep(NA, dim(covariates)[1])
  )

  pheno_consts <- list(
    N = dim(covariates)[1],
    mu_beta = beta_prior_mean,
    sigma_beta = beta_prior_sd,
    n_beta = length(beta_prior_mean),
    time_steps = dim(covariates)[2],
    ppcheck = any(c("mean","variance") %in% pp_checks),
    ppcheck_mean = "mean" %in% pp_checks,
    ppcheck_var = "variance" %in% pp_checks
  )

  pheno_inits <- list(
    Beta = beta_start, # only applies with no random effects
    z = ceiling(rowMeans(y))
  )

  ### Multiple random effects ###
  if (length(beta_random_indices) > 1 & correlated) {
    n_beta <- length(beta_prior_mean)
    n_random_betas <- length(beta_random_indices)
    beta_fixed_indices <- c(1:n_beta)[-beta_random_indices]
    n_rho <- ncol(combn(n_random_betas, 2))
    Rho_diag <- diag(n_beta)

    # This will be fed into NIMBLE for variable assignment within the Rho matrix
    rho_indices <- NULL
    for (p in 1:(n_random_betas - 1)) {
      for (q in (p + 1):n_random_betas) {
        rho_indices <- rbind(rho_indices, c(p, q))
      }
    }

    # Create inits for eps when random effects are modeled
    # (do this same thing with just one random effect as well)
    eps <- matrix(NA, ncol = length(unique(group_ids)), nrow = n_random_betas)

    for (group_j in seq_along(unique(group_ids))) {
      for (param_p in 1:n_random_betas) {
        eps[param_p, group_j] <- rnorm(1, 0, sd = sd_eps_start[param_p])
      }
    }

    pheno_inits <- c(pheno_inits, list(sd_eps = sd_eps_start,
                                       mean_beta = beta_start,
                                       rho = rep(0, n_rho),
                                       eps = eps,
                                       Rho = matrix(0,
                                                    nrow = n_random_betas,
                                                    ncol = n_random_betas),
                                       Sigma = diag(sd_eps_start)))
    pheno_inits$Beta <- NULL

    pheno_consts <- c(pheno_consts,
                      list(n_rho = n_rho, Rho_diag = Rho_diag,
                           n_group = length(unique(group_ids)),
                           rho_indices = rho_indices,
                           group_ids = group_ids,
                           mean_eps = rep(0, n_random_betas),
                           n_random_betas = n_random_betas,
                           n_fixed_betas = length(beta_fixed_indices),
                           beta_random_indices = beta_random_indices,
                           beta_fixed_indices = beta_fixed_indices))
  }

  ### Single random effect or multiple uncorrelated ###
  if (!any(is.na(beta_random_indices)) &
      (length(beta_random_indices) == 1 | !correlated)) {
    n_beta <- length(beta_prior_mean)
    n_random_betas <- length(beta_random_indices)

    beta_fixed_indices <- c(1:n_beta)[-beta_random_indices]

    eps <- matrix(NA, ncol = length(unique(group_ids)), nrow = n_random_betas)

    for (group_j in seq_along(unique(group_ids))) {
      for (param_p in 1:n_random_betas) {
        eps[param_p, group_j] <- rnorm(1, 0, sd = sd_eps_start[param_p])
      }
    }

    pheno_inits <- c(pheno_inits, list(sd_eps = sd_eps_start,
                                       mean_beta = beta_start,
                                       eps = eps))
    pheno_inits$Beta <- NULL
    pheno_consts <- c(pheno_consts,
                      list(n_group = length(unique(group_ids)),
                           group_ids = group_ids,
                           n_random_betas = n_random_betas,
                           n_fixed_betas = length(beta_fixed_indices),
                           beta_random_indices = beta_random_indices,
                           beta_fixed_indices = beta_fixed_indices))
  }

  return(list(
    pheno_data = pheno_data,
    pheno_consts = pheno_consts,
    pheno_inits = pheno_inits
  ))
}