tests/testthat/test-method_ebal.R

In WeightIt: Weighting for Covariate Balance in Observational Studies

test_that("Binary treatment", {
  skip_on_cran()
  skip_if_not_installed("rootSolve")
  skip_if_not_installed("cobalt")

  eps <- if (capabilities("long.double")) 1e-5 else 1e-3

  test_data <- readRDS(test_path("fixtures", "test_data.rds"))

  set.seed(123)
  base_weights <- runif(nrow(test_data))

  expect_no_condition({
    W0 <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                   data = test_data, method = "ebal", estimand = "ATE",
                   include.obj = TRUE, solver = "optim")
  })

  expect_M_parts_okay(W0, tolerance = eps)

  sw.opts <- c(FALSE, TRUE)
  bw.opts <- c(FALSE, TRUE)
  estimand.opts <- c("ATE", "ATT", "ATC")

  weight.mat <- matrix(nrow = nrow(test_data),
                       ncol = length(sw.opts) *
                         length(bw.opts) * length(estimand.opts))
  colnames(weight.mat) <- rep("", ncol(weight.mat))

  k <- 1

  for (sw in sw.opts) {
    for (bw in bw.opts) {
      for (estimand in estimand.opts) {
        test_that(sprintf("Ebal: sw = %s, bw = %s, estimand = %s", sw, bw, estimand), {
          W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                        data = test_data, method = "ebal", estimand = estimand,
                        s.weights = if (sw) "SW" else NULL,
                        base.weights = if (bw) base_weights else NULL,
                        include.obj = TRUE, solver = "multiroot")

          expect_M_parts_okay(W, tolerance = eps)
          expect_equal(cobalt::col_w_smd(W$covs, W$treat, W$weights,
                                         s.weights = W$s.weights),
                       0 * cobalt::col_w_smd(W$covs, W$treat,
                                             s.weights = W$s.weights),
                       expected.label = "all 0s",
                       tolerance = eps)

          expect_true(is_null(W$ps))
          expect_false(is_null(W$obj))

          if (estimand %in% c("ATT", "ATC")) {
            expect_ATT_weights_okay(W, tolerance = eps)
          }

          for (i in 0:1) {
            e <- {
              if (estimand == "ATT" && i == 1) expect_equal
              else if (estimand == "ATC" && i == 0) expect_equal
              else expect_not_equal
            }

            e(unname(W$weights[W$treat == i]),
              rep(1, sum(W$treat == i)),
              label = sprintf("%s weights", i),
              expected.label = "all 1s",
              tolerance = eps)
          }

          for (i in seq_len(k - 1)) {
            expect_not_equal(unname(W$weights), weight.mat[,i],
                             expected.label = sprintf("weights for %s", colnames(weight.mat)[i]),
                             tolerance = eps)
          }

          n <- sprintf("W_%s_%s_%s", sw, bw, estimand)
          colnames(weight.mat)[k] <<- n
          weight.mat[,k] <<- W$weights
          k <<- k + 1
        })
      }
    }
  }

  # Estimands
  expect_error({
    W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATO")
  }, "not an allowable estimand", ignore.case = TRUE)

  #Non-full rank
  expect_no_condition({
    W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9 +
                    I(1 - X5) + I(X9 * 2),
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, solver = "optim")
  })

  expect_M_parts_okay(W, tolerance = eps)
  expect_equal(W$weights, W0$weights, tolerance = eps)

  # All categorical covariates (issue #86)
  expect_no_condition({
    W <- weightit(A ~ cut(X1, 3) + cut(X2, 3) + cut(X3, 3),
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, solver = "optim", reltol = 1e-12)
  })

  expect_M_parts_okay(W, tolerance = eps)

  # tols > 0
  expect_no_condition({
    W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                   data = test_data, method = "ebal", estimand = "ATE",
                   include.obj = TRUE, tols = .05)
  })

  expect_not_equal(W$weights, W0$weights)

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) > eps)) #Some worse than 0

  expect_no_condition({
    W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, tols = .05, s.weights = "SW")
  })

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Diff.Adj) > eps)) #Some worse than 0

  #Should be equivalent to CBPS and IPT with logit link for ATT
  for (sw in sw.opts) {
    test_that(sprintf("Ebal matches CBPS/IPT for ATT: sw = %s", sw), {
      W <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                    data = test_data, method = "ebal", estimand = "ATT",
                    s.weights = if (sw) "SW" else NULL,
                    include.obj = TRUE, solver = "optim")

      Wcbps <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                        data = test_data, method = "cbps", estimand = "ATT",
                        s.weights = if (sw) "SW" else NULL,
                        link = "logit", solver = "multiroot",
                        include.obj = TRUE)

      expect_equal(ESS(W$weights[W$treat == 0] * W$s.weights[W$treat == 0]),
                   ESS(Wcbps$weights[Wcbps$treat == 0] * Wcbps$s.weights[Wcbps$treat == 0]),
                   expected.label = "ESS for CBPS",
                   tolerance = .01)

      Wipt <- weightit(A ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                       data = test_data, method = "ipt", estimand = "ATT",
                       s.weights = if (sw) "SW" else NULL,
                       link = "logit",
                       include.obj = TRUE)

      expect_equal(ESS(W$weights[W$treat == 0] * W$s.weights[W$treat == 0]),
                   ESS(Wipt$weights[Wipt$treat == 0] * Wipt$s.weights[Wipt$treat == 0]),
                   expected.label = "ESS for IPT",
                   tolerance = .01)
    })
  }
})

