Merge pull request #362 from stan-dev/rootogram

adding discrete option to ppc_rootogram

Author: behramulukir

NEWS.md

@@ -5,6 +5,7 @@
 * Add `ppc_dots()` and `ppd_dots()` by @behramulukir (#357)
 * Add `x` argument to `ppc_error_binned` by @behramulukir (#359)
 * Add `x` argument to `ppc_error_scatter_avg()` by @behramulukir (#367)
+* Add `discrete` style to `ppc_rootogram` by @behramulukir (#362)
 
 # bayesplot 1.13.0
 

R/ppc-discrete.R

@@ -21,9 +21,12 @@
 #' @param size,fatten,linewidth For bar plots, `size`, `fatten`, and `linewidth`
 #'   are passed to [ggplot2::geom_pointrange()] to control the appearance of the
 #'   `yrep` points and intervals. For rootograms `size` is passed to
-#'   [ggplot2::geom_line()].
+#'   [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)`,
+#'  if `TRUE` then the observed counts will be plotted with different shapes
+#'  depending on whether they are within the bounds of the `y` quantiles.
 #'
 #' @template return-ggplot-or-data
 #'
@@ -44,18 +47,29 @@
 #' }
 #' \item{`ppc_rootogram()`}{
 #'   Rootograms allow for diagnosing problems in count data models such as
-#'   overdispersion or excess zeros. They consist of a histogram of `y` with the
-#'   expected counts based on `yrep` overlaid as a line along with uncertainty
-#'   intervals. The y-axis represents the square roots of the counts to
-#'   approximately adjust for scale differences and thus ease comparison between
-#'   observed and expected counts. Using the `style` argument, the histogram
-#'   style can be adjusted to focus on different aspects of the data:
+#'   overdispersion or excess zeros. In `standing`, `hanging`, and `suspended`
+#'   styles, they consist of a histogram of `y` with the expected counts based on
+#'   `yrep` overlaid as a line along with uncertainty intervals.
+#'
+#'   Meanwhile, in `discrete` style, median counts based on `yrep` are laid
+#'   as a point range with uncertainty intervals along with dots
+#'   representing the `y`.
+#'
+#'   The y-axis represents the square roots of the counts to approximately
+#'   adjust for scale differences and thus ease comparison between observed
+#'   and expected counts. Using the `style` argument, the rootogram can be
+#'   adjusted to focus on different aspects of the data:
 #'   * _Standing_: basic histogram of observed counts with curve
 #'    showing expected counts.
-#'   * _Hanging_: observed counts counts hanging from the curve
+#'   * _Hanging_: observed counts hanging from the curve
 #'    representing expected counts.
 #'   * _Suspended_: histogram of the differences between expected and
 #'    observed counts.
+#'    * _Discrete_: a dot-and-whisker plot of the median counts and
+#'    dots representing observed counts.
+#'
+#'    As it emphasizes the discrete nature of the count data,
+#'    using `discrete` style is suggested.
 #'
 #'   **All of the rootograms are plotted on the square root scale**. See Kleiber
 #'   and Zeileis (2016) for advice on interpreting rootograms and selecting
@@ -198,22 +212,22 @@ ppc_bars_grouped <-
            fatten = 2.5,
            linewidth = 1,
            freq = TRUE) {
-  check_ignored_arguments(...)
-  call <- match.call(expand.dots = FALSE)
-  g <- eval(ungroup_call("ppc_bars", call), parent.frame())
-  if (fixed_y(facet_args)) {
-    g <- g + expand_limits(y = 1.05 * max(g$data[["h"]], na.rm = TRUE))
+    check_ignored_arguments(...)
+    call <- match.call(expand.dots = FALSE)
+    g <- eval(ungroup_call("ppc_bars", call), parent.frame())
+    if (fixed_y(facet_args)) {
+      g <- g + expand_limits(y = 1.05 * max(g$data[["h"]], na.rm = TRUE))
+    }
+    g +
+      bars_group_facets(facet_args) +
+      force_axes_in_facets()
   }
-  g +
-    bars_group_facets(facet_args) +
-    force_axes_in_facets()
-}
 
 
 #' @rdname PPC-discrete
 #' @export
 #' @param style For `ppc_rootogram`, a string specifying the rootogram
-#'   style. The options are `"standing"`, `"hanging"`, and
+#'   style. The options are `"discrete"`, `"standing"`, `"hanging"`, and
 #'   `"suspended"`. See the **Plot Descriptions** section, below, for
 #'   details on the different styles.
 #'
@@ -234,61 +248,34 @@ ppc_bars_grouped <-
 #'
 #' 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"),
+                          style = c("standing", "hanging", "suspended", "discrete"),
                           ...,
                           prob = 0.9,
-                          size = 1) {
+                          size = 1,
+                          bound_distinct = TRUE) {
   check_ignored_arguments(...)
   style <- match.arg(style)
-  y <- validate_y(y)
-  yrep <- validate_predictions(yrep, length(y))
-  if (!all_counts(y)) {
-    abort("ppc_rootogram expects counts as inputs to 'y'.")
-  }
-  if (!all_counts(yrep)) {
-    abort("ppc_rootogram expects counts as inputs to 'yrep'.")
-  }
 
-  alpha <- (1 - prob) / 2
-  probs <- c(alpha, 1 - alpha)
-  ymax <- max(y, yrep)
-  xpos <- 0L:ymax
-
-  # prepare a table for yrep
-  tyrep <- as.list(rep(NA, nrow(yrep)))
-  for (i in seq_along(tyrep)) {
-    tyrep[[i]] <- table(yrep[i,])
-    matches <- match(xpos, rownames(tyrep[[i]]))
-    tyrep[[i]] <- as.numeric(tyrep[[i]][matches])
-  }
-  tyrep <- do.call(rbind, tyrep)
-  tyrep[is.na(tyrep)] <- 0
-  tyexp <- sqrt(colMeans(tyrep))
-  tyquantile <- sqrt(t(apply(tyrep, 2, quantile, probs = probs)))
-  colnames(tyquantile) <- c("tylower", "tyupper")
-
-  # prepare a table for y
-  ty <- table(y)
-  ty <- sqrt(as.numeric(ty[match(xpos, rownames(ty))]))
-  if (style == "suspended") {
-    ty <- tyexp - ty
-  }
-  ty[is.na(ty)] <- 0
-  ypos <- ty / 2
-  if (style == "hanging") {
-    ypos <- tyexp - ypos
-  }
+  data <- .ppc_rootogram_data(
+    y = y,
+    yrep = yrep,
+    style = style,
+    prob = prob,
+    bound_distinct = bound_distinct
+  )
 
-  data <- data.frame(xpos, ypos, ty, tyexp, tyquantile)
-  graph <- ggplot(data) +
-    aes(
-      ymin = .data$tylower,
-      ymax = .data$tyupper,
-      height = .data$ty
-    ) +
+  # Building geoms for y and y_rep
+  geom_y <- if (style == "discrete") {
+    geom_point(
+      aes(y = .data$obs, shape = .data$obs_shape),
+      size = size * 1.5,
+      color = get_color("d"),
+      fill = get_color("d"))
+  } else {
     geom_tile(
       aes(
         x = .data$xpos,
@@ -298,34 +285,69 @@ ppc_rootogram <- function(y,
       color = get_color("lh"),
       linewidth = 0.25,
       width = 1
-    ) +
-    bayesplot_theme_get()
-
-  if (style != "standing") {
-    graph <- graph + hline_0(size = 0.4)
+    )
   }
 
-  graph <- graph +
+  geom_yrep <- if (style == "discrete") {
+    geom_pointrange(
+      aes(y = .data$pred_median, ymin = .data$lower, ymax = .data$upper, color = "y_rep"),
+      fill = get_color("lh"),
+      linewidth = size,
+      size = size,
+      fatten = 2,
+      alpha = 1
+    )
+  } else {
     geom_smooth(
-      aes(
-        x = .data$xpos,
-        y = .data$tyexp,
-        color = "Expected"
-      ),
+      aes(x = .data$xpos, y = .data$tyexp, color = "Expected"),
       fill = get_color("d"),
       linewidth = size,
       stat = "identity"
-    ) +
-    scale_fill_manual("", values = get_color("l")) +
-    scale_color_manual("", values = get_color("dh")) +
-    labs(x = expression(italic(y)),
-         y = expression(sqrt(Count)))
-
-  if (style == "standing") {
-    graph <- graph + dont_expand_y_axis()
+    )
   }
 
-  graph + reduce_legend_spacing(0.25)
+  # Creating the graph
+  graph <- ggplot(data)
+
+  if (style == "discrete") {
+    graph <- graph +
+      geom_yrep +
+      geom_y +
+      aes(x = .data$xpos) +
+      scale_y_sqrt() +
+      scale_fill_manual("", values = get_color("d"), guide = "none") +
+      scale_color_manual("", values = get_color("lh"), labels = yrep_label()) +
+      labs(x = expression(italic(y)), y = "Count") +
+      bayesplot_theme_get() +
+      reduce_legend_spacing(0.25) +
+      scale_shape_manual(values = c("In" = 22, "Out" = 23, "y" = 22), guide = "legend", labels = c("y" = expression(italic(y))))
+    if (bound_distinct) {
+      graph <- graph + guides(shape = guide_legend(expression(italic(y)~within~bounds)))
+    } else {
+      graph <- graph + guides(shape = guide_legend(" "))
+    }
+  } else {
+    graph <- graph +
+      geom_y +
+      geom_yrep +
+      aes(
+        ymin = .data$tylower,
+        ymax = .data$tyupper,
+        height = .data$ty
+      ) +
+      scale_fill_manual("", values = get_color("l")) +
+      scale_color_manual("", values = get_color("dh")) +
+      labs(x = expression(italic(y)), y = expression(sqrt(Count))) +
+      bayesplot_theme_get() +
+      reduce_legend_spacing(0.25)
+    if (style == "standing") {
+      graph <- graph + dont_expand_y_axis()
+    } else {
+      graph <- graph + hline_0(size = 0.4)
+    }
+  }
+
+  return(graph)
 }
 
 
@@ -395,7 +417,7 @@ ppc_bars_data <-
   data <-
     reshape2::melt(tmp_data, id.vars = "group") %>%
     count(.data$group, .data$value, .data$variable) %>%
-    tidyr::complete(.data$group, .data$value, .data$variable, fill = list(n = 0)) %>% 
+    tidyr::complete(.data$group, .data$value, .data$variable, fill = list(n = 0)) %>%
     group_by(.data$variable, .data$group) %>%
     mutate(proportion = .data$n / sum(.data$n)) %>%
     ungroup() %>%
@@ -440,3 +462,90 @@ bars_group_facets <- function(facet_args, scales_default = "fixed") {
 fixed_y <- function(facet_args) {
   !isTRUE(facet_args[["scales"]] %in% c("free", "free_y"))
 }
+
+#' Internal function for `ppc_rootogram()`
+#' @param y,yrep User's `y` and `yrep` arguments.
+#' @param style,prob,bound_distinct User's `style`, `prob`, and
+#' (if applicable) `bound_distinct` arguments.
+#' @noRd
+.ppc_rootogram_data <- function(y,
+                              yrep,
+                              style = c("standing", "hanging", "suspended", "discrete"),
+                              prob = 0.9,
+                              bound_distinct) {
+
+  y <- validate_y(y)
+  yrep <- validate_predictions(yrep, length(y))
+  if (!all_counts(y)) {
+    abort("ppc_rootogram expects counts as inputs to 'y'.")
+  }
+  if (!all_counts(yrep)) {
+    abort("ppc_rootogram expects counts as inputs to 'yrep'.")
+  }
+
+  alpha <- (1 - prob) / 2
+  probs <- c(alpha, 1 - alpha)
+  ymax <- max(y, yrep)
+  xpos <- 0L:ymax
+
+  # prepare a table for yrep
+  tyrep <- as.list(rep(NA, nrow(yrep)))
+  for (i in seq_along(tyrep)) {
+    tyrep[[i]] <- table(yrep[i,])
+    matches <- match(xpos, rownames(tyrep[[i]]))
+    tyrep[[i]] <- as.numeric(tyrep[[i]][matches])
+  }
+  tyrep <- do.call(rbind, tyrep)
+  tyrep[is.na(tyrep)] <- 0
+
+  # discrete style
+  if (style == "discrete"){
+    pred_median <- apply(tyrep, 2, median)
+    pred_quantile <- t(apply(tyrep, 2, quantile, probs = probs))
+    colnames(pred_quantile) <- c("lower", "upper")
+
+    # prepare a table for y
+    ty <- table(y)
+    y_count <- as.numeric(ty[match(xpos, rownames(ty))])
+    y_count[is.na(y_count)] <- 0
+
+    if (bound_distinct) {
+      # If the observed count is within the bounds of the predicted quantiles,
+      # use a different shape for the point
+      obs_shape <- obs_shape <- ifelse(y_count >= pred_quantile[, "lower"] & y_count <= pred_quantile[, "upper"], "In", "Out")
+    } else {
+      obs_shape <- rep("y", length(y_count)) # all points are the same shape for observed
+    }
+
+    data <- data.frame(
+      xpos = xpos,
+      obs = y_count,
+      pred_median = pred_median,
+      lower = pred_quantile[, "lower"],
+      upper = pred_quantile[, "upper"],
+      obs_shape = obs_shape
+    )
+  }
+  # standing, hanging, suspended styles
+  else {
+    tyexp <- sqrt(colMeans(tyrep))
+    tyquantile <- sqrt(t(apply(tyrep, 2, quantile, probs = probs)))
+    colnames(tyquantile) <- c("tylower", "tyupper")
+
+    # prepare a table for y
+    ty <- table(y)
+    ty <- sqrt(as.numeric(ty[match(xpos, rownames(ty))]))
+    if (style == "suspended") {
+      ty <- tyexp - ty
+    }
+    ty[is.na(ty)] <- 0
+    ypos <- ty / 2
+    if (style == "hanging") {
+      ypos <- tyexp - ypos
+    }
+
+    data <- data.frame(xpos, ypos, ty, tyexp, tyquantile)
+  }
+
+  return(data)
+}