## ----knitr-opts, include = FALSE----------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

## ----setup--------------------------------------------------------------------
library(deli)

## ----simulate-data------------------------------------------------------------
set.seed(42)
n <- 1000
W1 <- rnorm(n)
W2 <- rbinom(n, 1, 0.4)
A <- rbinom(n, 1, plogis(-0.5 + 0.5 * W1 + 0.3 * W2))
# True ATE = 1.5
Y <- 2 + 1.5 * A + W1 - 0.5 * W2 + rnorm(n)

## ----gformula-----------------------------------------------------------------
X <- cbind(1, A, W1, W2)      # Observed design matrix
X1 <- cbind(1, 1, W1, W2)     # Counterfactual: all treated
X0 <- cbind(1, 0, W1, W2)     # Counterfactual: all untreated

psi_gformula <- function(theta) {
  ee_gformula(theta, y = Y, X = X, X1 = X1, X0 = X0)
}

# theta: ACE, E[Y(1)], E[Y(0)], beta0, beta1, beta2, beta3
m_gformula <- m_estimate(
  stacked_equations = psi_gformula,
  init = c(0, 0, 0, 0, 0, 0, 0)
)

# ACE estimate (true = 1.5)
m_gformula@theta[1]

# With confidence interval
summary(m_gformula)

## ----ipw----------------------------------------------------------------------
W_ps <- cbind(1, W1, W2)  # Propensity score design matrix

psi_ipw <- function(theta) {
  ee_ipw(theta, y = Y, A = A, W = W_ps)
}

# theta: ACE, E[Y(1)], E[Y(0)], alpha0, alpha1, alpha2
m_ipw <- m_estimate(stacked_equations = psi_ipw, init = c(0, 0, 0, 0, 0, 0))

# ACE
m_ipw@theta[1]
summary(m_ipw)

## ----ps-range-----------------------------------------------------------------
# Propensity scores implied by the alphas in the fitted stack
ps <- plogis(drop(W_ps %*% m_ipw@theta[4:6]))
range(ps)

# How many scores a tight c(0.3, 0.7) range would clip, by tail
c(lower = sum(ps < 0.3), upper = sum(ps > 0.7))

## ----ipw-truncate-------------------------------------------------------------
psi_ipw_trunc <- function(theta) {
  ee_ipw(theta, y = Y, A = A, W = W_ps,
         truncate = c(0.3, 0.7))
}

m_ipw_trunc <- m_estimate(
  stacked_equations = psi_ipw_trunc,
  init = c(0, 0, 0, 0, 0, 0)
)
summary(m_ipw_trunc)

## ----compare-truncation-------------------------------------------------------
data.frame(
  fit = c("Untruncated", "Truncated"),
  estimate = round(c(m_ipw@theta[[1]], m_ipw_trunc@theta[[1]]), 4),
  std_err = round(
    sqrt(c(vcov(m_ipw)[1, 1], vcov(m_ipw_trunc)[1, 1])),
    4
  )
)

## ----aipw---------------------------------------------------------------------
psi_aipw <- function(theta) {
  ee_aipw(theta, y = Y, A = A,
          W = W_ps,             # Propensity score model
          X = X, X1 = X1, X0 = X0)  # Outcome model
}

# theta: ACE, E[Y(1)], E[Y(0)], alpha (3), beta (4)
m_aipw <- m_estimate(
  stacked_equations = psi_aipw,
  init = c(0, 0, 0, rep(0, 3), rep(0, 4))
)

summary(m_aipw)

## ----compare-methods----------------------------------------------------------
data.frame(
  method = c("G-formula", "IPW", "AIPW"),
  estimate = round(
    c(m_gformula@theta[[1]], m_ipw@theta[[1]], m_aipw@theta[[1]]),
    3
  ),
  true_ate = 1.5
)