Merge f339991 into 3965a5c

R/declare_sampling.R

@@ -38,11 +38,12 @@
 declare_sampling <- make_declarations(sampling_handler, "sampling")
 
 #' @param sampling_variable The prefix for the sampling inclusion probability variable.
+#' @param keep The set of sampling outcomes to retain. Use 0:1 to keep the entire population.
 #' @param data A data.frame.
 #' @importFrom rlang quos !!! call_modify eval_tidy quo
 #' @importFrom randomizr draw_rs obtain_inclusion_probabilities
 #' @rdname declare_sampling
-sampling_handler <- function(data, ..., sampling_variable = "S") {
+sampling_handler <- function(data, ..., sampling_variable = "S", keep = 1) {
   ## draw sample
 
   options <- quos(...)
@@ -63,7 +64,7 @@ sampling_handler <- function(data, ..., sampling_variable = "S") {
   S <- as.character(S)
 
   ## subset to the sampled observations
-  data[ data[[S]] %in% 1, names(data) != S, drop = FALSE]
+  data[ data[[S]] %in% keep,(length(keep) > 1) | names(data) != S, drop = FALSE]
 }
 
 validation_fn(sampling_handler) <- function(ret, dots, label) {

tests/testthat/test-sampling.R

@@ -103,3 +103,28 @@ test_that("Factor out declarations", {
   expect_true(inherits(attr(design[[3]], "dots")$declaration, "rs_complete"))
   expect_true(inherits(attr(design[[4]], "dots")$declaration, "ra_complete"))
 })
+
+
+
+
+
+test_that("Keep options on declare_sampling", {
+
+  N <- 500
+  n <- 2
+
+  design1 <- declare_population(N = N, noise = 1:N) + declare_sampling(n = n)
+  design2 <- declare_population(N = N, noise = 1:N) + declare_sampling(n = n, keep=0)
+  design3 <- declare_population(N = N, noise = 1:N) + declare_sampling(n = n, keep=0:1)
+
+
+  expect_equal(nrow(draw_data(design1)), n)
+  expect_equal(nrow(draw_data(design2)), N-n)
+
+  expect_equal(
+    table(draw_data(design3)$.__Sample),
+    structure(c(`0` = N-n, `1` = n), .Dim = 2L, .Dimnames = structure(list(
+      c("0", "1")), .Names = ""), class = "table")
+  )
+
+})