R/plot_build.R

In marginaleffects: Predictions, Comparisons, Slopes, Marginal Means, and Hypothesis Tests

plot_preprocess <- function(
    dat,
    v_x,
    v_color = NULL,
    v_facet_1 = NULL,
    v_facet_2 = NULL,
    condition = NULL,
    mfx = NULL
) {
    for (v in names(condition$condition)) {
        fun <- function(x, lab) {
            idx <- match(x, sort(unique(x)))
            factor(lab[idx], levels = lab)
        }
        if (identical(condition$condition[[v]], "threenum")) {
            dat[[v]] <- fun(dat[[v]], c("-SD", "Mean", "+SD"))
        } else if (identical(condition$condition[[v]], "fivenum")) {
            # Labels must be unique: `digits = 2` collapses close values
            # (e.g. 10.02 and 10.05) and `factor()` then fails.
            vals <- sort(unique(dat[[v]]))
            labs <- format(vals, digits = 2)
            digits <- 2
            while (anyDuplicated(labs) > 0 && digits < 15) {
                digits <- digits + 1
                labs <- format(vals, digits = digits)
            }
            labs <- make.unique(labs)
            dat[[v]] <- fun(dat[[v]], labs)
        } else if (identical(condition$condition[[v]], "minmax")) {
            dat[[v]] <- fun(dat[[v]], c("Min", "Max"))
        } else if (identical(condition$condition[[v]], "quartile")) {
            dat[[v]] <- fun(dat[[v]], c("Q1", "Q2", "Q3"))
        }
    }
    if (check_variable_class(mfx, v_x, "categorical")) {
        dat[[v_x]] <- factor(dat[[v_x]])
    }
    # colors, linetypes, and facets are categorical attributes
    if (isTRUE(v_color %in% colnames(dat))) {
        dat[[v_color]] <- factor(dat[[v_color]])
    }
    if (isTRUE(v_facet_1 %in% colnames(dat))) {
        dat[[v_facet_1]] <- factor(dat[[v_facet_1]])
    }
    if (isTRUE(v_facet_2 %in% colnames(dat))) {
        dat[[v_facet_2]] <- factor(dat[[v_facet_2]])
    }

    if (!is.null(mfx) && !is.null(mfx@modeldata)) {
        vars <- c(v_x, v_color, v_facet_1, v_facet_2)
        for (v in vars) {
            if (isTRUE(v %in% colnames(dat)) && isTRUE(v %in% colnames(mfx@modeldata))) {
                attr(dat[[v]], "label") <- attr(mfx@modeldata[[v]], "label")
            }
        }
    }
    return(dat)
}


