Feature/relationship matrix

.Rbuildignore

@@ -27,3 +27,4 @@ codemeta.json
 \.[Tt]xt$
 CLAUDE.md
 .claude
+Rplots.pdf

.gitignore

@@ -11,3 +11,4 @@ inst/doc
 /Meta/
 
 Rplots.pdf
+*.local.json

NAMESPACE

@@ -20,6 +20,7 @@ export(objective_function_piepho)
 export(objective_function_signature)
 export(optim_params)
 export(plot_progress)
+export(prep_relationship)
 export(speed)
 importFrom(farver,decode_colour)
 importFrom(ggplot2,aes)

R/calculate_adjacency_score.R

@@ -42,13 +42,77 @@ shift_pad <- function(m, dx, dy, fill = NA) {
   nr <- nrow(m)
   nc <- ncol(m)
   out <- matrix(fill, nr, nc)
+  # entire source falls off the grid when the shift meets or exceeds the
+  # corresponding dimension; `seq.int(1, nr - dy)` would otherwise produce
+  # a descending sequence and out-of-bounds indices
+  if (abs(dx) >= nc || abs(dy) >= nr) {
+    return(out)
+  }
 
   r_src <- if (dy >= 0) seq.int(1, nr - dy) else seq.int(1 - dy, nr)
   c_src <- if (dx >= 0) seq.int(1, nc - dx) else seq.int(1 - dx, nc)
   out[r_src + dy, c_src + dx] <- m[r_src, c_src]
   return(out)
 }
 
+#' Validate and Flatten a Relationship Matrix for Lookup
+#'
+#' @description
+#' Validates `relationship` and packs it into a flat numeric vector plus
+#' row/column level vectors that [adjacency_score_vec()] can index into in
+#' constant time per cell. [speed()] calls this once before the SA loop so
+#' the per-iteration cost stays in the lookup itself. Direct callers of
+#' [calculate_adjacency_score()] must wrap their matrix with this function
+#' before passing - the score functions consume the prepped form only.
+#'
+#' @param relationship A numeric matrix with rownames and colnames covering
+#'   every treatment value to be scored.
+#' @param treatments Optional character vector of treatment values that
+#'   `relationship` must cover. Skipped when `NULL`.
+#'
+#' @return A list with `flat` (column-major flattened matrix), `row_levels`,
+#'   `col_levels`, and `n_row`.
+#'
+#' @export
+prep_relationship <- function(relationship, treatments = NULL) {
+  if (!is.matrix(relationship) || !is.numeric(relationship)) {
+    stop("`relationship` must be a numeric matrix.", call. = FALSE)
+  }
+
+  if (is.null(rownames(relationship)) || is.null(colnames(relationship))) {
+    stop(
+      "`relationship` must have rownames and colnames covering all ",
+      "treatments to be scored.",
+      call. = FALSE
+    )
+  }
+
+  if (!is.null(treatments)) {
+    treatments <- as.character(treatments)
+    treatments <- unique(treatments[!is.na(treatments)])
+    missing <- unique(c(
+      setdiff(treatments, rownames(relationship)),
+      setdiff(treatments, colnames(relationship))
+    ))
+
+    if (length(missing)) {
+      stop(
+        "`relationship` is missing entries for treatments: ",
+        paste(missing, collapse = ", "),
+        ".",
+        call. = FALSE
+      )
+    }
+  }
+
+  return(list(
+    flat = as.vector(relationship),
+    row_levels = rownames(relationship),
+    col_levels = colnames(relationship),
+    n_row = nrow(relationship)
+  ))
+}
+
 #' Per-Cell Weighted Adjacency Score
 #'
 #' @description
