Skip to content

Instantly share code, notes, and snippets.

@andrewheiss
Last active March 20, 2018 23:04
Show Gist options
  • Save andrewheiss/a3134085e92c6607db39c5b14e1b879e to your computer and use it in GitHub Desktop.
Save andrewheiss/a3134085e92c6607db39c5b14e1b879e to your computer and use it in GitHub Desktop.
library(tidyverse)
library(Amelia)
library(broom)

# Use the africa dataset from Amelia
data(africa)
set.seed(1234)
imp_amelia <- amelia(x = africa, m = 5, cs = "country", ts = "year", logs = "gdp_pc", p2s = 0)

# Gather all the imputed datasets into one data frame and run a model on each
models_imputed_df <- bind_rows(unclass(imp_amelia$imputations), .id = "m") %>%
  group_by(m) %>%
  nest() %>% 
  mutate(model = data %>% map(~ lm(gdp_pc ~ trade + civlib, data = .)))
models_imputed_df
#> # A tibble: 5 x 3
#>   m     data               model   
#>   <chr> <list>             <list>  
#> 1 imp1  <tibble [120 × 7]> <S3: lm>
#> 2 imp2  <tibble [120 × 7]> <S3: lm>
#> 3 imp3  <tibble [120 × 7]> <S3: lm>
#> 4 imp4  <tibble [120 × 7]> <S3: lm>
#> 5 imp5  <tibble [120 × 7]> <S3: lm>


# We want to see how GDP per capita varies with changes in civil liberties, so
# we create a new data frame with values for each of the covariates in the
# model. We include the full range of civil liberties (from 0 to 1) and the mean
# of trade.
new_data <- data_frame(civlib = seq(0, 1, 0.1), 
                       trade = mean(africa$trade, na.rm = TRUE))
new_data
#> # A tibble: 11 x 2
#>    civlib trade
#>     <dbl> <dbl>
#>  1  0.     62.6
#>  2  0.100  62.6
#>  3  0.200  62.6
#>  4  0.300  62.6
#>  5  0.400  62.6
#>  6  0.500  62.6
#>  7  0.600  62.6
#>  8  0.700  62.6
#>  9  0.800  62.6
#> 10  0.900  62.6
#> 11  1.00   62.6


meld_predictions <- function(x) {
  # x is a data frame with m rows and two columns:
  #
  # m  .fitted  .se.fit
  # 1  1.05     0.34
  # 2  1.09     0.28
  # x  ...      ...
  
  # Meld the fitted values using Rubin's rules
  x_melded <- mi.meld(matrix(x$.fitted), matrix(x$.se.fit))
  
  data_frame(.fitted = as.numeric(x_melded$q.mi),
             .se.fit = as.numeric(x_melded$se.mi))
}

# We augment/predict using new_data in each of the imputed models, then we group
# by each of the values of civil liberties (so each value, like 0.1 and 0.2 has
# 5 values, 1 from each of the imputed models), and then we meld those 5
# predicted values into a single value with meld_predictions()
predict_melded <- data_frame(models = models_imputed_df$model) %>%
  mutate(m = 1:n(),
         fitted = models %>% map(~ augment(., newdata = new_data))) %>% 
  unnest(fitted) %>% 
  group_by(civlib) %>%  # Group by each of the variables that you vary
  nest(.fitted, .se.fit) %>% 
  mutate(fitted_melded = data %>% map(~ meld_predictions(.))) %>% 
  unnest(fitted_melded) %>% 
  mutate(ymin = .fitted + (qnorm(0.025) * .se.fit),
         ymax = .fitted + (qnorm(0.975) * .se.fit))

# Plot!
ggplot(predict_melded, aes(x = civlib, y = .fitted)) +
  geom_line(color = "blue") +
  geom_ribbon(aes(ymin = ymin, ymax = ymax), alpha = 0.2, fill = "blue")

# How does this compare with a single, non-imputed model?
model_simple <- lm(gdp_pc ~ trade + civlib, data = africa)
simple_predict <- augment(model_simple, newdata = new_data) %>% 
  mutate(ymin = .fitted + (qnorm(0.025) * .se.fit),
         ymax = .fitted + (qnorm(0.975) * .se.fit))

ggplot(predict_melded, aes(x = civlib, y = .fitted)) +
  geom_line(color = "blue") +
  geom_ribbon(aes(ymin = ymin, ymax = ymax), alpha = 0.2, fill = "blue") +
  geom_line(data = simple_predict, color = "red") +
  geom_ribbon(data = simple_predict, aes(ymin = ymin, ymax = ymax), alpha = 0.2, fill = "red")

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment