From 158b2e33c259a658afd02be4cd1ac9962e1635e8 Mon Sep 17 00:00:00 2001
From: Ben Schneider <benjamin.julius.schneider@gmail.com>
Date: Wed, 11 Oct 2023 12:08:02 -0400
Subject: [PATCH] Bug fix for random-groups jackknife with variance strata

---
 R/as_jackknife_design.R                      | 18 +++++++++++-------
 tests/testthat/test-random-group-jackknife.R | 20 +++++++++++++++++++-
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/R/as_jackknife_design.R b/R/as_jackknife_design.R
index c1c9efa..00c7ddd 100644
--- a/R/as_jackknife_design.R
+++ b/R/as_jackknife_design.R
@@ -229,18 +229,14 @@ as_random_group_jackknife_design.survey.design <- function(
     stop("`scale_method` must be either 'variance-stratum-psus' or 'variance-units'")
   }
 
-  # Extract necessary design information
+  # Begin extracting necessary design information
   n <- nrow(design)
-
   design_vars <- data.frame(
     'ROW_ID' = seq_len(n),
-    'STRATUM' = design$strata[,1,drop=TRUE],
-    'PSU' = interaction(design$strata[,1,drop=TRUE],
-                        design$cluster[,1,drop=TRUE],
-                        drop = TRUE) |> as.numeric(),
-    stringsAsFactors = FALSE
+    'STRATUM' = design$strata[,1,drop=TRUE]
   )
 
+  # Handle VAR_STRAT variable
   if (!is.null(var_strat)) {
     if (!(var_strat %in% colnames(design$variables))) {
       stop("`var_strat` must be either NULL or the name of a variable in the data.")
@@ -254,6 +250,13 @@ as_random_group_jackknife_design.survey.design <- function(
     design_vars[['VAR_STRAT']] <- 1
   }
 
+  # Create unique PSU IDs within
+  # combinations of variance strata and design strata
+  design_vars[['PSU']] <- interaction(design_vars[['VAR_STRAT']],
+                                      design$strata[,1,drop=TRUE],
+                                      design$cluster[,1,drop=TRUE],
+                                      drop = TRUE) |> as.numeric()
+
   # Warn the user about FPCs and check that `var_strat_frac` is valid
 
   if (!is.null(design$fpc$popsize)) {
@@ -308,6 +311,7 @@ as_random_group_jackknife_design.survey.design <- function(
   n_psus <- design_vars$PSU |> unique() |> length()
   psu_random_labels <- sample(x = n_psus, size = n_psus, replace = FALSE)
   design_vars[['RAND_PSU_ID']] <- interaction(
+    design_vars[['VAR_STRAT']],
     design_vars[['STRATUM']],
     psu_random_labels[design_vars[['PSU']]],
     drop = TRUE, lex.order = TRUE
diff --git a/tests/testthat/test-random-group-jackknife.R b/tests/testthat/test-random-group-jackknife.R
index 57644e6..f1a6a80 100644
--- a/tests/testthat/test-random-group-jackknife.R
+++ b/tests/testthat/test-random-group-jackknife.R
@@ -31,7 +31,7 @@ set.seed(2014)
 
 # Test for expected formation of random groups ----
 
-  test_that(desc = "Random groups formed correctly", {
+  test_that(desc = "Random groups (without var_strat) formed correctly", {
 
     ## Correct sizes for random groups
     jk_design <- as_random_group_jackknife_design(
@@ -62,6 +62,24 @@ set.seed(2014)
 
   })
 
+  test_that(desc = "Random groups (with var_strat) formed correctly", {
+
+    ## Correct sizes for random groups
+    jk_design <- as_random_group_jackknife_design(
+      design = dstrat_nofpc,
+      var_strat = "stype",
+      replicates = 5
+    )
+
+    expect_equal(
+      object = jk_design$variables |>
+        count(stype, .random_group) |>
+        pull("n"),
+      expected = c(rep(c(20, 10, 10), each = 5))
+    )
+
+  })
+
 # Given the random groups, replicates formed correctly ----
 
   test_that(desc = "Replicates formed correctly, given random groups", {