Document survey sampling terminology for the _rs functions (see #99)

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: randomizr
 Title: Easy-to-Use Tools for Common Forms of Random Assignment and Sampling
-Version: 2.0.1
+Version: 2.0.1.9000
 Authors@R: c(person("Alexander", "Coppock", email = "acoppock@gmail.com", role = c("aut", "cre"), comment = c(ORCID = "0000-0002-5733-2386")),
              person("Jasper", "Cooper", email = "jaspercooper@gmail.com", role = c("ctb"), comment = c(ORCID = "0000-0002-8639-3188")),
              person("Neal", "Fultz", email = "nfultz@gmail.com", role = c("ctb"), comment = "C version of restricted partitions"),

NEWS.md

@@ -1,3 +1,11 @@
+# randomizr (development version)
+
+* `complete_ra()`, `cluster_ra()`, `block_ra()`, and `block_and_cluster_ra()` now treat a lone unit (or lone cluster, or one-unit block) with its stated probability. Given `prob` or `block_prob` strictly between 0 and 1, such a unit was left in control on every draw, while the `_probabilities()` functions reported the stated probability. Designs with more than one unit per block keep their exact draws: the correcting coin is drawn after the existing ones, and only when the case arises.
+
+* `complete_rs()`, `cluster_rs()`, `strata_rs()`, and `strata_and_cluster_rs()` now reach the requested inclusion probability when rounding up would take the whole frame, stratum, or set of clusters. They had inherited `complete_ra()`'s rule of never rounding up to N, which keeps a control arm nonempty under assignment but has no purpose under sampling: `complete_rs(N = 10, prob = 0.95)` always sampled 9 units, for inclusion of 0.90, and a one-unit stratum with `strata_prob = 0.3` was never sampled. The `_probabilities()` functions report the corrected values. Draws for designs that never hit the case are unchanged (#99).
+
+* The documentation of `simple_rs()`, `complete_rs()`, `strata_rs()`, `cluster_rs()`, and `strata_and_cluster_rs()` gains a "Terminology" section giving the names the survey sampling literature uses for each design (Bernoulli and Poisson sampling, simple random sampling without replacement, stratified and cluster sampling). Function names are unchanged. See #99.
+
 # randomizr 2.0.1
 
 This release is a major internal restructuring plus one new capability, described here relative to randomizr 1.0.1, the version it replaces on CRAN. No exported function, argument, return type, S3 class, or object field carried by 1.0.1 is removed, renamed, or reordered. New arguments were added to `declare_ra()`, `conduct_ra()` and `obtain_condition_probabilities()` (`prob_unit_each` after `prob_each`; `ra_type` and `formula` before `permutation_matrix`; `data` last), so a caller passing `block_m` or later arguments by position must name them; `block_ra()` gains two arguments documented as internal (`.block_int`, `.N_per_block`). The package exports 31 objects against 1.0.1's 29, the two additions being `balanced_ra()` and `balanced_ra_probabilities()`.

R/block_ra.R

@@ -163,6 +163,15 @@ block_ra <- function(blocks = NULL,
       m_per_b <- integer(G)
     }
     raw        <- block_assign_cpp(block_int, m_per_b, prob_v, mode)
+    # The kernel follows complete_ra() in never rounding a block's count up to
+    # the whole block, so a one-unit block with block_prob strictly between 0
+    # and 1 was never treated. Its coin is drawn after the kernel's draws, so
+    # designs without such a block keep their streams.
+    if (mode == 2L) {
+      for (b in which(N_per_block == 1L & prob_v > 0 & prob_v < 1)) {
+        if (runif(1L) < prob_v[b]) raw[block_int == b] <- 1L
+      }
+    }
     cond       <- if (!is.null(conditions)) conditions else c(0L, 1L)
     assignment <- cond[raw + 1L]
     return(clean_condition_names(assignment, conditions))

R/cluster_rs.R

@@ -4,6 +4,9 @@
 #'
 #' By default the clusters are drawn by complete random sampling, so a fixed number of clusters is sampled on every draw. Setting \code{simple = TRUE} draws each cluster independently instead, using \code{\link{simple_rs}()}.
 #'
+#' @section Terminology:
+#' The name follows \code{\link{cluster_ra}()}, as the \code{_rs} functions follow their random assignment counterparts. In the survey sampling literature this design is one-stage cluster sampling: simple random sampling without replacement (SRSWOR) of clusters, with every unit in a sampled cluster included (Lohr 2022, \emph{Sampling: Design and Analysis}, 3rd ed.; Särndal, Swensson, and Wretman 1992, \emph{Model Assisted Survey Sampling}). With \code{simple = TRUE} the clusters are instead drawn by Bernoulli sampling (Poisson sampling if probabilities differ by cluster).
+#'
 #' @seealso \code{\link{complete_rs}()}, \code{\link{strata_and_cluster_rs}()}, \code{\link{cluster_ra}()}, \code{\link{cluster_rs_probabilities}()}
 #'
 #' @param clusters A vector of length N indicating which cluster each unit belongs to. (required)

R/complete_ra.R

@@ -151,6 +151,12 @@ complete_ra <- function(N,
       if (m_ceiling == N) {
         m <- m_floor
         assignment <- sample(rep(conditions, c(N - m, m)))
+        # Rounding down keeps a control unit when N > 1. With N = 1 it left the
+        # unit in control on every draw, although its probability is prob. The
+        # coin is drawn after the permutation, so no other design's stream moves.
+        if (N == 1 && m_floor < N * prob && runif(1L) < prob) {
+          assignment <- conditions[2]
+        }
         assignment <-
           clean_condition_names(assignment, conditions)
         return(assignment)

R/complete_rs.R

@@ -7,6 +7,9 @@
 #' If only \code{N} is specified, half the units are sampled. When \code{N} is odd, either \code{floor(N/2)} or \code{ceiling(N/2)} units are sampled.
 #'
 #'
+#' @section Terminology:
+#' The name follows \code{\link{complete_ra}()}, as the \code{_rs} functions follow their random assignment counterparts. With a whole-number \code{n}, every subset of size \code{n} is equally likely, which the survey sampling literature calls simple random sampling without replacement (SRSWOR; Lohr 2022, \emph{Sampling: Design and Analysis}, 3rd ed.; Särndal, Swensson, and Wretman 1992, \emph{Model Assisted Survey Sampling}). When \code{N*prob} is not a whole number, the design is a mixture of SRSWOR of size \code{floor(N*prob)} and SRSWOR of size \code{ceiling(N*prob)}, including when the ceiling is \code{N} itself.
+#'
 #' @seealso \code{\link{simple_rs}()}, \code{\link{strata_rs}()}, \code{\link{cluster_rs}()}, \code{\link{complete_ra}()}, \code{\link{complete_rs_probabilities}()}
 #'
 #' @param N The number of units in the sampling frame. Must be a positive integer. (required)
@@ -81,8 +84,10 @@ complete_rs <- function(N,
     n_dn <- floor(Np)
     n_up <- ceiling(Np)
 
-    # If rounding doesn't matter or rounds up to 100% use n_dn except when N=1
-    n <- if (n_up == n_dn || (N > 1 && n_up == N)) {
+    # Unlike complete_ra(), which keeps a unit out of treatment by never
+    # rounding up to N, sampling has no second arm to protect: rounding up to
+    # the whole frame is allowed, so inclusion is exactly prob.
+    n <- if (n_up == n_dn) {
       n_dn
     } else  {
       n_dn + sample(0:1, 1, prob = abs(1:0 - (Np - n_dn)))
@@ -148,8 +153,7 @@ complete_rs_probabilities <- function(N,
     #n / max(N,2) # 0,1=> 0, 1,1 => 1/2
     n / N
   } else if (is.numeric(prob)) {
-    ifelse(N > 1 && ceiling(N * prob) == N,
-           floor(N * prob) / N, prob)
+    prob
   } else {
     .5
   }

R/simple_rs.R

@@ -4,6 +4,9 @@
 #'
 #' If \code{prob} is not specified, each unit is sampled with probability 0.5.
 #'
+#' @section Terminology:
+#' The name follows \code{\link{simple_ra}()}, as the \code{_rs} functions follow their random assignment counterparts. In the survey sampling literature, independent draws like these are called Bernoulli sampling when \code{prob} is constant and Poisson sampling when \code{prob_unit} varies by unit (Lohr 2022, \emph{Sampling: Design and Analysis}, 3rd ed.; Särndal, Swensson, and Wretman 1992, \emph{Model Assisted Survey Sampling}). The sample size is random, and each unit's inclusion probability is exactly the probability supplied.
+#'
 #' @seealso \code{\link{complete_rs}()}, \code{\link{strata_rs}()}, \code{\link{simple_ra}()}, \code{\link{simple_rs_probabilities}()}
 #'
 #' @param N The number of units in the sampling frame. Must be a positive integer. (required)

R/strata_and_cluster_rs.R

@@ -10,6 +10,9 @@
 #' Clusters must nest within strata: every unit in a cluster has to belong to
 #' the same stratum.
 #'
+#' @section Terminology:
+#' The name follows \code{\link{block_and_cluster_ra}()}, as the \code{_rs} functions follow their random assignment counterparts. In the survey sampling literature this design is stratified one-stage cluster sampling: stratified simple random sampling without replacement (stratified SRSWOR) of clusters, with every unit in a sampled cluster included (Lohr 2022, \emph{Sampling: Design and Analysis}, 3rd ed.; Särndal, Swensson, and Wretman 1992, \emph{Model Assisted Survey Sampling}).
+#'
 #' @seealso \code{\link{cluster_rs}()}, \code{\link{strata_rs}()},
 #'   \code{\link{block_and_cluster_ra}()}
 #'

R/strata_rs.R

@@ -4,6 +4,9 @@
 #'
 #' The number sampled per stratum can be left to the function, set as a common count or probability across strata (\code{n}, \code{prob}), or set stratum by stratum (\code{strata_n}, \code{strata_prob}). When the probability varies across strata the sample is not self-weighting, and \code{\link{strata_rs_probabilities}()} gives the inclusion probabilities needed to weight it.
 #'
+#' @section Terminology:
+#' The name follows \code{\link{block_ra}()}, as the \code{_rs} functions follow their random assignment counterparts. In the survey sampling literature this design is stratified simple random sampling without replacement (stratified SRSWOR; Lohr 2022, \emph{Sampling: Design and Analysis}, 3rd ed.; Särndal, Swensson, and Wretman 1992, \emph{Model Assisted Survey Sampling}). Strata are sampled independently. When \code{N_stratum*prob} is not a whole number, a stratum's sample size is a mixture of the floor and the ceiling, as in \code{\link{complete_rs}()}.
+#'
 #' @seealso \code{\link{complete_rs}()}, \code{\link{strata_and_cluster_rs}()}, \code{\link{block_ra}()}, \code{\link{strata_rs_probabilities}()}
 #'
 #' @param strata A vector of length N indicating which stratum each unit belongs to. Can be a character, factor, or numeric vector. (required)
@@ -60,7 +63,7 @@ strata_rs <- function(strata = NULL,
                       strata_prob = NULL,
                       check_inputs = TRUE) {
   if (check_inputs) .invoke_check(check_samplr_arguments_new)
-  block_ra(
+  S <- block_ra(
     blocks       = strata,
     prob         = prob,
     prob_unit    = prob_unit,
@@ -72,6 +75,40 @@ strata_rs <- function(strata = NULL,
     num_arms     = 2L,
     check_inputs = FALSE
   )
+  if (.strata_rs_uses_block_prob(n, n_unit, strata_n, prob_unit, strata_prob)) {
+    r <- .strata_rs_rounded_down(strata, prob_unit, strata_prob)
+    for (s in which(r$rounded_down)) {
+      if (runif(1L) < r$frac[s]) S[r$stratum == s] <- 1
+    }
+  }
+  S
+}
+
+# block_ra() never rounds a block's count up to the whole block, which keeps
+# every block's control arm nonempty. Sampling has no second arm to protect,
+# so a stratum whose count block_ra() rounds down to N_s - 1 is instead taken
+# whole with probability equal to the fractional part of N_s * prob. Mixing
+# the whole stratum with block_ra()'s N_s - 1 draw makes inclusion exactly
+# prob, and the extra draw comes after block_ra()'s, so seeds for designs
+# that never hit the case are unchanged. block_ra() reaches the rule only
+# through block_prob, which prob_unit and strata_prob become; a scalar prob
+# takes another branch that already rounds up.
+.strata_rs_uses_block_prob <- function(n, n_unit, strata_n, prob_unit, strata_prob) {
+  is.null(n) && is.null(n_unit) && is.null(strata_n) &&
+    (!is.null(prob_unit) || is.numeric(strata_prob))
+}
+
+.strata_rs_rounded_down <- function(strata, prob_unit, strata_prob) {
+  strata <- droplevels(as.factor(strata))
+  prob <- if (!is.null(prob_unit)) as.numeric(tapply(prob_unit, strata, unique)) else strata_prob
+  N_s <- tabulate(as.integer(strata))
+  Np <- N_s * prob
+  list(
+    stratum = as.integer(strata),
+    prob = prob,
+    frac = Np - floor(Np),
+    rounded_down = N_s > 1 & ceiling(Np) == N_s & floor(Np) < Np
+  )
 }
 
 #' Inclusion probabilities: Stratified Random Sampling
@@ -131,5 +168,11 @@ strata_rs_probabilities <- function(strata = NULL,
     num_arms     = 2L,
     check_inputs = FALSE
   )
-  prob_mat[, "prob_1"]
+  probs <- prob_mat[, "prob_1"]
+  if (.strata_rs_uses_block_prob(n, n_unit, strata_n, prob_unit, strata_prob)) {
+    r <- .strata_rs_rounded_down(strata, prob_unit, strata_prob)
+    hit <- r$rounded_down[r$stratum]
+    probs[hit] <- r$prob[r$stratum][hit]
+  }
+  probs
 }

tests/testthat/test-block_ra.R

@@ -256,3 +256,22 @@ test_that("types OK",{
 })
 
 
+
+test_that("a one-unit block with block_prob strictly between 0 and 1 is treated with probability block_prob", {
+  blocks <- rep(c("a", "b"), c(1, 6))
+  set.seed(343)
+  Z <- replicate(10000, block_ra(blocks = blocks, block_prob = c(0.3, 0.5)))
+  expect_equal(mean(Z[1, ]), 0.3, tolerance = 0.03)
+  expect_true(all(colSums(Z[2:7, ]) == 3))
+  set.seed(343)
+  Z <- replicate(10000, block_and_cluster_ra(blocks = blocks, clusters = 1:7, block_prob = c(0.3, 0.5))[1])
+  expect_equal(mean(Z), 0.3, tolerance = 0.03)
+})
+
+test_that("the one-unit-block correction leaves other blocked designs' draws unchanged", {
+  blocks <- rep(c("a", "b", "c"), c(5, 8, 3))
+  set.seed(343)
+  Z <- block_ra(blocks = blocks, block_prob = c(0.3, 0.5, 0.9))
+  # Recorded under randomizr 2.0.1, before the correction.
+  expect_equal(Z, c(0L, 0L, 0L, 1L, 0L, 1L, 0L, 0L, 1L, 1L, 1L, 0L, 0L, 0L, 1L, 1L))
+})

tests/testthat/test-complete_ra.R

@@ -210,3 +210,13 @@ test_that("issue 90 is OK", {
 })
 
 
+
+test_that("a single unit with prob strictly between 0 and 1 is treated with probability prob", {
+  set.seed(343)
+  Z <- replicate(10000, complete_ra(N = 1, prob = 0.3))
+  expect_equal(mean(Z), 0.3, tolerance = 0.03)
+  expect_equal(unname(complete_ra_probabilities(N = 1, prob = 0.3)[1, "prob_1"]), 0.3)
+  set.seed(343)
+  Z <- replicate(10000, cluster_ra(clusters = rep(1, 4), prob = 0.3)[1])
+  expect_equal(mean(Z), 0.3, tolerance = 0.03)
+})

tests/testthat/test-complete_rs.R

@@ -84,3 +84,11 @@ test_that("_unit",{
   expect_error(complete_rs(100, n_unit = rep(c(40, 50), c(50, 50))))
 
 })
+
+test_that("complete_rs rounds up to the whole frame when the ceiling of N*prob is N", {
+  expect_equal(complete_rs_probabilities(N = 10, prob = 0.95), rep(0.95, 10))
+  set.seed(343)
+  sizes <- replicate(4000, sum(complete_rs(N = 10, prob = 0.95)))
+  expect_setequal(unique(sizes), c(9, 10))
+  expect_equal(mean(sizes == 10), 0.5, tolerance = 0.05)
+})

tests/testthat/test-strata_rs.R

@@ -103,3 +103,21 @@ test_that("prob_unit and n_unit", {
 
 })
 
+
+test_that("strata_rs rounds a stratum up to its whole size when the ceiling of N_s*prob is N_s", {
+  strata <- rep(c("a", "b", "c"), c(4, 10, 1))
+  probs <- c(0.9, 0.5, 0.3)
+  expect_equal(
+    strata_rs_probabilities(strata = strata, strata_prob = probs),
+    rep(probs, c(4, 10, 1))
+  )
+  expect_equal(
+    strata_rs_probabilities(strata = strata, prob_unit = rep(probs, c(4, 10, 1))),
+    rep(probs, c(4, 10, 1))
+  )
+  set.seed(343)
+  S <- replicate(10000, strata_rs(strata = strata, strata_prob = probs))
+  realized <- tapply(rowMeans(S), strata, mean)
+  expect_equal(as.numeric(realized), probs, tolerance = 0.02)
+  expect_setequal(unique(colSums(S[1:4, ])), c(3, 4))
+})