Handle factor random slopes and add tests

5 files changed, 111 insertions, 24 deletions

diff --git a/R/gdistremoval.R b/R/gdistremoval.R
index b253625..db5439b 100644
--- a/R/gdistremoval.R
+++ b/R/gdistremoval.R
@@ -199,19 +199,19 @@ setMethod("getDesign", "unmarkedFrameGDR",
 
   if(return.frames) return(list(sc=sc, ysc=ysc, oc=oc))
 
-  lam_fixed <- lme4::nobars(formula$lambdaformula)
+  lam_fixed <- remove_bars(formula$lambdaformula)
   Xlam <- model.matrix(lam_fixed,
             model.frame(lam_fixed, sc, na.action=NULL))
 
-  phi_fixed <- lme4::nobars(formula$phiformula)
+  phi_fixed <- remove_bars(formula$phiformula)
   Xphi <- model.matrix(phi_fixed,
             model.frame(phi_fixed, ysc, na.action=NULL))
 
-  dist_fixed <- lme4::nobars(formula$distanceformula)
+  dist_fixed <- remove_bars(formula$distanceformula)
   Xdist <- model.matrix(dist_fixed,
             model.frame(dist_fixed, ysc, na.action=NULL))
 
-  rem_fixed <- lme4::nobars(formula$removalformula)
+  rem_fixed <- remove_bars(formula$removalformula)
   Xrem <- model.matrix(rem_fixed,
             model.frame(rem_fixed, oc, na.action=NULL))
 
diff --git a/R/mixedModelTools.R b/R/mixedModelTools.R
index 4581f1e..d8dad59 100644
--- a/R/mixedModelTools.R
+++ b/R/mixedModelTools.R
@@ -28,11 +28,10 @@ get_Z <- function(formula, data, newdata=NULL){
   
   # Get new formula
   bars <- find_bars(formula)
-  new_form <- bars_to_formula(bars)
-
+  check_duplicate_terms(bars)
+  
   # Get partial Z for each bar expression
-  form_list <- lapply(1:length(bars), function(x) bars_to_formula(bars[x]))
-  Z_parts <- lapply(form_list, get_partial_Z, data=data, newdata=newdata)
+  Z_parts <- lapply(bars, get_partial_Z, data=data, newdata=newdata)
 
   # Create model frame
   #mf <- model.frame(new_form, data, na.action=stats::na.pass)
@@ -42,7 +41,7 @@ get_Z <- function(formula, data, newdata=NULL){
   #}
 
   # Create sparse Z matrix
-  #Z <- model.matrix(new_form, mf)
+  #Z <- model.matrix(new_form, mf, contrasts.arg=cont)
   Z <- do.call(cbind, Z_parts)
   colnames(Z) <- Z_colnames(formula, data)
   Matrix::Matrix(Z, sparse=TRUE)
@@ -56,26 +55,25 @@ has_random <- function(form){
 # Find 'bar' random effect parts of formula (i.e., the (x|y) structures)-------
 # Operates recursively
 find_bars <- function(form){
-  out <- NULL
+  if(is.null(form)) return(NULL)
   if(is.name(form)) return(NULL)
   if(form[[1]] == as.name("(")) return(form)
   if(is.call(form)){
-    out <- lapply(form, find_bars)
+    out <- return(unlist(lapply(form, find_bars)))
   }
-  unlist(out)
+  NULL
 }
 
 # Convert bar components into new formula--------------------------------------
-# E.g. ~x + (1|g) becomes ~g - 1
-bars_to_formula <- function(bars){
-  bar_terms <- lapply(bars, bars_to_terms)
-  check_duplicate_terms(bar_terms)
+# E.g. (1|g) becomes ~g - 1
+bar_to_formula <- function(bar){
+  bar_terms <- bar_to_terms(bar)
   as.formula(str2lang(paste("~", paste(bar_terms, collapse = " + "), "- 1")))
 }
 
 # Translate bar components into standard formula terms-------------------------
-bars_to_terms <- function(bars){
-  info <- get_bar_info(bars)
+bar_to_terms <- function(bar){
+  info <- get_bar_info(bar)
   new_terms <- sapply(info$LHS, function(x){
     if(x == "1") return(info$RHS)
     paste0(x, ":", info$RHS)
@@ -119,7 +117,8 @@ check_bar_info <- function(info){
 # Check terms to make sure there aren't any duplicates-------------------------
 # E.g. as a result of a formula like ~ (1|g) + (x||) where the intercept
 # is also implied in the second bar expression
-check_duplicate_terms <- function(bar_terms){
+check_duplicate_terms <- function(bars){
+  bar_terms <- lapply(bars, bar_to_terms)
   all_terms <- lapply(bar_terms, function(x){
     unlist(strsplit(x, " + ", fixed=TRUE))
   })
@@ -132,14 +131,22 @@ check_duplicate_terms <- function(bar_terms){
   invisible()
 }
 
+
 # Get partial Z for a given bar expression-------------------------------------
-get_partial_Z <- function(formula, data, newdata){
+get_partial_Z <- function(bar, data, newdata){
+  info <- get_bar_info(bar)
+  formula <- bar_to_formula(bar)
   mf <- model.frame(formula, data, na.action=stats::na.pass)
   if(!is.null(newdata)){
     mf <- model.frame(stats::terms(mf), newdata, na.action=stats::na.pass,
                       xlev=get_xlev(data, mf))
   }
-  model.matrix(formula, mf)
+  has_fac <- any(sapply(info$LHS[info$LHS != "1"], function(x) is.factor(data[[x]])))
+  if(has_fac){
+    stop("unmarked does not support random slopes for R factors.\n",
+         "Try converting the variable to a series of indicator variables instead.", call.=FALSE)
+  }
+  out <- model.matrix(formula, mf)
 }
 
 # Get levels of factor---------------------------------------------------------
diff --git a/R/power.R b/R/power.R
index 3240e8e..f0aa741 100644
--- a/R/power.R
+++ b/R/power.R
@@ -163,7 +163,7 @@ check_coefs <- function(coefs, fit, template=FALSE){
 
   # If there are random effects, adjust the expected coefficient names
   # to remove the b vector and add the grouping covariate name
-  rand <- lapply(formulas, lme4::findbars)
+  rand <- lapply(formulas, find_bars)
   if(!all(sapply(rand, is.null))){
     stopifnot(all(required_subs %in% names(formulas)))
     rvar <- lapply(rand, function(x) unlist(lapply(x, all.vars)))
diff --git a/R/simulate.R b/R/simulate.R
index a8887cb..efc372a 100644
--- a/R/simulate.R
+++ b/R/simulate.R
@@ -91,7 +91,7 @@ setMethod("simulate", "character",
 #replace_sigma <- function(coefs, fit){
 #  required_subs <- names(fit@estimates@estimates)
 #  formulas <- sapply(names(fit), function(x) get_formula(fit, x))
-#  rand <- lapply(formulas, lme4::findbars)
+#  rand <- lapply(formulas, find_bars)
 #  if(!all(sapply(rand, is.null))){
 #    rvar <- lapply(rand, function(x) unlist(lapply(x, all.vars)))
 #    for (i in required_subs){
@@ -108,7 +108,7 @@ setMethod("simulate", "character",
 generate_random_effects <- function(coefs, fit){
   required_subs <- names(fit@estimates@estimates)
   formulas <- sapply(names(fit), function(x) get_formula(fit, x))
-  rand <- lapply(formulas, lme4::findbars)
+  rand <- lapply(formulas, find_bars)
   if(!all(sapply(rand, is.null))){
     rvar <- lapply(rand, function(x) unlist(lapply(x, all.vars)))
     for (i in required_subs){
diff --git a/tests/testthat/test_mixed_models.R b/tests/testthat/test_mixed_models.R
new file mode 100644
index 0000000..f0f7d10
--- /dev/null
+++ b/tests/testthat/test_mixed_models.R
@@ -0,0 +1,80 @@
+context("mixed model tools")
+
+test_that("get_reTrms matches lme4::mkReTrms", {
+
+  skip_if(!requireNamespace("lme4", quietly=TRUE), 
+          "lme4 package unavailable")
+  
+  dat <- data.frame(x = rnorm(20), y = rnorm(20), z = factor(sample(letters[1:3], 20, replace=T)),
+                    group = factor(sample(letters[4:6], 20, replace=T)),
+                    id = factor(sample(letters[7:9], 20, replace=T)))
+
+  form1 <- ~x + (1|group)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~x + (x||group)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~x + (x||group) + (1|id)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~x + (x||group) + (y||id)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~x + (x*y||group) + (y||id)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~(1|group)
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+
+  form1 <- ~(1|group) + x
+  r1 <- lme4::mkReTrms(lme4::findbars(form1), dat)
+  r2 <- get_reTrms(form1, dat)
+  expect_identical(r2$Z, Matrix::t(r1$Zt))
+  expect_identical(r1$cnms, r2$cnms)
+  attributes(r1$flist)$assign <- NULL
+  expect_identical(r1$flist, r2$flist)
+})
+
+test_that("get_reTrms errors correctly", {
+  form1 <- ~x + (x+z||group) + (y||id)
+  expect_error(get_reTrms(form1, dat))
+
+  form1 <- ~x + (x|group:id)
+  expect_error(get_reTrms(form1, dat))
+
+  form1 <- ~x + (x|group/id)
+  expect_error(get_reTrms(form1, dat))
+
+  expect_identical(find_bars(NULL), NULL)
+})