## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  message = FALSE,
  warning = FALSE,
  dpi = 96,
  fig.width = 7,
  fig.height = 3.6,
  out.width = "100%"
)

## ----setup--------------------------------------------------------------------
library(dplyr)
library(estimatr)
library(ggplot2)
library(marginaleffects)
library(patchwork)
library(tidyr)
library(vayr)

theme_set(theme_bw())
theme_update(
  strip.background = element_blank(),
  legend.position = "none",
  plot.title = element_text(size = 10, face = "bold")
)

## ----two_arm, fig.alt = "a two-arm trial shown with data and model, and as a bar chart"----
summary_df <-
  two_arm_trial |>
  group_by(condition) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

good <-
  ggplot(two_arm_trial, aes(condition, Y)) +
  geom_point(position = position_sunflower(density = 60, aspect_ratio = 1.55),
             alpha = 0.2, stroke = 0) +
  geom_point(data = summary_df, size = 3) +
  geom_errorbar(data = summary_df, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  scale_y_continuous(breaks = seq(0, 1, 0.25)) +
  coord_cartesian(ylim = c(-0.15, 1.15)) +
  labs(title = "Shows the model and the data", x = NULL, y = "Outcome [1 = yes]")

bad <-
  ggplot(summary_df, aes(condition, Y)) +
  geom_col() +
  scale_y_continuous(breaks = seq(0, 1, 0.25)) +
  coord_cartesian(ylim = c(-0.15, 1.15)) +
  labs(title = "Shows the model only", x = NULL, y = "Outcome [1 = yes]")

good + bad

## ----blocked_weights, fig.alt = "a blocked experiment with and without inverse probability weights"----
weighted_df <-
  blocked_experiment |>
  group_by(condition) |>
  reframe(tidy(lm_robust(Y ~ 1, weights = 1 / Z_cond_prob))) |>
  mutate(Y = estimate)

unweighted_df <-
  blocked_experiment |>
  group_by(condition) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

good <-
  ggplot(blocked_experiment, aes(condition, Y)) +
  geom_point(aes(size = 1 / Z_cond_prob),
             position = position_circlepack(density = 0.03, aspect_ratio = 1),
             alpha = 0.2, stroke = 0) +
  geom_point(data = weighted_df, size = 4) +
  geom_errorbar(data = weighted_df, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  scale_size_continuous(range = c(1, 4)) +
  labs(title = "Weights by the probability of assignment", x = NULL, y = "Count outcome")

bad <-
  ggplot(blocked_experiment, aes(condition, Y)) +
  geom_point(position = position_sunflower(density = 1.5, aspect_ratio = 1),
             alpha = 0.2, stroke = 0) +
  geom_point(data = unweighted_df, size = 4) +
  geom_errorbar(data = unweighted_df, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  labs(title = "Ignores the probability of assignment", x = NULL, y = "Count outcome")

good + bad

## ----blocked_estimates--------------------------------------------------------
bind_rows(
  weighted = tidy(lm_robust(Y ~ Z, weights = 1 / Z_cond_prob, data = blocked_experiment)),
  unweighted = tidy(lm_robust(Y ~ Z, data = blocked_experiment)),
  .id = "estimator"
) |>
  filter(term == "Z") |>
  select(estimator, estimate, std.error, conf.low, conf.high)

## ----blocked_facets, fig.alt = "the same experiment faceted by block and faceted by treatment"----
blocked_labelled <-
  blocked_experiment |>
  mutate(neighborhood_long = paste("Neighborhood", neighborhood),
         neighborhood_short = paste0("N", neighborhood))

by_block <-
  blocked_labelled |>
  group_by(condition, neighborhood_long, neighborhood_short) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

good <-
  ggplot(blocked_labelled, aes(condition, Y)) +
  geom_point(position = position_sunflower(density = 1.5, aspect_ratio = 1),
             alpha = 0.2, stroke = 0) +
  geom_point(data = by_block, size = 3) +
  geom_errorbar(data = by_block, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  facet_wrap(~ neighborhood_long) +
  labs(title = "Compares across randomly formed groups", x = NULL, y = "Count outcome")

bad <-
  ggplot(blocked_labelled, aes(neighborhood_short, Y)) +
  geom_point(position = position_sunflower(density = 1.5, aspect_ratio = 1),
             alpha = 0.2, stroke = 0) +
  geom_point(data = by_block, size = 3) +
  geom_errorbar(data = by_block, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  facet_wrap(~ condition) +
  labs(title = "Compares across neighborhoods", x = NULL, y = "Count outcome")

good + bad

## ----clustered, fig.alt = "a cluster-randomized experiment shown at the cluster level and at the student level"----
class_level <-
  clustered_experiment |>
  group_by(class, condition, n_per_class) |>
  summarise(Y = mean(Y), .groups = "drop")

clustered_df <-
  clustered_experiment |>
  group_by(condition) |>
  reframe(tidy(lm_robust(Y ~ 1, clusters = class))) |>
  mutate(Y = estimate)

naive_df <-
  clustered_experiment |>
  group_by(condition) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

good <-
  ggplot(class_level, aes(condition, Y)) +
  geom_point(aes(size = n_per_class),
             position = position_jitter_ellipse(width = 0.2, height = 0, seed = 1),
             alpha = 0.5, stroke = 0) +
  geom_point(data = clustered_df, size = 3) +
  geom_errorbar(data = clustered_df, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  scale_size_continuous(range = c(1, 4)) +
  coord_cartesian(ylim = c(400, 1600)) +
  labs(title = "Clusters are the unit, and the size", x = NULL, y = "Class average score")

bad <-
  ggplot(clustered_experiment, aes(condition, Y)) +
  geom_point(position = position_jitter_ellipse(width = 0.25, height = 20, seed = 1),
             alpha = 0.2, stroke = 0) +
  geom_point(data = naive_df, size = 3) +
  geom_errorbar(data = naive_df, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  coord_cartesian(ylim = c(400, 1600)) +
  labs(title = "Students are the unit", x = NULL, y = "Student score")

good + bad

## ----clustered_estimates------------------------------------------------------
bind_rows(
  clustered = tidy(lm_robust(Y ~ condition, clusters = class, data = clustered_experiment)),
  naive = tidy(lm_robust(Y ~ condition, data = clustered_experiment)),
  .id = "estimator"
) |>
  filter(term == "conditionTreatment") |>
  select(estimator, estimate, std.error, conf.low, conf.high)

## ----covariate, fig.alt = "the same experiment before and after covariate adjustment"----
centred <- covariate_adjustment |> mutate(X_c = X - mean(X))

gg_df <-
  centred |>
  transmute(
    ID,
    Y_Adjusted = residuals(lm(Y ~ X_c + X_c:Z, data = centred)),
    Z_Adjusted = residuals(lm(Z ~ X_c, data = centred)),
    Y_Unadjusted = Y,
    Z_Unadjusted = Z
  ) |>
  pivot_longer(
    -ID,
    names_to = c("variable", "estimation"),
    names_sep = "_"
  ) |>
  pivot_wider(names_from = variable, values_from = value) |>
  mutate(estimation = factor(estimation, levels = c("Unadjusted", "Adjusted")))

ggplot(gg_df, aes(Z, Y)) +
  geom_point(alpha = 0.4, stroke = 0) +
  stat_smooth(method = "lm_robust", colour = "grey40") +
  facet_wrap(~ estimation, scales = "free") +
  labs(x = "Randomly assigned treatment", y = "Outcome")

## ----covariate_estimates------------------------------------------------------
bind_rows(
  lin = tidy(lm_lin(Y ~ Z, covariates = ~ X, data = covariate_adjustment)),
  residualized = tidy(lm_robust(Y ~ Z, data = filter(gg_df, estimation == "Adjusted"))),
  unadjusted = tidy(lm_robust(Y ~ Z, data = covariate_adjustment)),
  .id = "estimator"
) |>
  filter(term == "Z") |>
  select(estimator, estimate, std.error)

## ----interaction, fig.alt = "conditional effects shown in data-space and as a plot of estimates alone"----
fit <- lm_robust(Y ~ condition * X, data = continuous_interaction)

label_df <- data.frame(
  X = c(1.1, 0.4),
  Y = c(-3.2, 7.2),
  condition = c("Control", "Treatment"),
  label = c("Control", "Treated")
)

good <-
  ggplot(continuous_interaction, aes(X, Y, group = condition, shape = condition)) +
  geom_point(alpha = 0.2, stroke = 0) +
  stat_smooth(method = "lm_robust", fullrange = TRUE, colour = "black") +
  geom_label(data = label_df, aes(label = label), size = 3) +
  coord_cartesian(xlim = c(-2, 2), ylim = c(-5, 10)) +
  labs(title = "Shows the model in data-space",
       x = "Pretreatment covariate", y = "Outcome")

# The conditional effect at x is the gap between the two fitted lines, which is
# a contrast rather than a fitted value, so it is marginaleffects' job. The df
# argument makes the intervals t-based, matching what tidy(lm_robust()) reports
# everywhere else in this vignette; marginaleffects defaults to a normal
# approximation.
cate_df <- comparisons(
  fit,
  variables = "condition",
  newdata = datagrid(X = seq(-2, 2, by = 0.25)),
  df = df.residual(fit)
)

bad <-
  ggplot(cate_df, aes(X, estimate)) +
  geom_hline(yintercept = 0, linetype = "dashed") +
  geom_point() +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high), width = 0) +
  coord_cartesian(xlim = c(-2, 2), ylim = c(-5, 10)) +
  labs(title = "Shows the estimates only",
       x = "Pretreatment covariate", y = "Conditional effect")

good + bad

## ----noncompliance, fig.alt = "noncompliance shown by assignment and, wrongly, by treatment receipt"----
# The grouping column is named dv rather than outcome because estimatr's tidy()
# returns a column called outcome, which would collide and add a third facet.
long_df <-
  noncompliance_experiment |>
  pivot_longer(c(D, Y), names_to = "dv", values_to = "value") |>
  mutate(dv = factor(dv, c("D", "Y"), c("Treatment receipt", "Turnout")))

by_assignment <-
  long_df |>
  group_by(Z, dv) |>
  reframe(tidy(lm_robust(value ~ 1))) |>
  mutate(value = estimate)

good <-
  ggplot(long_df, aes(Z, value)) +
  geom_point(position = position_sunflower(density = 50, aspect_ratio = 3.3),
             alpha = 0.15, stroke = 0) +
  geom_point(data = by_assignment, size = 3) +
  geom_errorbar(data = by_assignment, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  facet_wrap(~ dv) +
  scale_y_continuous(breaks = seq(0, 1, 0.25)) +
  coord_cartesian(ylim = c(-0.15, 1.15)) +
  labs(title = "By random assignment", x = NULL, y = "Outcome [1 = yes]")

received <-
  noncompliance_experiment |>
  mutate(D = factor(D, 0:1, c("Did not receive", "Did receive")))

by_receipt <-
  received |>
  group_by(Z, D) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

bad <-
  ggplot(received, aes(Z, Y)) +
  geom_point(position = position_sunflower(density = 50, aspect_ratio = 3.3),
             alpha = 0.15, stroke = 0) +
  geom_point(data = by_receipt, size = 3) +
  geom_errorbar(data = by_receipt, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  facet_wrap(~ D) +
  scale_y_continuous(breaks = seq(0, 1, 0.25)) +
  coord_cartesian(ylim = c(-0.15, 1.15)) +
  labs(title = "By treatment receipt", x = NULL, y = "Turnout [1 = yes]")

good + bad

## ----attrition, fig.width = 6, fig.height = 4, fig.alt = "extreme value bounds under attrition, with imputed points marked"----
bounded <- impute_extreme_values(
  attrition_experiment,
  outcome = "Y",
  assignment = "Z",
  range = c(1, 7)
)

bounded <- bounded |> mutate(condition = if_else(Z == 1, "Treatment", "Control"))

bound_means <-
  bounded |>
  group_by(condition, scenario) |>
  reframe(tidy(lm_robust(Y ~ 1))) |>
  mutate(Y = estimate)

ggplot(bounded, aes(condition, Y)) +
  geom_point(aes(colour = imputed, shape = imputed),
             position = position_sunflower(density = 45, aspect_ratio = 0.34),
             alpha = 0.5, stroke = 0) +
  geom_point(data = bound_means, size = 3) +
  geom_errorbar(data = bound_means, aes(ymin = conf.low, ymax = conf.high), width = 0) +
  facet_wrap(~ scenario) +
  scale_colour_manual(values = c("#205C8A", "#C67800")) +
  scale_y_continuous(breaks = 1:7) +
  labs(x = NULL, y = "Outcome [1: strongly disagree, 7: strongly agree]") +
  theme(legend.position = "bottom", legend.title = element_blank())

## ----attrition_bounds---------------------------------------------------------
bounded |>
  group_by(scenario) |>
  reframe(tidy(lm_robust(Y ~ Z))) |>
  filter(term == "Z") |>
  select(scenario, estimate, std.error, conf.low, conf.high)