Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
AMEs for continuous variables and categorical outcomes
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment