## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.align = "center",
  fig.width = 7,
  fig.height = 4.5
)

## ----load-package-------------------------------------------------------------
library(sreg)

## ----generate-large-----------------------------------------------------------
set.seed(101)

large_data <- sreg.rgen(
  n = 600,
  tau.vec = c(0.3, 0.7),
  n.strata = 4,
  cluster = FALSE
)

head(large_data)
table(large_data$S, large_data$D)

## ----fit-large-unadjusted-----------------------------------------------------
fit_large_unadjusted <- sreg(
  Y = large_data$Y,
  S = large_data$S,
  D = large_data$D
)

fit_large_unadjusted$tau.hat
fit_large_unadjusted$se.rob

## ----fit-large-adjusted-------------------------------------------------------
X_large <- large_data[c("x_1", "x_2")]

fit_large <- sreg(
  Y = large_data$Y,
  S = large_data$S,
  D = large_data$D,
  X = X_large
)

print(fit_large)

## ----plot-large, fig.alt="Estimated treatment effects and 95 percent confidence intervals"----
plot(
  fit_large,
  treatment_labels = c("Program A", "Program B"),
  title = "Estimated treatment effects",
  x_axis_title = "ATE relative to control",
  bar_fill = c("#3B82F6", "#14B8A6"),
  point_fill = "white",
  zero_line = TRUE
)

## ----store-plot, eval=FALSE---------------------------------------------------
# p <- plot(fit_large, treatment_labels = c("Program A", "Program B"))
# p + ggplot2::theme_minimal()

## ----generate-small-----------------------------------------------------------
set.seed(102)

small_data <- sreg.rgen(
  n = 300,
  tau.vec = c(0.3, 0.7),
  cluster = FALSE,
  small.strata = TRUE,
  k = 3,
  treat.sizes = c(1, 1, 1)
)

table(table(small_data$S))

## ----fit-small----------------------------------------------------------------
fit_small <- sreg(
  Y = small_data$Y,
  S = small_data$S,
  D = small_data$D,
  X = small_data[c("x_1", "x_2")],
  small.strata = TRUE
)

print(fit_small)

## ----validate-small-k, eval=FALSE---------------------------------------------
# sreg(
#   Y = small_data$Y,
#   S = small_data$S,
#   D = small_data$D,
#   small.strata = TRUE,
#   k = 3
# )

## ----generate-mixed-----------------------------------------------------------
set.seed(103)

mixed_data <- sreg.rgen(
  n = 120,
  tau.vec = 0.5,
  cluster = FALSE,
  mixed.strata = TRUE,
  n.small = 80,
  k = 4,
  treat.sizes = c(2, 2),
  n.strata = 4
)

sort(table(mixed_data$S))

## ----fit-mixed----------------------------------------------------------------
fit_mixed <- suppressWarnings(sreg(
  Y = mixed_data$Y,
  S = mixed_data$S,
  D = mixed_data$D,
  small.strata = TRUE,
  k = 4
))

print(fit_mixed)

## ----generate-cluster---------------------------------------------------------
set.seed(104)

cluster_data <- sreg.rgen(
  n = 60,
  tau.vec = 0.5,
  n.strata = 4,
  cluster = TRUE
)

c(
  observations = nrow(cluster_data),
  clusters = length(unique(cluster_data$G.id)),
  strata = length(unique(cluster_data$S))
)

## ----fit-cluster--------------------------------------------------------------
fit_cluster <- sreg(
  Y = cluster_data$Y,
  S = cluster_data$S,
  D = cluster_data$D,
  G.id = cluster_data$G.id,
  Ng = cluster_data$Ng,
  X = cluster_data[c("x_1", "x_2")]
)

print(fit_cluster)

## ----no-stratification, eval=FALSE--------------------------------------------
# fit_unstratified <- sreg(
#   Y = Y,
#   S = NULL,
#   D = D,
#   X = X,
#   small.strata = FALSE
# )

## ----inspect-result-----------------------------------------------------------
names(fit_mixed)
fit_mixed$tau.hat
fit_mixed$se.rob
fit_mixed$mixed.design

## ----hc1, eval=FALSE----------------------------------------------------------
# fit_hc1 <- sreg(Y = Y, S = S, D = D, HC1 = TRUE)
# fit_no_hc1 <- sreg(Y = Y, S = S, D = D, HC1 = FALSE)

## ----empirical-data-----------------------------------------------------------
data("AEJapp", package = "sreg")

Y_emp <- as.numeric(unclass(AEJapp$gradesq34))
D_raw <- as.numeric(unclass(AEJapp$treatment))
D_emp <- ifelse(D_raw == 3, 0, D_raw)
S_emp <- as.numeric(unclass(AEJapp$class_level))

fit_empirical <- sreg(
  Y = Y_emp,
  S = S_emp,
  D = D_emp
)

fit_empirical$tau.hat
fit_empirical$se.rob