## ----echo=FALSE---------------------------------------------------------------
  set.seed(17760701)
  knitr::opts_chunk$set(
    collapse = TRUE,
    comment = "#>",
    message = FALSE,
    warning = FALSE,
    fig.width = 7,
    fig.height = 3.5
  )
  options(digits=2)

## -----------------------------------------------------------------------------
library(randomizr)
library(dplyr)
library(tidyr)
library(purrr)
library(ggplot2)
library(estimatr)

## -----------------------------------------------------------------------------
data(HairEyeColor)

HairEyeColor |>
  as.data.frame() |>
  as_tibble()

## -----------------------------------------------------------------------------
# uncount() repeats each row Freq times, which is exactly what we want
hec <-
  HairEyeColor |>
  as.data.frame() |>
  as_tibble() |>
  uncount(Freq) |>
  select(Hair, Eye, Sex)

N <- nrow(hec)

hec

## -----------------------------------------------------------------------------
# Set a seed for reproducibility
set.seed(343)

# Create untreated and treated outcomes for all subjects
hec <-
  hec |>
  mutate(
    Y0 = rnorm(n = N,
               mean = 2 * as.numeric(Hair) - 4 * as.numeric(Eye) - 6 * as.numeric(Sex),
               sd = 5),
    Y1 = Y0 + 6 * as.numeric(Hair) + 4 * as.numeric(Eye) + 2 * as.numeric(Sex)
  )

# Calculate true ATE
hec |> summarize(ATE = mean(Y1 - Y0))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- simple_ra(N = N)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- simple_ra(N = N, prob = 0.30)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- simple_ra(N = N, num_arms = 3)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- simple_ra(N = N, prob_each = c(0.2, 0.2, 0.6))

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- simple_ra(N = N,
               prob_each = c(0.2, 0.2, 0.6),
               conditions = c("control", "placebo", "treatment"))

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- complete_ra(N = N)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- complete_ra(N = N, m = 200)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- complete_ra(N = N, num_arms = 3)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- complete_ra(N = N, m_each = c(100, 200, 292))

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- complete_ra(N = N,
                 m_each = c(100, 200, 292),
                 conditions = c("control", "placebo", "treatment"))

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## -----------------------------------------------------------------------------
sims <- 1000

simulate_once <- function(i) {
  hec <-
    hec |>
    mutate(
      # Conduct both kinds of random assignment
      Z_simple = simple_ra(N = N),
      Z_complete = complete_ra(N = N),
      # Reveal observed potential outcomes
      Y_simple = if_else(Z_simple == 1, Y1, Y0),
      Y_complete = if_else(Z_complete == 1, Y1, Y0)
    )

  fit_simple <- difference_in_means(Y_simple ~ Z_simple, data = hec)
  fit_complete <- difference_in_means(Y_complete ~ Z_complete, data = hec)

  bind_rows(
    tidy(fit_simple) |> filter(term == "Z_simple") |> mutate(design = "Simple"),
    tidy(fit_complete) |> filter(term == "Z_complete") |> mutate(design = "Complete")
  )
}

estimates <- map(1:sims, simulate_once) |> list_rbind()

## -----------------------------------------------------------------------------
estimates |>
  group_by(design) |>
  summarize(empirical_se = sd(estimate))

## -----------------------------------------------------------------------------
gg_df <-
  estimates |>
  mutate(design = factor(design, levels = c("Simple", "Complete")))

ggplot(gg_df, aes(x = estimate)) +
  geom_histogram(bins = 40) +
  geom_vline(xintercept = mean(hec$Y1 - hec$Y0), linetype = "dashed") +
  facet_wrap(~design) +
  labs(x = "ATE estimate", y = "Count",
       title = "Sampling distributions under simple and complete random assignment",
       subtitle = "Dashed line is the true ATE") +
  theme_bw() +
  theme(legend.position = "none")

## ----echo=TRUE, results="hide"------------------------------------------------
# how many of 10 subjects end up treated, over many draws
m_simple <- replicate(sims, sum(simple_ra(N = 10)))
m_complete <- replicate(sims, sum(complete_ra(N = 10)))

table(m_simple)

table(m_complete)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(m_simple))))

knitr::kable(t(as.matrix(table(m_complete))))

## -----------------------------------------------------------------------------
set.seed(20260824)

sims_small <- 10000

hec_small <- hec |> slice(1:30)