test_that("Multi-category treatment", {
  skip_on_cran()
  skip_if_not_installed("rootSolve")
  skip_if_not_installed("cobalt")

  eps <- if (capabilities("long.double")) 1e-5 else 1e-3

  test_data <- readRDS(test_path("fixtures", "test_data.rds"))

  set.seed(123)
  base_weights <- runif(nrow(test_data))

  expect_no_condition({
    W0 <- weightit(Am ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                   data = test_data, method = "ebal", estimand = "ATE",
                   include.obj = TRUE, solver = "optim")
  })

  sw.opts <- c(FALSE, TRUE)
  bw.opts <- c(FALSE, TRUE)
  estimand.opts <- c("ATE", "ATT")

  weight.mat <- matrix(nrow = nrow(test_data),
                       ncol = length(sw.opts) *
                         length(bw.opts) * length(estimand.opts))
  colnames(weight.mat) <- rep("", ncol(weight.mat))

  k <- 1

  for (sw in sw.opts) {
    for (bw in bw.opts) {
      for (estimand in estimand.opts) {
        test_that(sprintf("Ebal: sw = %s, bw = %s, estimand = %s", sw, bw, estimand), {
          W <- weightit(Am ~ X1 + X2 + X3 + X4 + X5,
                        data = test_data, method = "ebal", estimand = estimand,
                        focal = if (estimand == "ATE") NULL else "T",
                        s.weights = if (sw) "SW" else NULL,
                        base.weights = if (bw) base_weights else NULL,
                        include.obj = TRUE, solver = "multiroot")

          expect_M_parts_okay(W, tolerance = eps)
          for (tt in combn(levels(W$treat), 2, simplify = FALSE)) {
            in_tt <- W$treat %in% tt
            expect_equal(cobalt::col_w_smd(W$covs[in_tt,], W$treat[in_tt], W$weights[in_tt],
                                           s.weights = W$s.weights[in_tt]),
                         0 * cobalt::col_w_smd(W$covs[in_tt,], W$treat[in_tt],
                                               s.weights = W$s.weights[in_tt]),
                         label = sprintf("SMDs for %s", paste(tt, collapse = " vs. ")),
                         expected.label = "all 0s",
                         tolerance = eps)
          }

          expect_true(is_null(W$ps))
          expect_false(is_null(W$obj))

          if (estimand %in% c("ATT", "ATC")) {
            expect_ATT_weights_okay(W, tolerance = eps)
          }

          for (i in levels(W$treat)) {
            e <- {
              if (estimand == "ATT" && i == W$focal) expect_equal
              else expect_not_equal
            }

            e(unname(W$weights[W$treat == i]),
              rep(1, sum(W$treat == i)),
              label = sprintf("%s weights", i),
              expected.label = "all 1s",
              tolerance = eps)
          }

          for (i in seq_len(k - 1)) {
            expect_not_equal(unname(W$weights), weight.mat[,i],
                             expected.label = sprintf("weights for %s", colnames(weight.mat)[i]),
                             tolerance = eps)
          }

          n <- sprintf("W_%s_%s_%s", sw, bw, estimand)
          colnames(weight.mat)[k] <<- n
          weight.mat[,k] <<- W$weights
          k <<- k + 1
        })
      }
    }
  }

  # tols > 0
  expect_no_condition({
    W <- weightit(Am ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, tols = .05)
  })

  expect_not_equal(W$weights, W0$weights)

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) > eps)) #Some worse than 0

  expect_no_condition({
    W <- weightit(Am ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, tols = .05, s.weights = "SW")
  })

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Max.Diff.Adj) > eps)) #Some worse than 0
})

