fix issue with poly() in re terms #1221

paul-buerkner · Nov 8, 2024 · bfc6398 · bfc6398
1 parent f7106d4
commit bfc6398
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 7 deletions.
diff --git a/R/brmsframe.R b/R/brmsframe.R
@@ -247,12 +247,16 @@ get_levels.list <- function(x, ...) {
   out <- vector("list", length(x))
   for (i in seq_along(out)) {
     levels <- get_levels(x[[i]], ...)
+    if (!length(levels)) {
+      next
+    }
     if (is.list(levels)) {
       stopifnot(!is.null(names(levels)))
       out[[i]] <- as.list(levels)
-    } else if (!is.null(levels)) {
-      stopifnot(isTRUE(nzchar(names(x)[i])))
-      out[[i]] <- setNames(list(levels), names(x)[[i]])
+    } else if (is.vector(levels)) {
+      name_i <- names(x)[i]
+      stopifnot(isTRUE(nzchar(name_i)))
+      out[[i]] <- setNames(list(levels), name_i)
     }
   }
   out <- unlist(out, recursive = FALSE)
@@ -267,7 +271,7 @@ get_levels.brmsterms <- function(x, data = NULL, ...) {
     return(out)
   }
   if (!is.null(data)) {
-    ls <- list(frame_re(x, data), frame_me(x, data))
+    ls <- list(frame_re_levels_only(x, data), frame_me(x, data))
     out <- get_levels(ls)
   }
   out

diff --git a/R/formula-re.R b/R/formula-re.R
@@ -708,6 +708,25 @@ frame_re <- function(bterms, data, old_levels = NULL) {
   out
 }
 
+# like frame_re but only returns its levels attribute
+# this avoids issue #1221 and likely some other edge cases
+frame_re_levels_only <- function(bterms, data) {
+  out <- empty_reframe()
+  data <- combine_groups(data, get_group_vars(bterms))
+  re <- get_re(bterms)
+  re <- re[!duplicated(re$group), ]
+  levels <- named_list(re$group)
+  for (i in seq_along(levels)) {
+    # combine levels of all grouping factors within one grouping term
+    levels[[i]] <- unique(ulapply(
+      re$gcall[[i]]$groups,
+      function(g) extract_levels(get(g, data))
+    ))
+  }
+  set_levels(out) <- levels
+  out
+}
+
 empty_reframe <- function() {
   out <- data.frame(
     id = numeric(0), group = character(0), gn = numeric(0),

diff --git a/R/formula-sp.R b/R/formula-sp.R
@@ -297,8 +297,7 @@ frame_me <- function(x, data, old_levels = NULL) {
     term = uni_me, xname = "", grname = "",
     stringsAsFactors = FALSE
   )
-  unique_grnames <- unique(out$grname)
-  levels <- named_list(unique_grnames)
+  levels <- list()
   for (i in seq_rows(out)) {
     tmp <- eval2(out$term[i])
     out$xname[i] <- tmp$term
@@ -313,7 +312,7 @@ frame_me <- function(x, data, old_levels = NULL) {
   out$cor <- isTRUE(x$mecor)
   if (!is.null(old_levels)) {
     # for newdata numeration has to depend on the original levels
-    set_levels(out) <- old_levels[[unique_grnames]]
+    set_levels(out) <- old_levels
     set_levels(out, "used") <- levels
   } else {
     set_levels(out) <- levels