simulate_once_small <- function(i) {
  hec_small <-
    hec_small |>
    mutate(
      # Conduct both kinds of random assignment
      Z_simple = simple_ra(N = 30),
      Z_complete = complete_ra(N = 30),
      # Reveal observed potential outcomes
      Y_simple = if_else(Z_simple == 1, Y1, Y0),
      Y_complete = if_else(Z_complete == 1, Y1, Y0)
    )

  fit_simple <- difference_in_means(Y_simple ~ Z_simple, data = hec_small)
  fit_complete <- difference_in_means(Y_complete ~ Z_complete, data = hec_small)

  bind_rows(
    tidy(fit_simple) |> filter(term == "Z_simple") |> mutate(design = "Simple"),
    tidy(fit_complete) |> filter(term == "Z_complete") |> mutate(design = "Complete")
  )
}

estimates_small <- map(1:sims_small, simulate_once_small) |> list_rbind()

estimates_small |>
  group_by(design) |>
  summarize(empirical_se = sd(estimate))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- block_ra(blocks = hec$Hair)

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- block_ra(blocks = hec$Hair, num_arms = 3)

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- block_ra(blocks = hec$Hair,
              conditions = c("Control", "Placebo", "Treatment"))

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- block_ra(blocks = hec$Hair, prob_each = c(0.3, 0.7))

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE----------------------------------------------------------------
sort(unique(hec$Hair))

block_m_each <- rbind(c(78, 30),
                      c(186, 100),
                      c(51, 20),
                      c(87, 40))

block_m_each

## ----echo=TRUE, results="hide"------------------------------------------------
Z <- block_ra(blocks = hec$Hair, block_m_each = block_m_each)

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE, results="hide"------------------------------------------------
declaration <-
  declare_ra(blocks = Hair, block_m_each = block_m_each, data = hec)

# show the probability that each unit is assigned to each condition
head(declaration$probabilities_matrix)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(declaration$probabilities_matrix))

## ----echo=TRUE, results="hide"------------------------------------------------
# Show that the probability of treatment is different within block
table(hec$Hair, round(declaration$probabilities_matrix[, 2], 3))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(hec$Hair, round(declaration$probabilities_matrix[, 2], 3)))

## ----echo=TRUE, results="hide"------------------------------------------------
hec <-
  hec |>
  mutate(
    Z_blocked = block_ra(blocks = Hair, block_m_each = block_m_each),
    Y_blocked = if_else(Z_blocked == 1, Y1, Y0),
    cond_prob = obtain_condition_probabilities(declaration, Z_blocked),
    IPW_weights = 1 / cond_prob
  )

fit_LSDV <- lm_robust(Y_blocked ~ Z_blocked + Hair, data = hec)
fit_IPW <- lm_robust(Y_blocked ~ Z_blocked, weights = IPW_weights, data = hec)

tidy(fit_LSDV)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(tidy(fit_LSDV))

## ----echo=TRUE, results="hide"------------------------------------------------
tidy(fit_IPW)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(tidy(fit_IPW))

## ----echo=TRUE, results="hide"------------------------------------------------
block_id <- paste(hec$Hair, hec$Eye, hec$Sex, sep = "_")

Z <- block_ra(blocks = block_id)

head(table(block_id, Z))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(block_id, Z)))

## ----eval=FALSE, results="hide"-----------------------------------------------
# library(blockTools)
# 
# # blockTools requires that all variables be numeric
# numeric_mat <- model.matrix(~ Hair + Eye + Sex, data = hec)[, -1]
# 
# # blockTools also requires an id variable
# df_forBT <- data.frame(id_var = 1:nrow(numeric_mat), numeric_mat)
# 
# # Conducting the actual blocking: let's make trios
# out <- blockTools::block(df_forBT,
#                          n.tr = 3,
#                          id.vars = "id_var",
#                          block.vars = colnames(df_forBT)[-1])
# 
# # Extract the block_ids
# hec <- hec |>
#   mutate(block_id = blockTools::createBlockIDs(out, df_forBT, id.var = "id_var"))
# 
# # Conduct actual random assignment with randomizr
# Z_blocked <- block_ra(blocks = hec$block_id, num_arms = 3)
# 
# head(table(hec$block_id, Z_blocked))

## ----echo=TRUE, results="hide"------------------------------------------------
blocks <- rep(1:2, each = 3)
Z <- balanced_ra(blocks = blocks)
table(blocks, Z)

