diff --git a/DESCRIPTION b/DESCRIPTION
index 16fc603..41daf8d 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -18,6 +18,7 @@ Depends:
     patchwork
 Imports:
     reshape2,
+    proto,
     tidyselect,
     stringr,
     stringi,
diff --git a/R/get-coefs.R b/R/get-coefs.R
index b2a3833..bcb3bed 100644
--- a/R/get-coefs.R
+++ b/R/get-coefs.R
@@ -61,30 +61,30 @@ get_coefs_raw <- function(fit, var = "area") {
   
   ps <- as_draws_df(x = fit, variable = var, regex = TRUE) %>%
     as.data.frame() %>%
-    mutate(iteration = .data$.draw) %>%
-    select(-.data$.iteration, -.data$.chain, -.data$.draw) %>%
+    mutate(iteration = .draw) %>%
+    select(-.iteration, -.chain, -.draw) %>%
     melt(id.vars = "iteration") %>%
-    # mutate(variable = gsub(".*\\[|\\]", "", .data$variable)) %>%
-    mutate(variable = gsub(",Intercept", "", .data$variable)) %>%
-    filter(!str_detect(.data$variable, "sd_"))
+    # mutate(variable = gsub(".*\\[|\\]", "", variable)) %>%
+    mutate(variable = gsub(",Intercept", "", variable)) %>%
+    filter(!str_detect(variable, "sd_"))
   
   # Remove interaction term(s) if not asked for
   if (!str_detect(var, ":")) {
     ps <- ps %>%
-      filter(!str_detect(.data$variable, ":"))
+      filter(!str_detect(variable, ":"))
   }
   
   # if (nrow(fit$ranef) > 0) {
   #   ps <- posterior_samples(fit, pars = paste0("r_", var)) %>%
   #     mutate(iteration = 1:n()) %>%
   #     melt(id.vars = "iteration") %>%
-  #     mutate(variable = gsub(".*\\[|\\]", "", .data$variable)) %>%
-  #     mutate(variable = gsub(",Intercept", "", .data$variable))
+  #     mutate(variable = gsub(".*\\[|\\]", "", variable)) %>%
+  #     mutate(variable = gsub(",Intercept", "", variable))
   # } else {
   #   ps <- posterior_samples(fit, pars = var) %>%
   #     mutate(iteration = 1:n()) %>%
   #     melt(id.vars = "iteration") %>%
-  #     mutate(variable = gsub("b_", "", .data$variable))
+  #     mutate(variable = gsub("b_", "", variable))
   # }
   
   # Get the missing variable and normalise
@@ -96,11 +96,11 @@ get_coefs_raw <- function(fit, var = "area") {
   #                     value = 0)
   #   ps1 <- rbind(ps0, ps)
   #   mean_coefs <- ps1 %>% 
-  #     group_by(.data$iteration) %>% 
-  #     summarise(mean_coef = mean(.data$value))
+  #     group_by(iteration) %>% 
+  #     summarise(mean_coef = mean(value))
   #   coefs <- left_join(ps1, mean_coefs, by = "iteration") %>%
-  #     mutate(value = .data$value - .data$mean_coef) %>%
-  #     select(-.data$mean_coef)
+  #     mutate(value = value - mean_coef) %>%
+  #     select(-mean_coef)
   # } else if (nrow(fit$ranef) == 0 & !normalise & !is_poly & length(unique(ps$variable)) != 1) {
   #   data <- fit$data
   #   data[,var] <- paste0(var, data[,var])
@@ -150,14 +150,14 @@ get_coefs <- function(fit, var = "area",
     #   mutate(variable = rownames(.))
     ps <- as_draws_df(x = fit, variable = paste0("r_", var), regex = TRUE) %>%
       as.data.frame() %>%
-      mutate(iteration = .data$.draw) %>%
-      select(-.data$.iteration, -.data$.chain, -.data$.draw) %>%
+      mutate(iteration = .draw) %>%
+      select(-.iteration, -.chain, -.draw) %>%
       melt(id.vars = "iteration") %>%
-      mutate(variable = gsub(".*\\[|\\]", "", .data$variable)) %>%
-      mutate(variable = gsub(",Intercept", "", .data$variable))
+      mutate(variable = gsub(".*\\[|\\]", "", variable)) %>%
+      mutate(variable = gsub(",Intercept", "", variable))
     if (!str_detect(var, ":")) {
       ps <- ps %>%
-        filter(!str_detect(.data$variable, ":"))
+        filter(!str_detect(variable, ":"))
     }
   } else {
     # Population-level effects
@@ -172,18 +172,18 @@ get_coefs <- function(fit, var = "area",
     # eff <- rbind(e2, eff)
     ps <- as_draws_df(x = fit, variable = var, regex = TRUE) %>%
       as.data.frame() %>%
-      mutate(iteration = .data$.draw) %>%
-      select(-.data$.iteration, -.data$.chain, -.data$.draw) %>%
+      mutate(iteration = .draw) %>%
+      select(-.iteration, -.chain, -.draw) %>%
       melt(id.vars = "iteration") %>%
-      mutate(variable = gsub("b_", "", .data$variable)) %>%
-      mutate(variable = gsub("r_", "", .data$variable)) %>%
-      filter(!str_detect(.data$variable, "sd_"))
+      mutate(variable = gsub("b_", "", variable)) %>%
+      mutate(variable = gsub("r_", "", variable)) %>%
+      filter(!str_detect(variable, "sd_"))
     if (!str_detect(var, ":")) {
       ps <- ps %>%
-        filter(!str_detect(.data$variable, ":"))
+        filter(!str_detect(variable, ":"))
     }
-      # mutate(variable = paste0(var, gregexpr("[[:digit:]]+", .data$variable)))
-      # mutate(variable = paste0(var, parse_number(as.character(.data$variable))))
+      # mutate(variable = paste0(var, gregexpr("[[:digit:]]+", variable)))
+      # mutate(variable = paste0(var, parse_number(as.character(variable))))
     # unique(ps$variable)
     # tail(ps)
   }
@@ -191,7 +191,7 @@ get_coefs <- function(fit, var = "area",
   # Check to see if this is a polynomial
   if (any(grepl("poly", ps$variable))) {
     # ps <- ps %>%
-    #   mutate(variable = gsub("poly", "", .data$variable))
+    #   mutate(variable = gsub("poly", "", variable))
     is_poly <- TRUE
     # order_poly <- unique(sub(".*?(\\d).*", "\\1", ps$variable))
   }
@@ -200,11 +200,11 @@ get_coefs <- function(fit, var = "area",
   if (any(grepl("hu", ps$variable))) {
     if (hurdle) {
       ps <- ps %>%
-        filter(grepl("hu", .data$variable)) %>%
-        mutate(variable = gsub("hu_", "", .data$variable))
+        filter(grepl("hu", variable)) %>%
+        mutate(variable = gsub("hu_", "", variable))
     } else {
       ps <- ps %>%
-        filter(!grepl("hu", .data$variable))
+        filter(!grepl("hu", variable))
     }
   }
   
@@ -219,11 +219,11 @@ get_coefs <- function(fit, var = "area",
                       value = 0)
     ps1 <- rbind(ps0, ps)
     mean_coefs <- ps1 %>% 
-      group_by(.data$iteration) %>% 
-      summarise(mean_coef = mean(.data$value))
+      group_by(iteration) %>% 
+      summarise(mean_coef = mean(value))
     coefs <- left_join(ps1, mean_coefs, by = "iteration") %>%
-      mutate(value = .data$value - .data$mean_coef) %>%
-      select(-.data$mean_coef)
+      mutate(value = value - mean_coef) %>%
+      select(-mean_coef)
   } else if (nrow(fit$ranef) == 0 & !normalise & !is_poly & length(unique(ps$variable)) != 1) {
     data <- fit$data
     data[,var] <- paste0(var, data[,var])
@@ -248,21 +248,21 @@ get_coefs <- function(fit, var = "area",
   # if (transform & !is_poly) {
   #   if (fit$family$family %in% c("lognormal", "hurdle_lognormal")) {
   #     if (fit$family$link == "identity") {
-  #       coefs <- coefs %>% mutate(value = exp(.data$value))
+  #       coefs <- coefs %>% mutate(value = exp(value))
   #     } else {
   #       stop("This link function for the lognormal family has not been coded in influ2 yet - please update the plot-cdi.R function.")
   #     }
   #   } else if (fit$family$family %in% c("gamma", "hurdle_gamma")) {
   #     if (fit$family$link == "inverse") {
-  #       coefs <- coefs %>% mutate(value = 1.0 / .data$value)
+  #       coefs <- coefs %>% mutate(value = 1.0 / value)
   #     } else if (fit$family$link == "identity") {
-  #       coefs <- coefs %>% mutate(value = .data$value)
+  #       coefs <- coefs %>% mutate(value = value)
   #     } else if (fit$family$link == "log") {
-  #       coefs <- coefs %>% mutate(value = exp(.data$value))
+  #       coefs <- coefs %>% mutate(value = exp(value))
   #     }
   #   } else if (fit$family$family %in% c("bernoulli")) {
   #     if (fit$family$link == "logit") {
-  #       coefs <- coefs %>% mutate(value = exp(.data$value) / (1.0 + exp(.data$value)))
+  #       coefs <- coefs %>% mutate(value = exp(value) / (1.0 + exp(value)))
   #     }
   #   } else {
   #     stop("This family has not been coded in influ2 yet - please update the plot-cdi.R function.")
diff --git a/tests/testthat/test-cdi.R b/tests/testthat/test-cdi.R
index e7168d6..d88b993 100644
--- a/tests/testthat/test-cdi.R
+++ b/tests/testthat/test-cdi.R
@@ -1,29 +1,33 @@
+library(brms)
+
+options(mc.cores = parallel::detectCores())
+
+data(lobsters_per_pot)
+
+brm1 <- brm(lobsters ~ year + month, data = lobsters_per_pot, family = poisson)
+glm1 <- glm(lobsters ~ year + month, data = lobsters_per_pot, family = poisson)
+brm2 <- brm(lobsters ~ year + (1 | month), data = lobsters_per_pot, family = negbinomial())
+
+
 context("CDI plot")
 
 test_that("summary gives the same thing as get_coefs for population-level effects", {
   
-  library(brms)
-  
-  data(iris)
-  iris <- iris %>% 
-    mutate(PetalLength = Petal.Length, SepalWidth = factor(round(Sepal.Width)))
-  
-  fit <- brm(PetalLength ~ SepalWidth, data = iris, family = lognormal())
-  
   # This is get_coefs from influ2
-  c1 <- get_coefs(fit = fit, var = "SepalWidth", normalise = FALSE) %>%
+  c1 <- get_coefs(fit = brm1, var = "year", normalise = FALSE) %>%
     group_by(variable) %>%
     summarise(Estimate = mean(value), Est.Error = sd(value), 
               Q5 = quantile(value, probs = 0.05), Q95 = quantile(value, probs = 0.95))
   
   # Here I use the fixef fuction from the nlme package
-  c2 <- fixef(fit, probs = c(0.05, 0.95)) %>%
+  c2 <- fixef(brm1, probs = c(0.05, 0.95)) %>%
     data.frame() %>%
+    filter(!grepl("month", rownames(.))) %>%
     mutate(variable = c1$variable)
   c2[1,1:4] <- 0
 
-  expect_is(fit, "brmsfit")
-  expect_error(ranef(fit))
+  expect_is(brm1, "brmsfit")
+  expect_error(ranef(brm1))
   expect_equal(c1$Estimate, c2$Estimate)
   expect_equal(c1$Est.Error, c2$Est.Error)
   expect_equal(as.numeric(c1$Q5), c2$Q5)
@@ -33,17 +37,8 @@ test_that("summary gives the same thing as get_coefs for population-level effect
 
 test_that("summary gives the same thing as get_coefs for group-level effects", {
 
-  library(brms)
-
-  data(iris)
-  iris <- iris %>%
-    mutate(PetalLength = Petal.Length,
-           SepalWidth = factor(round(Sepal.Width)))
-
-  fit <- brm(PetalLength ~ (1 | SepalWidth), data = iris, family = lognormal())
-
   # This is get_coefs from influ2
-  c1 <- get_coefs(fit = fit, var = "SepalWidth", normalise = FALSE) %>%
+  c1 <- get_coefs(fit = brm2, var = "month", normalise = FALSE) %>%
     group_by(variable) %>%
     summarise(Estimate = mean(value), 
               Est.Error = sd(value), 
@@ -51,11 +46,11 @@ test_that("summary gives the same thing as get_coefs for group-level effects", {
               Q95 = quantile(value, probs = 0.95))
 
   # Here I use the ranef fuction from the nlme package
-  c2 <- ranef(fit, groups = "SepalWidth", probs = c(0.05, 0.95))[[1]][,,1] %>%
+  c2 <- ranef(brm2, groups = "month", probs = c(0.05, 0.95))[[1]][,,1] %>%
     data.frame() %>%
     mutate(variable = rownames(.))
   
-  expect_is(fit, "brmsfit")
+  expect_is(brm2, "brmsfit")
   expect_equal(c1$Estimate, c2$Estimate)
   expect_equal(c1$Est.Error, c2$Est.Error)
   expect_equal(as.numeric(c1$Q5), c2$Q5)
@@ -65,22 +60,12 @@ test_that("summary gives the same thing as get_coefs for group-level effects", {
 
 test_that("this matches Nokome Bentley's influ package", {
   
-  library(brms)
-  # library(proto)
-  
-  data(iris)
-  iris <- iris %>% mutate(Sepal.Width = factor(Sepal.Width))
-
   # Test glm
-  fit1 <- glm(Petal.Length ~ Sepal.Width + Species, data = iris)
-  
-  # source("influ.R")
-  # source("tests/testthat/influ.R")
-  myInfl <- Influence$new(fit1)
+  myInfl <- Influence$new(glm1)
   c1a <- myInfl$coeffs()
   c1a <- c1a[2:length(c1a)]
-  c1b <- coef(fit1)
-  c1b <- c1b[grepl("Sepal.Width", names(c1b))]
+  c1b <- coef(glm1)
+  c1b <- c1b[grepl("year", names(c1b))]
   # myInfl$summary
   # myInfl$stanPlot()
   # myInfl$stepPlot()
@@ -89,10 +74,8 @@ test_that("this matches Nokome Bentley's influ package", {
   # myInfl$cdiPlotAll()
   
   # Test brms
-  fit2 <- brm(Petal.Length ~ Sepal.Width + Species, data = iris)
-  
-  c2a <- get_coefs(fit = fit2, var = "Sepal.Width", normalise = FALSE, transform = FALSE) %>%
-    filter(variable != "Sepal.Width2") %>%
+  c2a <- get_coefs(fit = brm1, var = "year", normalise = FALSE, transform = FALSE) %>%
+    filter(variable != "year2000") %>%
     group_by(variable) %>%
     summarise(Estimate = mean(value), 
               Est.Error = sd(value), 
@@ -100,10 +83,10 @@ test_that("this matches Nokome Bentley's influ package", {
               Q50 = quantile(value, probs = 0.5), 
               Q95 = quantile(value, probs = 0.95))
   
-  c2b <- fixef(fit2, probs = c(0.05, 0.95)) %>%
+  c2b <- fixef(object = brm1, probs = c(0.05, 0.95)) %>%
     data.frame() %>%
     mutate(variable = rownames(.)) %>%
-    filter(grepl("Sepal.Width", variable))
+    filter(grepl("year", variable))
   
   # Test Nokomes coeffs are the same as stats::coef function
   plot(c1a, type = "b")