R/rootogram.R

In gratia: Graceful 'ggplot'-Based Graphics and Other Functions for GAMs Fitted Using 'mgcv'

## Add rootogram support for suitable distributions

#' Rootograms to assess goodness of model fit
#'
#' A rootogram is a model diagnostic tool that assesses the goodness of fit of
#' a statistical model. The observed values of the response are compared with
#' those expected from the fitted model. For discrete, count responses, the
#' frequency of each count (0, 1, 2, etc) in the observed data and expected
#' from the conditional distribution of the response implied by the model are
#' compared. For continuous variables, the observed and expected frequencies
#' are obtained by grouping the data into bins. The rootogram is drawn using
#' [ggplot2::ggplot()] graphics. The design closely follows Kleiber & Zeileis
#' (2016).
#'
#' @param object an R object
#' @param max_count integer; the largest count to consider
#' @param breaks for continuous responses, how to group the response. Can be
#'   anything that is acceptable as the `breaks` argument of
#'   [graphics::hist.default()]
#' @param ... arguments passed to other methods
#'
#' @references Kleiber, C., Zeileis, A., (2016) Visualizing Count Data
#'   Regressions Using Rootograms. *Am. Stat.* **70**, 296–303.
#'   \doi{10.1080/00031305.2016.1173590}
#'
#' @export
#' @examples
#' load_mgcv()
#' \dontshow{op <- options(cli.unicode = FALSE, pillar.sigfig = 6)}
#' df <- data_sim("eg1", n = 1000, dist = "poisson", scale = 0.1, seed = 6)
#'
#' # A poisson example
#' m <- gam(y ~ s(x0, bs = "cr") + s(x1, bs = "cr") + s(x2, bs = "cr") +
#'          s(x3, bs = "cr"), family = poisson(), data = df, method = "REML")
#' rg <- rootogram(m)
#' rg
#' draw(rg) # plot the rootogram
#'
#' # A Gaussian example
#' df <- data_sim("eg1", dist = "normal", seed = 2)
#' m <- gam(y ~ s(x0) + s(x1) + s(x2) + s(x3), data = df, method = "REML")
#' draw(rootogram(m, breaks = "FD"), type = "suspended")
#' \dontshow{options(op)}
`rootogram` <- function(object, ...) {
    UseMethod("rootogram")
}

#' @rdname rootogram
#' @importFrom stringr str_detect str_to_title
#' @export
`rootogram.gam` <- function(object, max_count = NULL, breaks = "Sturges", ...) {
    ## limit this to a few distributions at the moment
    distrs <- c("poisson", "Negative Binomial", "gaussian")
    fname <- family_name(object)

    ## check model family is supported
    supported <- stringr::str_detect(fname, distrs)
    if (!any(supported)) {
        stop("Only <", paste(str_to_title(distrs), collapse = ", "),
             "> models supported.", call. = FALSE)
    }

    df <- switch(distrs[supported],
                 poisson = fitted_poisson(object, max_count = max_count),
                 `Negative Binomial` = fitted_neg_bin(object,
                                                      max_count = max_count),
                 gaussian = fitted_gaussian(object, breaks = breaks))
    class(df) <- append(class(df), "rootogram", after = 0L)
    attr(df, "distribution") <- stringr::str_to_title(distrs[supported])
    df
}

#' @importFrom purrr map_dfc
#' @importFrom stats dpois na.omit
#' @importFrom tibble tibble
#' @keywords internal
`fitted_poisson` <- function(model, max_count = NULL) {
    mf <- model.frame(model)
    y <- model.response(mf)
    if (is.null(max_count)) {
        max_count <- rootogram_max_count(y)
    }
    bin <- seq(0, max_count, by = 1L)
    names(bin) <- as.character(bin)
    obs <- table(factor(y, levels = bin))
    lambda <- predict(model, newdata = mf, na.action = na.omit,
                      type = "response")
    fitted <- purrr::map_dfc(bin, .f = dpois, lambda = lambda)
    fitted <- colSums(fitted)
    tibble(bin = bin,
           observed = as.integer(obs[seq_along(bin)]),
           fitted = fitted)
}

