grantmcdermott / etwfe

Extended two-way fixed effects
https://grantmcdermott.com/etwfe/
Other
50 stars 11 forks source link

Multiple xvar levels #34

Closed grantmcdermott closed 1 year ago

grantmcdermott commented 1 year ago

Closes #29.

Simulation example as a proof of concept (also added to tests):

devtools::load_all("~/Documents/Projects/etwfe")
#> ℹ Loading etwfe
library(data.table)

# 70 indivs
# 20 time periods
# staggered treat rollout at t = 11 and t = 16
# one control group (0), followed by tree equi-sized treatment groups 1:3
# (with each treatment group separated across rollout periods)

set.seed(1234L)

ids = 70
periods = 20

dat = CJ(id = 1:ids, period = 1:periods)

dat[
  , 
  x := 0.1*period + runif(n = .N, max = 0.1)
][
  ,
  te_grp := fcase(
    id <= 10, 0,
    id <= 20, 1,
    id <= 30, 2,
    id <= 40, 3,
    id <= 50, 1,
    id <= 60, 2,
    id <= 70, 3
  )
][
  ,
  first_treat := fcase(
    te_grp == 0, Inf,
    id <= 40, 11,
    id <= 70, 16
  )
][
  ,
  te := 0
][
  period >= first_treat,
  te := te_grp*(period-first_treat) + rnorm(.N, sd = 0.01) # add a little noise to the TEs
][
  ,
  te_grp := as.factor(te_grp)
][
  ,
  y := 1*x + te + rnorm(n = .N, sd = 0.1)
]

## show in plot form
# with(dat, plot(period, y, col = te_grp))

## etwfe
sim_mod = etwfe(
  y ~ x, tvar = period, gvar = first_treat, xvar = te_grp, data = dat
)
es = emfx(sim_mod, type = "event")

## known ATTs for the event study
truth = dat[
  period >= first_treat,
  .(ATE = mean(te)),
  by = .(te_grp, event = period - first_treat)
]

## plot comparison
with(truth, plot(event, ATE, col = te_grp, cex = 1.5, pch = 0))
with(es, points(event, estimate, col = te_grp, pch = 16))
legend("topleft", legend = c("Truth", "etwfe"), pch = c(0, 16))
title(main = "Heterogeneous ATEs (three groups), with staggered rollout")

Created on 2023-04-27 with reprex v2.0.2