Visualization: Collapse by [random] group (#597)

* fix

* fixes

* add for plotting

* fixes, docs

* desc, news

* vignette

* docs

* fix

* add snaps

* fix

* address comments

* fix

* fix pkgdown

Author: strengejacke

DESCRIPTION

@@ -1,7 +1,7 @@
 Type: Package
 Package: modelbased
 Title: Estimation of Model-Based Predictions, Contrasts and Means
-Version: 0.13.1.3
+Version: 0.13.1.4
 Authors@R:
     c(person(given = "Dominique",
              family = "Makowski",

NAMESPACE

@@ -83,6 +83,7 @@ S3method(visualisation_recipe,estimate_grouplevel)
 S3method(visualisation_recipe,estimate_means)
 S3method(visualisation_recipe,estimate_predicted)
 S3method(visualisation_recipe,estimate_slopes)
+export(collapse_by_group)
 export(describe_nonlinear)
 export(display)
 export(estimate_contrasts)

NEWS.md

@@ -1,5 +1,14 @@
 # modelbased (devel)
 
+## Changes
+
+* New function `collapse_by_group()`, which extracts the raw data points and
+  "averages" (i.e. "collapses") the response variable over the levels of the
+  grouping factor given in `collapse_by`. Only works with mixed models.
+  Additionally, the `plot()` and `visualization_recipe()` methods get a
+  `collapse_group` argument to use this feature when adding data points to plots
+  using `show_data` or `show_residuals`.
+
 ## Bug fixes
 
 * Fixed issue in `estimate_slope()` when `p_adjust = "esarey"`.

R/collapse_by_group.R

@@ -0,0 +1,105 @@
+#' @title Collapse raw data by random effect groups
+#' @name collapse_by_group
+#'
+#' @description This function extracts the raw data points (i.e. the data
+#'   that was used to fit the model) and "averages" (i.e. "collapses") the
+#'   response variable over the levels of the grouping factor given in
+#'   `collapse_by`. Only works with mixed models.
+#'
+#' @param collapse_by Name of the (random effects) grouping factor. Data is
+#'   collapsed by the levels of this factor.
+#' @param residuals Logical, if `TRUE`, collapsed partial residuals instead
+#'   of raw data by the levels of the grouping factor.
+#' @inheritParams residualize_over_grid
+#'
+#' @return A data frame with raw data points, averaged over the levels of
+#'   the given grouping factor from the random effects. The group level of
+#'   the random effect is saved in the column `"random"`.
+#'
+#' @examplesIf require("lme4", quietly = TRUE)
+#' data(efc, package = "modelbased")
+#' efc$e15relat <- as.factor(efc$e15relat)
+#' efc$c161sex <- as.factor(efc$c161sex)
+#' levels(efc$c161sex) <- c("male", "female")
+#' model <- lme4::lmer(neg_c_7 ~ c161sex + (1 | e15relat), data = efc)
+#' me <- estimate_means(model, "c161sex")
+#' head(efc)
+#' collapse_by_group(me, model, "e15relat")
+#'
+#' @export
+collapse_by_group <- function(grid, model, collapse_by = NULL, residuals = FALSE) {
+  if (!insight::is_mixed_model(model)) {
+    insight::format_error("This function only works with mixed effects models.")
+  }
+
+  model_data <- insight::get_data(model, source = "frame", verbose = FALSE)
+
+  if (is.null(collapse_by) || isTRUE(collapse_by) || isTRUE(residuals)) {
+    collapse_by <- insight::find_random(model, flatten = TRUE)
+  }
+
+  if (length(collapse_by) > 1) {
+    collapse_by <- collapse_by[1]
+    insight::format_alert(
+      "More than one random grouping variable found.",
+      paste0("Using `", collapse_by, "`.")
+    )
+  }
+
+  if (!collapse_by %in% colnames(model_data)) {
+    insight::format_error(paste0("Could not find `", collapse_by, "` column."))
+  }
+
+  if (residuals) {
+    rawdata <- residualize_over_grid(grid, model)
+    y_name <- "Mean"
+  } else {
+    rawdata <- insight::get_data(model, source = "environment", verbose = FALSE)
+    y_name <- insight::find_response(model)
+
+    # we need this column for labelling data points, but not for collapsing
+    rawdata$rowname <- NULL
+
+    predictor_names <- setdiff(colnames(rawdata), y_name)
+    if (any(vapply(rawdata[predictor_names], Negate(is.factor), logical(1)))) {
+      insight::format_alert(
+        "Collapsing usually not informative across a continuous variable."
+      )
+    }
+  }
+
+  if (is.factor(rawdata[[y_name]])) {
+    rawdata[[y_name]] <- as.numeric(rawdata[[y_name]])
+    if (insight::model_info(model)$is_binomial) {
+      rawdata[[y_name]] <- rawdata[[y_name]] - 1
+    } # else ordinal?
+  }
+
+  if (nrow(rawdata) == nrow(model_data)) {
+    rawdata$random <- factor(model_data[[collapse_by]])
+  } else if (collapse_by %in% colnames(rowdata)) {
+    rawdata$random <- factor(rawdata[[collapse_by]])
+  } else {
+    insight::format_error(paste0("Could not find `", collapse_by, "` column."))
+  }
+
+  agg_data <- stats::aggregate(
+    rawdata[[y_name]],
+    by = rawdata[colnames(rawdata) != y_name],
+    FUN = mean
+  )
+
+  if (residuals) {
+    y_name <- insight::find_response(model)
+  }
+
+  colnames(agg_data)[ncol(agg_data)] <- y_name
+  colnames(agg_data)[colnames(agg_data) == "group"] <- "group_col"
+
+  # sanity check, add dummy if not present
+  if (is.null(agg_data$group_col)) {
+    agg_data$group_col <- factor(1)
+  }
+
+  agg_data
+}

R/tinyplot.R

@@ -49,6 +49,7 @@ tinyplot.estimate_means <- function(
   type = NULL,
   dodge = NULL,
   show_data = FALSE,
+  collapse_group = NULL,
   numeric_as_discrete = NULL,
   ...
 ) {
@@ -191,7 +192,11 @@ tinyplot.estimate_means <- function(
   if (show_data) {
     # extract raw data from the model
     model <- attributes(x)$model
-    rawdata <- as.data.frame(insight::get_data(model, verbose = FALSE))
+    if (is.null(collapse_group)) {
+      rawdata <- as.data.frame(insight::get_data(model, verbose = FALSE))
+    } else {
+      rawdata <- collapse_by_group(x, model, collapse_by = collapse_group)
+    }
 
     # set alpha
     if (is.null(dots$alpha)) {

R/visualisation_recipe.R

@@ -23,8 +23,10 @@
 #'
 #' @param x A modelbased object.
 #' @param show_data Logical, if `TRUE`, display the "raw" data as a background
-#' to the model-based estimation. This argument will be ignored for plotting
-#' objects returned by `estimate_slopes()` or `estimate_grouplevel()`.
+#' to the model-based estimation. For mixed models, you can additionally use the
+#' `collapse_group` argument to "collapse" data points by random effects
+#' grouping factors. Argument `show_data` will be ignored for plotting objects
+#' returned by `estimate_slopes()` or `estimate_grouplevel()`.
 #' @param join_dots Logical, if `TRUE` (default) and for categorical focal terms
 #' in `by`, dots (estimates) are connected by lines, i.e. plots will be a
 #' combination of dots with error bars and connecting lines. If `FALSE`, only
@@ -39,9 +41,17 @@
 #' predictor. Use `FALSE` to always use continuous color scales for numeric
 #' predictors. It is possible to set a global default value using `options()`,
 #' e.g. `options(modelbased_numeric_as_discrete = 10)`.
-#' @param show_residuals Logical, if `TRUE`, display residuals of the model
-#' as a background to the model-based estimation. Residuals will be computed
-#' for the predictors in the data grid, using [`residualize_over_grid()`].
+#' @param show_residuals Logical, if `TRUE`, display residuals of the model as a
+#' background to the model-based estimation. Residuals will be computed for the
+#' predictors in the data grid, using [`residualize_over_grid()`]. For mixed
+#' models, you can additionally use the `collapse_group` argument to "collapse"
+#' data points from residuals by random effects grouping factors.
+#' @param collapse_group This argument only takes effect when either `show_data`
+#' or `show_residuals` is `TRUE`. For mixed effects models, name of the grouping
+#' variable of random effects. If `collapse_group = TRUE`, data points
+#' "collapsed" by the first random effect groups are added to the plot. Else, if
+#' `collapse_group` is a name of a group factor, data is collapsed by that
+#' specific random effect. See [`collapse_by_group()`] for further details.
 #' @param point,line,pointrange,ribbon,facet,grid Additional
 #' aesthetics and parameters for the geoms (see customization example).
 #' @param ... Arguments passed from `plot()` to `visualisation_recipe()`, or
@@ -74,7 +84,7 @@
 #'   will set a default value for the `dodge` argument (spacing between geoms)
 #'   when using `tinyplot::plt()`. Should be a number between `0` and `1`.
 #'
-#' @examplesIf all(insight::check_if_installed(c("marginaleffects", "see", "ggplot2"), quietly = TRUE)) && getRversion() >= "4.1.0"
+#' @examplesIf all(insight::check_if_installed(c("marginaleffects", "see", "ggplot2", "lme4"), quietly = TRUE)) && getRversion() >= "4.1.0"
 #' library(ggplot2)
 #' library(see)
 #' # ==============================================
@@ -152,25 +162,42 @@
 #' x <- estimate_means(model, by = c("cyl", "wt"))
 #' plot(x)
 #'
-#'
 #' # GLMs ---------------------
 #' data <- data.frame(vs = mtcars$vs, cyl = as.factor(mtcars$cyl))
 #' x <- estimate_means(glm(vs ~ cyl, data = data, family = "binomial"), by = c("cyl"))
 #' plot(x)
+#'
+#' # ==============================================
+#' # Adding original data to the plot
+#' # ==============================================
+#' data(efc, package = "modelbased")
+#' efc$e15relat <- as.factor(efc$e15relat)
+#' efc$c161sex <- as.factor(efc$c161sex)
+#' levels(efc$c161sex) <- c("male", "female")
+#' model <- lme4::lmer(neg_c_7 ~ c161sex + (1 | e15relat), data = efc)
+#'
+#' me <- estimate_means(model, "c161sex")
+#' plot(me, show_data = TRUE)
+#'
+#' # data points: collapse by / average over random effects groups -------
+#' plot(me, show_data = TRUE, collapse_group = "e15relat")
 #' }
 #' @export
-visualisation_recipe.estimate_predicted <- function(x,
-                                                    show_data = FALSE,
-                                                    show_residuals = FALSE,
-                                                    point = NULL,
-                                                    line = NULL,
-                                                    pointrange = NULL,
-                                                    ribbon = NULL,
-                                                    facet = NULL,
-                                                    grid = NULL,
-                                                    join_dots = NULL,
-                                                    numeric_as_discrete = NULL,
-                                                    ...) {
+visualisation_recipe.estimate_predicted <- function(
+  x,
+  show_data = FALSE,
+  show_residuals = FALSE,
+  collapse_group = NULL,
+  point = NULL,
+  line = NULL,
+  pointrange = NULL,
+  ribbon = NULL,
+  facet = NULL,
+  grid = NULL,
+  join_dots = NULL,
+  numeric_as_discrete = NULL,
+  ...
+) {
   # Process argument ---------------------------------------------------------
   # --------------------------------------------------------------------------
 
@@ -186,6 +213,7 @@ visualisation_recipe.estimate_predicted <- function(x,
     x,
     show_data = show_data,
     show_residuals = show_residuals,
+    collapse_by = collapse_group,
     point = point,
     line = line,
     pointrange = pointrange,
@@ -233,13 +261,15 @@ visualisation_recipe.estimate_means <- visualisation_recipe.estimate_predicted
 #' plot(visualisation_recipe(x))
 #' }
 #' @export
-visualisation_recipe.estimate_slopes <- function(x,
-                                                 line = NULL,
-                                                 pointrange = NULL,
-                                                 ribbon = NULL,
-                                                 facet = NULL,
-                                                 grid = NULL,
-                                                 ...) {
+visualisation_recipe.estimate_slopes <- function(
+  x,
+  line = NULL,
+  pointrange = NULL,
+  ribbon = NULL,
+  facet = NULL,
+  grid = NULL,
+  ...
+) {
   .visualization_recipe(
     x,
     show_data = FALSE,
@@ -285,13 +315,15 @@ visualisation_recipe.estimate_slopes <- function(x,
 #' plot(x)
 #' }
 #' @export
-visualisation_recipe.estimate_grouplevel <- function(x,
-                                                     line = NULL,
-                                                     pointrange = NULL,
-                                                     ribbon = NULL,
-                                                     facet = NULL,
-                                                     grid = NULL,
-                                                     ...) {
+visualisation_recipe.estimate_grouplevel <- function(
+  x,
+  line = NULL,
+  pointrange = NULL,
+  ribbon = NULL,
+  facet = NULL,
+  grid = NULL,
+  ...
+) {
   if (is.null(facet)) {
     facet <- list(scales = "free")
   } else {