#' @importFrom purrr map_dfc
#' @importFrom stats dnbinom na.omit
#' @importFrom tibble tibble
#' @keywords internal
`fitted_neg_bin` <- function(model, max_count = NULL) {
    mf <- model.frame(model)
    y <- model.response(mf)
    if (is.null(max_count)) {
        max_count <- rootogram_max_count(y)
    }
    bin <- seq(0, max_count, by = 1L)
    names(bin) <- as.character(bin)
    obs <- table(factor(y, levels = bin))
    mu <- predict(model, newdata = mf, na.action = na.omit,
                  type = "response")
    theta <- nb_theta(model)
    fitted <- purrr::map_dfc(bin, .f = dnbinom, mu = mu, size = theta)
    fitted <- colSums(fitted)
    df <- tibble(bin = bin,
                 observed = as.integer(obs[seq_along(bin)]),
                 fitted = fitted)
    attr(df, "theta") <- theta
    df
}

#' @importFrom purrr map_dfc
#' @importFrom stats pnorm na.omit
#' @importFrom graphics hist
#' @importFrom tibble tibble
#' @keywords internal
`fitted_gaussian` <- function(model, breaks = NULL) {
    if (is.null(breaks)) {
        breaks <- "Sturges" ## give same breaks as those of hist()
    }
    mf <- model.frame(model)
    y <- model.response(mf)
    mu <- predict(model, newdata = mf, na.action = na.omit)
    h <- hist(y, breaks = breaks, plot = FALSE)
    bin <- h$breaks
    names(bin) <- as.character(bin)
    obs <- h$counts
    mid <- h$mids
    sigma <- model$sig2
    pdf <- purrr::map_dfc(bin, .f = pnorm, mean = mu, sd = sigma)
    nc <- length(bin)
    fitted <- pdf[, seq(1, nc - 1, by = 1L)]
    fitted[] <- pdf[, seq(2, nc, by = 1L)] - pdf[, seq(1, nc - 1, by = 1L)]
    fitted <- colSums(fitted)
    df <- tibble::tibble(bin = bin[-nc],
                         observed = obs,
                         fitted = fitted)
    attr(df, "sigma") <- sigma
    attr(df, "mid") <- mid
    attr(df, "width") <- diff(h$breaks)
    df
}

## try to match countreg::rootogram
#' @noRd
#' @keywords internal
`rootogram_max_count` <- function(y) {
    max(1.5 * max(y), 20L)
}

