AMEs for continuous variables and categorical outcomes

gistfile1.txt

bayes_dydx.default_mn <- function(model, data = NULL, variable, stepsize = 1e-7, re_formula = NULL){
  
  # Check that everything is running properly and that the
  # user has provided all of the relevant information.
  
  if(is.null(model) == T){
    
    stop("Please provide a model to the function using the 'model =' argument (e.g. model = m1)")
    
  } else if(is.null(variable) == T){
    
    stop("Please provide a variable name to compute average marginal effects for using the 'variable =' argument (e.g. variable = 'x'")
    
  }
  
  # Get data from model where data = NULL
  if(is.null(data) == T){
    d <- model$data
  } else {
    d <- data
  }
  
  # Get outcome from model
  resp <- model$formula$resp
  
  # Omit outcome from data
  d <-
    d %>%
    select(-resp)
  
  # Omit random effects from the data if necessary
  if(is.null(re_formula) == F){
    
    # Get random effects
    rnfx <- unique(model$ranef$group)
    
    # Omit from data
    d <-
      d %>% 
      select(-rnfx)
  }
  
  
  # Calculate observed combinations and frequencies to reduce computation time
  d <-
    d %>% 
    group_by_all() %>% 
    count(name = "w") %>% 
    ungroup()
  
  # Create function to set "h" based on "eps" to deal with machine precision
  setstep <- function(x) {
    x + (max(abs(x), 1, na.rm = TRUE) * sqrt(stepsize)) - x
  }
  
  # Calculate numerical derivative
  d1 <- d0 <- d
  
  d0[[variable]] <- d0[[variable]] - setstep(d0[[variable]])
  d1[[variable]] <- d1[[variable]] + setstep(d1[[variable]])
  
  # Add fitted draws
  f0 <- 
    d0 %>%
    add_fitted_draws(model = model,
                     re_formula = re_formula,
                     value = paste0(variable, "_d0"))
  
  f1 <- 
    d1 %>%
    add_fitted_draws(model = model,
                     re_formula = re_formula,
                     value = paste0(variable, "_d1"))
  
  # Calculate average marginal effect
  out <-
    f0 %>% 
    ungroup() %>% 
    mutate(
      me = (f1[[paste0(variable, "_d1")]] - f0[[paste0(variable, "_d0")]])/(d1[[variable]] - d0[[variable]])
    ) %>% 
    group_by(.draw, .category) %>% 
    summarise(ame = sum(me * w)/sum(w)) %>% 
    ungroup()
  
  # Return AME
  out
  
}