Feature/optimize swapping

R/design_utils.R

@@ -38,6 +38,22 @@ generate_neighbour <- function(design,
   }
 }
 
+#' Plot Rows in Each Group
+#'
+#' @description
+#' Rows with a treatment, split by group and fixed for a whole level.
+#'
+#' @inheritParams generate_neighbour
+#'
+#' @return A named list of row positions, one element per level of `swap_within`.
+#'
+#' @keywords internal
+plots_by_group <- function(design, swap, swap_within) {
+  groups <- design[[swap_within]]
+  plots <- which(!is.na(groups) & !is.na(design[[swap]]))
+  return(split(plots, groups[plots]))
+}
+
 #' Exchange Linked Columns Between Two Sets of Plots
 #'
 #' @description
@@ -84,14 +100,12 @@ exchange_linked <- function(design, linked_cols, plots_1, plots_2) {
 #'
 #' @keywords internal
 swappable_groups <- function(design, swap, swap_within, swap_all) {
-  groups <- design[[swap_within]]
-  treatments <- as.character(design[[swap]])
-  keep <- !is.na(groups) & !is.na(treatments)
-
   # One pass over the design rather than a scan per group. Splitting on the
   # factor keeps every level, so a level the data no longer uses arrives empty
   # and falls out below rather than being dropped silently.
-  by_group <- split(treatments[keep], groups[keep])
+  all_plots <- plots_by_group(design, swap, swap_within)
+  treatments <- as.character(design[[swap]])
+  by_group <- lapply(all_plots, function(plots) return(treatments[plots]))
   counts <- lapply(by_group, function(x) return(as.integer(table(x))))
 
   # Two distinct treatments are the minimum for any exchange, which also rules
@@ -105,7 +119,7 @@ swappable_groups <- function(design, swap, swap_within, swap_all) {
     !vapply(counts, function(x) return(any(duplicated(x))), logical(1))
 
   return(list(
-    swappable = names(by_group)[exchangeable & !unequal],
+    swappable = all_plots[exchangeable & !unequal],
     unequal_replication = names(by_group)[unequal]
   ))
 }
@@ -157,103 +171,81 @@ swappable_groups <- function(design, swap, swap_within, swap_all) {
 # fmt: skip
 generate_single_swap_neighbour <- function(design, swap, swap_within, swap_count, swap_all_blocks,
                                            linked_cols = NULL, swappable = NULL) {
-  new_design <- design
-
   # Only groups a swap can be proposed in, so no iteration is spent on one that
   # cannot move. Settled once per level by `swappable_groups()`.
-  all_blocks <- design[[swap_within]]
-  blocks <- swappable %||% levels(all_blocks)
-
-  if (swap_all_blocks) {
-    # Swap in all blocks
-    blocks_to_swap <- blocks
-  } else {
-    # Pick a random block
-    blocks_to_swap <- sample(blocks, 1)
+  swappable <- swappable %||% plots_by_group(design, swap, swap_within)
+  if (!swap_all_blocks) {
+    swappable <- swappable[sample.int(length(swappable), 1)]
   }
 
-  swapped_idx <- 1
-  swapped_items <- character(2 * swap_count * length(blocks_to_swap))
-
-  # Perform swaps in selected blocks
-  for (block in blocks_to_swap) {
-    # Get indices of plots in this block
-    block_indices <- which(all_blocks == block & !is.na(all_blocks) & !is.na(design[[swap]]))
-
-    if (length(block_indices) >= 2) {
-      # Need at least 2 plots to swap
-      for (i in 1:swap_count) {
-        # Select two random plots in this block
-        swap_pair <- sample(block_indices, 2)
-        # `as.character()` because the swap column is a factor, whose values
-        # would otherwise reach `swapped_items` as their integer codes
-        to_be_swapped <- as.character(new_design[[swap]][swap_pair])
-
-        # If both plots have the same treatment, try to find a different one
-        if (to_be_swapped[1] == to_be_swapped[2]) {
-          different_indices <- block_indices[new_design[[swap]][block_indices] != to_be_swapped[1]]
-
-          # Only proceed with swap if different treatments are available
-          if (length(different_indices) > 0) {
-            swap_pair[[2]] <- sample(different_indices, 1)
-            to_be_swapped[2] <- as.character(new_design[[swap]][[swap_pair[[2]]]])
-          } else {
-            # Skip this swap - no different treatments available
-            to_be_swapped <- NULL
-          }
-        }
+  # Swapped as integer codes, sparing the label matching of factor `[<-` and `==`
+  items <- design[[swap]]
+  codes <- unclass(items)
+  swapped <- integer(2 * swap_count * length(swappable))
+  n_swapped <- 0L
 
-        # Perform the swap only if we have valid treatments to swap
-        if (!is.null(to_be_swapped)) {
-          new_design[[swap]][rev(swap_pair)] <- to_be_swapped
-          new_design <- exchange_linked(new_design, linked_cols, swap_pair[1], swap_pair[2])
-          swapped_items[swapped_idx:(swapped_idx + 1)] <- to_be_swapped
-          swapped_idx <- swapped_idx + 2
+  for (plots in swappable) {
+    # Need at least 2 plots to swap
+    if (length(plots) < 2) {
+      next
+    }
+
+    for (i in seq_len(swap_count)) {
+      pair <- plots[sample.int(length(plots), 2)]
+      code_1 <- codes[pair[1]]
+      code_2 <- codes[pair[2]]
+
+      # If both plots have the same treatment, try to find a different one
+      if (code_1 == code_2) {
+        different <- plots[codes[plots] != code_1]
+        if (length(different) == 0) {
+          next
         }
+
+        pair[2] <- different[sample.int(length(different), 1)]
+        code_2 <- codes[pair[2]]
       }
+
+      codes[pair] <- c(code_2, code_1)
+      design <- exchange_linked(design, linked_cols, pair[1], pair[2])
+      swapped[n_swapped + 1:2] <- c(code_1, code_2)
+      n_swapped <- n_swapped + 2L
     }
   }
 
-  return(list(design = new_design, swapped_items = swapped_items))
+  attributes(codes) <- attributes(items)
+  design[[swap]] <- codes
+
+  return(list(design = design, swapped_items = levels(items)[swapped[seq_len(n_swapped)]]))
 }
 
 #' Generate neighbour for sequential or hierarchical designs
 #' @keywords internal
 # fmt: skip
 generate_multi_swap_neighbour <- function(design, swap, swap_within, swap_count, swap_all_blocks,
                                           linked_cols = NULL, swappable = NULL) {
-  new_design <- design
-
   # Settled once per level by `swappable_groups()`
-  groups <- swappable %||% levels(design[[swap_within]])
-
-  if (swap_all_blocks) {
-    # Swap in all groups
-    groups_to_swap <- groups
-  } else {
-    # Pick a random group
-    groups_to_swap <- sample(groups, 1)
+  swappable <- swappable %||% plots_by_group(design, swap, swap_within)
+  if (!swap_all_blocks) {
+    swappable <- swappable[sample.int(length(swappable), 1)]
   }
 
-  swapped_idx <- 1
-  swapped_items <- character(2 * swap_count * length(groups_to_swap))
+  # Swapped as integer codes, sparing the label matching of factor `[<-` and `==`
+  items <- design[[swap]]
+  codes <- unclass(items)
+  swapped <- integer(2 * swap_count * length(swappable))
+  n_swapped <- 0L
 
-  # Perform swaps in selected groups
-  for (group in groups_to_swap) {
-    # Get unique treatments within this group
-    group_filter <- new_design[[swap_within]] == group & !is.na(new_design[[swap_within]])
-    group_data <- new_design[group_filter & !is.na(new_design[[swap]]), ]
-    # `as.character()` so everything taken from it - `eligible`, `swap_pair`,
-    # and in turn `swapped_items` - is a treatment label rather than the integer
-    # code a factor contributes to a character vector
-    group_treatments <- as.character(unique(group_data[[swap]]))
+  for (plots in swappable) {
+    group_codes <- codes[plots]
+    group_treatments <- unique(group_codes)
 
     # Counted once: every swap below exchanges equally replicated treatments, so
     # these counts are unaffected by them
-    group_counts <- tabulate(match(group_data[[swap]], group_treatments), length(group_treatments))
+    group_counts <- tabulate(match(group_codes, group_treatments), length(group_treatments))
 
-    if (nrow(group_data) >= 2) {
-      for (i in 1:swap_count) {
+    if (length(plots) >= 2) {
+      for (i in seq_len(swap_count)) {
         # Only proceed if there are at least 2 different treatments
         if (length(group_treatments) < 2) {
           # Skip this swap - only one treatment in this group
@@ -283,27 +275,28 @@ generate_multi_swap_neighbour <- function(design, swap, swap_within, swap_count,
           eligible <- group_treatments[group_counts == as.integer(chosen)]
         }
 
-        # Select two different treatments
-        # Use sample with replace=FALSE to ensure they're different
-        swap_pair <- sample(eligible, 2, replace = FALSE)
+        # Two different treatments; `eligible` always holds at least two here
+        swap_pair <- eligible[sample.int(length(eligible), 2)]
 
-        # Find all plots with these treatments in this group
-        plots_1 <- which(group_filter & new_design[[swap]] == swap_pair[1])
-        plots_2 <- which(group_filter & new_design[[swap]] == swap_pair[2])
+        # Every plot of each treatment in this group
+        group_codes <- codes[plots]
+        plots_1 <- plots[group_codes == swap_pair[1]]
+        plots_2 <- plots[group_codes == swap_pair[2]]
 
-        # Swap all instances of these treatments
-        new_design[[swap]][plots_1] <- swap_pair[2]
-        new_design[[swap]][plots_2] <- swap_pair[1]
-        new_design <- exchange_linked(new_design, linked_cols, plots_1, plots_2)
+        codes[plots_1] <- swap_pair[2]
+        codes[plots_2] <- swap_pair[1]
+        design <- exchange_linked(design, linked_cols, plots_1, plots_2)
 
-        swapped_items[swapped_idx] <- swap_pair[1]
-        swapped_items[swapped_idx + 1] <- swap_pair[2]
-        swapped_idx <- swapped_idx + 2
+        swapped[n_swapped + 1:2] <- swap_pair
+        n_swapped <- n_swapped + 2L
       }
     }
   }
 
-  return(list(design = new_design, swapped_items = swapped_items[1:(swapped_idx - 1)]))
+  attributes(codes) <- attributes(items)
+  design[[swap]] <- codes
+
+  return(list(design = design, swapped_items = levels(items)[swapped[seq_len(n_swapped)]]))
 }
 
 #' Infer 'row' and 'col' with Patterns

tests/testthat/test-design_utils.R

@@ -136,13 +136,13 @@ test_that("swappable_groups separates unequal replication from other blockers",
   )
 
   all_swap <- swappable_groups(design, "treatment", "block", swap_all = TRUE)
-  expect_equal(all_swap$swappable, "g2")
+  expect_named(all_swap$swappable, "g2")
   expect_equal(all_swap$unequal_replication, "g1")
 
   # Without `swap_all` a single pair of plots moves, so replication is irrelevant
   # and only the single-treatment group is stuck
   single <- swappable_groups(design, "treatment", "block", swap_all = FALSE)
-  expect_equal(single$swappable, c("g1", "g2"))
+  expect_identical(single$swappable, list(g1 = 1:6, g2 = 7:12))
   expect_length(single$unequal_replication, 0)
 })
 
@@ -154,10 +154,32 @@ test_that("swappable_groups counts a level with no plots as unswappable", {
   )
 
   result <- swappable_groups(design, "treatment", "site", swap_all = TRUE)
-  expect_equal(result$swappable, c("a", "b"))
+  expect_named(result$swappable, c("a", "b"))
   expect_length(result$unequal_replication, 0)
 })
 
+test_that("plots_by_group gives each group's rows, leaving out unswappable plots", {
+  design <- data.frame(
+    block = factor(c("b2", "b1", NA, "b2", "b1", "b1")),
+    treatment = factor(c("A", "B", "A", NA, "A", "C"))
+  )
+
+  expect_identical(
+    plots_by_group(design, "treatment", "block"),
+    list(b1 = c(2L, 5L, 6L), b2 = 1L)
+  )
+})
+
+test_that("plots_by_group keeps a level with no plots as an empty group", {
+  design <- data.frame(
+    site = factor(c("a", "a", "b"), levels = c("a", "b", "c")),
+    treatment = factor(c("A", "B", "A"))
+  )
+
+  result <- plots_by_group(design, "treatment", "site")
+  expect_named(result, c("a", "b", "c"))
+  expect_identical(result$c, integer(0))
+})
 test_that("speed() warns when a swap group is left frozen mid-search", {
   # A level 1 swap can leave a site with no two treatments sharing a replication
   # count. The search cannot move anything there, which should be reported rather