#' Draw a rootogram
#'
#' A rootogram is a model diagnostic tool that assesses the goodness of fit of
#' a statistical model. The observed values of the response are compared with
#' those expected from the fitted model. For discrete, count responses, the
#' frequency of each count (0, 1, 2, etc) in the observed data and expected
#' from the conditional distribution of the response implied by the model are
#' compared. For continuous variables, the observed and expected frequencies
#' are obtained by grouping the data into bins. The rootogram is drawn using
#' [ggplot2::ggplot()] graphics. The design closely follows Kleiber & Zeileis
#' (2016).
#'
#' @param type character; the type of rootogram to draw.
#' @param sqrt logical; show the observed and fitted frequencies
#' @param ref_line logical; draw a reference line at zero?
#' @param warn_limits logical; draw Tukey's warning limit lines at +/- 1?
#' @param fitted_colour,bar_colour,bar_fill,ref_line_colour,warn_line_colour
#'   colours used to draw the respective element of the rootogram.
#' @param xlab,ylab character; labels for the x and y axis of the rootogram.
#'   May be missing (`NULL`), in which case suitable labels will be used.
#''
#' @inheritParams draw
#'
#' @references Kleiber, C., Zeileis, A., (2016) Visualizing Count Data
#'   Regressions Using Rootograms. *Am. Stat.* **70**, 296–303.
#'   \doi{10.1080/00031305.2016.1173590}
#'
#' @return A 'ggplot' object.
#'
#' @family draw methods
#' @seealso [rootogram()] to compute the data for the rootogram.
#'
#' @export
#' @importFrom rlang .data
#' @importFrom dplyr mutate
#' @importFrom ggplot2 ggplot aes geom_rect geom_point geom_line labs geom_hline
#'
#' @examples
#' load_mgcv()
#' df <- data_sim("eg1", n = 1000, dist = "poisson", scale = 0.1, seed = 6)
#'
#' # A poisson example
#' m <- gam(y ~ s(x0, bs = "cr") + s(x1, bs = "cr") + s(x2, bs = "cr") +
#'          s(x3, bs = "cr"), family = poisson(), data = df, method = "REML")
#' rg <- rootogram(m)
#'
#' # plot the rootogram
#' draw(rg)
#'
#' # change the type of rootogram
#' draw(rg, type = "suspended")
`draw.rootogram` <- function(object,
                             type = c("hanging", "standing", "suspended"),
                             sqrt = TRUE,
                             ref_line = TRUE,
                             warn_limits = TRUE,
                             fitted_colour = "steelblue",
                             bar_colour = NA, #"grey",
                             bar_fill = "grey",
                             ref_line_colour = "black",
                             warn_line_colour = "black",
                             ylab = NULL, xlab = NULL, ...) {
    type <- match.arg(type)

    ## Create a suitable axis label for y axis if none supplied
    if (is.null(ylab)) {
        ylab <- if (as.logical(sqrt)) {
            object <- mutate(object,
                             observed = sqrt(.data$observed),
                             fitted = sqrt(.data$fitted))
            expression(sqrt(Frequency))
        } else {
            "Frequency"
        }
    }
    ## ...and same for the x axis
    if (is.null(xlab)) {
        xlab <- "Response"
    }

    ## bin width
    distr <- attr(object, "distribution")
    ## bins for counts are 1 unit wide so make bins 5% narrower on either side
    width <- if (distr %in% c("Poisson", "Negative Binomial")) {
        0.45
    } else {
        ## but if a continuous distrib then we capture the width attribute and
        ## scale it by 90 (dividing by 2 as we +/- half the width to observed
        ## value)
        w <- attr(object, "width")
        (w * 0.9) / 2
    }

    object <- mutate(object,
                     x_low = .data$bin - width,
                     x_high = .data$bin + width)

    nr <- nrow(object)
    object <- if (type == "hanging") {
        mutate(object,
               y_bot = .data$fitted - .data$observed,
               y_top = .data$y_bot + .data$observed)
    } else if (type == "suspended") {
        mutate(object,
               y_bot = rep(0, nr),
               y_top = .data$fitted - .data$observed)
    } else {
        mutate(object,
               y_bot = rep(0, nr),
               y_top = .data$y_bot + .data$observed)
    }

    plt <- ggplot(object) +
        geom_rect(aes(xmin = .data$x_low,
                      xmax = .data$x_high,
                      ymin = .data$y_bot,
                      ymax = .data$y_top),
                  fill = bar_fill, col = bar_colour) +
        geom_line(aes(x = .data$bin,
                      y = .data$fitted),
                  colour = fitted_colour, linewidth = 1) +
        geom_point(aes(x = .data$bin,
                       y = .data$fitted),
                   colour = fitted_colour, size = 2.5)

    if (as.logical(ref_line) & type != "standing") {
        plt <- plt + geom_hline(yintercept = 0, colour = ref_line_colour)
    }

    if (as.logical(warn_limits) & type != "standing") {
        plt <- plt + geom_hline(yintercept = c(-1, 1), lty = "dashed",
                                colour = warn_line_colour)
    }

    plt <- plt + labs(y = ylab, x = xlab)

    plt
}