## ----setup, include = FALSE-------------------------------------------------------------
knitr::opts_chunk$set(collapse = TRUE, comment = "#>", message = FALSE)
options(digits = 4, width = 90)

## ----warning = FALSE--------------------------------------------------------------------
library(estimatr)
library(randomizr)
library(dplyr)

## ---------------------------------------------------------------------------------------
set.seed(343)
N <- 100

dat <- tibble(
  X = runif(N),
  cluster = rep(letters[1:20], each = 5),
  block = cut(X, breaks = quantile(X, seq(0, 1, 0.1)), labels = FALSE, include.lowest = TRUE),
  W = runif(N, 0.5, 1.5), # sampling weights, used further down
  Y_Z_0 = rnorm(N, mean = X) + rep(rnorm(20, sd = 0.5), each = 5), # control potential outcome, with a shock shared within each cluster
  Y_Z_1 = Y_Z_0 + 0.35 # treatment potential outcome (constant effects)
)

print(dat, n = 5)

## ---------------------------------------------------------------------------------------
dat <- dat |>
  mutate(
    Z = complete_ra(N, m = N / 2),
    Y = if_else(Z == 1, Y_Z_1, Y_Z_0)
  )

fit <- lm_robust(Y ~ Z + X, data = dat)
fit

## ---------------------------------------------------------------------------------------
tidy(fit)

## ---------------------------------------------------------------------------------------
dat <- dat |>
  mutate(
    Z_cl = cluster_ra(clusters = cluster, m = 10),
    Y_cl = if_else(Z_cl == 1, Y_Z_1, Y_Z_0)
  )

## ---------------------------------------------------------------------------------------
# Ignoring the clustered assignment understates the standard error
fit_naive <-
  lm_robust(Y_cl ~ Z_cl + X, data = dat)
tidy(fit_naive)

# Accounting for it
fit_cluster_aware <-
  lm_robust(Y_cl ~ Z_cl + X, data = dat, clusters = cluster)
tidy(fit_cluster_aware)

## ---------------------------------------------------------------------------------------
fit_HC2 <- lm_robust(Y ~ Z + X, data = dat, se_type = "HC2")

fit_classical <- lm_robust(Y ~ Z + X, data = dat, se_type = "classical")

fit_stata <- lm_robust(Y ~ Z + X, data = dat, se_type = "stata")

# The standard error on Z under each
c(
  HC2 = fit_HC2$std.error[["Z"]],
  classical = fit_classical$std.error[["Z"]],
  stata = fit_stata$std.error[["Z"]]
)

## ---------------------------------------------------------------------------------------
dat <- dat |>
  mutate(
    Z_bl = block_ra(blocks = block, prob = 0.5),
    Y_bl = if_else(Z_bl == 1, Y_Z_1, Y_Z_0)
  )

## ---------------------------------------------------------------------------------------
lm_robust(Y_bl ~ Z_bl, data = dat, fixed_effects = ~ block)

# The same numbers, using the slower dummy matrix way.
lm_robust(Y_bl ~ Z_bl + factor(block), data = dat)

## ---------------------------------------------------------------------------------------
weighted_fit <- lm_robust(Y ~ Z + X, data = dat, weights = W)
tidy(weighted_fit)

## ---------------------------------------------------------------------------------------
lin_fit <- lm_lin(Y ~ Z, covariates = ~ X, data = dat)
tidy(lin_fit)

## ---------------------------------------------------------------------------------------
lin_fit$scaled_center

## ---------------------------------------------------------------------------------------
dat <- dat |>
  mutate(
    Y_D_0 = Y_Z_0, # potential outcome if treatment is not received
    Y_D_1 = Y_D_0 + 0.35, # potential outcome if treatment is received
    complier = rbinom(N, size = 1, prob = 0.6),
    D = Z * complier,
    Y_nc = if_else(D == 1, Y_D_1, Y_D_0)
  )

## ---------------------------------------------------------------------------------------
iv_fit <- iv_robust(Y_nc ~ D + X | Z + X, data = dat)
summary(iv_fit)

## ---------------------------------------------------------------------------------------
dim_fit <- difference_in_means(Y ~ Z, data = dat)
dim_fit
dim_fit$design

## ---------------------------------------------------------------------------------------
dim_cl <- difference_in_means(Y_cl ~ Z_cl, data = dat, clusters = cluster)
dim_cl$design
tidy(dim_cl)

## ---------------------------------------------------------------------------------------
# The blocked version of the experiment
dim_bl <- difference_in_means(Y_bl ~ Z_bl, data = dat, blocks = block)
dim_bl$design
tidy(dim_bl)

# R's sleep data, ten patients each measured under two drugs, is a
# matched-pairs design, and is recognised as such
difference_in_means(extra ~ group, data = sleep, blocks = ID)$design

## ---------------------------------------------------------------------------------------
horvitz_thompson(Y ~ Z, data = dat, condition_prs = c("0" = 0.5, "1" = 0.5))

## ---------------------------------------------------------------------------------------
# declare_ra is from the randomizr package
decl <- declare_ra(N = N, prob = 0.5, simple = FALSE)
horvitz_thompson(Y ~ Z, data = dat, condition_prs = decl)

## ---------------------------------------------------------------------------------------
lh_fit <- lh_robust(Y ~ Z + X, data = dat, linear_hypothesis = "Z + 2*X = 0")
tidy(lh_fit)

## ---------------------------------------------------------------------------------------
joint <- lh_robust(Y ~ Z + X, data = dat, linear_hypothesis = c("Z = 0", "X = 0"))
joint$joint_hypothesis

## ---------------------------------------------------------------------------------------
glance(fit)
confint(fit)

## ----results = "asis"-------------------------------------------------------------------
texreg::htmlreg(list(fit, lin_fit), include.ci = FALSE, caption = "")

## ----eval = requireNamespace("emmeans", quietly = TRUE)---------------------------------
emmeans::emmeans(fit, "Z", at = list(Z = 0:1))

## ----eval = requireNamespace("DeclareDesign", quietly = TRUE) && requireNamespace("future.apply", quietly = TRUE)----
library(DeclareDesign)

declaration_lin <-
  declare_model(N = 100, 
                X = runif(N), 
                U = rnorm(N, mean = X), 
                potential_outcomes(Y ~ 0.35 * Z + U)) +
  declare_inquiry(ATE = mean(Y_Z_1 - Y_Z_0)) +
  declare_assignment(Z = complete_ra(N)) +
  declare_measurement(Y = reveal_outcomes(Y ~ Z)) +
  declare_estimator(Y ~ Z, covariates = ~ X, 
                    .method = lm_lin, inquiry = "ATE")

set.seed(343)
diagnose_design(declaration_lin, sims = 500)