## ----echo=FALSE---------------------------------------------------------------
blocks <- rep(1:2, each = 3)
Z <- balanced_ra(blocks = blocks)
knitr::kable(table(blocks, Z))

## ----echo=TRUE, results="hide"------------------------------------------------
hec <-
  hec |>
  mutate(cluster_id = paste(Hair, Eye, Sex, sep = "_"))

Z_clust <- cluster_ra(clusters = hec$cluster_id)

head(table(hec$cluster_id, Z_clust))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(hec$cluster_id, Z_clust)))

## ----echo=TRUE, results="hide"------------------------------------------------
Z_clust <- cluster_ra(clusters = hec$cluster_id, num_arms = 3)

head(table(hec$cluster_id, Z_clust))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(hec$cluster_id, Z_clust)))

## ----echo=TRUE, results="hide"------------------------------------------------
Z_clust <- cluster_ra(clusters = hec$cluster_id,
                      conditions = c("Control", "Placebo", "Treatment"))

head(table(hec$cluster_id, Z_clust))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(hec$cluster_id, Z_clust)))

## ----echo=TRUE, results="hide"------------------------------------------------
Z_clust <- cluster_ra(clusters = hec$cluster_id, m_each = c(5, 15, 12))

head(table(hec$cluster_id, Z_clust))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(hec$cluster_id, Z_clust)))

## ----echo=TRUE, results="hide"------------------------------------------------
cluster_level_df <-
  hec |>
  group_by(cluster_id) |>
  summarize(cluster_size = n()) |>
  arrange(cluster_size) |>
  mutate(block_id = paste0("block_", sprintf("%02d", rep(1:16, each = 2))))

hec <- left_join(hec, cluster_level_df, by = "cluster_id")

Z <- block_and_cluster_ra(clusters = hec$cluster_id, blocks = hec$block_id)

head(table(hec$cluster_id, Z))

head(table(hec$block_id, Z))

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(table(hec$cluster_id, Z)))

knitr::kable(head(table(hec$block_id, Z)))

## ----echo=TRUE, results="hide"------------------------------------------------
block_m_each <- rbind(c(78, 30),
                      c(186, 100),
                      c(51, 20),
                      c(87, 40))

Z <- block_ra(blocks = hec$Hair, block_m_each = block_m_each)

table(Z, hec$Hair)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(Z, hec$Hair))

## ----echo=TRUE, results="hide"------------------------------------------------
declaration <-
  declare_ra(blocks = Hair, block_m_each = block_m_each, data = hec)

prob_mat <- declaration$probabilities_matrix

head(prob_mat)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(head(prob_mat))

## ----echo=TRUE, results="hide"------------------------------------------------
cond_prob <- obtain_condition_probabilities(declaration, Z)

table(round(cond_prob, 2), Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(table(round(cond_prob, 2), Z))

## ----echo=TRUE, results="hide"------------------------------------------------
# 400 families have 1 child in the lottery, 100 families have 2
family_id <- c(sprintf("%03d", 1:500), sprintf("%03d", 1:100))

school_ra <- function(m) {
  N <- length(family_id)
  random_number <- sample(1:N, replace = FALSE)
  Z <- rep(0, N)
  i <- 1
  while (sum(Z) < m) {
    Z[family_id == family_id[random_number[i]]] <- 1
    i <- i + 1
  }
  return(Z)
}

Z <- school_ra(200)

table(Z)

## ----echo=FALSE---------------------------------------------------------------
knitr::kable(t(as.matrix(table(Z))))

## -----------------------------------------------------------------------------
Z_matrix <- replicate(1000, school_ra(200))

gg_df <- tibble(student = seq_len(nrow(Z_matrix)),
                prob = rowMeans(Z_matrix))

ggplot(gg_df, aes(x = student, y = prob)) +
  geom_point(size = 0.8) +
  labs(x = "Student", y = "Estimated probability of assignment") +
  theme_bw()

## ----eval=FALSE---------------------------------------------------------------
# hec <-
#   hec |>
#   mutate(
#     Z_complete = complete_ra(N = N,
#                              m_each = c(100, 200, 292),
#                              conditions = c("control", "placebo", "treatment")),
#     id_var = row_number()
#   )
# 
# hec |>
#   select(id_var, Z_complete) |>
#   readr::write_csv("MyRandomAssignment.csv")