adding ppc_rootogram_grouped function

NAMESPACE

@@ -155,6 +155,7 @@ export(ppc_ribbon)
 export(ppc_ribbon_data)
 export(ppc_ribbon_grouped)
 export(ppc_rootogram)
+export(ppc_rootogram_grouped)
 export(ppc_scatter)
 export(ppc_scatter_avg)
 export(ppc_scatter_avg_data)

R/ppc-discrete.R

@@ -24,7 +24,7 @@
 #'   [ggplot2::geom_line()] and [ggplot2::geom_pointrange()].
 #' @param freq For bar plots only, if `TRUE` (the default) the y-axis will
 #'   display counts. Setting `freq=FALSE` will put proportions on the y-axis.
-#' @param bound_distinct For `ppc_rootogram(style = "discrete)`,
+#' @param bound_distinct For `ppc_rootogram(style = "discrete)` and `ppc_rootogram_grouped(style = "discrete)`,
 #'  if `TRUE` then the observed counts will be plotted with different shapes
 #'  depending on whether they are within the bounds of the `y` quantiles.
 #'
@@ -75,6 +75,10 @@
 #'   and Zeileis (2016) for advice on interpreting rootograms and selecting
 #'   among the different styles.
 #' }
+#' \item{`ppc_rootogram_grouped()`}{
+#'   Same as `ppc_rootogram()` but a separate plot (facet) is generated for each
+#'   level of a grouping variable.
+#' }
 #' }
 #'
 #' @section Related functions:
@@ -147,6 +151,24 @@
 #' )
 #' }
 #'
+#' # rootograms for counts
+#' y <- rpois(100, 20)
+#' yrep <- matrix(rpois(10000, 20), ncol = 100)
+#'
+#' color_scheme_set("brightblue")
+#' ppc_rootogram(y, yrep)
+#' ppc_rootogram(y, yrep, prob = 0)
+#'
+#' ppc_rootogram(y, yrep, style = "hanging", prob = 0.8)
+#' ppc_rootogram(y, yrep, style = "suspended")
+#' ppc_rootogram(y, yrep, style = "discrete")
+#'
+#' # rootograms for counts with groups
+#' group <- gl(2, 50, length = 100, labels = c("GroupA", "GroupB"))
+#' ppc_rootogram_grouped(y, yrep, group)
+#' ppc_rootogram_grouped(y, yrep, group, style = "hanging", facet_args = list(nrow = 2))
+#' ppc_rootogram_grouped(y, yrep, group, style = "discrete", prob = 0.5)
+#'
 NULL
 
 #' @rdname PPC-discrete
@@ -253,32 +275,26 @@ ppc_bars_grouped <-
 #' *The American Statistician*. 70(3): 296--303.
 #' <https://arxiv.org/abs/1605.01311>.
 #'
-#' @examples
-#' # rootograms for counts
-#' y <- rpois(100, 20)
-#' yrep <- matrix(rpois(10000, 20), ncol = 100)
-#'
-#' color_scheme_set("brightblue")
-#' ppc_rootogram(y, yrep)
-#' ppc_rootogram(y, yrep, prob = 0)
-#'
-#' ppc_rootogram(y, yrep, style = "hanging", prob = 0.8)
-#' ppc_rootogram(y, yrep, style = "suspended")
-#' ppc_rootogram(y, yrep, style = "discrete")
-#'
 ppc_rootogram <- function(y,
                           yrep,
                           style = c("standing", "hanging", "suspended", "discrete"),
                           ...,
                           prob = 0.9,
                           size = 1,
                           bound_distinct = TRUE) {
-  check_ignored_arguments(...)
+
+  dots <- list(...)
+  if (!from_grouped(dots)) {
+    check_ignored_arguments(...)
+    dots$group <- NULL
+  }
+
   style <- match.arg(style)
 
   data <- .ppc_rootogram_data(
     y = y,
     yrep = yrep,
+    group = dots$group,
     style = style,
     prob = prob,
     bound_distinct = bound_distinct
@@ -367,6 +383,33 @@ ppc_rootogram <- function(y,
 }
 
 
+#' @rdname PPC-discrete
+#' @export
+ppc_rootogram_grouped <- 
+  function(y,
+           yrep,
+           group,
+           style = c("standing", "hanging", "suspended", "discrete"),
+           ...,
+           facet_args = list(),
+           prob = 0.9,
+           size = 1,
+           bound_distinct = TRUE) {
+    check_ignored_arguments(...)
+    style <- match.arg(style)
+    call <- match.call(expand.dots = FALSE)
+    g <- eval(ungroup_call("ppc_rootogram", call), parent.frame())
+
+    # In style = discrete, scale_y_sqrt() can't handle -Inf values in axis segments
+    # so force_axes_in_facets() results in errors
+    if (style != "discrete") {
+      g <- g + force_axes_in_facets()
+    }
+
+    g + bars_group_facets(facet_args)
+  }
+
+
 #' @rdname PPC-discrete
 #' @export
 ppc_bars_data <-
@@ -481,11 +524,13 @@ fixed_y <- function(facet_args) {
 
 #' Internal function for `ppc_rootogram()`
 #' @param y,yrep User's `y` and `yrep` arguments.
+#' @param group User's `group` argument (can be NULL).
 #' @param style,prob,bound_distinct User's `style`, `prob`, and
 #' (if applicable) `bound_distinct` arguments.
 #' @noRd
 .ppc_rootogram_data <- function(y,
                               yrep,
+                              group = NULL,
                               style = c("standing", "hanging", "suspended", "discrete"),
                               prob = 0.9,
                               bound_distinct) {
@@ -498,7 +543,41 @@ fixed_y <- function(facet_args) {
   if (!all_counts(yrep)) {
     abort("ppc_rootogram expects counts as inputs to 'yrep'.")
   }
+
+  if (!is.null(group)) {
+    group <- validate_group(group, length(y))
+  }
+
+  # Handle grouped data
+  if (!is.null(group)) {
+    group_levels <- unique(group)
+    all_data <- list()
+
+    for (g in seq_along(group_levels)) {
+      grp <- group_levels[g]
+      idx <- which(group == grp)
+
+      y_g <- y[idx]
+      yrep_g <- yrep[, idx, drop = FALSE]
+
+      # Call this function recursively without group
+      data_g <- .ppc_rootogram_data(
+        y = y_g,
+        yrep = yrep_g,
+        group = NULL,
+        style = style,
+        prob = prob,
+        bound_distinct = bound_distinct
+      )
+
+      data_g$group <- grp
+      all_data[[g]] <- data_g
+    }
+
+    return(do.call(rbind, all_data))
+  }
 
+  # Ungrouped data processing
   alpha <- (1 - prob) / 2
   probs <- c(alpha, 1 - alpha)
   ymax <- max(y, yrep)