R/model_functions.R

In cfhammill/hierarchyTrees: Code for testing taxonomy tree models

#' Fit a flat hierarchical model
#'
#' Fits `scaled_vol ~ group + (group | p0) + (1 | ID)`
#' with [stan_lmer]
#'
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
h0 <- function(data, cores = 4){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  stan_lmer(scaled_vol ~ group + (group | p0)
            + (1 | ID)
          , data = data$vol_frame
          , open_progress = FALSE) %>%
      over(class_l, ~ c("h0_fit", .))
}

#' Fit a one-parent hierarchical model
#'
#' Fits `scaled_vol ~ group + (group | p0) + (group | p1) + (1 | ID)`
#' with [stan_lmer]
#'
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
h1 <- function(data, cores = 4){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  stan_lmer(scaled_vol ~ group + (group | p0)
            + (group | p1)
            + (1 | ID)
          , data = data$vol_frame
          , open_progress = FALSE) %>%
      over(class_l, ~ c("h1_fit", .))
}

#' Fit a two-parent hierarchical model
#'
#' Fits `scaled_vol ~ group + (group | p0) + (group | p1) + (group | p2) + (1 | ID)`
#' with [stan_lmer]
#'
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
h2 <- function(data, cores = 4){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  stan_lmer(scaled_vol ~ group + (group | p0)
            + (group | p1)
            + (group | p2)
            + (1 | ID)
          , data = data$vol_frame
          , open_progress = FALSE) %>%
    over(class_l, ~ c("h2_fit", .))
}

#' Fit the effect diffusion tree
#'
#' Fits the effect diffusion tree model
#'
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
edt <- function(data, cores = 4){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  stan(file = system.file(file.path("models", "effect_diffusion_tree.stan")
                        , package = "hierarchyTrees")
     , control = list(max_treedepth = 15, adapt_delta = .99)
     , data = data$stan_list
     , open_progress = FALSE) %>%
    list(model = ., data = data$stan_list) %>%
    set(class_l, "edt_fit")
}

#' Fit the no-pooling model
#'
#' Fits a massively univariate model using [stan_glm]
#' with `scaled_vol ~ group`. This runs faster single core.
#' 
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
np <- function(data, cores = 1){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  data$vol_frame %>%
    group_by(name) %>%
    do(fitted = stan_glm(scaled_vol ~ group, data = .
                       , open_progress = FALSE)) %>%
    over(class_l, ~ c("np_fit", .))
}

#' Fit the no-pooling model with brain volume covaried
#'
#' Fits a massively univariate model using [stan_glm]
#' with `scaled_vol ~ group + bv`. This runs faster single core.
#' 
#' @param data A data object generated by `tree_to_edt_data`
#' @param cores number of cores to use, default 4
#' @export
npcv <- function(data, cores = 1){
  op <- options()
  on.exit({ options(op) })
  options(mc.cores = cores)
  
  data$vol_frame  %>%
    group_by(p0) %>%
    do(fitted = stan_glm(scaled_vol ~ group + bv, data = .
                       , open_progress = FALSE)) %>%
      over(class_l, ~ c("np_fit", .))
}