Merge 779812e into 8fb3f0a

DeclareDesign · Apr 11, 2018 · ab4f635 · ab4f635
2 parents 8fb3f0a + 779812e
commit ab4f635
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 32 deletions.
diff --git a/R/estimatr_horvitz_thompson.R b/R/estimatr_horvitz_thompson.R
@@ -283,25 +283,22 @@ horvitz_thompson <- function(formula,
     }
 
     if (!is.null(condition2)) {
-      treatnum <-
-        which(ra_declaration$condition_names == condition2)
+      treatnum <- match(condition2, ra_declaration$conditions)
 
-      if (!(condition2 %in% ra_declaration$condition_names)) {
+      if (is.na(treatnum)) {
         stop(
           "If `condition2` and `ra_declaration` are both specified, ",
           "`condition2` must match the condition_names in `ra_declaration`.",
           "\n`condition2`: ", condition2, "\n`condition_names`: ",
           paste0(
-            ra_declaration$condition_names,
+            ra_declaration$conditions,
             collapse = ", "
           )
         )
       }
 
-      treatment_prob <- obtain(
-        ra_declaration,
-        condition2
-      )
+      treatment_prob <- obtain(ra_declaration,condition2)
+
     } else {
       # assuming treatment is second column
       treatment_prob <- ra_declaration$probabilities_matrix[, 2]
@@ -370,7 +367,7 @@ horvitz_thompson <- function(formula,
     }
     # If simple, just use condition probabilities shortcut
     # Same if se not needed
-    if ("ra_simple" %in% class(ra_declaration) || se_type == "none") {
+    if (inherits(ra_declaration, "ra_simple") || se_type == "none") {
       condition_pr_mat <- NULL
     } else {
       # TODO to allow for declaration with multiple arms, get probability matrix

diff --git a/R/helper_condition_pr_matrix.R b/R/helper_condition_pr_matrix.R
@@ -114,8 +114,8 @@ declaration_to_condition_pr_mat <- function(ra_declaration,
   }
 
   if (is.null(condition1) && is.null(condition2)) {
-    condition1 <- ra_declaration$condition_names[1]
-    condition2 <- ra_declaration$condition_names[2]
+    condition1 <- ra_declaration$conditions[1]
+    condition2 <- ra_declaration$conditions[2]
   } else if (is.null(condition1) && !is.null(condition2)) {
     stop(
       "Cannot have `condition1 == NULL` and `condition2 != NULL`"

diff --git a/tests/testthat/test-condition-pr-matrix.R b/tests/testthat/test-condition-pr-matrix.R
@@ -1,15 +1,18 @@
 context("Helper - HT condition_pr_matrix")
+n <- 5
 
-cond_pr_mat_tests <- function() {
 
-  # Errors appropriately
-  expect_error(
-    declaration_to_condition_pr_mat(rbinom(5, 1, 0.5)),
-    "`ra_declaration` must be an object of class 'ra_declaration'"
-  )
+test_that("Checks class", {
+    # Errors appropriately
+    expect_error(
+      declaration_to_condition_pr_mat(rbinom(5, 1, 0.5)),
+      "`ra_declaration` must be an object of class 'ra_declaration'"
+    )
+})
+
+test_that("Complete randomization", {
 
-  # Complete randomization
-  n <- 5
+  #
   prs <- rep(0.4, times = n)
   comp_ra <- randomizr::declare_ra(N = n, prob = prs[1])
   perms <- randomizr::obtain_permutation_matrix(comp_ra)
@@ -18,7 +21,10 @@ cond_pr_mat_tests <- function() {
     permutations_to_condition_pr_mat(perms)
   )
 
-  # declaration to condition_pr_mat errors
+})
+
+test_that("declaration to condition_pr_mat errors", {
+
   expect_error(
     declaration_to_condition_pr_mat(randomizr::declare_ra(N = n), 1, NULL),
     "Cannot have `condition2 == NULL`"
@@ -31,6 +37,9 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(rbinom(5, 1, 0.5)),
     "`ra_declaration` must be an object of class 'ra_declaration'"
   )
+})
+
+test_that("condition args work properly", {
 
   # Condition args work properly
   mat01 <- declaration_to_condition_pr_mat(
@@ -48,8 +57,11 @@ cond_pr_mat_tests <- function() {
   # colnames(mat01)
   # colnames(mat10)
   expect_equal(mat01, mat10[rownames(mat01), colnames(mat01)])
+})
+
+test_that("Complete randomization with number of treated units not fixed", {
 
-  # Complete randomization with number of treated units not fixed
+  #
   comp_odd_ra <- randomizr::declare_ra(N = 3, prob = 0.5)
   perms <- randomizr::obtain_permutation_matrix(comp_odd_ra)
 
@@ -67,7 +79,9 @@ cond_pr_mat_tests <- function() {
     permutations_to_condition_pr_mat(perms)
   )
 
-  # Complete randomization with non 0.5 as remainder
+})
+
+test_that("Complete randomization with non 0.5 as remainder", {
   comp_odd_ra <- randomizr::declare_ra(N = 3, prob = 0.4)
   decl_cond_pr_mat <- declaration_to_condition_pr_mat(comp_odd_ra)
 
@@ -78,6 +92,8 @@ cond_pr_mat_tests <- function() {
     permutations_to_condition_pr_mat(get_perms),
     tolerance = 0.01
   )
+})
+test_that("Simple ra", {
 
   # Simple randomization
   prs <- rep(0.4, times = n)
@@ -94,6 +110,8 @@ cond_pr_mat_tests <- function() {
     permutations_to_condition_pr_mat(perms),
     tolerance = 0.02
   )
+})
+test_that("Blocked complete ra", {
 
   # Blocked case
   dat <- data.frame(
@@ -108,9 +126,15 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(bl_ra),
     permutations_to_condition_pr_mat(bl_perms)
   )
+})
+test_that("Blocked complete ra with remainder", {
+  dat <- data.frame(
+    bl = c("A", "B", "A", "B", "B", "B"),
+    pr = c(0.5, 0.25, 0.5, 0.25, 0.25, 0.25)
+  )
 
   # with remainder
-  bl <- c("A", "B", "A", "A", "B", "B")
+  bl <- c("A", "B", "A", "A", "B", "B") # Is this used anywhere?
 
   bl_ra <- randomizr::declare_ra(blocks = dat$bl, prob = 0.4)
   bl_perms <- replicate(5000, randomizr::conduct_ra(bl_ra))
@@ -120,6 +144,8 @@ cond_pr_mat_tests <- function() {
     permutations_to_condition_pr_mat(bl_perms),
     tolerance = 0.02
   )
+})
+test_that("Clustered complete ra", {
 
   # Cluster complete case
   dat <- data.frame(
@@ -143,8 +169,15 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(cl_ra),
     permutations_to_condition_pr_mat(cl_perms)
   )
+})
+
+test_that("Clustered ra", {
+
+  # Cluster simple ? Should this be simple or no? --NJF
+  dat <- data.frame(
+    cl = c("A", "B", "A", "C", "A", "B")
+  )
 
-  # Cluster simple
   dat$prs <- 0.3
   cl_simp_ra <- randomizr::declare_ra(clusters = dat$cl, prob = dat$prs[1])
   cl_simp_perms <- randomizr::obtain_permutation_matrix(cl_simp_ra)
@@ -168,6 +201,9 @@ cond_pr_mat_tests <- function() {
     tolerance = 0.01
   )
 
+})
+
+test_that("Blocked and Clustered ra", {
 
   # Blocked and clustered
   dat <- data.frame(
@@ -182,6 +218,9 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(bl_cl_ra),
     permutations_to_condition_pr_mat(bl_cl_perms)
   )
+})
+
+test_that("Blocked and clusted ra with remainder", {
 
   # with remainder
   dat <- data.frame(
@@ -196,17 +235,21 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(bl_cl_ra),
     permutations_to_condition_pr_mat(bl_cl_perms)
   )
+})
 
-  # Custom case
+test_that("Custom ra", {
   cust_perms <- cbind(c(1, 0, 1, 0), c(1, 1, 0, 0))
   cust_ra <- randomizr::declare_ra(permutation_matrix = cust_perms)
 
   expect_equal(
     declaration_to_condition_pr_mat(cust_ra),
     permutations_to_condition_pr_mat(cust_perms)
   )
+})
+
+test_that("Errors for things that we can't support", {
 
-  # Errors for things that we can't support
+  #
   # multiple armed experiments
   mult_ra <- randomizr::declare_ra(N = 10, prob_each = c(0.2, 0.2, 0.6))
   expect_error(
@@ -229,8 +272,9 @@ cond_pr_mat_tests <- function() {
     ),
     "Treatment probabilities cannot vary within blocks"
   )
+})
 
-  # probability not fixed within blocks
+test_that("probability not fixed within blocks", {
   bl_small <- randomizr::declare_ra(
     blocks = c(1, 1, 2, 2),
     prob = 0.4
@@ -248,6 +292,9 @@ cond_pr_mat_tests <- function() {
     declaration_to_condition_pr_mat(bl_small),
     "Treatment probabilities must be fixed within blocks for block randomized"
   )
+})
+
+test_that("N=2, m=1", {
 
   comp <- randomizr::declare_ra(N = 2, m = 1)
   assign(
@@ -269,9 +316,6 @@ cond_pr_mat_tests <- function() {
     estimatr:::gen_pr_matrix_block(c(1, 2), c(1, 2)),
     "Must specify one of `t`, `p2`, or `p1`"
   )
-}
 
-test_that("condition_pr_matrix behaves as expected w/ randomizr", {
-  cond_pr_mat_tests()
-})
 
+})
diff --git a/tests/testthat/test-horvitz-thompson.R b/tests/testthat/test-horvitz-thompson.R
@@ -273,7 +273,7 @@ test_that("Horvitz-Thompson works with clustered data", {
     y = rnorm(12)
   )
   # pr = 0.25 in first, 0.5 in second
-  blcl_ra <- randomizr::declare_ra(blocks = clbl_dat$bl, clusters = clbl_dat$cl_new, m = c(1, 2, 1))
+  blcl_ra <- randomizr::declare_ra(blocks = clbl_dat$bl, clusters = clbl_dat$cl_new, block_m = c(1, 2, 1))
   clbl_dat$z_clbl <- randomizr::conduct_ra(blcl_ra)
   expect_equivalent(
     horvitz_thompson(y ~ z_clbl, data = clbl_dat, ra_declaration = blcl_ra),