test_that("Continuous treatment", {
  skip_on_cran()
  skip_if_not_installed("rootSolve")
  skip_if_not_installed("cobalt")

  eps <- if (capabilities("long.double")) 1e-5 else 1e-3

  test_data <- readRDS(test_path("fixtures", "test_data.rds"))

  set.seed(123)
  base_weights <- runif(nrow(test_data))

  expect_no_condition({
    W0 <- weightit(Ac ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                   data = test_data, method = "ebal",
                   include.obj = TRUE, solver = "optim")
  })

  expect_M_parts_okay(W0, tolerance = eps)

  sw.opts <- c(FALSE, TRUE)
  bw.opts <- c(FALSE, TRUE)
  d.moments.opts <- c(1, 3)

  weight.mat <- matrix(nrow = nrow(test_data),
                       ncol = length(sw.opts) *
                         length(bw.opts) * length(d.moments.opts))
  colnames(weight.mat) <- rep("", ncol(weight.mat))

  k <- 1

  for (sw in sw.opts) {
    for (bw in bw.opts) {
      for (d.moments in d.moments.opts) {
        test_that(sprintf("Ebal: sw = %s, bw = %s, d.moments = %s", sw, bw, d.moments), {
          W <- weightit(Ac ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                        data = test_data, method = "ebal",
                        d.moments = d.moments,
                        s.weights = if (sw) "SW" else NULL,
                        base.weights = if (bw) base_weights else NULL,
                        include.obj = TRUE, solver = "multiroot")

          expect_M_parts_okay(W, tolerance = eps)
          expect_equal(cobalt::col_w_cov(W$covs, W$treat, W$weights, std = TRUE,
                                         s.weights = W$s.weights),
                       0 * cobalt::col_w_cov(W$covs, W$treat, std = TRUE,
                                             s.weights = W$s.weights),
                       expected.label = "all 0s",
                       tolerance = eps)

          expect_equal(cobalt::col_w_mean(cbind(poly(W$treat, d.moments), W$covs), W$weights,
                                          s.weights = W$s.weights),
                       cobalt::col_w_mean(cbind(poly(W$treat, d.moments), W$covs),
                                              s.weights = W$s.weights),
                       expected.label = "unweighted means",
                       tolerance = eps)

          expect_true(is_null(W$ps))
          expect_false(is_null(W$obj))

          for (i in seq_len(k - 1)) {
            expect_not_equal(unname(W$weights), weight.mat[,i],
                             expected.label = sprintf("weights for %s", colnames(weight.mat)[i]),
                             tolerance = eps)
          }

          n <- sprintf("W_%s_%s_%s", sw, bw, d.moments)
          colnames(weight.mat)[k] <<- n
          weight.mat[,k] <<- W$weights
          k <<- k + 1
        })
      }
    }
  }

  #Non-full rank
  expect_no_condition({
    W <- weightit(Ac ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9 +
                    I(1 - X5) + I(X9 * 2),
                  data = test_data, method = "ebal",
                  include.obj = TRUE, solver = "optim")
  })

  expect_M_parts_okay(W, tolerance = eps)
  expect_equal(W$weights, W0$weights, tolerance = eps)

  # tols > 0
  expect_no_condition({
    W <- weightit(Ac ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, tols = .05)
  })

  expect_not_equal(W$weights, W0$weights)

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) > eps)) #Some worse than 0

  expect_no_condition({
    W <- weightit(Ac ~ X1 + X2 + X3 + X4 + X5 + X6 + X7 + X8 + X9,
                  data = test_data, method = "ebal", estimand = "ATE",
                  include.obj = TRUE, tols = .05, s.weights = "SW")
  })

  expect_true(all(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) <= .05 + eps)) #None worse than tols
  expect_true(any(abs(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) - .05) <= eps)) #Some exactly tols
  expect_true(any(abs(cobalt::bal.tab(W)$Balance$Corr.Adj) > eps)) #Some worse than 0
})

test_that("Continuous treatment: M-estimation parts assemble into a variance", {
  skip_on_cran()
  skip_if_not_installed("rootSolve")

  eps <- if (capabilities("long.double")) 1e-5 else 1e-3

  test_data <- readRDS(test_path("fixtures", "test_data.rds"))

  # The continuous parts supply `dw_dBtreat` but no `hess_treat`, so assembling them
  # takes the mixed analytic/numeric branch in `.get_glm_weightit_vcov()`. The parts
  # themselves were checked above; nothing verified that the assembly produces a usable
  # variance, which is where a wrong standard error would appear without any error.
  W <- weightit(Ac ~ X1 + X2 + X3, data = test_data, method = "ebal")

  expect_M_parts_okay(W, tolerance = eps)
  expect_false(is_null(attr(W, "Mparts")$dw_dBtreat))
  expect_null(attr(W, "Mparts")$hess_treat)

  fit <- lm_weightit(Y_C ~ Ac, data = test_data, weightit = W)

  se_asympt <- sqrt(diag(vcov(fit)))
  se_hc0 <- sqrt(diag(vcov(fit, vcov = "HC0")))

  expect_true(all(is.finite(se_asympt)))
  expect_true(all(se_asympt > 0))

  # Accounting for estimation of the weights changes the answer, so the parts are
  # actually contributing rather than being dropped
  expect_not_equal(unname(se_asympt), unname(se_hc0))
})