@@ -68,7 +132,8 @@ shift_pad <- function(m, dx, dy, fill = NA) {
 adjacency_score_vec <- function(design_matrix,
                                 dists = c(1, 2),
                                 weights = c(1, 2),
-                                ring_type = c("manhattan", "chebyshev")) {
+                                ring_type = c("manhattan", "chebyshev"),
+                                relationship = NULL) {
   ring_type <- match.arg(ring_type)
   stopifnot(length(dists) == length(weights))
   nr <- nrow(design_matrix)
@@ -91,10 +156,19 @@ adjacency_score_vec <- function(design_matrix,
   # force 3D even when n_offsets == 1, where simplify2array returns a matrix
   dim(shifted_stack) <- c(nr, nc, n_offsets)
 
-  matches <- sweep(shifted_stack, c(1, 2), design_matrix, FUN = "==")
-  matches[is.na(matches)] <- FALSE
+  if (is.null(relationship)) {
+    pair_values <- sweep(shifted_stack, c(1, 2), design_matrix, FUN = "==")
+    pair_values[is.na(pair_values)] <- FALSE
+  } else {
+    orig_idx <- match(as.character(design_matrix), relationship$row_levels)
+    shifted_idx <- match(as.character(shifted_stack), relationship$col_levels)
+    flat_idx <- rep(orig_idx, n_offsets) + relationship$n_row * (shifted_idx - 1L)
+    pair_values <- relationship$flat[flat_idx]
+    pair_values[is.na(pair_values)] <- 0
+    dim(pair_values) <- c(nr, nc, n_offsets)
+  }
 
-  weighted <- sweep(matches, 3, offset_weights, FUN = "*")
+  weighted <- sweep(pair_values, 3, offset_weights, FUN = "*")
   score <- rowSums(weighted, dims = 2)
   return(score)
 }
@@ -109,6 +183,9 @@ adjacency_score_vec <- function(design_matrix,
 #'
 #' Internally this is a thin wrapper around [adjacency_score_vec()].
 #'
+#' Pass a `relationship` matrix to score neighbour pairs by a graded
+#' similarity (e.g. genetic relatedness) instead of a strict identity match.
+#'
 #' @param layout_df A data frame containing the design.
 #' @param swap Column name of the treatments to be scored.
 #' @param row_column Name of the column representing rows (default `"row"`).
@@ -120,6 +197,13 @@ adjacency_score_vec <- function(design_matrix,
 #'   `1`).
 #' @param ring_type Ring shape: `"manhattan"` (default; diamond ring)
 #'   or `"chebyshev"` (square ring). See [ring_offsets()].
+#' @param relationship Optional pairwise-relationship lookup produced by
+#'   [prep_relationship()]. When supplied, each neighbour pair contributes
+#'   `relationship[cell, neighbour]` rather than `1` for matches and `0`
+#'   otherwise. NA-padded cells off the design edge contribute `0`. Defaults
+#'   to `NULL`, which keeps the strict identity match. Pass the raw matrix
+#'   through `prep_relationship()` first; the score functions consume only
+#'   the prepped form.
 #'
 #' @return A non-negative numeric value: the number of like-treatment edges
 #'   in the row/column adjacency graph.
@@ -141,6 +225,15 @@ adjacency_score_vec <- function(design_matrix,
 #' )
 #' calculate_adjacency_score(design_with_adj, "treatment") # 6
 #'
+#' # Example 3: graded relationship between A and B
+#' rel <- prep_relationship(matrix(
+#'   c(1, 0.3, 0.3, 1),
+#'   nrow = 2,
+#'   dimnames = list(c("A", "B"), c("A", "B"))
+#' ))
+#' calculate_adjacency_score(design_no_adj, "treatment", relationship = rel)
+#' # 3.6: each of the 12 A-B edges contributes 0.3
+#'
 #' @seealso [adjacency_score_vec()]
 #'
 #' @export
@@ -150,8 +243,10 @@ calculate_adjacency_score <- function(layout_df,
                                       col_column = "col",
                                       ring_dists = 1,
                                       ring_weights = 1,
-                                      ring_type = c("manhattan", "chebyshev")) {
+                                      ring_type = c("manhattan", "chebyshev"),
+                                      relationship = NULL) {
   ring_type <- match.arg(ring_type)
+
   design_matrix <- matrix(
     layout_df[[swap]],
     nrow = max(as_numeric_factor(layout_df[[row_column]]), na.rm = TRUE),
@@ -163,7 +258,8 @@ calculate_adjacency_score <- function(layout_df,
     design_matrix,
     dists = ring_dists,
     weights = ring_weights,
-    ring_type = ring_type
+    ring_type = ring_type,
+    relationship = relationship
   )
   return(sum(per_cell) / 2)
 }

R/metrics.R

@@ -58,7 +58,7 @@ objective_function <- function(layout_df,
   ring_args <- list(...)
   ring_args <- ring_args[intersect(
     names(ring_args),
-    c("ring_dists", "ring_weights", "ring_type")
+    c("ring_dists", "ring_weights", "ring_type", "relationship")
   )]
   adj_score <- ifelse(adj_weight != 0,
     do.call(

R/speed.R

@@ -208,8 +208,15 @@ speed <- function(data,
     }
   }
 
-  design <- speed_hierarchical(data, optimise, quiet, seed, row_column = row_column,
-                               col_column = col_column, ...)
+  dots <- list(...)
+  .reject_optim_params_in_dots(dots)
+  dots <- .prep_dots(dots, optimise, data)
+
+  design <- do.call(speed_hierarchical, c(
+    list(data = data, optimise = optimise, quiet = quiet, seed = seed,
+         row_column = row_column, col_column = col_column),
+    dots
+  ))
   design$design_df[[dummy_group]] <- NULL
   design$design_df <- to_types(design$design_df, factored$input_types)
 
@@ -288,7 +295,7 @@ speed_hierarchical <- function(data, optimise, quiet, seed, ...) {
       # Calculate new score
       new_score_obj <- opt$obj_function(new_design$design,opt$swap, spatial_cols, adj_weight = adj_weight,
                                         bal_weight = bal_weight, current_score_obj = current_score_obj,
-                                        swapped_items = new_design$swapped_items,...)
+                                        swapped_items = new_design$swapped_items, ...)
       new_score <- new_score_obj$score
 
       # Decide whether to accept the new design
@@ -424,3 +431,39 @@ print.design <- function(x, ...) {
 
   return(invisible(x))
 }
+
+#' Reject `...` arguments that must travel through [optim_params()].
+#'
+#' Stops with a message pointing the user at `optimise_params = optim_params(...)`
+#' when any of the listed names is found in `dots`.
+#'
+#' @param dots A named list captured from `...`.
+#' @return `NULL`, invisibly. Called for its side effect.
+#' @keywords internal
+.reject_optim_params_in_dots <- function(dots) {
+  forbidden <- intersect(names(dots), c("adj_weight", "bal_weight"))
+  if (length(forbidden) == 0) return(invisible(NULL))
+  stop(
+    "Argument(s) ", paste(sprintf("`%s`", forbidden), collapse = ", "),
+    " must be passed via `optim_params()`, not directly to `speed()`. ",
+    "For example: `optimise_params = optim_params(",
+    paste0(forbidden[1], " = ..."), ")`.",
+    call. = FALSE
+  )
+}
+
+#' Prep `dots$relationship` once with the union of treatments seen at
+#' every swap level.
+#'
+#' @param dots A named list captured from `...`.
+#' @param optimise Per-level `optimise` list as built by [create_speed_input()].
+#' @param data The (factor-converted) design data frame.
+#' @return `dots`, with `relationship` replaced by the prepped form when present.
+#' @keywords internal
+.prep_dots <- function(dots, optimise, data) {
+  if (is.null(dots$relationship)) return(dots)
+  swap_cols <- unique(vapply(optimise, function(o) o$swap, character(1)))
+  treatments <- unlist(lapply(swap_cols, function(s) as.character(data[[s]])))
+  dots$relationship <- prep_relationship(dots$relationship, treatments)
+  dots
+}

_pkgdown.yml

@@ -30,6 +30,7 @@ reference:
       - create_pair_mapping
       - initialise_design_df
       - add_buffers
+      - prep_relationship
   - title: Advanced options
   - contents:
       - optim_params