diff --git a/DESCRIPTION b/DESCRIPTION
index 5008f7bb..465c07ea 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: pammtools
 Title: Piece-Wise Exponential Additive Mixed Modeling Tools
-Version: 0.1.10
-Date: 2019-03-14
+Version: 0.1.11
+Date: 2019-03-28
 Authors@R: c(
     person("Andreas", "Bender", , "andreas.bender@stat.uni-muenchen.de", role = c("aut", "cre"), comment=c(ORCID = "0000-0001-5628-8611")),
     person("Fabian", "Scheipl", , "fabian.scheipl@stat.uni-muenchen.de", role = c("aut")))
diff --git a/NAMESPACE b/NAMESPACE
index d231492b..ba487548 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -194,6 +194,7 @@ importFrom(purrr,transpose)
 importFrom(rlang,":=")
 importFrom(rlang,.data)
 importFrom(rlang,UQ)
+importFrom(rlang,enquo)
 importFrom(rlang,eval_tidy)
 importFrom(rlang,quo_name)
 importFrom(rlang,quos)
diff --git a/R/cumulative-coefficient.R b/R/cumulative-coefficient.R
index 7aa5d337..b66287af 100644
--- a/R/cumulative-coefficient.R
+++ b/R/cumulative-coefficient.R
@@ -40,14 +40,17 @@ get_cumu_coef.gam <- function(model, data, terms, ...) {
 #' @export
 get_cumu_coef.aalen <- function(model, data = NULL, terms, ci = TRUE, ...) {
 
-    cumu_coef <- model$cum %>% as_tibble() %>%
-      select(map_int(c("time", terms),
-        ~which(grepl(., colnames(model$cum))))) %>%
-      gather("variable", "cumu_hazard", -.data$time)
-    cumu_var <- model[["var.cum"]] %>% as_tibble() %>%
-     select(map_int(c("time", terms),
-        ~which(grepl(., colnames(model$cum))))) %>%
-      gather("variable", "cumu_var", -.data$time)
+  terms <- map(c("time", terms),
+      ~grep(.x, colnames(model$cum), value = TRUE)) %>%
+    reduce(union)
+  cumu_coef <- model[["cum"]] %>%
+    as_tibble() %>%
+    select(one_of(terms)) %>%
+    gather("variable", "cumu_hazard", -.data$time)
+  cumu_var <- model[["var.cum"]] %>%
+    as_tibble() %>%
+    select(terms) %>%
+    gather("variable", "cumu_var", -.data$time)
 
     suppressMessages(
       left_join(cumu_coef, cumu_var) %>%
@@ -82,7 +85,7 @@ get_cumu_diff <- function(d1, d2, model, nsim = 100L) {
 
 #' @inheritParams get_cumu_coef
 #' @import dplyr purrr
-#' @importFrom rlang sym
+#' @importFrom rlang sym enquo quo_name
 #' @keywords internal
 cumu_coef <- function(data, model, term, ...) {
 
diff --git a/tests/testthat/test-cumulative-coefficients.R b/tests/testthat/test-cumulative-coefficients.R
index 11c0b3f6..d918fe04 100644
--- a/tests/testthat/test-cumulative-coefficients.R
+++ b/tests/testthat/test-cumulative-coefficients.R
@@ -2,32 +2,25 @@ context("Test cumulative coefficients functionality")
 
 test_that("Cumulative coefficients work", {
 
-  data("tumor")
-
-  # PAMs
-  tumor <- tumor %>% dplyr::slice(1:100)
-  tumor_ped <- tumor %>% as_ped(
-    formula = Surv(days, status) ~ age + complications)
-  pam <- mgcv::gam(
-    formula = ped_status ~ s(tend, by = as.ordered(complications)) +
-      s(tend, by = age),
-      data = tumor_ped, family = poisson(), offset = offset)
-  cumu_coef <- get_cumu_coef(pam, tumor_ped,
-    terms = c("age", "complications"))
-  expect_data_frame(cumu_coef, nrows = 82L, ncols = 6L)
-  cumu_coef_pam <- get_cumu_coef(pam, tumor_ped,
-    terms = c("(Intercept)", "age"))
-  expect_data_frame(cumu_coef_pam, nrows = 82L, ncols = 6L)
+  df <- tumor[1:30, c("days", "status", "age")]
+  df$x1 <- as.factor(rep(letters[1:3], each = nrow(df) / 3L))
 
   ## aalen model
   library(timereg)
-  aalen <- aalen(Surv(days, status)~ age + complications, data = tumor)
-  cumu_coef_aalen <- get_cumu_coef(aalen, terms = c("age", "complications"))
-  expect_data_frame(cumu_coef_aalen, nrows = 86L, ncols = 6L)
-  # cox aalen
-  cox.aalen <- cox.aalen(Surv(days, status) ~ age + prop(complications),
-    data = tumor)
-  cumu_coef_cox.aalen <- get_cumu_coef(cox.aalen, terms = c("age"))
-  expect_data_frame(cumu_coef_cox.aalen, nrows = 43, ncols = 6L)
+  mod       <- aalen(Surv(days, status) ~ x1 + age, data = df)
+  cumu_coef_aalen <- get_cumu_coef(
+    mod,
+    df,
+    terms = c("(Intercept)", "x1"))
+  expect_data_frame(cumu_coef_aalen, nrows = 42L, ncols = 6L)
+  expect_equal(unique(cumu_coef_aalen$variable), c("(Intercept)", "x1b", "x1c"))
+
+  ## pam
+  ped <- as_ped(df, formula = Surv(days, status)~ x1 + age)
+  pam <- mgcv::gam(ped_status ~ x1 + age, data = ped, family = poisson(),
+    offset = offset)
+  cumu_coef_pam <- get_cumu_coef(pam, ped, terms = c("age", "x1"), nsim = 20L)
+  expect_data_frame(cumu_coef_pam, nrows = 36L, ncols = 6L)
+  expect_equal(unique(cumu_coef_pam$variable), c("age", "x1 (b)", "x1 (c)"))
 
 })