mean_se_cluster.R

mean_se_cluster <- function (x, mult = 1, cluster = NULL) 
{
    x_na <- is.na(x)
    x <- x[!x_na]
    cluster <- cluster[!x_na]
    stopifnot(!is.null(cluster))
    mod <- lme4::lmer(x ~ 1 + (1 | cluster))
    intercept <- broom.mixed::tidy(mod, effects = "fixed")
    se <- mult * intercept$std.error
    mean <- intercept$estimate
    data.frame(list(y = mean, ymin = mean - se, ymax = mean + 
                            se), n = 1)
}

ggplot(mtcars, aes(factor(cyl), mpg)) +
    geom_pointrange(fun.data = 'mean_se', stat = 'summary')

# doesn't work
ggplot(mtcars, aes(factor(cyl), mpg)) +
    geom_pointrange(fun.data = 'mean_se_cluster', stat = 'summary', fun.args = list(cluster = cyl))

# works
ggplot(mtcars, aes(1, mpg)) +
    geom_pointrange(fun.data = 'mean_se_cluster', stat = 'summary', fun.args = list(cluster = mtcars$cyl))



mean_se_cluster_lme4 <- function (df, mult = 1)
{
  stopifnot(!is.null(df$cluster))
  y_na <- is.na(df$y)
  df <- df[!y_na, ]
  mod <- lme4::lmer(y ~ 1 + (1 | cluster), data = df)
  intercept <- broom.mixed::tidy(mod, effects = "fixed")
  se <- mult * intercept$std.error
  mean <- intercept$estimate
  data.frame(list(y = mean, ymin = mean - se, ymax = mean +
                    se), n = 1)
}


StatClusterSummary <- ggproto("StatClusterSummary", StatSummary,
                    compute_panel = function(data, scales, fun.data = NULL, fun = NULL,
                                             fun.max = NULL, fun.min = NULL, fun.args = list(),
                                             na.rm = FALSE, flipped_aes = FALSE) {
                      data <- flip_data(data, flipped_aes)
                      force(fun.data)
                      fun <- function(df) { fun.data(df) }
                      summarised <- ggplot2:::summarise_by_x(data, fun)
                      summarised$flipped_aes <- flipped_aes
                      flip_data(summarised, flipped_aes)
                    },
                     required_aes = c("x", "y", "cluster")
)

stat_cluster_summary <- function(mapping = NULL, data = NULL, geom = "pointrange",
                                 position = "identity", na.rm = FALSE, show.legend = NA,
                                 inherit.aes = TRUE, fun.data = mean_se_cluster_lme4, ...) {
  layer(
    stat = StatClusterSummary, data = data, mapping = mapping, geom = geom,
    position = position, show.legend = show.legend, inherit.aes = inherit.aes,
    params = list(fun.data = fun.data, na.rm = na.rm, ...)
  )
}

ggplot(mtcars, aes(factor(vs), mpg, cluster = gear)) +
  geom_pointrange(stat = 'cluster_summary', fun.data = mean_se_cluster_lme4)
ggplot(mtcars, aes(factor(vs), mpg, cluster = gear)) +
  geom_pointrange(stat = 'summary')