plot_build <- function(
    dat,
    v_x,
    v_color = NULL,
    v_facet_1 = NULL,
    v_facet_2 = NULL,
    dv = NULL,
    points = 0,
    rug = FALSE,
    gray = FALSE,
    mfx = NULL
) {
    checkmate::assert_flag(rug)
    checkmate::assert_flag(gray)

    is_discrete <- function(df, var, mfx = NULL) {
        if (is.null(var) || !isTRUE(var %in% colnames(df))) {
            return(FALSE)
        }
        out <- is.factor(df[[var]]) ||
            is.character(df[[var]]) ||
            is.logical(df[[var]])
        if (!is.null(mfx)) {
            out <- out || check_variable_class(mfx, var, "categorical")
        }
        return(out)
    }

    # create index before building ggplot to make sure it is available
    dat$marginaleffects_term_index <- get_unique_index(dat)
    multi_variables <- isTRUE(length(unique(dat$marginaleffects_term_index)) > 1)
    x_is_discrete <- is_discrete(dat, v_x, mfx)

    # The stored response may be a combined expression such as "cbind(s, f)",
    # which is not a column of `modeldata`. Plot the first raw response
    # variable instead, and skip points/rug when no column matches.
    dv_plot <- if (length(dv) > 0 && !is.na(dv[1]) && nzchar(dv[1])) dv[1] else NULL
    if (!is.null(dv_plot) && !dv_plot %in% colnames(mfx@modeldata)) {
        dv_raw <- hush(insight::find_response(mfx@model, combine = FALSE, component = "all", flatten = TRUE))
        dv_raw <- intersect(dv_raw, colnames(mfx@modeldata))
        dv_plot <- if (length(dv_raw) > 0) dv_raw[1] else NULL
    }
    if (is.null(dv_plot)) {
        points <- 0
    }

    p <- ggplot2::ggplot(data = dat)

    if (
        points > 0 &&
            !x_is_discrete &&
            !check_variable_class(mfx, dv, "categorical")
    ) {
        dat_points <- mfx@modeldata
        if (!is.null(dat_points)) {
            if (is_discrete(dat_points, v_x, mfx)) {
                dat_points[[v_x]] <- factor(dat_points[[v_x]])
            }
            if (is_discrete(dat_points, v_color, mfx)) {
                dat_points[[v_color]] <- factor(dat_points[[v_color]])
            }
        }
        if (is_discrete(dat_points, v_color, mfx)) {
            if (isTRUE(gray)) {
                p <- p +
                    ggplot2::geom_point(
                        data = dat_points,
                        alpha = points,
                        ggplot2::aes(
                            x = .data[[v_x]],
                            y = .data[[dv_plot]],
                            shape = .data[[v_color]]
                        )
                    )
            } else {
                p <- p +
                    ggplot2::geom_point(
                        data = dat_points,
                        alpha = points,
                        ggplot2::aes(
                            x = .data[[v_x]],
                            y = .data[[dv_plot]],
                            color = .data[[v_color]]
                        )
                    )
            }
        } else {
            p <- p +
                ggplot2::geom_point(
                    data = dat_points,
                    alpha = points,
                    ggplot2::aes(x = .data[[v_x]], y = .data[[dv_plot]])
                )
        }
    }

    if (isTRUE(rug)) {
        p <- p + ggplot2::geom_rug(data = mfx@modeldata, ggplot2::aes(x = .data[[v_x]]))
        if (!is.null(dv_plot)) {
            p <- p +
                ggplot2::geom_rug(data = mfx@modeldata, ggplot2::aes(y = .data[[dv_plot]]))
        }
    }

    aes_args <- list(
        x = substitute(.data[[v_x]]),
        y = substitute(estimate)
    )

    if ("conf.low" %in% colnames(dat)) {
        aes_args$ymin <- substitute(conf.low)
        aes_args$ymax <- substitute(conf.high)
    }

    aes_args_ribbon <- aes_args
    aes_args_ribbon$fill <- aes_args$color
    aes_args_ribbon$color <- NULL

    # discrete x-axis
    if (x_is_discrete) {
        if (!is.null(v_color)) {
            if (gray) {
                aes_args$shape <- substitute(.data[[v_color]])
            } else {
                aes_args$color <- substitute(.data[[v_color]])
            }
        }
        aes_obj <- do.call(ggplot2::aes, aes_args)
        if ("conf.low" %in% colnames(dat)) {
            p <- p +
                ggplot2::geom_pointrange(
                    mapping = aes_obj,
                    position = ggplot2::position_dodge(0.15)
                )
        } else {
            p <- p +
                ggplot2::geom_point(
                    mapping = aes_obj,
                    position = ggplot2::position_dodge(0.15)
                )
        }

        # continuous x-axis
    } else {
        if (!is.null(v_color)) {
            if (gray) {
                aes_args$linetype <- substitute(.data[[v_color]])
                aes_args_ribbon$linetype <- substitute(.data[[v_color]])
            } else {
                aes_args$color <- substitute(.data[[v_color]])
                aes_args_ribbon$fill <- substitute(.data[[v_color]])
            }
        }
        aes_obj_ribbon <- do.call(ggplot2::aes, aes_args_ribbon)
        aes_args$ymin <- aes_args$ymax <- NULL
        aes_obj <- do.call(ggplot2::aes, aes_args)
        if ("conf.low" %in% colnames(dat)) {
            p <- p + ggplot2::geom_ribbon(aes_obj_ribbon, alpha = 0.1)
        }
        p <- p + ggplot2::geom_line(aes_obj)
    }

    # facets: 3rd and 4th variable and/or multiple effects
    ## If pass two facets then make facet grid
    if (!is.null(v_facet_1) && !is.null(v_facet_2)) {
        fo <- stats::as.formula(paste(
            v_facet_2,
            "~",
            ifelse(multi_variables, "marginaleffects_term_index +", ""),
            v_facet_1
        ))
        p <- p +
            ggplot2::facet_grid(fo, scales = "free", labeller = function(x) {
                lapply(
                    ggplot2::label_both(x),
                    gsub,
                    pattern = "marginaleffects_term_index: ",
                    replacement = ""
                )
            })
        ## if pass only 1 facet then facet_wrap
    } else if (!is.null(v_facet_1) && is.null(v_facet_2)) {
        fo <- stats::as.formula(paste(
            "~",
            ifelse(multi_variables, "marginaleffects_term_index +", ""),
            v_facet_1
        ))
        p <- p + ggplot2::facet_wrap(fo, scales = "free")
    } else if (multi_variables) {
        fo <- stats::as.formula("~ marginaleffects_term_index")
        p <- p + ggplot2::facet_wrap(fo, scales = "free")
    }

    return(p)
}