AMEs for Continuous Variables

gistfile1.txt

# Install and load packages
librarian::shelf(brms, tidybayes, tidyverse, magrittr, margins)

# Create "posterior_ame" function
posterior_ame <- function(model, variable, eps = 1e-7){
  
  # Get model outcome
  resp <- model$formula$resp
  
  # Omit outcome variable from data
  d <-
    model$data %>% 
    select(-resp)
  
  # Calculate observed combinations and their frequencies
  # to reduce computation time where n is very large
  d <-
    d %>% 
    group_by_all() %>% 
    count(name = "w") %>% 
    ungroup()
  
  # Create data frames with which to calculate numerical derivatives
  d0 <- d1 <- d
  
  # 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(eps)) - x
  }
  
  # Calculate numerical derivative
  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 = NA,
                     value = paste0(variable, "_d0"))
  
  f1 <- 
    d1 %>%
    add_fitted_draws(model = model,
                     re_formula = NA,
                     value = paste0(variable, "_d1"))
  
  # Calculate average marginal effect
  fit <-
    f0 %>% 
    ungroup() %>% 
    mutate(
      me =
        f1[[paste0(variable, "_d1")]] %>%
        subtract(f0[[paste0(variable, "_d0")]]) %>% 
        divide_by(d1[[variable]] - d0[[variable]])
    ) %>% 
    mutate(me_w = me * w) %>% 
    group_by_at(".draw") %>% 
    summarise(ame = sum(me_w)/sum(w)) %>% 
    select(ame)
  
  # Return AME
  fit$ame
  
}

# Run simple Frequentist model and get AMEs
freq_m <- lm(mpg ~ 1 + wt, data = mtcars)
freq_ame <- freq_m %>% margins()
summary(freq_ame)

# Run simple Bayesian model and get AMEs
bayes_m <- brm(mpg ~ 1 + wt,
               data = mtcars,
               chains = 2,
               cores = 2)
bayes_ame <- posterior_ame(bayes_m, "wt")
median(bayes_ame)
sd(bayes_ame)