From 58befdab1119fb30c5f194e1dcdf574163167b6a Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Thu, 13 Aug 2026 15:04:39 -0400
Subject: [PATCH 01/10] Base function modified from Aaron's script

---
 NAMESPACE                              |   1 +
 NEWS.md                                |   4 +
 R/util_k_marginals.R                   | 142 ++++++++++++++++++++
 man/dot-util_k_marginals.Rd            |  29 +++++
 man/util_k_marginals.Rd                |  24 ++++
 tests/testthat/test-util_k_marginals.R | 171 +++++++++++++++++++++++++
 6 files changed, 371 insertions(+)
 create mode 100644 R/util_k_marginals.R
 create mode 100644 man/dot-util_k_marginals.Rd
 create mode 100644 man/util_k_marginals.Rd
 create mode 100644 tests/testthat/test-util_k_marginals.R

diff --git a/NAMESPACE b/NAMESPACE
index 2067fa9..2251a8a 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -26,6 +26,7 @@ export(prep_discrete_eval_data)
 export(util_ci_overlap)
 export(util_co_occurrence)
 export(util_corr_fit)
+export(util_k_marginals)
 export(util_ks_distance)
 export(util_moments)
 export(util_percentiles)
diff --git a/NEWS.md b/NEWS.md
index d9262f1..1a321c5 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,7 @@
+# syntheval 0.1.0
+
+* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals.
+
 # syntheval 0.0.5
 
 * Remove `util_tails()`
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
new file mode 100644
index 0000000..16f77a1
--- /dev/null
+++ b/R/util_k_marginals.R
@@ -0,0 +1,142 @@
+#' @title Worker function for the k-marginals metric
+#'
+#' @description This worker function takes a specified k-marginal and calculates
+#' all unique k-combinations of variables shared by the input data. For each
+#' combination, tibbles containing unique combinations of observed value levels
+#' are created with the marginal probabilities for each cell for both
+#' synthetic and confidential data included. The mean absolute
+#' difference (MabsDD) between the synthetic and confidential data's marginal
+#' probabilities is computed for each k-combination. The MabsDDs are
+#' averaged into a scalar before being rescaled (1 - mean) * 1000.
+#'
+#' @param synth_data A tibble with synthetic data.
+#' @param conf_data A tibble with confidential data.
+#' @param k Scalar order of the k-marginal (valid range = 1:3).
+#'
+#' @return A value in range [0, 1000] where a higher value denotes lower MabsDDs
+#' and consequently greater similarity between confidential and synthetic data.
+#'
+.util_k_marginals <- function(synth_data, conf_data, k) {
+
+  if (!(is.numeric(k) && length(k) == 1 && k %in% 1:3)) {
+
+    stop("`k` must be a single integer between 1 and 3")
+
+  }
+
+  stopifnot(inherits(synth_data, "data.frame"))
+  stopifnot(inherits(conf_data, "data.frame"))
+
+  if (nrow(synth_data) == 0 || nrow(conf_data) == 0) {
+
+    stop("`synth_data` and `conf_data` must each contain at least one row")
+
+  }
+
+  # only variables present in both datasets contribute marginals
+  shared_vars <- intersect(names(synth_data), names(conf_data))
+
+  if (length(shared_vars) < k) {
+
+    stop("`k` cannot exceed the number of variables shared by both datasets")
+
+  }
+
+  kmarginals_vars <- t(utils::combn(x = shared_vars, m = k))
+
+  # cell proportions for one dataset over one set of variables
+  process_data <- function(data, vars, prop_name) {
+
+    props <- data |>
+      dplyr::select(dplyr::all_of(vars)) |>
+      dplyr::group_by_all() |>
+      dplyr::count() |>
+      dplyr::ungroup() |>
+      dplyr::mutate("{prop_name}" := .data$n / sum(.data$n)) |>
+      dplyr::select(-"n")
+
+    return(props)
+
+  }
+  
+  # MabsDD for one set of variables; cells absent from one dataset count as 0
+  madd <- function(vars) {
+
+    combined_data <- dplyr::full_join(
+      process_data(data = synth_data, vars = vars, prop_name = "prop_synth"),
+      process_data(data = conf_data, vars = vars, prop_name = "prop_conf"),
+      by = vars
+    ) |>
+      tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0))
+
+    madd <- combined_data |>
+      dplyr::summarize(
+        madd = mean(abs(.data$prop_synth - .data$prop_conf))
+      ) |>
+      dplyr::pull("madd")
+
+    return(madd)
+
+  }
+
+  # iterate over all k-way marginals
+  madds <- purrr::map_dbl(
+    .x = seq_len(nrow(kmarginals_vars)),
+    .f = \(i) madd(vars = kmarginals_vars[i, ])
+  )
+
+  # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000]
+  return((1 - mean(madds)) * 1000)
+
+}
+
+#' @title Calculate the k-marginals metric
+#'
+#' @description For each unique k-combination of variables shared by the
+#' synthetic and confidential data, the mean absolute difference (MabsDD)
+#' between the two datasets' marginal cell probabilities is computed. The
+#' MabsDDs are averaged across combinations and rescaled as (1 - mean) * 1000.
+#'
+#' @param eval_data An `eval_data` object.
+#' @param k Scalar order of the k-marginal (valid range = 1:3).
+#'
+#' @return A value in range [0, 1000] where a higher value denotes lower MabsDDs
+#' and consequently greater similarity between confidential and synthetic data.
+#' For multiple replicates, a list of such values, one per replicate.
+#'
+#' @export
+#'
+util_k_marginals <- function(eval_data, k = 3) {
+
+  stopifnot(is_eval_data(eval_data))
+
+  if (eval_data$n_rep == 1) {
+
+    return(
+      .util_k_marginals(
+        synth_data = eval_data$synth_data,
+        conf_data = eval_data$conf_data,
+        k = k
+      )
+    )
+
+  } else {
+
+    result <- purrr::map(
+      .x = eval_data$synth_data,
+      .f = \(sd) {
+
+        .util_k_marginals(
+          synth_data = sd,
+          conf_data = eval_data$conf_data,
+          k = k
+        )
+
+      }
+    )
+
+    return(result)
+
+  }
+
+}
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
new file mode 100644
index 0000000..1fed0ec
--- /dev/null
+++ b/man/dot-util_k_marginals.Rd
@@ -0,0 +1,29 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/util_k_marginals.R
+\name{.util_k_marginals}
+\alias{.util_k_marginals}
+\title{Worker function for the k-marginals metric}
+\usage{
+.util_k_marginals(synth_data, conf_data, k)
+}
+\arguments{
+\item{synth_data}{A tibble with synthetic data.}
+
+\item{conf_data}{A tibble with confidential data.}
+
+\item{k}{Scalar order of the k-marginal (valid range = 1:3).}
+}
+\value{
+A value in range \link{0, 1000} where a higher value denotes lower MabsDDs
+and consequently greater similarity between confidential and synthetic data.
+}
+\description{
+This worker function takes a specified k-marginal and calculates
+all unique k-combinations of variables shared by the input data. For each
+combination, tibbles containing unique combinations of observed value levels
+are created with the marginal probabilities for each cell for both
+synthetic and confidential data included. The mean absolute
+difference (MabsDD) between the synthetic and confidential data's marginal
+probabilities is computed for each k-combination. The MabsDDs are
+averaged into a scalar before being rescaled (1 - mean) * 1000.
+}
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
new file mode 100644
index 0000000..dc71a2f
--- /dev/null
+++ b/man/util_k_marginals.Rd
@@ -0,0 +1,24 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/util_k_marginals.R
+\name{util_k_marginals}
+\alias{util_k_marginals}
+\title{Calculate the k-marginals metric}
+\usage{
+util_k_marginals(eval_data, k = 3)
+}
+\arguments{
+\item{eval_data}{An \code{eval_data} object.}
+
+\item{k}{Scalar order of the k-marginal (valid range = 1:3).}
+}
+\value{
+A value in range \link{0, 1000} where a higher value denotes lower MabsDDs
+and consequently greater similarity between confidential and synthetic data.
+For multiple replicates, a list of such values, one per replicate.
+}
+\description{
+For each unique k-combination of variables shared by the
+synthetic and confidential data, the mean absolute difference (MabsDD)
+between the two datasets' marginal cell probabilities is computed. The
+MabsDDs are averaged across combinations and rescaled as (1 - mean) * 1000.
+}
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
new file mode 100644
index 0000000..f26cf45
--- /dev/null
+++ b/tests/testthat/test-util_k_marginals.R
@@ -0,0 +1,171 @@
+# hand-computable example
+#
+# conf_data          synth_data
+#   a  b               a  b
+#   x  p               x  p
+#   x  p               y  q
+#   y  p               y  q
+#   y  q               y  p
+#
+# k = 1
+#   marginal a: conf (x = 0.50, y = 0.50), synth (x = 0.25, y = 0.75)
+#     MabsDD = mean(|0.25 - 0.50|, |0.75 - 0.50|) = 0.25
+#   marginal b: conf (p = 0.75, q = 0.25), synth (p = 0.50, q = 0.50)
+#     MabsDD = mean(|0.50 - 0.75|, |0.50 - 0.25|) = 0.25
+#   score = (1 - mean(0.25, 0.25)) * 1000 = 750
+#
+# k = 2 (single combination: a x b)
+#   cells: conf (x,p = 0.50, y,p = 0.25, y,q = 0.25)
+#          synth (x,p = 0.25, y,p = 0.25, y,q = 0.50)
+#     per-cell |synth - conf|: (x,p) 0.25, (y,p) 0, (y,q) 0.25
+#     MabsDD = mean(0.25, 0, 0.25) = 1/6
+#   score = (1 - 1/6) * 1000 = 5000/6
+
+conf <- tibble::tibble(
+  a = c("x", "x", "y", "y"),
+  b = c("p", "p", "p", "q")
+)
+
+synth <- tibble::tibble(
+  a = c("x", "y", "y", "y"),
+  b = c("p", "q", "q", "p")
+)
+
+test_that("k = 1 score matches hand-computed value", {
+
+  expect_equal(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1),
+    750
+  )
+
+})
+
+test_that("k = 2 score matches hand-computed value", {
+
+  expect_equal(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 2),
+    5000 / 6
+  )
+
+})
+
+test_that("identical data scores exactly 1000 for every k", {
+
+  for (k in 1:2) {
+
+    expect_equal(
+      .util_k_marginals(synth_data = conf, conf_data = conf, k = k),
+      1000
+    )
+
+  }
+
+})
+
+test_that("cells absent from one dataset count as proportion zero", {
+
+  # conf has level y that synth lacks; synth is all x
+  # marginal a: conf (x = 0.5, y = 0.5), synth (x = 1, y = 0)
+  # MabsDD = mean(0.5, 0.5) = 0.5 -> score 500
+  conf_gap <- tibble::tibble(a = c("x", "y"))
+  synth_gap <- tibble::tibble(a = c("x", "x"))
+
+  expect_equal(
+    .util_k_marginals(synth_data = synth_gap, conf_data = conf_gap, k = 1),
+    500
+  )
+
+})
+
+test_that("k outside 1:3 throws an error", {
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 4),
+    regexp = "`k` must be a single integer between 1 and 3"
+  )
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 0),
+    regexp = "`k` must be a single integer between 1 and 3"
+  )
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = c(1, 2)),
+    regexp = "`k` must be a single integer between 1 and 3"
+  )
+
+  # %in% coerces, so non-numeric scalars need an explicit type check
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = TRUE),
+    regexp = "`k` must be a single integer between 1 and 3"
+  )
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = "1"),
+    regexp = "`k` must be a single integer between 1 and 3"
+  )
+
+})
+
+test_that("zero-row inputs throw an error instead of returning NaN", {
+
+  empty <- conf[0, ]
+
+  expect_error(
+    .util_k_marginals(synth_data = empty, conf_data = conf, k = 1),
+    regexp = "at least one row"
+  )
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = empty, k = 1),
+    regexp = "at least one row"
+  )
+
+})
+
+test_that("k exceeding the number of shared variables throws an error", {
+
+  expect_error(
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 3),
+    regexp = "shared by both datasets"
+  )
+
+})
+
+test_that("variables not shared by both datasets are ignored", {
+
+  # extra synth-only column must not create combinations
+  synth_extra <- dplyr::mutate(synth, c = c("m", "m", "n", "n"))
+
+  expect_equal(
+    .util_k_marginals(synth_data = synth_extra, conf_data = conf, k = 1),
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+  )
+
+})
+
+test_that("util_k_marginals accepts an eval_data object", {
+
+  ed <- eval_data(conf_data = conf, synth_data = synth)
+
+  expect_equal(util_k_marginals(eval_data = ed, k = 1), 750)
+
+  expect_equal(util_k_marginals(eval_data = ed, k = 2), 5000 / 6)
+
+})
+
+test_that("util_k_marginals maps over replicates", {
+
+  ed <- eval_data(conf_data = conf, synth_data = list(synth, conf))
+
+  result <- util_k_marginals(eval_data = ed, k = 1)
+
+  expect_equal(result, list(750, 1000))
+
+})
+
+test_that("util_k_marginals rejects non-eval_data input", {
+
+  expect_error(util_k_marginals(eval_data = synth, k = 1))
+
+})

From 493eda6c71ad2d5d6be78c54623682b2ef0d00fe Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Thu, 13 Aug 2026 15:45:03 -0400
Subject: [PATCH 02/10] Return worst marginals and cells from
 util_k_marginals()

---
 NAMESPACE                              |   1 +
 R/util_k_marginals.R                   | 143 +++++++++++---
 man/dot-util_k_marginals.Rd            |  25 ++-
 man/print.k_marginals.Rd               |  21 +++
 man/util_k_marginals.Rd                |  21 ++-
 tests/testthat/test-util_k_marginals.R | 248 ++++++++++++++++++++++++-
 6 files changed, 418 insertions(+), 41 deletions(-)
 create mode 100644 man/print.k_marginals.Rd

diff --git a/NAMESPACE b/NAMESPACE
index 2251a8a..9528f53 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -1,6 +1,7 @@
 # Generated by roxygen2: do not edit by hand
 
 S3method(print,eval_data)
+S3method(print,k_marginals)
 export("%>%")
 export(add_discriminator_auc)
 export(add_pmse)
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 16f77a1..2cac50f 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -12,11 +12,40 @@
 #' @param synth_data A tibble with synthetic data.
 #' @param conf_data A tibble with confidential data.
 #' @param k Scalar order of the k-marginal (valid range = 1:3).
+#' @param keep_marginals Single integer number of worst marginals to retain
+#' in the output.
+#' Defaults to `Inf`, which retains all of them.
+#' @param keep_cells Single integer number of worst cells to retain in the
+#' output. Defaults to `Inf`, which retains all of them.
 #'
-#' @return A value in range [0, 1000] where a higher value denotes lower MabsDDs
-#' and consequently greater similarity between confidential and synthetic data.
+#' @return A `k_marginals` object with three elements: `score`, a value in
+#' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
+#' greater similarity between confidential and synthetic data; `marginals`, a
+#' tibble with the MabsDD for each combination of variables, worst first; and
+#' `cells`, a tibble with the synthetic and confidential proportions and their
+#' absolute difference for every cell, worst first. `score` is always computed
+#' from all marginals, even when `keep_marginals` or `keep_cells` truncate the
+#' detail tables.
 #'
-.util_k_marginals <- function(synth_data, conf_data, k) {
+.util_k_marginals <- function(
+    synth_data,
+    conf_data,
+    k,
+    keep_marginals = Inf,
+    keep_cells = Inf) {
+
+  for (keep in list(keep_marginals, keep_cells)) {
+
+    if (!(is.numeric(keep) && length(keep) == 1 && !is.na(keep) &&
+          keep >= 1 && keep == floor(keep))) {
+
+      stop(
+        "`keep_marginals` and `keep_cells` must be single integers >= 1 or Inf"
+      )
+
+    }
+
+  }
 
   if (!(is.numeric(k) && length(k) == 1 && k %in% 1:3)) {
 
@@ -59,34 +88,79 @@
 
   }
   
-  # MabsDD for one set of variables; cells absent from one dataset count as 0
-  madd <- function(vars) {
+  # per-cell differences for one set of variables; cells absent from one
+  # dataset count as 0
+  marginal_cells <- function(vars) {
 
-    combined_data <- dplyr::full_join(
+    cells <- dplyr::full_join(
       process_data(data = synth_data, vars = vars, prop_name = "prop_synth"),
       process_data(data = conf_data, vars = vars, prop_name = "prop_conf"),
       by = vars
     ) |>
-      tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0))
-
-    madd <- combined_data |>
-      dplyr::summarize(
-        madd = mean(abs(.data$prop_synth - .data$prop_conf))
+      tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0)) |>
+      tidyr::unite(col = "cell", dplyr::all_of(vars), sep = ", ") |>
+      dplyr::mutate(
+        variables = paste(vars, collapse = ", "),
+        abs_diff = abs(.data$prop_synth - .data$prop_conf)
       ) |>
-      dplyr::pull("madd")
-
-    return(madd)
+      dplyr::select(
+        "variables", "cell", "prop_synth", "prop_conf", "abs_diff"
+      )
+    # variables disambiguates cells across combinations and drives the
+    # per-combination summary; the prop columns show the direction of the
+    # discrepancy, not just its size
+    return(cells)
 
   }
 
-  # iterate over all k-way marginals
-  madds <- purrr::map_dbl(
+  # per-cell differences across all k-way marginals, worst cells first
+  cells <- purrr::map(
     .x = seq_len(nrow(kmarginals_vars)),
-    .f = \(i) madd(vars = kmarginals_vars[i, ])
+    .f = \(i) marginal_cells(vars = kmarginals_vars[i, ])
+  ) |>
+    purrr::list_rbind() |>
+    dplyr::arrange(dplyr::desc(.data$abs_diff))
+
+  # MabsDD per combination, worst marginals first
+  marginals <- cells |>
+    dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables") |>
+    dplyr::arrange(dplyr::desc(.data$madd))
+
+  # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000];
+  # computed from all marginals before any truncation
+  score <- (1 - mean(marginals$madd)) * 1000
+
+  result <- structure(
+    list(
+      score = score,
+      marginals = utils::head(marginals, n = keep_marginals),
+      cells = utils::head(cells, n = keep_cells)
+    ),
+    class = "k_marginals"
   )
 
-  # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000]
-  return((1 - mean(madds)) * 1000)
+  return(result)
+
+}
+
+#' @title Print a k_marginals object
+#'
+#' @param x A `k_marginals` object from [util_k_marginals()].
+#' @param n Number of worst marginals to display.
+#' @param ... Additional arguments passed to methods (unused).
+#'
+#' @return `x`, invisibly.
+#'
+#' @export
+#'
+print.k_marginals <- function(x, n = 5, ...) {
+
+  cat("k-marginals score:", round(x$score, digits = 2), "\n\n")
+
+  cat("Worst marginals:\n")
+  print(utils::head(x$marginals, n = n))
+
+  return(invisible(x))
 
 }
 
@@ -99,14 +173,29 @@
 #'
 #' @param eval_data An `eval_data` object.
 #' @param k Scalar order of the k-marginal (valid range = 1:3).
+#' @param keep_marginals Single integer number of worst marginals to retain
+#' in the output.
+#' Defaults to `Inf`, which retains all of them.
+#' @param keep_cells Single integer number of worst cells to retain in the
+#' output. Defaults to `Inf`, which retains all of them.
 #'
-#' @return A value in range [0, 1000] where a higher value denotes lower MabsDDs
-#' and consequently greater similarity between confidential and synthetic data.
-#' For multiple replicates, a list of such values, one per replicate.
+#' @return A `k_marginals` object with three elements: `score`, a value in
+#' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
+#' greater similarity between confidential and synthetic data; `marginals`, a
+#' tibble with the MabsDD for each combination of variables, worst first; and
+#' `cells`, a tibble with the synthetic and confidential proportions and their
+#' absolute difference for every cell, worst first. `score` is always computed
+#' from all marginals, even when `keep_marginals` or `keep_cells` truncate the
+#' detail tables. For multiple replicates, a list of such objects, one per
+#' replicate.
 #'
 #' @export
 #'
-util_k_marginals <- function(eval_data, k = 3) {
+util_k_marginals <- function(
+    eval_data,
+    k = 3,
+    keep_marginals = Inf,
+    keep_cells = Inf) {
 
   stopifnot(is_eval_data(eval_data))
 
@@ -116,7 +205,9 @@ util_k_marginals <- function(eval_data, k = 3) {
       .util_k_marginals(
         synth_data = eval_data$synth_data,
         conf_data = eval_data$conf_data,
-        k = k
+        k = k,
+        keep_marginals = keep_marginals,
+        keep_cells = keep_cells
       )
     )
 
@@ -129,7 +220,9 @@ util_k_marginals <- function(eval_data, k = 3) {
         .util_k_marginals(
           synth_data = sd,
           conf_data = eval_data$conf_data,
-          k = k
+          k = k,
+          keep_marginals = keep_marginals,
+          keep_cells = keep_cells
         )
 
       }
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 1fed0ec..9c91a15 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -4,7 +4,13 @@
 \alias{.util_k_marginals}
 \title{Worker function for the k-marginals metric}
 \usage{
-.util_k_marginals(synth_data, conf_data, k)
+.util_k_marginals(
+  synth_data,
+  conf_data,
+  k,
+  keep_marginals = Inf,
+  keep_cells = Inf
+)
 }
 \arguments{
 \item{synth_data}{A tibble with synthetic data.}
@@ -12,10 +18,23 @@
 \item{conf_data}{A tibble with confidential data.}
 
 \item{k}{Scalar order of the k-marginal (valid range = 1:3).}
+
+\item{keep_marginals}{Single integer number of worst marginals to retain
+in the output.
+Defaults to \code{Inf}, which retains all of them.}
+
+\item{keep_cells}{Single integer number of worst cells to retain in the
+output. Defaults to \code{Inf}, which retains all of them.}
 }
 \value{
-A value in range \link{0, 1000} where a higher value denotes lower MabsDDs
-and consequently greater similarity between confidential and synthetic data.
+A \code{k_marginals} object with three elements: \code{score}, a value in
+range \link{0, 1000} where a higher value denotes lower MabsDDs and consequently
+greater similarity between confidential and synthetic data; \code{marginals}, a
+tibble with the MabsDD for each combination of variables, worst first; and
+\code{cells}, a tibble with the synthetic and confidential proportions and their
+absolute difference for every cell, worst first. \code{score} is always computed
+from all marginals, even when \code{keep_marginals} or \code{keep_cells} truncate the
+detail tables.
 }
 \description{
 This worker function takes a specified k-marginal and calculates
diff --git a/man/print.k_marginals.Rd b/man/print.k_marginals.Rd
new file mode 100644
index 0000000..ab3b920
--- /dev/null
+++ b/man/print.k_marginals.Rd
@@ -0,0 +1,21 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/util_k_marginals.R
+\name{print.k_marginals}
+\alias{print.k_marginals}
+\title{Print a k_marginals object}
+\usage{
+\method{print}{k_marginals}(x, n = 5, ...)
+}
+\arguments{
+\item{x}{A \code{k_marginals} object from \code{\link[=util_k_marginals]{util_k_marginals()}}.}
+
+\item{n}{Number of worst marginals to display.}
+
+\item{...}{Additional arguments passed to methods (unused).}
+}
+\value{
+\code{x}, invisibly.
+}
+\description{
+Print a k_marginals object
+}
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index dc71a2f..7f100e6 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -4,17 +4,30 @@
 \alias{util_k_marginals}
 \title{Calculate the k-marginals metric}
 \usage{
-util_k_marginals(eval_data, k = 3)
+util_k_marginals(eval_data, k = 3, keep_marginals = Inf, keep_cells = Inf)
 }
 \arguments{
 \item{eval_data}{An \code{eval_data} object.}
 
 \item{k}{Scalar order of the k-marginal (valid range = 1:3).}
+
+\item{keep_marginals}{Single integer number of worst marginals to retain
+in the output.
+Defaults to \code{Inf}, which retains all of them.}
+
+\item{keep_cells}{Single integer number of worst cells to retain in the
+output. Defaults to \code{Inf}, which retains all of them.}
 }
 \value{
-A value in range \link{0, 1000} where a higher value denotes lower MabsDDs
-and consequently greater similarity between confidential and synthetic data.
-For multiple replicates, a list of such values, one per replicate.
+A \code{k_marginals} object with three elements: \code{score}, a value in
+range \link{0, 1000} where a higher value denotes lower MabsDDs and consequently
+greater similarity between confidential and synthetic data; \code{marginals}, a
+tibble with the MabsDD for each combination of variables, worst first; and
+\code{cells}, a tibble with the synthetic and confidential proportions and their
+absolute difference for every cell, worst first. \code{score} is always computed
+from all marginals, even when \code{keep_marginals} or \code{keep_cells} truncate the
+detail tables. For multiple replicates, a list of such objects, one per
+replicate.
 }
 \description{
 For each unique k-combination of variables shared by the
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index f26cf45..d8c6cd7 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -34,7 +34,7 @@ synth <- tibble::tibble(
 test_that("k = 1 score matches hand-computed value", {
 
   expect_equal(
-    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1),
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score,
     750
   )
 
@@ -43,7 +43,7 @@ test_that("k = 1 score matches hand-computed value", {
 test_that("k = 2 score matches hand-computed value", {
 
   expect_equal(
-    .util_k_marginals(synth_data = synth, conf_data = conf, k = 2),
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 2)$score,
     5000 / 6
   )
 
@@ -54,7 +54,7 @@ test_that("identical data scores exactly 1000 for every k", {
   for (k in 1:2) {
 
     expect_equal(
-      .util_k_marginals(synth_data = conf, conf_data = conf, k = k),
+      .util_k_marginals(synth_data = conf, conf_data = conf, k = k)$score,
       1000
     )
 
@@ -71,7 +71,7 @@ test_that("cells absent from one dataset count as proportion zero", {
   synth_gap <- tibble::tibble(a = c("x", "x"))
 
   expect_equal(
-    .util_k_marginals(synth_data = synth_gap, conf_data = conf_gap, k = 1),
+    .util_k_marginals(synth_data = synth_gap, conf_data = conf_gap, k = 1)$score,
     500
   )
 
@@ -138,8 +138,8 @@ test_that("variables not shared by both datasets are ignored", {
   synth_extra <- dplyr::mutate(synth, c = c("m", "m", "n", "n"))
 
   expect_equal(
-    .util_k_marginals(synth_data = synth_extra, conf_data = conf, k = 1),
-    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+    .util_k_marginals(synth_data = synth_extra, conf_data = conf, k = 1)$score,
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
   )
 
 })
@@ -148,9 +148,9 @@ test_that("util_k_marginals accepts an eval_data object", {
 
   ed <- eval_data(conf_data = conf, synth_data = synth)
 
-  expect_equal(util_k_marginals(eval_data = ed, k = 1), 750)
+  expect_equal(util_k_marginals(eval_data = ed, k = 1)$score, 750)
 
-  expect_equal(util_k_marginals(eval_data = ed, k = 2), 5000 / 6)
+  expect_equal(util_k_marginals(eval_data = ed, k = 2)$score, 5000 / 6)
 
 })
 
@@ -160,7 +160,7 @@ test_that("util_k_marginals maps over replicates", {
 
   result <- util_k_marginals(eval_data = ed, k = 1)
 
-  expect_equal(result, list(750, 1000))
+  expect_equal(purrr::map_dbl(result, "score"), c(750, 1000))
 
 })
 
@@ -169,3 +169,233 @@ test_that("util_k_marginals rejects non-eval_data input", {
   expect_error(util_k_marginals(eval_data = synth, k = 1))
 
 })
+
+test_that("marginals and cells report worst-first detail", {
+
+  result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 2)
+
+  expect_s3_class(result, "k_marginals")
+
+  # single a x b combination with MabsDD = 1/6
+  expect_equal(result$marginals$variables, "a, b")
+  expect_equal(result$marginals$madd, 1 / 6)
+
+  # three observed cells sorted by descending absolute difference
+  expect_equal(
+    names(result$cells),
+    c("variables", "cell", "prop_synth", "prop_conf", "abs_diff")
+  )
+  expect_equal(result$cells$abs_diff, c(0.25, 0.25, 0))
+  expect_equal(result$cells$cell[3], "y, p")
+
+})
+
+test_that("cells absent from the synthetic data appear with proportion zero", {
+
+  conf_gap <- tibble::tibble(a = c("x", "y"))
+  synth_gap <- tibble::tibble(a = c("x", "x"))
+
+  result <- .util_k_marginals(synth_data = synth_gap, conf_data = conf_gap, k = 1)
+
+  y_cell <- dplyr::filter(result$cells, .data$cell == "y")
+
+  expect_equal(y_cell$prop_synth, 0)
+  expect_equal(y_cell$prop_conf, 0.5)
+
+})
+
+test_that("print method reports the score and worst marginals", {
+
+  result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+
+  expect_output(print(result), regexp = "k-marginals score: 750")
+  expect_output(print(result), regexp = "Worst marginals:")
+
+})
+
+test_that("keep_marginals and keep_cells truncate the detail tables", {
+
+  full <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+
+  kept <- .util_k_marginals(
+    synth_data = synth,
+    conf_data = conf,
+    k = 1,
+    keep_marginals = 1,
+    keep_cells = 2
+  )
+
+  expect_equal(nrow(kept$marginals), 1)
+  expect_equal(nrow(kept$cells), 2)
+
+  # retained rows are the worst ones from the full tables
+  expect_equal(kept$marginals, full$marginals[1, ])
+  expect_equal(kept$cells, full$cells[1:2, ])
+
+})
+
+test_that("retention keeps the highest abs_diff cells", {
+
+  # three levels with a strict worst cell: y has the largest abs_diff
+  # conf: x = 0.500, y = 0.250, z = 0.250
+  # synth: x = 0.250, y = 0.625, z = 0.125
+  # abs_diff: x = 0.250, y = 0.375, z = 0.125
+  conf_tri <- tibble::tibble(a = c(rep("x", 4), rep("y", 2), rep("z", 2)))
+  synth_tri <- tibble::tibble(a = c(rep("x", 2), rep("y", 5), "z"))
+
+  kept <- .util_k_marginals(
+    synth_data = synth_tri,
+    conf_data = conf_tri,
+    k = 1,
+    keep_cells = 1
+  )
+
+  expect_equal(kept$cells$cell, "y")
+  expect_equal(kept$cells$abs_diff, 0.375)
+
+})
+
+test_that("score is computed from all marginals, not the retained subset", {
+
+  full <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+
+  kept <- .util_k_marginals(
+    synth_data = synth,
+    conf_data = conf,
+    k = 1,
+    keep_marginals = 1,
+    keep_cells = 1
+  )
+
+  expect_equal(kept$score, full$score)
+
+})
+
+test_that("invalid retention arguments throw an error", {
+
+  for (bad_keep in list(0, "5", 1.5, NA_real_, NaN, c(1, 2))) {
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth, conf_data = conf, k = 1, keep_marginals = bad_keep
+      ),
+      regexp = "must be single integers >= 1 or Inf"
+    )
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth, conf_data = conf, k = 1, keep_cells = bad_keep
+      ),
+      regexp = "must be single integers >= 1 or Inf"
+    )
+
+  }
+
+  # Inf remains valid: it is the documented keep-everything default
+  expect_equal(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1,
+      keep_marginals = Inf, keep_cells = Inf
+    )$score,
+    750
+  )
+
+})
+
+test_that("util_k_marginals passes retention arguments through", {
+
+  ed <- eval_data(conf_data = conf, synth_data = synth)
+
+  result <- util_k_marginals(
+    eval_data = ed, k = 1, keep_marginals = 1, keep_cells = 2
+  )
+
+  expect_equal(nrow(result$marginals), 1)
+  expect_equal(nrow(result$cells), 2)
+  expect_equal(result$score, 750)
+
+})
+
+test_that("non-integer and non-finite k values throw an error", {
+
+  for (bad_k in list(1.5, NA_real_, NaN, Inf)) {
+
+    expect_error(
+      .util_k_marginals(synth_data = synth, conf_data = conf, k = bad_k),
+      regexp = "`k` must be a single integer between 1 and 3"
+    )
+
+  }
+
+})
+
+test_that("marginals are sorted by descending madd across combinations", {
+
+  # third shared variable, identical in conf but perturbed in synth
+  # pair madds: (a, c) = 1/4, (a, b) = 1/6, (b, c) = 1/6
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  result <- .util_k_marginals(synth_data = synth_3, conf_data = conf_3, k = 2)
+
+  expect_equal(result$marginals$variables[1], "a, c")
+  expect_equal(result$marginals$madd, c(1 / 4, 1 / 6, 1 / 6))
+
+  # global cell ordering holds across cells from different combinations
+  expect_equal(
+    result$cells$abs_diff,
+    sort(result$cells$abs_diff, decreasing = TRUE)
+  )
+
+})
+
+test_that("each replicate result is a complete k_marginals object", {
+
+  ed <- eval_data(conf_data = conf, synth_data = list(synth, conf))
+
+  result <- util_k_marginals(eval_data = ed, k = 1)
+
+  for (rep in result) {
+
+    expect_s3_class(rep, "k_marginals")
+    expect_named(rep, c("score", "marginals", "cells"))
+    expect_gt(nrow(rep$marginals), 0)
+    expect_gt(nrow(rep$cells), 0)
+
+  }
+
+})
+
+test_that("conf-only extra columns are ignored", {
+
+  conf_extra <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+
+  expect_equal(
+    .util_k_marginals(synth_data = synth, conf_data = conf_extra, k = 1)$score,
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
+  )
+
+})
+
+test_that("one-row datasets produce a valid result", {
+
+  one_row <- tibble::tibble(a = "x", b = "p")
+
+  result <- .util_k_marginals(synth_data = one_row, conf_data = one_row, k = 1)
+
+  expect_equal(result$score, 1000)
+  expect_equal(nrow(result$cells), 2)
+
+})
+
+test_that("print truncates the marginals display to n rows", {
+
+  result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+
+  out <- utils::capture.output(print(result, n = 1))
+
+  # tibble rows print with a leading row number: row 1 only, no row 2
+  expect_true(any(grepl("^1 ", out)))
+  expect_false(any(grepl("^2 ", out)))
+
+})

From 10058ae0411418cf0df86e2ce5cd4a91a0a903b1 Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Thu, 13 Aug 2026 16:03:13 -0400
Subject: [PATCH 03/10] Add sampling and priority variables to
 util_k_marginals()

---
 R/util_k_marginals.R                   |  83 ++++++++--
 man/dot-util_k_marginals.Rd            |  19 ++-
 man/util_k_marginals.Rd                |  26 +++-
 tests/testthat/test-util_k_marginals.R | 204 +++++++++++++++++++++++++
 4 files changed, 314 insertions(+), 18 deletions(-)

diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 2cac50f..d6853c7 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -17,6 +17,15 @@
 #' Defaults to `Inf`, which retains all of them.
 #' @param keep_cells Single integer number of worst cells to retain in the
 #' output. Defaults to `Inf`, which retains all of them.
+#' @param n_marginals Single integer cap on the number of variable
+#' combinations to evaluate. When the number of possible combinations exceeds
+#' the cap, a random subset is sampled; set a seed before calling for
+#' reproducible results. Defaults to `Inf`, which evaluates all combinations.
+#' @param priority_vars Optional character vector of variable names. Every
+#' combination containing at least one of these variables is always evaluated;
+#' sampling under `n_marginals` only applies to the remaining combinations. If
+#' the priority combinations alone exceed `n_marginals`, all of them are still
+#' evaluated. Defaults to `NULL`.
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -24,23 +33,26 @@
 #' tibble with the MabsDD for each combination of variables, worst first; and
 #' `cells`, a tibble with the synthetic and confidential proportions and their
 #' absolute difference for every cell, worst first. `score` is always computed
-#' from all marginals, even when `keep_marginals` or `keep_cells` truncate the
-#' detail tables.
+#' from all evaluated marginals, even when `keep_marginals` or `keep_cells`
+#' truncate the detail tables.
 #'
 .util_k_marginals <- function(
     synth_data,
     conf_data,
     k,
     keep_marginals = Inf,
-    keep_cells = Inf) {
+    keep_cells = Inf,
+    n_marginals = Inf,
+    priority_vars = NULL) {
 
-  for (keep in list(keep_marginals, keep_cells)) {
+  for (keep in list(keep_marginals, keep_cells, n_marginals)) {
 
     if (!(is.numeric(keep) && length(keep) == 1 && !is.na(keep) &&
           keep >= 1 && keep == floor(keep))) {
 
       stop(
-        "`keep_marginals` and `keep_cells` must be single integers >= 1 or Inf"
+        "`keep_marginals`, `keep_cells`, and `n_marginals` must be single ",
+        "integers >= 1 or Inf"
       )
 
     }
@@ -71,8 +83,42 @@
 
   }
 
+  if (!is.null(priority_vars)) {
+
+    if (!(is.character(priority_vars) &&
+          all(priority_vars %in% shared_vars))) {
+
+      stop(
+        "`priority_vars` must be a character vector of variables shared by ",
+        "both datasets"
+      )
+
+    }
+
+  }
+
   kmarginals_vars <- t(utils::combn(x = shared_vars, m = k))
 
+  # sample combinations down to n_marginals, always keeping combinations that
+  # contain a priority variable
+  if (nrow(kmarginals_vars) > n_marginals) {
+
+    is_priority <- apply(
+      X = kmarginals_vars,
+      MARGIN = 1,
+      FUN = \(vars) any(vars %in% priority_vars)
+    )
+
+    n_sampled <- min(max(n_marginals - sum(is_priority), 0), sum(!is_priority))
+
+    sampled_rows <- sample(x = which(!is_priority), size = n_sampled)
+
+    kmarginals_vars <- kmarginals_vars[
+      sort(c(which(is_priority), sampled_rows)), , drop = FALSE
+    ]
+
+  }
+
   # cell proportions for one dataset over one set of variables
   process_data <- function(data, vars, prop_name) {
 
@@ -178,6 +224,15 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' Defaults to `Inf`, which retains all of them.
 #' @param keep_cells Single integer number of worst cells to retain in the
 #' output. Defaults to `Inf`, which retains all of them.
+#' @param n_marginals Single integer cap on the number of variable
+#' combinations to evaluate. When the number of possible combinations exceeds
+#' the cap, a random subset is sampled; set a seed before calling for
+#' reproducible results. Defaults to `Inf`, which evaluates all combinations.
+#' @param priority_vars Optional character vector of variable names. Every
+#' combination containing at least one of these variables is always evaluated;
+#' sampling under `n_marginals` only applies to the remaining combinations. If
+#' the priority combinations alone exceed `n_marginals`, all of them are still
+#' evaluated. Defaults to `NULL`.
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -185,9 +240,9 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' tibble with the MabsDD for each combination of variables, worst first; and
 #' `cells`, a tibble with the synthetic and confidential proportions and their
 #' absolute difference for every cell, worst first. `score` is always computed
-#' from all marginals, even when `keep_marginals` or `keep_cells` truncate the
-#' detail tables. For multiple replicates, a list of such objects, one per
-#' replicate.
+#' from all evaluated marginals, even when `keep_marginals` or `keep_cells`
+#' truncate the detail tables. For multiple replicates, a list of such
+#' objects, one per replicate.
 #'
 #' @export
 #'
@@ -195,7 +250,9 @@ util_k_marginals <- function(
     eval_data,
     k = 3,
     keep_marginals = Inf,
-    keep_cells = Inf) {
+    keep_cells = Inf,
+    n_marginals = Inf,
+    priority_vars = NULL) {
 
   stopifnot(is_eval_data(eval_data))
 
@@ -207,7 +264,9 @@ util_k_marginals <- function(
         conf_data = eval_data$conf_data,
         k = k,
         keep_marginals = keep_marginals,
-        keep_cells = keep_cells
+        keep_cells = keep_cells,
+        n_marginals = n_marginals,
+        priority_vars = priority_vars
       )
     )
 
@@ -222,7 +281,9 @@ util_k_marginals <- function(
           conf_data = eval_data$conf_data,
           k = k,
           keep_marginals = keep_marginals,
-          keep_cells = keep_cells
+          keep_cells = keep_cells,
+          n_marginals = n_marginals,
+          priority_vars = priority_vars
         )
 
       }
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 9c91a15..12fcf8b 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -9,7 +9,9 @@
   conf_data,
   k,
   keep_marginals = Inf,
-  keep_cells = Inf
+  keep_cells = Inf,
+  n_marginals = Inf,
+  priority_vars = NULL
 )
 }
 \arguments{
@@ -25,6 +27,17 @@ Defaults to \code{Inf}, which retains all of them.}
 
 \item{keep_cells}{Single integer number of worst cells to retain in the
 output. Defaults to \code{Inf}, which retains all of them.}
+
+\item{n_marginals}{Single integer cap on the number of variable
+combinations to evaluate. When the number of possible combinations exceeds
+the cap, a random subset is sampled; set a seed before calling for
+reproducible results. Defaults to \code{Inf}, which evaluates all combinations.}
+
+\item{priority_vars}{Optional character vector of variable names. Every
+combination containing at least one of these variables is always evaluated;
+sampling under \code{n_marginals} only applies to the remaining combinations. If
+the priority combinations alone exceed \code{n_marginals}, all of them are still
+evaluated. Defaults to \code{NULL}.}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
@@ -33,8 +46,8 @@ greater similarity between confidential and synthetic data; \code{marginals}, a
 tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their
 absolute difference for every cell, worst first. \code{score} is always computed
-from all marginals, even when \code{keep_marginals} or \code{keep_cells} truncate the
-detail tables.
+from all evaluated marginals, even when \code{keep_marginals} or \code{keep_cells}
+truncate the detail tables.
 }
 \description{
 This worker function takes a specified k-marginal and calculates
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 7f100e6..014190c 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -4,7 +4,14 @@
 \alias{util_k_marginals}
 \title{Calculate the k-marginals metric}
 \usage{
-util_k_marginals(eval_data, k = 3, keep_marginals = Inf, keep_cells = Inf)
+util_k_marginals(
+  eval_data,
+  k = 3,
+  keep_marginals = Inf,
+  keep_cells = Inf,
+  n_marginals = Inf,
+  priority_vars = NULL
+)
 }
 \arguments{
 \item{eval_data}{An \code{eval_data} object.}
@@ -17,6 +24,17 @@ Defaults to \code{Inf}, which retains all of them.}
 
 \item{keep_cells}{Single integer number of worst cells to retain in the
 output. Defaults to \code{Inf}, which retains all of them.}
+
+\item{n_marginals}{Single integer cap on the number of variable
+combinations to evaluate. When the number of possible combinations exceeds
+the cap, a random subset is sampled; set a seed before calling for
+reproducible results. Defaults to \code{Inf}, which evaluates all combinations.}
+
+\item{priority_vars}{Optional character vector of variable names. Every
+combination containing at least one of these variables is always evaluated;
+sampling under \code{n_marginals} only applies to the remaining combinations. If
+the priority combinations alone exceed \code{n_marginals}, all of them are still
+evaluated. Defaults to \code{NULL}.}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
@@ -25,9 +43,9 @@ greater similarity between confidential and synthetic data; \code{marginals}, a
 tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their
 absolute difference for every cell, worst first. \code{score} is always computed
-from all marginals, even when \code{keep_marginals} or \code{keep_cells} truncate the
-detail tables. For multiple replicates, a list of such objects, one per
-replicate.
+from all evaluated marginals, even when \code{keep_marginals} or \code{keep_cells}
+truncate the detail tables. For multiple replicates, a list of such
+objects, one per replicate.
 }
 \description{
 For each unique k-combination of variables shared by the
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index d8c6cd7..e800435 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -399,3 +399,207 @@ test_that("print truncates the marginals display to n rows", {
   expect_false(any(grepl("^2 ", out)))
 
 })
+
+test_that("n_marginals caps the number of evaluated combinations", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  set.seed(20250813)
+
+  result <- .util_k_marginals(
+    synth_data = synth_3, conf_data = conf_3, k = 2, n_marginals = 2
+  )
+
+  expect_equal(nrow(result$marginals), 2)
+  expect_true(result$score >= 0 && result$score <= 1000)
+
+})
+
+test_that("sampling is reproducible given a seed", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  set.seed(1)
+  first <- .util_k_marginals(
+    synth_data = synth_3, conf_data = conf_3, k = 2, n_marginals = 1
+  )
+
+  set.seed(1)
+  second <- .util_k_marginals(
+    synth_data = synth_3, conf_data = conf_3, k = 2, n_marginals = 1
+  )
+
+  expect_equal(first, second)
+
+})
+
+test_that("priority_vars combinations are always evaluated", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  # the two combinations containing a fill the cap exactly, so the
+  # selection is deterministic despite sampling
+  result <- .util_k_marginals(
+    synth_data = synth_3,
+    conf_data = conf_3,
+    k = 2,
+    n_marginals = 2,
+    priority_vars = "a"
+  )
+
+  expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
+
+})
+
+test_that("priority combinations exceeding n_marginals are all kept", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  result <- .util_k_marginals(
+    synth_data = synth_3,
+    conf_data = conf_3,
+    k = 2,
+    n_marginals = 1,
+    priority_vars = "a"
+  )
+
+  expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
+
+})
+
+test_that("n_marginals at or above the combination count changes nothing", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  full <- .util_k_marginals(synth_data = synth_3, conf_data = conf_3, k = 2)
+
+  capped <- .util_k_marginals(
+    synth_data = synth_3, conf_data = conf_3, k = 2, n_marginals = 3
+  )
+
+  expect_equal(capped, full)
+
+})
+
+test_that("invalid sampling arguments throw an error", {
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1, n_marginals = 1.5
+    ),
+    regexp = "must be single integers >= 1 or Inf"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1, priority_vars = "zzz"
+    ),
+    regexp = "`priority_vars` must be a character vector"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1, priority_vars = 1
+    ),
+    regexp = "`priority_vars` must be a character vector"
+  )
+
+})
+
+test_that("util_k_marginals passes sampling arguments through", {
+
+  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
+  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
+
+  ed <- eval_data(conf_data = conf_3, synth_data = synth_3)
+
+  result <- util_k_marginals(
+    eval_data = ed, k = 2, n_marginals = 2, priority_vars = "a"
+  )
+
+  expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
+
+})
+
+test_that("sampling fills remaining slots after priority combinations", {
+
+  # 4 shared variables, k = 2: 6 combinations, 3 containing a
+  conf_4 <- dplyr::mutate(
+    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
+  )
+  synth_4 <- dplyr::mutate(
+    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
+  )
+
+  set.seed(20250813)
+
+  result <- .util_k_marginals(
+    synth_data = synth_4,
+    conf_data = conf_4,
+    k = 2,
+    n_marginals = 4,
+    priority_vars = "a"
+  )
+
+  # cap respected exactly: all 3 priority combos plus 1 sampled non-priority
+  expect_equal(nrow(result$marginals), 4)
+
+  has_a <- grepl("a", result$marginals$variables)
+  expect_equal(sum(has_a), 3)
+  expect_equal(sum(!has_a), 1)
+  expect_true(
+    all(result$marginals$variables[!has_a] %in% c("b, c", "b, d", "c, d"))
+  )
+
+})
+
+test_that("n_marginals caps k = 3 combinations", {
+
+  # 4 shared variables, k = 3: 4 combinations
+  conf_4 <- dplyr::mutate(
+    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
+  )
+  synth_4 <- dplyr::mutate(
+    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
+  )
+
+  set.seed(20250813)
+
+  result <- .util_k_marginals(
+    synth_data = synth_4, conf_data = conf_4, k = 3, n_marginals = 2
+  )
+
+  expect_equal(nrow(result$marginals), 2)
+  expect_true(result$score >= 0 && result$score <= 1000)
+
+})
+
+test_that("priority_vars applies to k = 3 combinations", {
+
+  # priority a appears in 3 of the 4 triples; cap of 3 keeps exactly those
+  conf_4 <- dplyr::mutate(
+    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
+  )
+  synth_4 <- dplyr::mutate(
+    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_4,
+    conf_data = conf_4,
+    k = 3,
+    n_marginals = 3,
+    priority_vars = "a"
+  )
+
+  expect_equal(
+    sort(result$marginals$variables),
+    c("a, b, c", "a, b, d", "a, c, d")
+  )
+
+})

From f27cd5c48a7554df8dee07f0b6e09c3dc14cff2b Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Thu, 13 Aug 2026 16:12:33 -0400
Subject: [PATCH 04/10] Added sample weighting

---
 R/util_k_marginals.R                   |  92 ++++++++++++++--
 man/dot-util_k_marginals.Rd            |   9 +-
 man/util_k_marginals.Rd                |   9 +-
 tests/testthat/test-util_k_marginals.R | 145 +++++++++++++++++++++++++
 4 files changed, 241 insertions(+), 14 deletions(-)

diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index d6853c7..9bc7761 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -26,6 +26,11 @@
 #' sampling under `n_marginals` only applies to the remaining combinations. If
 #' the priority combinations alone exceed `n_marginals`, all of them are still
 #' evaluated. Defaults to `NULL`.
+#' @param weight_var Optional character name of a numeric sample-weight
+#' column present in both datasets. When set, cell proportions are weight
+#' shares instead of row shares, and the weight column is excluded from the
+#' marginals. Weights must be finite and non-negative with a positive total.
+#' Defaults to `NULL` (unweighted).
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -43,7 +48,8 @@
     keep_marginals = Inf,
     keep_cells = Inf,
     n_marginals = Inf,
-    priority_vars = NULL) {
+    priority_vars = NULL,
+    weight_var = NULL) {
 
   for (keep in list(keep_marginals, keep_cells, n_marginals)) {
 
@@ -74,8 +80,51 @@
 
   }
 
-  # only variables present in both datasets contribute marginals
-  shared_vars <- intersect(names(synth_data), names(conf_data))
+  if (!is.null(weight_var)) {
+
+    if (!(is.character(weight_var) && length(weight_var) == 1)) {
+
+      stop("`weight_var` must be a single character string")
+
+    }
+
+    if (!(weight_var %in% names(synth_data) &&
+          weight_var %in% names(conf_data))) {
+
+      stop("`weight_var` must be a column in both datasets")
+
+    }
+
+    if (!(is.numeric(synth_data[[weight_var]]) &&
+          is.numeric(conf_data[[weight_var]]))) {
+
+      stop("`weight_var` must be a numeric column in both datasets")
+
+    }
+
+    # invalid weights break the probability interpretation of proportions
+    for (weights in list(synth_data[[weight_var]], conf_data[[weight_var]])) {
+
+      if (!all(is.finite(weights)) || any(weights < 0) ||
+          sum(weights) <= 0) {
+
+        stop(
+          "`weight_var` values must be finite and non-negative with a ",
+          "positive total in both datasets"
+        )
+
+      }
+
+    }
+
+  }
+
+  # only variables present in both datasets contribute marginals; the weight
+  # column is never itself a marginal
+  shared_vars <- setdiff(
+    intersect(names(synth_data), names(conf_data)),
+    weight_var
+  )
 
   if (length(shared_vars) < k) {
 
@@ -119,14 +168,25 @@
 
   }
 
-  # cell proportions for one dataset over one set of variables
+  # cell proportions for one dataset over one set of variables; weighted
+  # proportions are weight shares instead of row shares
   process_data <- function(data, vars, prop_name) {
 
-    props <- data |>
-      dplyr::select(dplyr::all_of(vars)) |>
-      dplyr::group_by_all() |>
-      dplyr::count() |>
-      dplyr::ungroup() |>
+    if (is.null(weight_var)) {
+
+      counts <- dplyr::count(data, dplyr::across(dplyr::all_of(vars)))
+
+    } else {
+
+      counts <- dplyr::count(
+        data,
+        dplyr::across(dplyr::all_of(vars)),
+        wt = .data[[weight_var]]
+      )
+
+    }
+
+    props <- counts |>
       dplyr::mutate("{prop_name}" := .data$n / sum(.data$n)) |>
       dplyr::select(-"n")
 
@@ -233,6 +293,11 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' sampling under `n_marginals` only applies to the remaining combinations. If
 #' the priority combinations alone exceed `n_marginals`, all of them are still
 #' evaluated. Defaults to `NULL`.
+#' @param weight_var Optional character name of a numeric sample-weight
+#' column present in both datasets. When set, cell proportions are weight
+#' shares instead of row shares, and the weight column is excluded from the
+#' marginals. Weights must be finite and non-negative with a positive total.
+#' Defaults to `NULL` (unweighted).
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -252,7 +317,8 @@ util_k_marginals <- function(
     keep_marginals = Inf,
     keep_cells = Inf,
     n_marginals = Inf,
-    priority_vars = NULL) {
+    priority_vars = NULL,
+    weight_var = NULL) {
 
   stopifnot(is_eval_data(eval_data))
 
@@ -266,7 +332,8 @@ util_k_marginals <- function(
         keep_marginals = keep_marginals,
         keep_cells = keep_cells,
         n_marginals = n_marginals,
-        priority_vars = priority_vars
+        priority_vars = priority_vars,
+        weight_var = weight_var
       )
     )
 
@@ -283,7 +350,8 @@ util_k_marginals <- function(
           keep_marginals = keep_marginals,
           keep_cells = keep_cells,
           n_marginals = n_marginals,
-          priority_vars = priority_vars
+          priority_vars = priority_vars,
+          weight_var = weight_var
         )
 
       }
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 12fcf8b..9e132d3 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -11,7 +11,8 @@
   keep_marginals = Inf,
   keep_cells = Inf,
   n_marginals = Inf,
-  priority_vars = NULL
+  priority_vars = NULL,
+  weight_var = NULL
 )
 }
 \arguments{
@@ -38,6 +39,12 @@ combination containing at least one of these variables is always evaluated;
 sampling under \code{n_marginals} only applies to the remaining combinations. If
 the priority combinations alone exceed \code{n_marginals}, all of them are still
 evaluated. Defaults to \code{NULL}.}
+
+\item{weight_var}{Optional character name of a numeric sample-weight
+column present in both datasets. When set, cell proportions are weight
+shares instead of row shares, and the weight column is excluded from the
+marginals. Weights must be finite and non-negative with a positive total.
+Defaults to \code{NULL} (unweighted).}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 014190c..69f74ed 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -10,7 +10,8 @@ util_k_marginals(
   keep_marginals = Inf,
   keep_cells = Inf,
   n_marginals = Inf,
-  priority_vars = NULL
+  priority_vars = NULL,
+  weight_var = NULL
 )
 }
 \arguments{
@@ -35,6 +36,12 @@ combination containing at least one of these variables is always evaluated;
 sampling under \code{n_marginals} only applies to the remaining combinations. If
 the priority combinations alone exceed \code{n_marginals}, all of them are still
 evaluated. Defaults to \code{NULL}.}
+
+\item{weight_var}{Optional character name of a numeric sample-weight
+column present in both datasets. When set, cell proportions are weight
+shares instead of row shares, and the weight column is excluded from the
+marginals. Weights must be finite and non-negative with a positive total.
+Defaults to \code{NULL} (unweighted).}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index e800435..2811163 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -603,3 +603,148 @@ test_that("priority_vars applies to k = 3 combinations", {
   )
 
 })
+
+test_that("weighted proportions match hand-computed values", {
+
+  # conf: weight shares x = 3/4, y = 1/4; synth: x = 1/2, y = 1/2
+  # MabsDD = mean(0.25, 0.25) = 0.25 -> score 750
+  conf_w <- tibble::tibble(a = c("x", "y"), w = c(3, 1))
+  synth_w <- tibble::tibble(a = c("x", "y"), w = c(1, 1))
+
+  result <- .util_k_marginals(
+    synth_data = synth_w, conf_data = conf_w, k = 1, weight_var = "w"
+  )
+
+  expect_equal(result$score, 750)
+  expect_equal(
+    dplyr::filter(result$cells, .data$cell == "x")$prop_conf,
+    0.75
+  )
+
+})
+
+test_that("unit weights reproduce the unweighted result", {
+
+  conf_w <- dplyr::mutate(conf, w = 1)
+  synth_w <- dplyr::mutate(synth, w = 1)
+
+  weighted <- .util_k_marginals(
+    synth_data = synth_w, conf_data = conf_w, k = 1, weight_var = "w"
+  )
+
+  unweighted <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
+
+  expect_equal(weighted, unweighted)
+
+})
+
+test_that("the weight column is never a marginal", {
+
+  conf_w <- dplyr::mutate(conf, w = 1)
+  synth_w <- dplyr::mutate(synth, w = 1)
+
+  result <- .util_k_marginals(
+    synth_data = synth_w, conf_data = conf_w, k = 1, weight_var = "w"
+  )
+
+  expect_equal(sort(result$marginals$variables), c("a", "b"))
+
+})
+
+test_that("invalid weight_var throws an error", {
+
+  conf_w <- dplyr::mutate(conf, w = 1)
+  synth_w <- dplyr::mutate(synth, w = 1)
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_w, conf_data = conf_w, k = 1, weight_var = 1
+    ),
+    regexp = "`weight_var` must be a single character string"
+  )
+
+  # column missing from one dataset
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf_w, k = 1, weight_var = "w"
+    ),
+    regexp = "`weight_var` must be a column in both datasets"
+  )
+
+  # non-numeric weight column
+  conf_chr <- dplyr::mutate(conf, w = "1")
+  synth_chr <- dplyr::mutate(synth, w = "1")
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_chr, conf_data = conf_chr, k = 1, weight_var = "w"
+    ),
+    regexp = "`weight_var` must be a numeric column in both datasets"
+  )
+
+})
+
+test_that("util_k_marginals passes weight_var through", {
+
+  conf_w <- tibble::tibble(a = c("x", "y"), w = c(3, 1))
+  synth_w <- tibble::tibble(a = c("x", "y"), w = c(1, 1))
+
+  ed <- eval_data(conf_data = conf_w, synth_data = synth_w)
+
+  expect_equal(
+    util_k_marginals(eval_data = ed, k = 1, weight_var = "w")$score,
+    750
+  )
+
+})
+
+test_that("invalid weight values throw an error", {
+
+  synth_w <- dplyr::mutate(synth, w = 1)
+
+  bad_weights <- list(
+    c(1, 1, 1, -1),   # negative
+    c(1, 1, 1, NA),   # missing
+    c(1, 1, 1, Inf),  # non-finite
+    c(0, 0, 0, 0)     # zero total
+  )
+
+  for (bad_w in bad_weights) {
+
+    conf_bad <- dplyr::mutate(conf, w = bad_w)
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_w, conf_data = conf_bad, k = 1, weight_var = "w"
+      ),
+      regexp = "finite and non-negative with a positive total"
+    )
+
+    # symmetric: same weights are rejected on the synthetic side
+    synth_bad <- dplyr::mutate(synth, w = bad_w)
+    conf_w <- dplyr::mutate(conf, w = 1)
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_bad, conf_data = conf_w, k = 1, weight_var = "w"
+      ),
+      regexp = "finite and non-negative with a positive total"
+    )
+
+  }
+
+})
+
+test_that("zero weights are valid when the total is positive", {
+
+  # zero-weight rows drop out: conf weight shares x = 1, y = 0
+  conf_w <- tibble::tibble(a = c("x", "y"), w = c(2, 0))
+  synth_w <- tibble::tibble(a = c("x", "y"), w = c(1, 1))
+
+  result <- .util_k_marginals(
+    synth_data = synth_w, conf_data = conf_w, k = 1, weight_var = "w"
+  )
+
+  expect_equal(result$score, 500)
+
+})

From 29c5f9fd025cc2f7ac6fe4d931ef1c5d7ad16e60 Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Thu, 13 Aug 2026 16:38:41 -0400
Subject: [PATCH 05/10] Add discretization of numeric variables

---
 NEWS.md                                |   2 +-
 R/util_k_marginals.R                   | 129 ++++++++++++++-
 man/dot-util_k_marginals.Rd            |  16 +-
 man/util_k_marginals.Rd                |  16 +-
 tests/testthat/test-util_k_marginals.R | 213 +++++++++++++++++++++++++
 5 files changed, 369 insertions(+), 7 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 1a321c5..a3e2769 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,6 @@
 # syntheval 0.1.0
 
-* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals.
+* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, and discretization of numeric variables.
 
 # syntheval 0.0.5
 
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 9bc7761..05bb1ae 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -31,6 +31,16 @@
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param bins Optional single integer >= 2. When set, every numeric shared
+#' variable is discretized into this many bins (fewer, with a warning, if
+#' tied quantile cut points collapse) with breaks derived from the
+#' confidential data and applied to both datasets; the outer bins extend to
+#' +/-Inf so synthetic values outside the confidential range land in edge
+#' bins. Defaults to `NULL` (no discretization).
+#' @param discretize_method Method used to place bin breaks when `bins` is
+#' set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
+#' for univariate k-means clustering (set a seed before calling for
+#' reproducible clusters). Defaults to "width".
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -49,7 +59,11 @@
     keep_cells = Inf,
     n_marginals = Inf,
     priority_vars = NULL,
-    weight_var = NULL) {
+    weight_var = NULL,
+    bins = NULL,
+    discretize_method = c("width", "ntile", "cluster")) {
+
+  discretize_method <- match.arg(discretize_method)
 
   for (keep in list(keep_marginals, keep_cells, n_marginals)) {
 
@@ -126,6 +140,85 @@
     weight_var
   )
 
+  if (!is.null(bins)) {
+
+    if (!(is.numeric(bins) && length(bins) == 1 && !is.na(bins) &&
+          bins >= 2 && bins == floor(bins))) {
+
+      stop("`bins` must be a single integer >= 2")
+
+    }
+
+    numeric_vars <- shared_vars[
+      purrr::map_lgl(.x = shared_vars, .f = \(v) is.numeric(conf_data[[v]]))
+    ]
+
+    for (var in numeric_vars) {
+
+      conf_values <- conf_data[[var]]
+
+      if (!all(is.finite(conf_values))) {
+
+        stop(
+          "numeric variables must be finite to discretize; `", var,
+          "` is not"
+        )
+
+      }
+
+      if (dplyr::n_distinct(conf_values) < bins) {
+
+        stop(
+          "`", var, "` has fewer distinct confidential values than `bins`"
+        )
+
+      }
+
+      # interior cut points always derive from the confidential data; each
+      # method yields bins - 1 of them
+      cut_points <- switch(
+        EXPR = discretize_method,
+        width = seq(
+          from = min(conf_values),
+          to = max(conf_values),
+          length.out = bins + 1
+        )[2:bins],
+        ntile = stats::quantile(
+          x = conf_values,
+          probs = seq(from = 0, to = 1, length.out = bins + 1),
+          names = FALSE
+        )[2:bins],
+        cluster = {
+
+          centers <- sort(
+            stats::kmeans(x = conf_values, centers = bins)$centers[, 1]
+          )
+
+          (centers[-1] + centers[-length(centers)]) / 2
+
+        }
+      )
+
+      # outer bins extend to +/-Inf so out-of-range synthetic values land in
+      # edge bins; unique() collapses ties from skewed quantiles
+      breaks <- unique(c(-Inf, cut_points, Inf))
+
+      if (length(breaks) - 1 < bins) {
+
+        warning(
+          "`", var, "` was discretized into ", length(breaks) - 1,
+          " bins instead of ", bins, " because of tied cut points"
+        )
+
+      }
+
+      synth_data[[var]] <- cut(x = synth_data[[var]], breaks = breaks)
+      conf_data[[var]] <- cut(x = conf_values, breaks = breaks)
+
+    }
+
+  }
+
   if (length(shared_vars) < k) {
 
     stop("`k` cannot exceed the number of variables shared by both datasets")
@@ -298,6 +391,16 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param bins Optional single integer >= 2. When set, every numeric shared
+#' variable is discretized into this many bins (fewer, with a warning, if
+#' tied quantile cut points collapse) with breaks derived from the
+#' confidential data and applied to both datasets; the outer bins extend to
+#' +/-Inf so synthetic values outside the confidential range land in edge
+#' bins. Defaults to `NULL` (no discretization).
+#' @param discretize_method Method used to place bin breaks when `bins` is
+#' set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
+#' for univariate k-means clustering (set a seed before calling for
+#' reproducible clusters). Defaults to "width".
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
 #' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
@@ -318,10 +421,24 @@ util_k_marginals <- function(
     keep_cells = Inf,
     n_marginals = Inf,
     priority_vars = NULL,
-    weight_var = NULL) {
+    weight_var = NULL,
+    bins = NULL,
+    discretize_method = c("width", "ntile", "cluster")) {
 
   stopifnot(is_eval_data(eval_data))
 
+  discretize_method <- match.arg(discretize_method)
+
+  # surface the resolved method so a forgotten discretize_method is visible
+  if (!is.null(bins)) {
+
+    message(
+      "Discretizing numeric variables into ", bins,
+      " bins using the '", discretize_method, "' method"
+    )
+
+  }
+
   if (eval_data$n_rep == 1) {
 
     return(
@@ -333,7 +450,9 @@ util_k_marginals <- function(
         keep_cells = keep_cells,
         n_marginals = n_marginals,
         priority_vars = priority_vars,
-        weight_var = weight_var
+        weight_var = weight_var,
+        bins = bins,
+        discretize_method = discretize_method
       )
     )
 
@@ -351,7 +470,9 @@ util_k_marginals <- function(
           keep_cells = keep_cells,
           n_marginals = n_marginals,
           priority_vars = priority_vars,
-          weight_var = weight_var
+          weight_var = weight_var,
+          bins = bins,
+          discretize_method = discretize_method
         )
 
       }
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 9e132d3..8278d9d 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -12,7 +12,9 @@
   keep_cells = Inf,
   n_marginals = Inf,
   priority_vars = NULL,
-  weight_var = NULL
+  weight_var = NULL,
+  bins = NULL,
+  discretize_method = c("width", "ntile", "cluster")
 )
 }
 \arguments{
@@ -45,6 +47,18 @@ column present in both datasets. When set, cell proportions are weight
 shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
+
+\item{bins}{Optional single integer >= 2. When set, every numeric shared
+variable is discretized into this many bins (fewer, with a warning, if
+tied quantile cut points collapse) with breaks derived from the
+confidential data and applied to both datasets; the outer bins extend to
++/-Inf so synthetic values outside the confidential range land in edge
+bins. Defaults to \code{NULL} (no discretization).}
+
+\item{discretize_method}{Method used to place bin breaks when \code{bins} is
+set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
+for univariate k-means clustering (set a seed before calling for
+reproducible clusters). Defaults to "width".}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 69f74ed..61ffed5 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -11,7 +11,9 @@ util_k_marginals(
   keep_cells = Inf,
   n_marginals = Inf,
   priority_vars = NULL,
-  weight_var = NULL
+  weight_var = NULL,
+  bins = NULL,
+  discretize_method = c("width", "ntile", "cluster")
 )
 }
 \arguments{
@@ -42,6 +44,18 @@ column present in both datasets. When set, cell proportions are weight
 shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
+
+\item{bins}{Optional single integer >= 2. When set, every numeric shared
+variable is discretized into this many bins (fewer, with a warning, if
+tied quantile cut points collapse) with breaks derived from the
+confidential data and applied to both datasets; the outer bins extend to
++/-Inf so synthetic values outside the confidential range land in edge
+bins. Defaults to \code{NULL} (no discretization).}
+
+\item{discretize_method}{Method used to place bin breaks when \code{bins} is
+set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
+for univariate k-means clustering (set a seed before calling for
+reproducible clusters). Defaults to "width".}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index 2811163..746b0ba 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -748,3 +748,216 @@ test_that("zero weights are valid when the total is positive", {
   expect_equal(result$score, 500)
 
 })
+
+test_that("width discretization matches hand-computed values", {
+
+  # conf 1:4 with 2 bins: interior cut at 2.5, so low = {1, 2}, high = {3, 4}
+  # conf shares (0.5, 0.5); synth c(1, 1, 1, 4) shares (0.75, 0.25)
+  # MabsDD = mean(0.25, 0.25) = 0.25 -> score 750
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  result <- .util_k_marginals(
+    synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2
+  )
+
+  expect_equal(result$score, 750)
+  expect_equal(nrow(result$cells), 2)
+
+})
+
+test_that("ntile discretization uses confidential quantiles", {
+
+  # conf quartile cut points at 25/50/75th percentiles of 1:8
+  # 4 bins of 2 values each: conf shares 0.25 apiece
+  # synth all in the lowest bin: shares (1, 0, 0, 0)
+  # MabsDD = mean(0.75, 0.25, 0.25, 0.25) = 0.375 -> score 625
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4, 5, 6, 7, 8))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 1, 1, 1, 1, 1))
+
+  result <- .util_k_marginals(
+    synth_data = synth_num,
+    conf_data = conf_num,
+    k = 1,
+    bins = 4,
+    discretize_method = "ntile"
+  )
+
+  expect_equal(result$score, 625)
+
+})
+
+test_that("cluster discretization separates well-separated groups", {
+
+  # two tight clusters around 1 and 10: the midpoint break lands between
+  # them, so conf shares (0.5, 0.5) and synth (1, 0) -> score 500
+  conf_num <- tibble::tibble(v = c(1, 1.1, 10, 10.1))
+  synth_num <- tibble::tibble(v = c(1, 1, 1.1, 1.1))
+
+  set.seed(20250813)
+
+  result <- .util_k_marginals(
+    synth_data = synth_num,
+    conf_data = conf_num,
+    k = 1,
+    bins = 2,
+    discretize_method = "cluster"
+  )
+
+  expect_equal(result$score, 500)
+
+})
+
+test_that("synthetic values outside the confidential range land in edge bins", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(-100, -100, 100, 100))
+
+  result <- .util_k_marginals(
+    synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2
+  )
+
+  # extremes split evenly across the two edge bins, matching conf shares
+  expect_equal(result$score, 1000)
+
+})
+
+test_that("non-numeric variables are untouched by discretization", {
+
+  conf_mix <- dplyr::mutate(conf, v = c(1, 2, 3, 4))
+  synth_mix <- dplyr::mutate(synth, v = c(1, 2, 3, 4))
+
+  result <- .util_k_marginals(
+    synth_data = synth_mix, conf_data = conf_mix, k = 1, bins = 2
+  )
+
+  # categorical marginals a and b keep their original levels
+  a_cells <- dplyr::filter(result$cells, .data$variables == "a")
+  expect_equal(sort(a_cells$cell), c("x", "y"))
+
+})
+
+test_that("bins = NULL leaves numeric variables as-is", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  result <- .util_k_marginals(
+    synth_data = synth_num, conf_data = conf_num, k = 1
+  )
+
+  # every distinct value is its own cell
+  expect_equal(nrow(result$cells), 4)
+
+})
+
+test_that("invalid discretization arguments throw an error", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  for (bad_bins in list(1, 2.5, "2", NA_real_, c(2, 3))) {
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_num, conf_data = conf_num, k = 1, bins = bad_bins
+      ),
+      regexp = "`bins` must be a single integer >= 2"
+    )
+
+  }
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2,
+      discretize_method = "magic"
+    )
+  )
+
+  # non-finite confidential values cannot be discretized
+  conf_inf <- tibble::tibble(v = c(1, 2, 3, Inf))
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_num, conf_data = conf_inf, k = 1, bins = 2
+    ),
+    regexp = "must be finite to discretize"
+  )
+
+})
+
+test_that("util_k_marginals passes discretization arguments through", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  ed <- eval_data(conf_data = conf_num, synth_data = synth_num)
+
+  expect_equal(
+    suppressMessages(util_k_marginals(eval_data = ed, k = 1, bins = 2))$score,
+    750
+  )
+
+})
+
+test_that("util_k_marginals messages the resolved discretization method", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  ed <- eval_data(conf_data = conf_num, synth_data = synth_num)
+
+  expect_message(
+    util_k_marginals(eval_data = ed, k = 1, bins = 2),
+    regexp = "2 bins using the 'width' method"
+  )
+
+  expect_message(
+    util_k_marginals(
+      eval_data = ed, k = 1, bins = 2, discretize_method = "ntile"
+    ),
+    regexp = "2 bins using the 'ntile' method"
+  )
+
+  # no discretization, no message
+  expect_no_message(util_k_marginals(eval_data = ed, k = 1))
+
+})
+
+test_that("too few distinct confidential values throw an error", {
+
+  conf_const <- tibble::tibble(v = c(2, 2, 2, 2))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  for (method in c("width", "ntile", "cluster")) {
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_num, conf_data = conf_const, k = 1, bins = 2,
+        discretize_method = method
+      ),
+      regexp = "fewer distinct confidential values than `bins`"
+    )
+
+  }
+
+})
+
+test_that("tied quantile cut points collapse bins with a warning", {
+
+  # heavily tied data passes the distinct-value pre-check, but the 25th and
+  # 50th percentiles coincide at 1, collapsing a quantile bin
+  conf_ties <- tibble::tibble(v = c(1, 1, 1, 1, 1, 1, 2, 3, 4, 5))
+  synth_num <- tibble::tibble(v = c(1, 2, 3, 4, 5))
+
+  expect_warning(
+    result <- .util_k_marginals(
+      synth_data = synth_num, conf_data = conf_ties, k = 1, bins = 4,
+      discretize_method = "ntile"
+    ),
+    regexp = "bins instead of 4 because of tied cut points"
+  )
+
+  expect_lt(nrow(result$cells), 5)
+
+})

From 5cd495fd5eea2ee65976868208d7299209b4d61e Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Fri, 14 Aug 2026 14:52:23 -0400
Subject: [PATCH 06/10] Add synth_vars flag to restrict marginals to
 synthesized variables

---
 NEWS.md                                |   2 +-
 R/util_k_marginals.R                   |  73 ++++++-
 man/dot-util_k_marginals.Rd            |   6 +
 man/util_k_marginals.Rd                |   6 +
 tests/testthat/test-util_k_marginals.R | 287 ++++++++++++++++++++++++-
 5 files changed, 369 insertions(+), 5 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index a3e2769..4369626 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,6 @@
 # syntheval 0.1.0
 
-* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, and discretization of numeric variables.
+* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, a `synth_vars` flag to restrict marginals to synthesized variables, and discretization of numeric variables.
 
 # syntheval 0.0.5
 
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 05bb1ae..11e8aeb 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -31,6 +31,10 @@
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param synth_varnames Optional character vector of synthesized variable
+#' names. When set, only these variables (intersected with the variables
+#' shared by both datasets) contribute marginals. Defaults to `NULL`, which
+#' places no restriction on the shared variables.
 #' @param bins Optional single integer >= 2. When set, every numeric shared
 #' variable is discretized into this many bins (fewer, with a warning, if
 #' tied quantile cut points collapse) with breaks derived from the
@@ -60,6 +64,7 @@
     n_marginals = Inf,
     priority_vars = NULL,
     weight_var = NULL,
+    synth_varnames = NULL,
     bins = NULL,
     discretize_method = c("width", "ntile", "cluster")) {
 
@@ -140,6 +145,33 @@
     weight_var
   )
 
+  if (!is.null(synth_varnames)) {
+
+    if (!(is.character(synth_varnames) && length(synth_varnames) >= 1 &&
+          !anyNA(synth_varnames))) {
+
+      stop(
+        "`synth_varnames` must be a non-empty character vector without ",
+        "missing values, or NULL"
+      )
+
+    }
+
+    shared_vars <- intersect(shared_vars, synth_varnames)
+
+    # fail here rather than at the later k check, whose message would point
+    # away from the real problem
+    if (length(shared_vars) == 0) {
+
+      stop(
+        "`synth_varnames` shares no variables with both datasets; no ",
+        "marginals can be computed"
+      )
+
+    }
+
+  }
+
   if (!is.null(bins)) {
 
     if (!(is.numeric(bins) && length(bins) == 1 && !is.na(bins) &&
@@ -231,8 +263,9 @@
           all(priority_vars %in% shared_vars))) {
 
       stop(
-        "`priority_vars` must be a character vector of variables shared by ",
-        "both datasets"
+        "`priority_vars` must be a character vector of variables available ",
+        "for marginals after applying shared-variable and `synth_varnames` ",
+        "filtering"
       )
 
     }
@@ -391,6 +424,10 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param synth_vars A logical for if only synthesized variables should
+#' contribute marginals. Only meaningful when the `eval_data` records which
+#' variables were synthesized (i.e., was built from a `postsynth`); for plain
+#' data frames all shared variables are used regardless. Defaults to `TRUE`.
 #' @param bins Optional single integer >= 2. When set, every numeric shared
 #' variable is discretized into this many bins (fewer, with a warning, if
 #' tied quantile cut points collapse) with breaks derived from the
@@ -422,6 +459,7 @@ util_k_marginals <- function(
     n_marginals = Inf,
     priority_vars = NULL,
     weight_var = NULL,
+    synth_vars = TRUE,
     bins = NULL,
     discretize_method = c("width", "ntile", "cluster")) {
 
@@ -429,6 +467,35 @@ util_k_marginals <- function(
 
   discretize_method <- match.arg(discretize_method)
 
+  if (!(is.logical(synth_vars) && length(synth_vars) == 1 &&
+        !is.na(synth_vars))) {
+
+    stop("`synth_vars` must be a single TRUE or FALSE")
+
+  }
+
+  # NULL for plain data frame eval_data, so the worker applies no restriction
+  synth_varnames <- if (synth_vars) {
+
+    eval_data$synth_vars
+
+  } else {
+
+    NULL
+
+  }
+
+  # empty metadata can only come from a user-supplied eval_data(synth_vars =)
+  # argument; fail with a message in terms of this function's arguments
+  if (!is.null(synth_varnames) && length(synth_varnames) == 0) {
+
+    stop(
+      "`eval_data` records no synthesized variables; use `synth_vars = ",
+      "FALSE` to evaluate all shared variables"
+    )
+
+  }
+
   # surface the resolved method so a forgotten discretize_method is visible
   if (!is.null(bins)) {
 
@@ -451,6 +518,7 @@ util_k_marginals <- function(
         n_marginals = n_marginals,
         priority_vars = priority_vars,
         weight_var = weight_var,
+        synth_varnames = synth_varnames,
         bins = bins,
         discretize_method = discretize_method
       )
@@ -471,6 +539,7 @@ util_k_marginals <- function(
           n_marginals = n_marginals,
           priority_vars = priority_vars,
           weight_var = weight_var,
+          synth_varnames = synth_varnames,
           bins = bins,
           discretize_method = discretize_method
         )
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 8278d9d..49cf2dc 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -13,6 +13,7 @@
   n_marginals = Inf,
   priority_vars = NULL,
   weight_var = NULL,
+  synth_varnames = NULL,
   bins = NULL,
   discretize_method = c("width", "ntile", "cluster")
 )
@@ -48,6 +49,11 @@ shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
 
+\item{synth_varnames}{Optional character vector of synthesized variable
+names. When set, only these variables (intersected with the variables
+shared by both datasets) contribute marginals. Defaults to \code{NULL}, which
+places no restriction on the shared variables.}
+
 \item{bins}{Optional single integer >= 2. When set, every numeric shared
 variable is discretized into this many bins (fewer, with a warning, if
 tied quantile cut points collapse) with breaks derived from the
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 61ffed5..9bf52c1 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -12,6 +12,7 @@ util_k_marginals(
   n_marginals = Inf,
   priority_vars = NULL,
   weight_var = NULL,
+  synth_vars = TRUE,
   bins = NULL,
   discretize_method = c("width", "ntile", "cluster")
 )
@@ -45,6 +46,11 @@ shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
 
+\item{synth_vars}{A logical for if only synthesized variables should
+contribute marginals. Only meaningful when the \code{eval_data} records which
+variables were synthesized (i.e., was built from a \code{postsynth}); for plain
+data frames all shared variables are used regardless. Defaults to \code{TRUE}.}
+
 \item{bins}{Optional single integer >= 2. When set, every numeric shared
 variable is discretized into this many bins (fewer, with a warning, if
 tied quantile cut points collapse) with breaks derived from the
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index 746b0ba..f2007db 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -499,14 +499,14 @@ test_that("invalid sampling arguments throw an error", {
     .util_k_marginals(
       synth_data = synth, conf_data = conf, k = 1, priority_vars = "zzz"
     ),
-    regexp = "`priority_vars` must be a character vector"
+    regexp = "`priority_vars` must be a character vector of variables available"
   )
 
   expect_error(
     .util_k_marginals(
       synth_data = synth, conf_data = conf, k = 1, priority_vars = 1
     ),
-    regexp = "`priority_vars` must be a character vector"
+    regexp = "`priority_vars` must be a character vector of variables available"
   )
 
 })
@@ -961,3 +961,286 @@ test_that("tied quantile cut points collapse bins with a warning", {
   expect_lt(nrow(result$cells), 5)
 
 })
+
+test_that("synth_varnames restricts the worker's variable universe", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q"),
+    c = c("m", "n", "m", "n")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p"),
+    c = c("m", "n", "m", "n")
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_sv,
+    conf_data = conf_sv,
+    k = 2,
+    synth_varnames = c("a", "b")
+  )
+
+  expect_equal(result$marginals$variables, "a, b")
+
+  # NULL means no restriction
+  result_all <- .util_k_marginals(
+    synth_data = synth_sv,
+    conf_data = conf_sv,
+    k = 2,
+    synth_varnames = NULL
+  )
+
+  expect_setequal(
+    result_all$marginals$variables,
+    c("a, b", "a, c", "b, c")
+  )
+
+})
+
+test_that("synth_vars = TRUE keeps only synthesized variables for postsynth", {
+
+  ed <- eval_data(
+    conf_data = penguins_conf,
+    synth_data = penguins_postsynth
+  )
+
+  result <- util_k_marginals(eval_data = ed, k = 1, synth_vars = TRUE)
+
+  # species and island are carried over from start_data, not synthesized
+  expect_setequal(result$marginals$variables, ed$synth_vars)
+
+})
+
+test_that("synth_vars = FALSE includes carried-over variables", {
+
+  ed <- eval_data(
+    conf_data = penguins_conf,
+    synth_data = penguins_postsynth
+  )
+
+  result <- util_k_marginals(eval_data = ed, k = 1, synth_vars = FALSE)
+
+  expect_setequal(
+    result$marginals$variables,
+    intersect(names(ed$conf_data), names(ed$synth_data))
+  )
+
+  # penguins_postsynth's start data was sampled, so the carried-over
+  # variables have their own discrepancies; including them must change the
+  # score, guarding against the flag being silently ignored
+  result_synth_only <- util_k_marginals(
+    eval_data = ed, k = 1, synth_vars = TRUE
+  )
+
+  expect_false(isTRUE(all.equal(result$score, result_synth_only$score)))
+
+})
+
+test_that("synth_vars = TRUE is a no-op for plain data frame eval_data", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  ed <- eval_data(conf_data = conf_sv, synth_data = synth_sv)
+
+  result <- util_k_marginals(eval_data = ed, k = 1, synth_vars = TRUE)
+
+  expect_setequal(result$marginals$variables, c("a", "b"))
+
+})
+
+test_that("priority_vars excluded by synth_varnames error informatively", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q"),
+    c = c("m", "n", "m", "n")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p"),
+    c = c("m", "n", "m", "n")
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_sv,
+      conf_data = conf_sv,
+      k = 1,
+      priority_vars = "c",
+      synth_varnames = c("a", "b")
+    ),
+    regexp = "`priority_vars` must be a character vector of variables available"
+  )
+
+})
+
+test_that("k is validated against the restricted variable universe", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q"),
+    c = c("m", "n", "m", "n")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p"),
+    c = c("m", "n", "m", "n")
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_sv,
+      conf_data = conf_sv,
+      k = 3,
+      synth_varnames = c("a", "b")
+    ),
+    regexp = "`k` cannot exceed"
+  )
+
+})
+
+test_that("invalid synth_vars values error", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  ed <- eval_data(conf_data = conf_sv, synth_data = synth_sv)
+
+  for (bad in list("x", c(TRUE, FALSE), NA, 1)) {
+
+    expect_error(
+      util_k_marginals(eval_data = ed, k = 1, synth_vars = bad),
+      regexp = "`synth_vars` must be a single TRUE or FALSE"
+    )
+
+  }
+
+})
+
+test_that("invalid synth_varnames values error", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  for (bad in list(character(0), NA_character_, c("a", NA), 1)) {
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_sv,
+        conf_data = conf_sv,
+        k = 1,
+        synth_varnames = bad
+      ),
+      regexp = "`synth_varnames` must be a non-empty character vector"
+    )
+
+  }
+
+})
+
+test_that("synth_varnames with no shared variables errors informatively", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_sv,
+      conf_data = conf_sv,
+      k = 1,
+      synth_varnames = "zzz"
+    ),
+    regexp = "shares no variables with both datasets"
+  )
+
+})
+
+test_that("empty synthesized-variable metadata errors informatively", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  ed <- eval_data(
+    conf_data = conf_sv,
+    synth_data = synth_sv,
+    synth_vars = character(0)
+  )
+
+  expect_error(
+    util_k_marginals(eval_data = ed, k = 1, synth_vars = TRUE),
+    regexp = "records no synthesized variables"
+  )
+
+  # synth_vars = FALSE ignores the empty metadata
+  expect_no_error(util_k_marginals(eval_data = ed, k = 1, synth_vars = FALSE))
+
+})
+
+test_that("wrapper restriction bounds k by the synthesized-variable set", {
+
+  conf_sv <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  synth_sv <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  ed <- eval_data(
+    conf_data = conf_sv,
+    synth_data = synth_sv,
+    synth_vars = "a"
+  )
+
+  expect_error(
+    util_k_marginals(eval_data = ed, k = 2, synth_vars = TRUE),
+    regexp = "`k` cannot exceed"
+  )
+
+  # the same call succeeds once the restriction is lifted
+  expect_no_error(util_k_marginals(eval_data = ed, k = 2, synth_vars = FALSE))
+
+})

From ea4760954f68841daf4d1a10efd7ac79ff36f3de Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Fri, 14 Aug 2026 15:20:24 -0400
Subject: [PATCH 07/10] Add na.rm handling of missing values

---
 NEWS.md                                |   2 +-
 R/util_k_marginals.R                   |  91 ++++++-
 man/dot-util_k_marginals.Rd            |  15 +-
 man/util_k_marginals.Rd                |  15 +-
 tests/testthat/test-util_k_marginals.R | 334 +++++++++++++++++++++++++
 5 files changed, 440 insertions(+), 17 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 4369626..c557aa2 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,6 @@
 # syntheval 0.1.0
 
-* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, a `synth_vars` flag to restrict marginals to synthesized variables, and discretization of numeric variables.
+* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, a `synth_vars` flag to restrict marginals to synthesized variables, `na.rm` handling of missing values, and discretization of numeric variables.
 
 # syntheval 0.0.5
 
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 11e8aeb..99dee6e 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -35,12 +35,19 @@
 #' names. When set, only these variables (intersected with the variables
 #' shared by both datasets) contribute marginals. Defaults to `NULL`, which
 #' places no restriction on the shared variables.
+#' @param na.rm A logical for ignoring `NA` values in proportion
+#' calculations. When `FALSE`, missing values form their own `"NA"` level in
+#' each marginal (and a message lists the affected variables); when `TRUE`,
+#' rows with a missing value are dropped from each marginal that uses the
+#' affected variable, leaving marginals of complete variables untouched.
+#' Defaults to `FALSE`.
 #' @param bins Optional single integer >= 2. When set, every numeric shared
 #' variable is discretized into this many bins (fewer, with a warning, if
 #' tied quantile cut points collapse) with breaks derived from the
-#' confidential data and applied to both datasets; the outer bins extend to
-#' +/-Inf so synthetic values outside the confidential range land in edge
-#' bins. Defaults to `NULL` (no discretization).
+#' observed (non-missing) confidential values and applied to both datasets;
+#' the outer bins extend to +/-Inf so synthetic values outside the
+#' confidential range land in edge bins, and missing values follow `na.rm`
+#' like any other variable. Defaults to `NULL` (no discretization).
 #' @param discretize_method Method used to place bin breaks when `bins` is
 #' set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
 #' for univariate k-means clustering (set a seed before calling for
@@ -65,11 +72,18 @@
     priority_vars = NULL,
     weight_var = NULL,
     synth_varnames = NULL,
+    na.rm = FALSE,
     bins = NULL,
     discretize_method = c("width", "ntile", "cluster")) {
 
   discretize_method <- match.arg(discretize_method)
 
+  if (!(is.logical(na.rm) && length(na.rm) == 1 && !is.na(na.rm))) {
+
+    stop("`na.rm` must be a single TRUE or FALSE")
+
+  }
+
   for (keep in list(keep_marginals, keep_cells, n_marginals)) {
 
     if (!(is.numeric(keep) && length(keep) == 1 && !is.na(keep) &&
@@ -187,12 +201,14 @@
 
     for (var in numeric_vars) {
 
-      conf_values <- conf_data[[var]]
+      # cut points derive from observed values; missing values follow na.rm
+      # like every other variable, becoming an NA bin or dropped rows
+      conf_values <- conf_data[[var]][!is.na(conf_data[[var]])]
 
-      if (!all(is.finite(conf_values))) {
+      if (!all(is.finite(conf_values)) || length(conf_values) == 0) {
 
         stop(
-          "numeric variables must be finite to discretize; `", var,
+          "observed numeric values must be finite to discretize; `", var,
           "` is not"
         )
 
@@ -245,7 +261,7 @@
       }
 
       synth_data[[var]] <- cut(x = synth_data[[var]], breaks = breaks)
-      conf_data[[var]] <- cut(x = conf_values, breaks = breaks)
+      conf_data[[var]] <- cut(x = conf_data[[var]], breaks = breaks)
 
     }
 
@@ -257,6 +273,32 @@
 
   }
 
+  if (!na.rm) {
+
+    na_vars <- shared_vars[
+      purrr::map_lgl(
+        .x = shared_vars,
+        .f = \(v) anyNA(synth_data[[v]]) || anyNA(conf_data[[v]])
+      )
+    ]
+
+    if (length(na_vars) > 0) {
+
+      message(
+        "Some variables contain missing data: ",
+        paste(na_vars, collapse = ", ")
+      )
+
+    }
+
+    # missing values become their own "NA" level so they participate in
+    # marginals; numeric variables without bins keep NA, which count() still
+    # groups separately
+    synth_data[shared_vars] <- convert_na_to_level(synth_data[shared_vars])
+    conf_data[shared_vars] <- convert_na_to_level(conf_data[shared_vars])
+
+  }
+
   if (!is.null(priority_vars)) {
 
     if (!(is.character(priority_vars) &&
@@ -298,6 +340,25 @@
   # proportions are weight shares instead of row shares
   process_data <- function(data, vars, prop_name) {
 
+    if (na.rm) {
+
+      data <- dplyr::filter(
+        data,
+        !dplyr::if_any(.cols = dplyr::all_of(vars), .fns = is.na)
+      )
+
+      if (nrow(data) == 0) {
+
+        stop(
+          "no rows remain for the marginal over ",
+          paste(vars, collapse = ", "),
+          " after removing missing values"
+        )
+
+      }
+
+    }
+
     if (is.null(weight_var)) {
 
       counts <- dplyr::count(data, dplyr::across(dplyr::all_of(vars)))
@@ -428,12 +489,19 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' contribute marginals. Only meaningful when the `eval_data` records which
 #' variables were synthesized (i.e., was built from a `postsynth`); for plain
 #' data frames all shared variables are used regardless. Defaults to `TRUE`.
+#' @param na.rm A logical for ignoring `NA` values in proportion
+#' calculations. When `FALSE`, missing values form their own `"NA"` level in
+#' each marginal (and a message lists the affected variables); when `TRUE`,
+#' rows with a missing value are dropped from each marginal that uses the
+#' affected variable, leaving marginals of complete variables untouched.
+#' Defaults to `FALSE`.
 #' @param bins Optional single integer >= 2. When set, every numeric shared
 #' variable is discretized into this many bins (fewer, with a warning, if
 #' tied quantile cut points collapse) with breaks derived from the
-#' confidential data and applied to both datasets; the outer bins extend to
-#' +/-Inf so synthetic values outside the confidential range land in edge
-#' bins. Defaults to `NULL` (no discretization).
+#' observed (non-missing) confidential values and applied to both datasets;
+#' the outer bins extend to +/-Inf so synthetic values outside the
+#' confidential range land in edge bins, and missing values follow `na.rm`
+#' like any other variable. Defaults to `NULL` (no discretization).
 #' @param discretize_method Method used to place bin breaks when `bins` is
 #' set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
 #' for univariate k-means clustering (set a seed before calling for
@@ -460,6 +528,7 @@ util_k_marginals <- function(
     priority_vars = NULL,
     weight_var = NULL,
     synth_vars = TRUE,
+    na.rm = FALSE,
     bins = NULL,
     discretize_method = c("width", "ntile", "cluster")) {
 
@@ -519,6 +588,7 @@ util_k_marginals <- function(
         priority_vars = priority_vars,
         weight_var = weight_var,
         synth_varnames = synth_varnames,
+        na.rm = na.rm,
         bins = bins,
         discretize_method = discretize_method
       )
@@ -540,6 +610,7 @@ util_k_marginals <- function(
           priority_vars = priority_vars,
           weight_var = weight_var,
           synth_varnames = synth_varnames,
+          na.rm = na.rm,
           bins = bins,
           discretize_method = discretize_method
         )
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 49cf2dc..c568de4 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -14,6 +14,7 @@
   priority_vars = NULL,
   weight_var = NULL,
   synth_varnames = NULL,
+  na.rm = FALSE,
   bins = NULL,
   discretize_method = c("width", "ntile", "cluster")
 )
@@ -54,12 +55,20 @@ names. When set, only these variables (intersected with the variables
 shared by both datasets) contribute marginals. Defaults to \code{NULL}, which
 places no restriction on the shared variables.}
 
+\item{na.rm}{A logical for ignoring \code{NA} values in proportion
+calculations. When \code{FALSE}, missing values form their own \code{"NA"} level in
+each marginal (and a message lists the affected variables); when \code{TRUE},
+rows with a missing value are dropped from each marginal that uses the
+affected variable, leaving marginals of complete variables untouched.
+Defaults to \code{FALSE}.}
+
 \item{bins}{Optional single integer >= 2. When set, every numeric shared
 variable is discretized into this many bins (fewer, with a warning, if
 tied quantile cut points collapse) with breaks derived from the
-confidential data and applied to both datasets; the outer bins extend to
-+/-Inf so synthetic values outside the confidential range land in edge
-bins. Defaults to \code{NULL} (no discretization).}
+observed (non-missing) confidential values and applied to both datasets;
+the outer bins extend to +/-Inf so synthetic values outside the
+confidential range land in edge bins, and missing values follow \code{na.rm}
+like any other variable. Defaults to \code{NULL} (no discretization).}
 
 \item{discretize_method}{Method used to place bin breaks when \code{bins} is
 set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 9bf52c1..5e5fe51 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -13,6 +13,7 @@ util_k_marginals(
   priority_vars = NULL,
   weight_var = NULL,
   synth_vars = TRUE,
+  na.rm = FALSE,
   bins = NULL,
   discretize_method = c("width", "ntile", "cluster")
 )
@@ -51,12 +52,20 @@ contribute marginals. Only meaningful when the \code{eval_data} records which
 variables were synthesized (i.e., was built from a \code{postsynth}); for plain
 data frames all shared variables are used regardless. Defaults to \code{TRUE}.}
 
+\item{na.rm}{A logical for ignoring \code{NA} values in proportion
+calculations. When \code{FALSE}, missing values form their own \code{"NA"} level in
+each marginal (and a message lists the affected variables); when \code{TRUE},
+rows with a missing value are dropped from each marginal that uses the
+affected variable, leaving marginals of complete variables untouched.
+Defaults to \code{FALSE}.}
+
 \item{bins}{Optional single integer >= 2. When set, every numeric shared
 variable is discretized into this many bins (fewer, with a warning, if
 tied quantile cut points collapse) with breaks derived from the
-confidential data and applied to both datasets; the outer bins extend to
-+/-Inf so synthetic values outside the confidential range land in edge
-bins. Defaults to \code{NULL} (no discretization).}
+observed (non-missing) confidential values and applied to both datasets;
+the outer bins extend to +/-Inf so synthetic values outside the
+confidential range land in edge bins, and missing values follow \code{na.rm}
+like any other variable. Defaults to \code{NULL} (no discretization).}
 
 \item{discretize_method}{Method used to place bin breaks when \code{bins} is
 set: "width" for fixed binwidths, "ntile" for quantile bins, or "cluster"
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index f2007db..ca22df4 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -1244,3 +1244,337 @@ test_that("wrapper restriction bounds k by the synthesized-variable set", {
   expect_no_error(util_k_marginals(eval_data = ed, k = 2, synth_vars = FALSE))
 
 })
+
+# NA handling
+#
+# conf_na            synth_na
+#   a   b              a   b
+#   x   p              x   p
+#   x   p              y   q
+#   y   p              y   q
+#   NA  q              NA  p
+#
+# na.rm = FALSE, k = 1 (NA is its own level)
+#   marginal a: conf (x = 0.50, y = 0.25, NA = 0.25),
+#               synth (x = 0.25, y = 0.50, NA = 0.25)
+#     MabsDD = mean(0.25, 0.25, 0) = 1/6
+#   marginal b: MabsDD = mean(0.25, 0.25) = 0.25
+#   score = (1 - mean(1/6, 1/4)) * 1000 = 19000/24
+#
+# na.rm = TRUE, k = 1 (rows dropped per marginal)
+#   marginal a (3 rows each): conf (x = 2/3, y = 1/3),
+#                             synth (x = 1/3, y = 2/3)
+#     MabsDD = mean(1/3, 1/3) = 1/3
+#   marginal b (all 4 rows): MabsDD = 0.25
+#   score = (1 - mean(1/3, 1/4)) * 1000 = 17000/24
+
+conf_na <- tibble::tibble(
+  a = c("x", "x", "y", NA),
+  b = c("p", "p", "p", "q")
+)
+
+synth_na <- tibble::tibble(
+  a = c("x", "y", "y", NA),
+  b = c("p", "q", "q", "p")
+)
+
+test_that("NA values form their own level by default", {
+
+  result <- suppressMessages(
+    .util_k_marginals(synth_data = synth_na, conf_data = conf_na, k = 1)
+  )
+
+  expect_equal(result$score, 19000 / 24)
+
+  expect_true("NA" %in% result$cells$cell)
+
+})
+
+test_that("na.rm = TRUE drops missing values per marginal", {
+
+  result <- .util_k_marginals(
+    synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = TRUE
+  )
+
+  expect_equal(result$score, 17000 / 24)
+
+  expect_false("NA" %in% result$cells$cell)
+
+})
+
+test_that("missing data triggers a message when na.rm = FALSE", {
+
+  expect_message(
+    .util_k_marginals(synth_data = synth_na, conf_data = conf_na, k = 1),
+    regexp = "contain missing data: a"
+  )
+
+  expect_no_message(
+    .util_k_marginals(
+      synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = TRUE
+    )
+  )
+
+})
+
+test_that("invalid na.rm values error", {
+
+  for (bad in list("x", c(TRUE, FALSE), NA, 1)) {
+
+    expect_error(
+      .util_k_marginals(
+        synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = bad
+      ),
+      regexp = "`na.rm` must be a single TRUE or FALSE"
+    )
+
+  }
+
+})
+
+test_that("a literal 'NA' level alongside true NA values errors", {
+
+  conf_lit <- tibble::tibble(a = c("NA", "x", NA))
+  synth_lit <- tibble::tibble(a = c("x", "x", "x"))
+
+  expect_error(
+    suppressMessages(
+      .util_k_marginals(synth_data = synth_lit, conf_data = conf_lit, k = 1)
+    ),
+    regexp = "'NA' already exists"
+  )
+
+})
+
+test_that("numeric NA values land in an NA bin or are dropped", {
+
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+  synth_num <- tibble::tibble(v = c(1, 4, NA, NA))
+
+  kept <- suppressMessages(
+    .util_k_marginals(
+      synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2
+    )
+  )
+
+  # cut() maps synthetic NAs to an NA bin that becomes its own level
+  expect_true("NA" %in% kept$cells$cell)
+
+  dropped <- .util_k_marginals(
+    synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2,
+    na.rm = TRUE
+  )
+
+  expect_false("NA" %in% dropped$cells$cell)
+
+  # with the NAs dropped, the two synthetic values split evenly like the
+  # confidential data, so the marginal matches exactly
+  expect_equal(dropped$score, 1000)
+
+})
+
+test_that("a marginal with no complete rows errors under na.rm = TRUE", {
+
+  conf_all_na <- tibble::tibble(a = c(NA_character_, NA_character_))
+  synth_ok <- tibble::tibble(a = c("x", "y"))
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_ok, conf_data = conf_all_na, k = 1, na.rm = TRUE
+    ),
+    regexp = "no rows remain"
+  )
+
+})
+
+test_that("util_k_marginals passes na.rm through", {
+
+  ed <- eval_data(conf_data = conf_na, synth_data = synth_na)
+
+  expect_equal(
+    util_k_marginals(eval_data = ed, k = 1, na.rm = TRUE)$score,
+    17000 / 24
+  )
+
+})
+
+test_that("na.rm = TRUE drops rows per combination, not globally", {
+
+  # a's NA sits on a different row in each dataset; b and c are identical
+  # across datasets
+  #
+  # pairwise deletion, k = 2:
+  #   (b, c): all 4 rows, identical -> MabsDD = 0
+  #   (a, b): conf drops row 4 -> (x,p), (x,q), (y,p) each 1/3
+  #           synth drops row 3 -> (x,p), (x,q), (y,q) each 1/3
+  #           MabsDD = mean(0, 0, 1/3, 1/3) = 1/6
+  #   (a, c): conf -> (x,m) = 2/3, (y,n) = 1/3; synth identical -> MabsDD = 0
+  #   score = (1 - mean(0, 1/6, 0)) * 1000 = 17000/18
+  #
+  # global (listwise) deletion would filter different rows from each dataset
+  # and corrupt the complete (b, c) marginal, so madd = 0 for (b, c) is the
+  # discriminating assertion
+  conf_pair <- tibble::tibble(
+    a = c("x", "x", "y", NA),
+    b = c("p", "q", "p", "q"),
+    c = c("m", "m", "n", "n")
+  )
+
+  synth_pair <- tibble::tibble(
+    a = c("x", "x", NA, "y"),
+    b = c("p", "q", "p", "q"),
+    c = c("m", "m", "n", "n")
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_pair, conf_data = conf_pair, k = 2, na.rm = TRUE
+  )
+
+  expect_equal(result$score, 17000 / 18)
+
+  bc <- dplyr::filter(result$marginals, .data$variables == "b, c")
+
+  expect_equal(bc$madd, 0)
+
+})
+
+test_that("the missing-data message lists every affected variable", {
+
+  conf_two <- tibble::tibble(
+    a = c("x", NA),
+    b = c(NA, "q"),
+    c = c("m", "n")
+  )
+
+  synth_two <- tibble::tibble(
+    a = c("x", "y"),
+    b = c("p", "q"),
+    c = c("m", "n")
+  )
+
+  expect_message(
+    .util_k_marginals(synth_data = synth_two, conf_data = conf_two, k = 1),
+    regexp = "contain missing data: a, b"
+  )
+
+})
+
+test_that("a literal 'NA' level in the synthetic data also errors", {
+
+  conf_lit <- tibble::tibble(a = c("x", "x", "x"))
+  synth_lit <- tibble::tibble(a = c("NA", "x", NA))
+
+  expect_error(
+    suppressMessages(
+      .util_k_marginals(synth_data = synth_lit, conf_data = conf_lit, k = 1)
+    ),
+    regexp = "'NA' already exists"
+  )
+
+  # both datasets carrying the collision still errors
+  expect_error(
+    suppressMessages(
+      .util_k_marginals(synth_data = synth_lit, conf_data = synth_lit, k = 1)
+    ),
+    regexp = "'NA' already exists"
+  )
+
+})
+
+test_that("variables excluded by synth_varnames do not drive NA handling", {
+
+  # a has missing values but is filtered out, so no message and no NA cells
+  conf_excl <- tibble::tibble(
+    a = c("x", NA),
+    b = c("p", "q")
+  )
+
+  synth_excl <- tibble::tibble(
+    a = c(NA, "y"),
+    b = c("p", "p")
+  )
+
+  expect_no_message(
+    result <- .util_k_marginals(
+      synth_data = synth_excl,
+      conf_data = conf_excl,
+      k = 1,
+      synth_varnames = "b"
+    )
+  )
+
+  expect_false("NA" %in% result$cells$cell)
+
+})
+
+test_that("weighted proportions drop missing rows before computing shares", {
+
+  # na.rm = TRUE drops each dataset's NA row, and weight shares are computed
+  # from the surviving rows' weights:
+  #   conf keeps weights 1, 1, 2 -> x = 2/4, y = 2/4
+  #   synth keeps weights 1, 3 -> x = 1/4, y = 3/4
+  #   MabsDD = mean(0.25, 0.25) = 0.25 -> score 750
+  # dividing by the full weight total (including the dropped 10s) would give
+  # a different score, so 750 pins down the recomputation
+  conf_w <- tibble::tibble(
+    a = c("x", "x", "y", NA),
+    w = c(1, 1, 2, 10)
+  )
+
+  synth_w <- tibble::tibble(
+    a = c("x", "y", NA),
+    w = c(1, 3, 10)
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_w,
+    conf_data = conf_w,
+    k = 1,
+    weight_var = "w",
+    na.rm = TRUE
+  )
+
+  expect_equal(result$score, 750)
+
+})
+
+test_that("confidential numeric NA values discretize under both na.rm modes", {
+
+  # breaks derive from the observed confidential values (1:4, cut at 2.5)
+  conf_num_na <- tibble::tibble(v = c(1, 2, 3, 4, NA))
+  synth_num_na <- tibble::tibble(v = c(1, 2, 4, 4, NA))
+
+  # na.rm = FALSE: both datasets bin as low 2/5, high 2/5, NA 1/5
+  kept <- suppressMessages(
+    .util_k_marginals(
+      synth_data = synth_num_na, conf_data = conf_num_na, k = 1, bins = 2
+    )
+  )
+
+  expect_true("NA" %in% kept$cells$cell)
+  expect_equal(kept$score, 1000)
+
+  # na.rm = TRUE: observed values bin as low 2/4, high 2/4 in both datasets
+  dropped <- .util_k_marginals(
+    synth_data = synth_num_na, conf_data = conf_num_na, k = 1, bins = 2,
+    na.rm = TRUE
+  )
+
+  expect_false("NA" %in% dropped$cells$cell)
+  expect_equal(dropped$score, 1000)
+
+})
+
+test_that("infinite confidential values still refuse to discretize", {
+
+  conf_inf <- tibble::tibble(v = c(1, 2, 3, Inf))
+  synth_num <- tibble::tibble(v = c(1, 2, 3, 3))
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_num, conf_data = conf_inf, k = 1, bins = 2
+    ),
+    regexp = "must be finite to discretize"
+  )
+
+})

From a67d366e01cc3ecfdc459b5396e3c2b0bfad400a Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Fri, 14 Aug 2026 16:20:55 -0400
Subject: [PATCH 08/10] Add group_by stratification with per-stratum scores

---
 NEWS.md                                |   2 +-
 R/util_k_marginals.R                   | 378 +++++++++++-------
 man/dot-util_k_marginals.Rd            |  12 +-
 man/util_k_marginals.Rd                |  12 +-
 tests/testthat/test-util_k_marginals.R | 522 ++++++++++++++++++++++++-
 5 files changed, 773 insertions(+), 153 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index c557aa2..76c0d02 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,6 @@
 # syntheval 0.1.0
 
-* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, a `synth_vars` flag to restrict marginals to synthesized variables, `na.rm` handling of missing values, and discretization of numeric variables.
+* Add `util_k_marginals()` to calculate the k-marginals metric for 1-, 2-, and 3-way marginals, with worst-marginal and worst-cell output, marginal sampling with priority variables, sample weights, a `synth_vars` flag to restrict marginals to synthesized variables, `na.rm` handling of missing values, `group_by` stratification with per-stratum scores, and discretization of numeric variables.
 
 # syntheval 0.0.5
 
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 99dee6e..0b6e066 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -31,6 +31,12 @@
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param group_by Optional character vector of grouping variable names
+#' present in both datasets. When set, the grouping variables are excluded
+#' from the marginals and the metric is computed within each stratum observed
+#' in the confidential data; the headline score is the mean of the per-stratum
+#' scores weighted by each stratum's confidential share (weight share when
+#' `weight_var` is set). Defaults to `NULL` (no stratification).
 #' @param synth_varnames Optional character vector of synthesized variable
 #' names. When set, only these variables (intersected with the variables
 #' shared by both datasets) contribute marginals. Defaults to `NULL`, which
@@ -60,115 +66,113 @@
 #' `cells`, a tibble with the synthetic and confidential proportions and their
 #' absolute difference for every cell, worst first. `score` is always computed
 #' from all evaluated marginals, even when `keep_marginals` or `keep_cells`
-#' truncate the detail tables.
+#' truncate the detail tables. When `group_by` is set, `marginals` and
+#' `cells` gain the grouping columns and a fourth element `group_scores`
+#' reports each stratum's share and score, worst first.
 #'
 .util_k_marginals <- function(
-    synth_data,
-    conf_data,
-    k,
-    keep_marginals = Inf,
-    keep_cells = Inf,
-    n_marginals = Inf,
-    priority_vars = NULL,
-    weight_var = NULL,
-    synth_varnames = NULL,
-    na.rm = FALSE,
-    bins = NULL,
-    discretize_method = c("width", "ntile", "cluster")) {
-
+  synth_data,
+  conf_data,
+  k,
+  keep_marginals = Inf,
+  keep_cells = Inf,
+  n_marginals = Inf,
+  priority_vars = NULL,
+  weight_var = NULL,
+  group_by = NULL,
+  synth_varnames = NULL,
+  na.rm = FALSE,
+  bins = NULL,
+  discretize_method = c("width", "ntile", "cluster")
+) {
   discretize_method <- match.arg(discretize_method)
 
   if (!(is.logical(na.rm) && length(na.rm) == 1 && !is.na(na.rm))) {
-
     stop("`na.rm` must be a single TRUE or FALSE")
-
   }
 
   for (keep in list(keep_marginals, keep_cells, n_marginals)) {
-
     if (!(is.numeric(keep) && length(keep) == 1 && !is.na(keep) &&
-          keep >= 1 && keep == floor(keep))) {
-
+      keep >= 1 && keep == floor(keep))) {
       stop(
         "`keep_marginals`, `keep_cells`, and `n_marginals` must be single ",
         "integers >= 1 or Inf"
       )
-
     }
-
   }
 
   if (!(is.numeric(k) && length(k) == 1 && k %in% 1:3)) {
-
     stop("`k` must be a single integer between 1 and 3")
-
   }
 
   stopifnot(inherits(synth_data, "data.frame"))
   stopifnot(inherits(conf_data, "data.frame"))
 
   if (nrow(synth_data) == 0 || nrow(conf_data) == 0) {
-
     stop("`synth_data` and `conf_data` must each contain at least one row")
-
   }
 
   if (!is.null(weight_var)) {
-
     if (!(is.character(weight_var) && length(weight_var) == 1)) {
-
       stop("`weight_var` must be a single character string")
-
     }
 
     if (!(weight_var %in% names(synth_data) &&
-          weight_var %in% names(conf_data))) {
-
+      weight_var %in% names(conf_data))) {
       stop("`weight_var` must be a column in both datasets")
-
     }
 
     if (!(is.numeric(synth_data[[weight_var]]) &&
-          is.numeric(conf_data[[weight_var]]))) {
-
+      is.numeric(conf_data[[weight_var]]))) {
       stop("`weight_var` must be a numeric column in both datasets")
-
     }
 
     # invalid weights break the probability interpretation of proportions
     for (weights in list(synth_data[[weight_var]], conf_data[[weight_var]])) {
-
       if (!all(is.finite(weights)) || any(weights < 0) ||
-          sum(weights) <= 0) {
-
+        sum(weights) <= 0) {
         stop(
           "`weight_var` values must be finite and non-negative with a ",
           "positive total in both datasets"
         )
-
       }
+    }
+  }
+
+  if (!is.null(group_by)) {
+    if (!(is.character(group_by) && length(group_by) >= 1 &&
+      !anyNA(group_by) &&
+      all(group_by %in% names(synth_data)) &&
+      all(group_by %in% names(conf_data)))) {
+      stop(
+        "`group_by` must be a character vector of variables present in ",
+        "both datasets"
+      )
+    }
 
+    if (anyDuplicated(group_by) > 0) {
+      stop("`group_by` must not contain duplicate variable names")
     }
 
+    if (!is.null(weight_var) && weight_var %in% group_by) {
+      stop("`group_by` cannot include `weight_var`")
+    }
   }
 
   # only variables present in both datasets contribute marginals; the weight
-  # column is never itself a marginal
+  # column and grouping variables are never themselves marginals
   shared_vars <- setdiff(
     intersect(names(synth_data), names(conf_data)),
-    weight_var
+    c(weight_var, group_by)
   )
 
   if (!is.null(synth_varnames)) {
-
     if (!(is.character(synth_varnames) && length(synth_varnames) >= 1 &&
-          !anyNA(synth_varnames))) {
-
+      !anyNA(synth_varnames))) {
       stop(
         "`synth_varnames` must be a non-empty character vector without ",
         "missing values, or NULL"
       )
-
     }
 
     shared_vars <- intersect(shared_vars, synth_varnames)
@@ -176,23 +180,17 @@
     # fail here rather than at the later k check, whose message would point
     # away from the real problem
     if (length(shared_vars) == 0) {
-
       stop(
-        "`synth_varnames` shares no variables with both datasets; no ",
-        "marginals can be computed"
+        "`synth_varnames` matches no variables available for marginals ",
+        "after shared-variable, `group_by`, and `weight_var` filtering"
       )
-
     }
-
   }
 
   if (!is.null(bins)) {
-
     if (!(is.numeric(bins) && length(bins) == 1 && !is.na(bins) &&
-          bins >= 2 && bins == floor(bins))) {
-
+      bins >= 2 && bins == floor(bins))) {
       stop("`bins` must be a single integer >= 2")
-
     }
 
     numeric_vars <- shared_vars[
@@ -200,26 +198,21 @@
     ]
 
     for (var in numeric_vars) {
-
       # cut points derive from observed values; missing values follow na.rm
       # like every other variable, becoming an NA bin or dropped rows
       conf_values <- conf_data[[var]][!is.na(conf_data[[var]])]
 
       if (!all(is.finite(conf_values)) || length(conf_values) == 0) {
-
         stop(
           "observed numeric values must be finite to discretize; `", var,
           "` is not"
         )
-
       }
 
       if (dplyr::n_distinct(conf_values) < bins) {
-
         stop(
           "`", var, "` has fewer distinct confidential values than `bins`"
         )
-
       }
 
       # interior cut points always derive from the confidential data; each
@@ -237,13 +230,11 @@
           names = FALSE
         )[2:bins],
         cluster = {
-
           centers <- sort(
             stats::kmeans(x = conf_values, centers = bins)$centers[, 1]
           )
 
           (centers[-1] + centers[-length(centers)]) / 2
-
         }
       )
 
@@ -252,66 +243,77 @@
       breaks <- unique(c(-Inf, cut_points, Inf))
 
       if (length(breaks) - 1 < bins) {
-
         warning(
           "`", var, "` was discretized into ", length(breaks) - 1,
           " bins instead of ", bins, " because of tied cut points"
         )
-
       }
 
       synth_data[[var]] <- cut(x = synth_data[[var]], breaks = breaks)
       conf_data[[var]] <- cut(x = conf_data[[var]], breaks = breaks)
-
     }
-
   }
 
   if (length(shared_vars) < k) {
-
-    stop("`k` cannot exceed the number of variables shared by both datasets")
-
+    stop(
+      "`k` cannot exceed the number of variables available for marginals ",
+      "after shared-variable, `group_by`, `weight_var`, and ",
+      "`synth_varnames` filtering"
+    )
   }
 
-  if (!na.rm) {
+  # grouping variables get the same NA treatment as marginal variables: an
+  # "NA" stratum by default, or their incomplete rows dropped entirely
+  na_vars_scope <- c(shared_vars, group_by)
 
-    na_vars <- shared_vars[
+  if (!na.rm) {
+    na_vars <- na_vars_scope[
       purrr::map_lgl(
-        .x = shared_vars,
+        .x = na_vars_scope,
         .f = \(v) anyNA(synth_data[[v]]) || anyNA(conf_data[[v]])
       )
     ]
 
     if (length(na_vars) > 0) {
-
       message(
         "Some variables contain missing data: ",
         paste(na_vars, collapse = ", ")
       )
-
     }
 
     # missing values become their own "NA" level so they participate in
     # marginals; numeric variables without bins keep NA, which count() still
     # groups separately
-    synth_data[shared_vars] <- convert_na_to_level(synth_data[shared_vars])
-    conf_data[shared_vars] <- convert_na_to_level(conf_data[shared_vars])
+    synth_data[na_vars_scope] <- convert_na_to_level(synth_data[na_vars_scope])
+    conf_data[na_vars_scope] <- convert_na_to_level(conf_data[na_vars_scope])
+  } else if (!is.null(group_by)) {
+    # rows without a stratum cannot enter any stratified marginal
+    synth_data <- dplyr::filter(
+      synth_data,
+      !dplyr::if_any(.cols = dplyr::all_of(group_by), .fns = is.na)
+    )
+
+    conf_data <- dplyr::filter(
+      conf_data,
+      !dplyr::if_any(.cols = dplyr::all_of(group_by), .fns = is.na)
+    )
 
+    if (nrow(conf_data) == 0) {
+      stop(
+        "no confidential rows remain after removing missing `group_by` values"
+      )
+    }
   }
 
   if (!is.null(priority_vars)) {
-
     if (!(is.character(priority_vars) &&
-          all(priority_vars %in% shared_vars))) {
-
+      all(priority_vars %in% shared_vars))) {
       stop(
         "`priority_vars` must be a character vector of variables available ",
-        "for marginals after applying shared-variable and `synth_varnames` ",
-        "filtering"
+        "for marginals after shared-variable, `group_by`, `weight_var`, ",
+        "and `synth_varnames` filtering"
       )
-
     }
-
   }
 
   kmarginals_vars <- t(utils::combn(x = shared_vars, m = k))
@@ -319,7 +321,6 @@
   # sample combinations down to n_marginals, always keeping combinations that
   # contain a priority variable
   if (nrow(kmarginals_vars) > n_marginals) {
-
     is_priority <- apply(
       X = kmarginals_vars,
       MARGIN = 1,
@@ -331,46 +332,37 @@
     sampled_rows <- sample(x = which(!is_priority), size = n_sampled)
 
     kmarginals_vars <- kmarginals_vars[
-      sort(c(which(is_priority), sampled_rows)), , drop = FALSE
+      sort(c(which(is_priority), sampled_rows)), ,
+      drop = FALSE
     ]
-
   }
 
   # cell proportions for one dataset over one set of variables; weighted
   # proportions are weight shares instead of row shares
-  process_data <- function(data, vars, prop_name) {
-
+  process_data <- function(data, vars, prop_name, allow_empty) {
     if (na.rm) {
-
       data <- dplyr::filter(
         data,
         !dplyr::if_any(.cols = dplyr::all_of(vars), .fns = is.na)
       )
 
-      if (nrow(data) == 0) {
-
+      if (nrow(data) == 0 && !allow_empty) {
         stop(
           "no rows remain for the marginal over ",
           paste(vars, collapse = ", "),
           " after removing missing values"
         )
-
       }
-
     }
 
     if (is.null(weight_var)) {
-
       counts <- dplyr::count(data, dplyr::across(dplyr::all_of(vars)))
-
     } else {
-
       counts <- dplyr::count(
         data,
         dplyr::across(dplyr::all_of(vars)),
         wt = .data[[weight_var]]
       )
-
     }
 
     props <- counts |>
@@ -378,16 +370,23 @@
       dplyr::select(-"n")
 
     return(props)
-
   }
-  
+
   # per-cell differences for one set of variables; cells absent from one
   # dataset count as 0
-  marginal_cells <- function(vars) {
-
+  marginal_cells <- function(vars, synth_d, conf_d, allow_empty_synth) {
+    # only the synthetic side may be empty (a stratum the synthesis never
+    # produced); a confidential marginal with no rows has nothing to score
+    # against and errors inside process_data
     cells <- dplyr::full_join(
-      process_data(data = synth_data, vars = vars, prop_name = "prop_synth"),
-      process_data(data = conf_data, vars = vars, prop_name = "prop_conf"),
+      process_data(
+        data = synth_d, vars = vars, prop_name = "prop_synth",
+        allow_empty = allow_empty_synth
+      ),
+      process_data(
+        data = conf_d, vars = vars, prop_name = "prop_conf",
+        allow_empty = FALSE
+      ),
       by = vars
     ) |>
       tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0)) |>
@@ -403,37 +402,123 @@
     # per-combination summary; the prop columns show the direction of the
     # discrepancy, not just its size
     return(cells)
+  }
 
+  # all per-cell differences for one (synth, conf) pair of datasets
+  compute_cells <- function(synth_d, conf_d, allow_empty_synth) {
+    cells <- purrr::map(
+      .x = seq_len(nrow(kmarginals_vars)),
+      .f = \(i) marginal_cells(
+        vars = kmarginals_vars[i, ],
+        synth_d = synth_d,
+        conf_d = conf_d,
+        allow_empty_synth = allow_empty_synth
+      )
+    ) |>
+      purrr::list_rbind()
+
+    return(cells)
   }
 
-  # per-cell differences across all k-way marginals, worst cells first
-  cells <- purrr::map(
-    .x = seq_len(nrow(kmarginals_vars)),
-    .f = \(i) marginal_cells(vars = kmarginals_vars[i, ])
-  ) |>
-    purrr::list_rbind() |>
+  if (is.null(group_by)) {
+    cells <- compute_cells(
+      synth_d = synth_data, conf_d = conf_data, allow_empty_synth = FALSE
+    ) |>
+      dplyr::arrange(dplyr::desc(.data$abs_diff))
+
+    # MabsDD per combination, worst marginals first
+    marginals <- cells |>
+      dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables") |>
+      dplyr::arrange(dplyr::desc(.data$madd))
+
+    # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000];
+    # computed from all marginals before any truncation
+    score <- (1 - mean(marginals$madd)) * 1000
+
+    result <- structure(
+      list(
+        score = score,
+        marginals = utils::head(marginals, n = keep_marginals),
+        cells = utils::head(cells, n = keep_cells)
+      ),
+      class = "k_marginals"
+    )
+
+    return(result)
+  }
+
+  # strata are defined by the confidential data; a stratum with no synthetic
+  # rows scores against all-zero synthetic proportions. Shares are the
+  # confidential row (or weight) share of each stratum, computed once
+  conf_totals <- if (is.null(weight_var)) {
+    rep(1, nrow(conf_data))
+  } else {
+    conf_data[[weight_var]]
+  }
+
+  strata <- conf_data |>
+    dplyr::mutate(.stratum_total = conf_totals) |>
+    dplyr::summarize(
+      .share = sum(.data$.stratum_total),
+      .by = dplyr::all_of(group_by)
+    ) |>
+    dplyr::mutate(.share = .data$.share / sum(.data$.share))
+
+  per_stratum <- purrr::map(
+    .x = seq_len(nrow(strata)),
+    .f = \(i) {
+      stratum <- strata[i, group_by, drop = FALSE]
+
+      synth_g <- dplyr::semi_join(synth_data, stratum, by = group_by)
+      conf_g <- dplyr::semi_join(conf_data, stratum, by = group_by)
+
+      cells_g <- compute_cells(
+        synth_d = synth_g, conf_d = conf_g, allow_empty_synth = TRUE
+      )
+
+      marginals_g <- cells_g |>
+        dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables")
+
+      list(
+        cells = dplyr::bind_cols(stratum, cells_g),
+        marginals = dplyr::bind_cols(stratum, marginals_g),
+        group_scores = dplyr::bind_cols(
+          stratum,
+          tibble::tibble(
+            share = strata$.share[i],
+            score = (1 - mean(marginals_g$madd)) * 1000
+          )
+        )
+      )
+    }
+  )
+
+  cells <- purrr::list_rbind(purrr::map(per_stratum, "cells")) |>
     dplyr::arrange(dplyr::desc(.data$abs_diff))
 
-  # MabsDD per combination, worst marginals first
-  marginals <- cells |>
-    dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables") |>
+  marginals <- purrr::list_rbind(purrr::map(per_stratum, "marginals")) |>
     dplyr::arrange(dplyr::desc(.data$madd))
 
-  # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000];
-  # computed from all marginals before any truncation
-  score <- (1 - mean(marginals$madd)) * 1000
+  group_scores <- purrr::list_rbind(
+    purrr::map(per_stratum, "group_scores")
+  ) |>
+    dplyr::arrange(.data$score)
+
+  # per-stratum scores roll up weighted by confidential shares, so small
+  # strata surface in group_scores without dominating the headline
+  score <- sum(group_scores$share * group_scores$score)
 
   result <- structure(
     list(
       score = score,
       marginals = utils::head(marginals, n = keep_marginals),
-      cells = utils::head(cells, n = keep_cells)
+      cells = utils::head(cells, n = keep_cells),
+      group_scores = group_scores
     ),
     class = "k_marginals"
   )
 
   return(result)
-
 }
 
 #' @title Print a k_marginals object
@@ -447,14 +532,17 @@
 #' @export
 #'
 print.k_marginals <- function(x, n = 5, ...) {
-
   cat("k-marginals score:", round(x$score, digits = 2), "\n\n")
 
   cat("Worst marginals:\n")
   print(utils::head(x$marginals, n = n))
 
-  return(invisible(x))
+  if (!is.null(x$group_scores)) {
+    cat("\nWorst groups:\n")
+    print(utils::head(x$group_scores, n = n))
+  }
 
+  return(invisible(x))
 }
 
 #' @title Calculate the k-marginals metric
@@ -485,6 +573,12 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' shares instead of row shares, and the weight column is excluded from the
 #' marginals. Weights must be finite and non-negative with a positive total.
 #' Defaults to `NULL` (unweighted).
+#' @param group_by Optional character vector of grouping variable names
+#' present in both datasets. When set, the grouping variables are excluded
+#' from the marginals and the metric is computed within each stratum observed
+#' in the confidential data; the headline score is the mean of the per-stratum
+#' scores weighted by each stratum's confidential share (weight share when
+#' `weight_var` is set). Defaults to `NULL` (no stratification).
 #' @param synth_vars A logical for if only synthesized variables should
 #' contribute marginals. Only meaningful when the `eval_data` records which
 #' variables were synthesized (i.e., was built from a `postsynth`); for plain
@@ -514,69 +608,61 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' `cells`, a tibble with the synthetic and confidential proportions and their
 #' absolute difference for every cell, worst first. `score` is always computed
 #' from all evaluated marginals, even when `keep_marginals` or `keep_cells`
-#' truncate the detail tables. For multiple replicates, a list of such
+#' truncate the detail tables. When `group_by` is set, `marginals` and
+#' `cells` gain the grouping columns and a fourth element `group_scores`
+#' reports each stratum's share and score, worst first. For multiple replicates, a list of such
 #' objects, one per replicate.
 #'
 #' @export
 #'
 util_k_marginals <- function(
-    eval_data,
-    k = 3,
-    keep_marginals = Inf,
-    keep_cells = Inf,
-    n_marginals = Inf,
-    priority_vars = NULL,
-    weight_var = NULL,
-    synth_vars = TRUE,
-    na.rm = FALSE,
-    bins = NULL,
-    discretize_method = c("width", "ntile", "cluster")) {
-
+  eval_data,
+  k = 3,
+  keep_marginals = Inf,
+  keep_cells = Inf,
+  n_marginals = Inf,
+  priority_vars = NULL,
+  weight_var = NULL,
+  group_by = NULL,
+  synth_vars = TRUE,
+  na.rm = FALSE,
+  bins = NULL,
+  discretize_method = c("width", "ntile", "cluster")
+) {
   stopifnot(is_eval_data(eval_data))
 
   discretize_method <- match.arg(discretize_method)
 
   if (!(is.logical(synth_vars) && length(synth_vars) == 1 &&
-        !is.na(synth_vars))) {
-
+    !is.na(synth_vars))) {
     stop("`synth_vars` must be a single TRUE or FALSE")
-
   }
 
   # NULL for plain data frame eval_data, so the worker applies no restriction
   synth_varnames <- if (synth_vars) {
-
     eval_data$synth_vars
-
   } else {
-
     NULL
-
   }
 
   # empty metadata can only come from a user-supplied eval_data(synth_vars =)
   # argument; fail with a message in terms of this function's arguments
   if (!is.null(synth_varnames) && length(synth_varnames) == 0) {
-
     stop(
       "`eval_data` records no synthesized variables; use `synth_vars = ",
       "FALSE` to evaluate all shared variables"
     )
-
   }
 
   # surface the resolved method so a forgotten discretize_method is visible
   if (!is.null(bins)) {
-
     message(
       "Discretizing numeric variables into ", bins,
       " bins using the '", discretize_method, "' method"
     )
-
   }
 
   if (eval_data$n_rep == 1) {
-
     return(
       .util_k_marginals(
         synth_data = eval_data$synth_data,
@@ -587,19 +673,17 @@ util_k_marginals <- function(
         n_marginals = n_marginals,
         priority_vars = priority_vars,
         weight_var = weight_var,
+        group_by = group_by,
         synth_varnames = synth_varnames,
         na.rm = na.rm,
         bins = bins,
         discretize_method = discretize_method
       )
     )
-
   } else {
-
     result <- purrr::map(
       .x = eval_data$synth_data,
       .f = \(sd) {
-
         .util_k_marginals(
           synth_data = sd,
           conf_data = eval_data$conf_data,
@@ -609,17 +693,15 @@ util_k_marginals <- function(
           n_marginals = n_marginals,
           priority_vars = priority_vars,
           weight_var = weight_var,
+          group_by = group_by,
           synth_varnames = synth_varnames,
           na.rm = na.rm,
           bins = bins,
           discretize_method = discretize_method
         )
-
       }
     )
 
     return(result)
-
   }
-
 }
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index c568de4..651c67b 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -13,6 +13,7 @@
   n_marginals = Inf,
   priority_vars = NULL,
   weight_var = NULL,
+  group_by = NULL,
   synth_varnames = NULL,
   na.rm = FALSE,
   bins = NULL,
@@ -50,6 +51,13 @@ shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
 
+\item{group_by}{Optional character vector of grouping variable names
+present in both datasets. When set, the grouping variables are excluded
+from the marginals and the metric is computed within each stratum observed
+in the confidential data; the headline score is the mean of the per-stratum
+scores weighted by each stratum's confidential share (weight share when
+\code{weight_var} is set). Defaults to \code{NULL} (no stratification).}
+
 \item{synth_varnames}{Optional character vector of synthesized variable
 names. When set, only these variables (intersected with the variables
 shared by both datasets) contribute marginals. Defaults to \code{NULL}, which
@@ -83,7 +91,9 @@ tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their
 absolute difference for every cell, worst first. \code{score} is always computed
 from all evaluated marginals, even when \code{keep_marginals} or \code{keep_cells}
-truncate the detail tables.
+truncate the detail tables. When \code{group_by} is set, \code{marginals} and
+\code{cells} gain the grouping columns and a fourth element \code{group_scores}
+reports each stratum's share and score, worst first.
 }
 \description{
 This worker function takes a specified k-marginal and calculates
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 5e5fe51..9e1f0d1 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -12,6 +12,7 @@ util_k_marginals(
   n_marginals = Inf,
   priority_vars = NULL,
   weight_var = NULL,
+  group_by = NULL,
   synth_vars = TRUE,
   na.rm = FALSE,
   bins = NULL,
@@ -47,6 +48,13 @@ shares instead of row shares, and the weight column is excluded from the
 marginals. Weights must be finite and non-negative with a positive total.
 Defaults to \code{NULL} (unweighted).}
 
+\item{group_by}{Optional character vector of grouping variable names
+present in both datasets. When set, the grouping variables are excluded
+from the marginals and the metric is computed within each stratum observed
+in the confidential data; the headline score is the mean of the per-stratum
+scores weighted by each stratum's confidential share (weight share when
+\code{weight_var} is set). Defaults to \code{NULL} (no stratification).}
+
 \item{synth_vars}{A logical for if only synthesized variables should
 contribute marginals. Only meaningful when the \code{eval_data} records which
 variables were synthesized (i.e., was built from a \code{postsynth}); for plain
@@ -80,7 +88,9 @@ tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their
 absolute difference for every cell, worst first. \code{score} is always computed
 from all evaluated marginals, even when \code{keep_marginals} or \code{keep_cells}
-truncate the detail tables. For multiple replicates, a list of such
+truncate the detail tables. When \code{group_by} is set, \code{marginals} and
+\code{cells} gain the grouping columns and a fourth element \code{group_scores}
+reports each stratum's share and score, worst first. For multiple replicates, a list of such
 objects, one per replicate.
 }
 \description{
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index ca22df4..1bba2ea 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -127,7 +127,7 @@ test_that("k exceeding the number of shared variables throws an error", {
 
   expect_error(
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 3),
-    regexp = "shared by both datasets"
+    regexp = "`k` cannot exceed the number of variables available"
   )
 
 })
@@ -1184,7 +1184,7 @@ test_that("synth_varnames with no shared variables errors informatively", {
       k = 1,
       synth_varnames = "zzz"
     ),
-    regexp = "shares no variables with both datasets"
+    regexp = "`synth_varnames` matches no variables available"
   )
 
 })
@@ -1578,3 +1578,521 @@ test_that("infinite confidential values still refuse to discretize", {
   )
 
 })
+
+# group_by stratification
+#
+# conf_g             synth_g
+#   g  a               g  a
+#   A  x               A  x
+#   A  y               A  x
+#   B  x               B  x
+#   B  y               B  y
+#
+# universe = {a}; g stratifies and never marginalizes
+# stratum A (conf share 0.5): conf (x = 0.5, y = 0.5), synth (x = 1)
+#   MabsDD = mean(0.5, 0.5) = 0.5 -> score 500
+# stratum B (conf share 0.5): identical -> MabsDD = 0 -> score 1000
+# overall = 0.5 * 500 + 0.5 * 1000 = 750
+
+conf_g <- tibble::tibble(
+  g = c("A", "A", "B", "B"),
+  a = c("x", "y", "x", "y")
+)
+
+synth_g <- tibble::tibble(
+  g = c("A", "A", "B", "B"),
+  a = c("x", "x", "x", "y")
+)
+
+test_that("group_by stratifies the score by confidential shares", {
+
+  result <- .util_k_marginals(
+    synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
+  )
+
+  expect_equal(result$score, 750)
+
+})
+
+test_that("grouped output gains group columns and group_scores", {
+
+  result <- .util_k_marginals(
+    synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
+  )
+
+  expect_named(result$group_scores, c("g", "share", "score"))
+
+  # worst stratum first
+  expect_equal(result$group_scores$g, c("A", "B"))
+  expect_equal(result$group_scores$share, c(0.5, 0.5))
+  expect_equal(result$group_scores$score, c(500, 1000))
+
+  expect_true("g" %in% names(result$marginals))
+  expect_true("g" %in% names(result$cells))
+
+  # worst-first ordering across strata
+  expect_equal(result$marginals$g[1], "A")
+
+  # ungrouped results carry no group_scores element
+  ungrouped <- .util_k_marginals(synth_data = synth_g, conf_data = conf_g, k = 1)
+
+  expect_null(ungrouped$group_scores)
+
+})
+
+test_that("an empty synthetic stratum scores against zero proportions", {
+
+  # synth has no B rows: stratum B conf cells (x = 0.5, y = 0.5) face
+  # synthetic proportions of 0 -> MabsDD = 0.5 -> score 500
+  # stratum A: conf (x = 0.5, y = 0.5), synth (x = 0.5, y = 0.5) -> 1000
+  # overall = 0.5 * 1000 + 0.5 * 500 = 750
+  synth_a_only <- tibble::tibble(
+    g = c("A", "A", "A", "A"),
+    a = c("x", "x", "y", "y")
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_a_only, conf_data = conf_g, k = 1, group_by = "g"
+  )
+
+  expect_equal(result$score, 750)
+
+})
+
+test_that("group shares use weights when weight_var is set", {
+
+  # conf weight shares: A = 2/4, B = 2/4 (row shares would be 2/3, 1/3)
+  # stratum A: conf (x = 0.5, y = 0.5), synth (x = 1) -> score 500
+  # stratum B: conf (x = 1), synth (x = 1) -> score 1000
+  # overall = 0.5 * 500 + 0.5 * 1000 = 750; row shares would give 2000/3
+  conf_gw <- tibble::tibble(
+    g = c("A", "A", "B"),
+    a = c("x", "y", "x"),
+    w = c(1, 1, 2)
+  )
+
+  synth_gw <- tibble::tibble(
+    g = c("A", "A", "B"),
+    a = c("x", "x", "x"),
+    w = c(1, 1, 1)
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_gw,
+    conf_data = conf_gw,
+    k = 1,
+    group_by = "g",
+    weight_var = "w"
+  )
+
+  expect_equal(result$score, 750)
+
+})
+
+test_that("invalid group_by values error", {
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 1, group_by = 1
+    ),
+    regexp = "`group_by` must be a character vector"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "zzz"
+    ),
+    regexp = "`group_by` must be a character vector"
+  )
+
+  conf_w <- dplyr::mutate(conf_g, w = 1)
+  synth_w <- dplyr::mutate(synth_g, w = 1)
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_w, conf_data = conf_w, k = 1,
+      group_by = "w", weight_var = "w"
+    ),
+    regexp = "`group_by` cannot include `weight_var`"
+  )
+
+})
+
+test_that("group variables are excluded from the marginal universe", {
+
+  result <- .util_k_marginals(
+    synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
+  )
+
+  expect_false(any(result$marginals$variables == "g"))
+
+  # k is checked against the universe without the group variables
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 2, group_by = "g"
+    ),
+    regexp = "`k` cannot exceed"
+  )
+
+})
+
+test_that("grouped print shows group scores", {
+
+  result <- .util_k_marginals(
+    synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
+  )
+
+  expect_output(print(result), regexp = "Worst groups:")
+
+})
+
+test_that("util_k_marginals passes group_by through", {
+
+  ed <- eval_data(conf_data = conf_g, synth_data = synth_g)
+
+  expect_equal(
+    util_k_marginals(eval_data = ed, k = 1, group_by = "g")$score,
+    750
+  )
+
+})
+
+test_that("missing group values follow na.rm", {
+
+  conf_gna <- tibble::tibble(
+    g = c("A", "A", NA, NA),
+    a = c("x", "y", "x", "y")
+  )
+
+  synth_gna <- tibble::tibble(
+    g = c("A", "A", NA, NA),
+    a = c("x", "x", "x", "y")
+  )
+
+  # na.rm = FALSE: NA forms its own stratum (perfect match -> 1000);
+  # stratum A scores 500 -> overall 750
+  kept <- suppressMessages(
+    .util_k_marginals(
+      synth_data = synth_gna, conf_data = conf_gna, k = 1, group_by = "g"
+    )
+  )
+
+  expect_equal(kept$score, 750)
+  expect_true("NA" %in% kept$group_scores$g)
+
+  # na.rm = TRUE: NA-group rows drop entirely, leaving only stratum A
+  dropped <- .util_k_marginals(
+    synth_data = synth_gna, conf_data = conf_gna, k = 1, group_by = "g",
+    na.rm = TRUE
+  )
+
+  expect_equal(dropped$score, 500)
+  expect_equal(nrow(dropped$group_scores), 1)
+
+})
+
+test_that("multi-column group_by stratifies by joint combinations", {
+
+  # four joint strata of two rows each; only stratum (A, p) diverges:
+  #   conf (x = 0.5, y = 0.5), synth (x = 1) -> MabsDD = 0.5 -> score 500
+  # the other three strata are identical -> 1000
+  # overall = 0.25 * 500 + 0.75 * 1000 = 875
+  conf_g2 <- tibble::tibble(
+    g1 = c("A", "A", "A", "A", "B", "B", "B", "B"),
+    g2 = c("p", "p", "q", "q", "p", "p", "q", "q"),
+    a = c("x", "y", "x", "y", "x", "y", "x", "y")
+  )
+
+  synth_g2 <- dplyr::mutate(
+    conf_g2,
+    a = c("x", "x", "x", "y", "x", "y", "x", "y")
+  )
+
+  result <- .util_k_marginals(
+    synth_data = synth_g2, conf_data = conf_g2, k = 1,
+    group_by = c("g1", "g2")
+  )
+
+  expect_equal(result$score, 875)
+
+  # both grouping columns ride along in every output
+  expect_named(result$group_scores, c("g1", "g2", "share", "score"))
+  expect_true(all(c("g1", "g2") %in% names(result$marginals)))
+  expect_true(all(c("g1", "g2") %in% names(result$cells)))
+
+  expect_equal(result$group_scores$share, rep(0.25, 4))
+
+  # worst joint stratum first
+  expect_equal(result$group_scores$g1[1], "A")
+  expect_equal(result$group_scores$g2[1], "p")
+
+  # neither grouping variable enters the marginal universe
+  expect_false(any(result$marginals$variables %in% c("g1", "g2")))
+
+})
+
+test_that("partially missing joint strata follow na.rm", {
+
+  # g2 is missing on rows 3-4; only the (A, NA) stratum diverges:
+  #   conf (x = 0.5, y = 0.5), synth (x = 1) -> score 500
+  conf_gpart <- tibble::tibble(
+    g1 = c("A", "A", "A", "A"),
+    g2 = c("p", "p", NA, NA),
+    a = c("x", "y", "x", "y")
+  )
+
+  synth_gpart <- dplyr::mutate(conf_gpart, a = c("x", "y", "x", "x"))
+
+  # na.rm = FALSE: the partial combination becomes an (A, "NA") stratum
+  # overall = 0.5 * 1000 + 0.5 * 500 = 750
+  kept <- suppressMessages(
+    .util_k_marginals(
+      synth_data = synth_gpart, conf_data = conf_gpart, k = 1,
+      group_by = c("g1", "g2")
+    )
+  )
+
+  expect_equal(kept$score, 750)
+  expect_true("NA" %in% kept$group_scores$g2)
+
+  # na.rm = TRUE: rows missing any grouping value drop, leaving (A, p) only
+  dropped <- .util_k_marginals(
+    synth_data = synth_gpart, conf_data = conf_gpart, k = 1,
+    group_by = c("g1", "g2"), na.rm = TRUE
+  )
+
+  expect_equal(dropped$score, 1000)
+  expect_equal(nrow(dropped$group_scores), 1)
+
+})
+
+test_that("empty, missing, and duplicate group_by values error", {
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 1,
+      group_by = character(0)
+    ),
+    regexp = "`group_by` must be a character vector"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 1,
+      group_by = c("g", NA)
+    ),
+    regexp = "`group_by` must be a character vector"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_g, conf_data = conf_g, k = 1,
+      group_by = c("g", "g")
+    ),
+    regexp = "must not contain duplicate"
+  )
+
+})
+
+test_that("group_by composes with synth_varnames", {
+
+  # b is shared but unsynthesized; g stratifies; the universe is {a} only
+  conf_gsv <- tibble::tibble(
+    g = c("A", "A", "B", "B"),
+    a = c("x", "y", "x", "y"),
+    b = c("m", "m", "n", "n")
+  )
+
+  synth_gsv <- dplyr::mutate(conf_gsv, a = c("x", "x", "x", "y"))
+
+  result <- .util_k_marginals(
+    synth_data = synth_gsv, conf_data = conf_gsv, k = 1,
+    group_by = "g", synth_varnames = "a"
+  )
+
+  expect_equal(result$score, 750)
+  expect_equal(unique(result$marginals$variables), "a")
+
+  # a grouping variable named in synth_varnames still stratifies rather
+  # than marginalizing
+  result_gname <- .util_k_marginals(
+    synth_data = synth_gsv, conf_data = conf_gsv, k = 1,
+    group_by = "g", synth_varnames = c("g", "a")
+  )
+
+  expect_equal(unique(result_gname$marginals$variables), "a")
+
+})
+
+test_that("literal 'NA' strings in grouping columns collide with true NA", {
+
+  # single grouping column carrying both a literal "NA" level and a true NA
+  conf_collide <- tibble::tibble(
+    g = c("NA", "A", NA),
+    a = c("x", "x", "x")
+  )
+
+  synth_ok <- tibble::tibble(
+    g = c("A", "A", "A"),
+    a = c("x", "x", "x")
+  )
+
+  expect_error(
+    suppressMessages(
+      .util_k_marginals(
+        synth_data = synth_ok, conf_data = conf_collide, k = 1,
+        group_by = "g"
+      )
+    ),
+    regexp = "'NA' already exists"
+  )
+
+  # multi-column grouping where only one key collides
+  conf_multi <- tibble::tibble(
+    g1 = c("A", "A", "A"),
+    g2 = c("NA", "p", NA),
+    a = c("x", "x", "x")
+  )
+
+  synth_multi <- tibble::tibble(
+    g1 = c("A", "A", "A"),
+    g2 = c("p", "p", "p"),
+    a = c("x", "x", "x")
+  )
+
+  expect_error(
+    suppressMessages(
+      .util_k_marginals(
+        synth_data = synth_multi, conf_data = conf_multi, k = 1,
+        group_by = c("g1", "g2")
+      )
+    ),
+    regexp = "'NA' already exists"
+  )
+
+  # convert_na_to_level() rejects a literal "NA" level even without true
+  # missing values, so such data must be scored with na.rm = TRUE
+  conf_legit <- tibble::tibble(
+    g = c("NA", "NA", "A", "A"),
+    a = c("x", "y", "x", "y")
+  )
+
+  synth_legit <- dplyr::mutate(conf_legit, a = c("x", "x", "x", "y"))
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_legit, conf_data = conf_legit, k = 1, group_by = "g"
+    ),
+    regexp = "'NA' already exists"
+  )
+
+  # na.rm = TRUE bypasses the conversion, so the "NA" stratum scores
+  # normally: stratum "NA" diverges (500), stratum A matches (1000)
+  result <- .util_k_marginals(
+    synth_data = synth_legit, conf_data = conf_legit, k = 1, group_by = "g",
+    na.rm = TRUE
+  )
+
+  expect_equal(result$score, 750)
+
+})
+
+test_that("grouped results map over replicates with the same structure", {
+
+  ed <- eval_data(
+    conf_data = conf_g,
+    synth_data = list(synth_g, conf_g)
+  )
+
+  result <- util_k_marginals(
+    eval_data = ed, k = 1, group_by = "g"
+  )
+
+  expect_length(result, 2)
+
+  for (rep in result) {
+
+    expect_s3_class(rep, "k_marginals")
+    expect_named(rep$group_scores, c("g", "share", "score"))
+    expect_true("g" %in% names(rep$marginals))
+    expect_true("g" %in% names(rep$cells))
+
+  }
+
+  # first replicate diverges in stratum A, second is identical data
+  expect_equal(purrr::map_dbl(result, "score"), c(750, 1000))
+
+})
+
+test_that("a confidential stratum emptied by na.rm errors instead of NaN", {
+
+  # stratum B's confidential rows are all missing on a, so per-marginal NA
+  # removal leaves nothing to score against
+  conf_gna2 <- tibble::tibble(
+    g = c("A", "A", "B", "B"),
+    a = c("x", "y", NA, NA)
+  )
+
+  synth_gna2 <- tibble::tibble(
+    g = c("A", "A", "B", "B"),
+    a = c("x", "y", "x", "y")
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth_gna2, conf_data = conf_gna2, k = 1, group_by = "g",
+      na.rm = TRUE
+    ),
+    regexp = "no rows remain"
+  )
+
+  # the synthetic side emptying in a stratum is still allowed
+  conf_swap <- synth_gna2
+  synth_swap <- conf_gna2
+
+  result <- .util_k_marginals(
+    synth_data = synth_swap, conf_data = conf_swap, k = 1, group_by = "g",
+    na.rm = TRUE
+  )
+
+  # stratum B scores conf (x = 0.5, y = 0.5) against zero synth -> 500
+  expect_equal(result$score, 750)
+
+})
+
+test_that("priority_vars with NA entries hits the intended error", {
+
+  # %in% never propagates NA, so the membership test is FALSE, not NA, and
+  # the package error fires rather than a base R condition failure
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1,
+      priority_vars = NA_character_
+    ),
+    regexp = "`priority_vars` must be a character vector of variables available"
+  )
+
+  expect_error(
+    .util_k_marginals(
+      synth_data = synth, conf_data = conf, k = 1,
+      priority_vars = c("a", NA)
+    ),
+    regexp = "`priority_vars` must be a character vector of variables available"
+  )
+
+})
+
+test_that("priority_vars = character(0) behaves like NULL", {
+
+  # an empty priority set passes validation and simply guarantees nothing
+  result <- .util_k_marginals(
+    synth_data = synth, conf_data = conf, k = 1,
+    priority_vars = character(0)
+  )
+
+  expect_equal(
+    result$score,
+    .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
+  )
+
+})

From 322cc176316386b99cb7dc125d93a12f9736169b Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Fri, 14 Aug 2026 16:38:53 -0400
Subject: [PATCH 09/10] Refactor util_k_marginals into internal helpers

---
 R/compute_marginal_cells.R                    | 105 +++++
 R/discretize_k_marginal_vars.R                |  94 +++++
 R/select_k_marginal_combos.R                  |  47 +++
 R/stratify_k_marginals.R                      | 105 +++++
 R/util_k_marginals.R                          | 268 ++-----------
 man/dot-compute_marginal_cells.Rd             |  43 ++
 man/dot-discretize_k_marginal_vars.Rd         |  39 ++
 man/dot-select_k_marginal_combos.Rd           |  30 ++
 man/dot-stratify_k_marginals.Rd               |  43 ++
 tests/testthat/test-compute_marginal_cells.R  |  59 +++
 .../test-discretize_k_marginal_vars.R         |  90 +++++
 .../testthat/test-select_k_marginal_combos.R  |  55 +++
 tests/testthat/test-stratify_k_marginals.R    |  86 ++++
 tests/testthat/test-util_k_marginals.R        | 370 +-----------------
 14 files changed, 851 insertions(+), 583 deletions(-)
 create mode 100644 R/compute_marginal_cells.R
 create mode 100644 R/discretize_k_marginal_vars.R
 create mode 100644 R/select_k_marginal_combos.R
 create mode 100644 R/stratify_k_marginals.R
 create mode 100644 man/dot-compute_marginal_cells.Rd
 create mode 100644 man/dot-discretize_k_marginal_vars.Rd
 create mode 100644 man/dot-select_k_marginal_combos.Rd
 create mode 100644 man/dot-stratify_k_marginals.Rd
 create mode 100644 tests/testthat/test-compute_marginal_cells.R
 create mode 100644 tests/testthat/test-discretize_k_marginal_vars.R
 create mode 100644 tests/testthat/test-select_k_marginal_combos.R
 create mode 100644 tests/testthat/test-stratify_k_marginals.R

diff --git a/R/compute_marginal_cells.R b/R/compute_marginal_cells.R
new file mode 100644
index 0000000..d814b08
--- /dev/null
+++ b/R/compute_marginal_cells.R
@@ -0,0 +1,105 @@
+#' @title Compute per-cell proportion differences for the k-marginals metric
+#'
+#' @description For each supplied variable combination, computes the marginal
+#' cell proportions in the synthetic and confidential data and their absolute
+#' differences. Cells absent from one dataset count as proportion zero.
+#'
+#' @param synth_data A tibble with synthetic data.
+#' @param conf_data A tibble with confidential data.
+#' @param combos A character matrix with one row per variable combination.
+#' @param weight_var Optional character name of a numeric sample-weight
+#' column; when set, cell proportions are weight shares instead of row
+#' shares. Defaults to `NULL` (unweighted).
+#' @param na.rm A logical for dropping rows with a missing value from each
+#' marginal that uses the affected variable.
+#' @param allow_empty_synth A logical for permitting the synthetic data to
+#' contribute no rows to a marginal (a stratum the synthesis never produced),
+#' in which case its proportions are zero. A confidential marginal with no
+#' rows always errors, since there is nothing to score against.
+#'
+#' @return A tibble with one row per cell: `variables`, `cell`,
+#' `prop_synth`, `prop_conf`, and `abs_diff`.
+#'
+.compute_marginal_cells <- function(
+  synth_data,
+  conf_data,
+  combos,
+  weight_var = NULL,
+  na.rm = FALSE,
+  allow_empty_synth = FALSE
+) {
+  # cell proportions for one dataset over one set of variables; weighted
+  # proportions are weight shares instead of row shares
+  process_data <- function(data, vars, prop_name, allow_empty) {
+    if (na.rm) {
+      data <- dplyr::filter(
+        data,
+        !dplyr::if_any(.cols = dplyr::all_of(vars), .fns = is.na)
+      )
+
+      if (nrow(data) == 0 && !allow_empty) {
+        stop(
+          "no rows remain for the marginal over ",
+          paste(vars, collapse = ", "),
+          " after removing missing values"
+        )
+      }
+    }
+
+    if (is.null(weight_var)) {
+      counts <- dplyr::count(data, dplyr::across(dplyr::all_of(vars)))
+    } else {
+      counts <- dplyr::count(
+        data,
+        dplyr::across(dplyr::all_of(vars)),
+        wt = .data[[weight_var]]
+      )
+    }
+
+    props <- counts |>
+      dplyr::mutate("{prop_name}" := .data$n / sum(.data$n)) |>
+      dplyr::select(-"n")
+
+    return(props)
+  }
+
+  # per-cell differences for one set of variables; cells absent from one
+  # dataset count as 0
+  marginal_cells <- function(vars) {
+    # only the synthetic side may be empty (a stratum the synthesis never
+    # produced); a confidential marginal with no rows has nothing to score
+    # against and errors inside process_data
+    cells <- dplyr::full_join(
+      process_data(
+        data = synth_data, vars = vars, prop_name = "prop_synth",
+        allow_empty = allow_empty_synth
+      ),
+      process_data(
+        data = conf_data, vars = vars, prop_name = "prop_conf",
+        allow_empty = FALSE
+      ),
+      by = vars
+    ) |>
+      tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0)) |>
+      tidyr::unite(col = "cell", dplyr::all_of(vars), sep = ", ") |>
+      dplyr::mutate(
+        variables = paste(vars, collapse = ", "),
+        abs_diff = abs(.data$prop_synth - .data$prop_conf)
+      ) |>
+      dplyr::select(
+        "variables", "cell", "prop_synth", "prop_conf", "abs_diff"
+      )
+    # variables disambiguates cells across combinations and drives the
+    # per-combination summary; the prop columns show the direction of the
+    # discrepancy, not just its size
+    return(cells)
+  }
+
+  cells <- purrr::map(
+    .x = seq_len(nrow(combos)),
+    .f = \(i) marginal_cells(vars = combos[i, ])
+  ) |>
+    purrr::list_rbind()
+
+  return(cells)
+}
diff --git a/R/discretize_k_marginal_vars.R b/R/discretize_k_marginal_vars.R
new file mode 100644
index 0000000..96ac9f8
--- /dev/null
+++ b/R/discretize_k_marginal_vars.R
@@ -0,0 +1,94 @@
+#' @title Discretize numeric variables for the k-marginals metric
+#'
+#' @description Discretizes every numeric variable among `vars` into `bins`
+#' bins with interior cut points derived from the observed (non-missing)
+#' confidential values and applied to both datasets. The outer bins extend to
+#' +/-Inf so synthetic values outside the confidential range land in edge
+#' bins.
+#'
+#' @param synth_data A tibble with synthetic data.
+#' @param conf_data A tibble with confidential data.
+#' @param vars Character vector of candidate variables; only numeric ones are
+#' discretized.
+#' @param bins Single integer >= 2 giving the number of bins. Fewer bins are
+#' produced, with a warning, if tied cut points collapse.
+#' @param discretize_method Method used to place bin breaks: "width" for
+#' fixed binwidths, "ntile" for quantile bins, or "cluster" for univariate
+#' k-means clustering (set a seed before calling for reproducible clusters).
+#'
+#' @return A list with the discretized `synth_data` and `conf_data`.
+#'
+.discretize_k_marginal_vars <- function(
+  synth_data,
+  conf_data,
+  vars,
+  bins,
+  discretize_method
+) {
+  if (!(is.numeric(bins) && length(bins) == 1 && !is.na(bins) &&
+    bins >= 2 && bins == floor(bins))) {
+    stop("`bins` must be a single integer >= 2")
+  }
+
+  numeric_vars <- vars[
+    purrr::map_lgl(.x = vars, .f = \(v) is.numeric(conf_data[[v]]))
+  ]
+
+  for (var in numeric_vars) {
+    # cut points derive from observed values; missing values follow na.rm
+    # like every other variable, becoming an NA bin or dropped rows
+    conf_values <- conf_data[[var]][!is.na(conf_data[[var]])]
+
+    if (!all(is.finite(conf_values)) || length(conf_values) == 0) {
+      stop(
+        "observed numeric values must be finite to discretize; `", var,
+        "` is not"
+      )
+    }
+
+    if (dplyr::n_distinct(conf_values) < bins) {
+      stop(
+        "`", var, "` has fewer distinct confidential values than `bins`"
+      )
+    }
+
+    # interior cut points always derive from the confidential data; each
+    # method yields bins - 1 of them
+    cut_points <- switch(
+      EXPR = discretize_method,
+      width = seq(
+        from = min(conf_values),
+        to = max(conf_values),
+        length.out = bins + 1
+      )[2:bins],
+      ntile = stats::quantile(
+        x = conf_values,
+        probs = seq(from = 0, to = 1, length.out = bins + 1),
+        names = FALSE
+      )[2:bins],
+      cluster = {
+        centers <- sort(
+          stats::kmeans(x = conf_values, centers = bins)$centers[, 1]
+        )
+
+        (centers[-1] + centers[-length(centers)]) / 2
+      }
+    )
+
+    # outer bins extend to +/-Inf so out-of-range synthetic values land in
+    # edge bins; unique() collapses ties from skewed quantiles
+    breaks <- unique(c(-Inf, cut_points, Inf))
+
+    if (length(breaks) - 1 < bins) {
+      warning(
+        "`", var, "` was discretized into ", length(breaks) - 1,
+        " bins instead of ", bins, " because of tied cut points"
+      )
+    }
+
+    synth_data[[var]] <- cut(x = synth_data[[var]], breaks = breaks)
+    conf_data[[var]] <- cut(x = conf_data[[var]], breaks = breaks)
+  }
+
+  return(list(synth_data = synth_data, conf_data = conf_data))
+}
diff --git a/R/select_k_marginal_combos.R b/R/select_k_marginal_combos.R
new file mode 100644
index 0000000..6379d75
--- /dev/null
+++ b/R/select_k_marginal_combos.R
@@ -0,0 +1,47 @@
+#' @title Select the variable combinations for the k-marginals metric
+#'
+#' @description Enumerates all unique k-combinations of the supplied
+#' variables and, when the total exceeds `n_marginals`, samples the
+#' combinations down to the cap while always retaining combinations that
+#' contain a priority variable.
+#'
+#' @param shared_vars Character vector of variables available for marginals.
+#' @param k Scalar order of the k-marginal.
+#' @param n_marginals Single integer target maximum for the number of
+#' variable combinations, or `Inf` for no cap. All priority combinations are
+#' retained even when they alone exceed the target.
+#' @param priority_vars Optional character vector of variable names whose
+#' combinations always survive sampling. Defaults to `NULL`.
+#'
+#' @return A character matrix with one row per selected combination and `k`
+#' columns.
+#'
+.select_k_marginal_combos <- function(
+  shared_vars,
+  k,
+  n_marginals,
+  priority_vars = NULL
+) {
+  kmarginals_vars <- t(utils::combn(x = shared_vars, m = k))
+
+  # sample combinations down to n_marginals, always keeping combinations that
+  # contain a priority variable
+  if (nrow(kmarginals_vars) > n_marginals) {
+    is_priority <- apply(
+      X = kmarginals_vars,
+      MARGIN = 1,
+      FUN = \(vars) any(vars %in% priority_vars)
+    )
+
+    n_sampled <- min(max(n_marginals - sum(is_priority), 0), sum(!is_priority))
+
+    sampled_rows <- sample(x = which(!is_priority), size = n_sampled)
+
+    kmarginals_vars <- kmarginals_vars[
+      sort(c(which(is_priority), sampled_rows)), ,
+      drop = FALSE
+    ]
+  }
+
+  return(kmarginals_vars)
+}
diff --git a/R/stratify_k_marginals.R b/R/stratify_k_marginals.R
new file mode 100644
index 0000000..7ea0932
--- /dev/null
+++ b/R/stratify_k_marginals.R
@@ -0,0 +1,105 @@
+#' @title Compute stratified k-marginals results
+#'
+#' @description Splits the synthetic and confidential data into the strata
+#' observed in the confidential data, computes per-cell differences and
+#' per-combination MabsDDs within each stratum, and rolls the per-stratum
+#' scores up weighted by each stratum's confidential share.
+#'
+#' @param synth_data A tibble with synthetic data.
+#' @param conf_data A tibble with confidential data.
+#' @param combos A character matrix with one row per variable combination.
+#' @param group_by Character vector of grouping variable names.
+#' @param weight_var Optional character name of a numeric sample-weight
+#' column; when set, cell proportions and stratum shares are weight shares
+#' instead of row shares. Defaults to `NULL` (unweighted).
+#' @param na.rm A logical for dropping rows with a missing value from each
+#' marginal that uses the affected variable.
+#'
+#' @return A list with `score` (the share-weighted mean of per-stratum
+#' scores), `marginals` and `cells` (stacked per-stratum tables with the
+#' grouping columns, worst first), and `group_scores` (one row per stratum
+#' with its share and score, worst first).
+#'
+.stratify_k_marginals <- function(
+  synth_data,
+  conf_data,
+  combos,
+  group_by,
+  weight_var = NULL,
+  na.rm = FALSE
+) {
+  # strata are defined by the confidential data; a stratum with no synthetic
+  # rows scores against all-zero synthetic proportions. Shares are the
+  # confidential row (or weight) share of each stratum, computed once
+  conf_totals <- if (is.null(weight_var)) {
+    rep(1, nrow(conf_data))
+  } else {
+    conf_data[[weight_var]]
+  }
+
+  strata <- conf_data |>
+    dplyr::mutate(.stratum_total = conf_totals) |>
+    dplyr::summarize(
+      .share = sum(.data$.stratum_total),
+      .by = dplyr::all_of(group_by)
+    ) |>
+    dplyr::mutate(.share = .data$.share / sum(.data$.share))
+
+  per_stratum <- purrr::map(
+    .x = seq_len(nrow(strata)),
+    .f = \(i) {
+      stratum <- strata[i, group_by, drop = FALSE]
+
+      synth_g <- dplyr::semi_join(synth_data, stratum, by = group_by)
+      conf_g <- dplyr::semi_join(conf_data, stratum, by = group_by)
+
+      cells_g <- .compute_marginal_cells(
+        synth_data = synth_g,
+        conf_data = conf_g,
+        combos = combos,
+        weight_var = weight_var,
+        na.rm = na.rm,
+        allow_empty_synth = TRUE
+      )
+
+      marginals_g <- cells_g |>
+        dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables")
+
+      list(
+        cells = dplyr::bind_cols(stratum, cells_g),
+        marginals = dplyr::bind_cols(stratum, marginals_g),
+        group_scores = dplyr::bind_cols(
+          stratum,
+          tibble::tibble(
+            share = strata$.share[i],
+            score = (1 - mean(marginals_g$madd)) * 1000
+          )
+        )
+      )
+    }
+  )
+
+  cells <- purrr::list_rbind(purrr::map(per_stratum, "cells")) |>
+    dplyr::arrange(dplyr::desc(.data$abs_diff))
+
+  marginals <- purrr::list_rbind(purrr::map(per_stratum, "marginals")) |>
+    dplyr::arrange(dplyr::desc(.data$madd))
+
+  group_scores <- purrr::list_rbind(
+    purrr::map(per_stratum, "group_scores")
+  ) |>
+    dplyr::arrange(.data$score)
+
+  # per-stratum scores roll up weighted by confidential shares, so small
+  # strata surface in group_scores without dominating the headline
+  score <- sum(group_scores$share * group_scores$score)
+
+  return(
+    list(
+      score = score,
+      marginals = marginals,
+      cells = cells,
+      group_scores = group_scores
+    )
+  )
+}
diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index 0b6e066..a6f1a90 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -188,70 +188,16 @@
   }
 
   if (!is.null(bins)) {
-    if (!(is.numeric(bins) && length(bins) == 1 && !is.na(bins) &&
-      bins >= 2 && bins == floor(bins))) {
-      stop("`bins` must be a single integer >= 2")
-    }
-
-    numeric_vars <- shared_vars[
-      purrr::map_lgl(.x = shared_vars, .f = \(v) is.numeric(conf_data[[v]]))
-    ]
-
-    for (var in numeric_vars) {
-      # cut points derive from observed values; missing values follow na.rm
-      # like every other variable, becoming an NA bin or dropped rows
-      conf_values <- conf_data[[var]][!is.na(conf_data[[var]])]
-
-      if (!all(is.finite(conf_values)) || length(conf_values) == 0) {
-        stop(
-          "observed numeric values must be finite to discretize; `", var,
-          "` is not"
-        )
-      }
-
-      if (dplyr::n_distinct(conf_values) < bins) {
-        stop(
-          "`", var, "` has fewer distinct confidential values than `bins`"
-        )
-      }
-
-      # interior cut points always derive from the confidential data; each
-      # method yields bins - 1 of them
-      cut_points <- switch(
-        EXPR = discretize_method,
-        width = seq(
-          from = min(conf_values),
-          to = max(conf_values),
-          length.out = bins + 1
-        )[2:bins],
-        ntile = stats::quantile(
-          x = conf_values,
-          probs = seq(from = 0, to = 1, length.out = bins + 1),
-          names = FALSE
-        )[2:bins],
-        cluster = {
-          centers <- sort(
-            stats::kmeans(x = conf_values, centers = bins)$centers[, 1]
-          )
-
-          (centers[-1] + centers[-length(centers)]) / 2
-        }
-      )
-
-      # outer bins extend to +/-Inf so out-of-range synthetic values land in
-      # edge bins; unique() collapses ties from skewed quantiles
-      breaks <- unique(c(-Inf, cut_points, Inf))
-
-      if (length(breaks) - 1 < bins) {
-        warning(
-          "`", var, "` was discretized into ", length(breaks) - 1,
-          " bins instead of ", bins, " because of tied cut points"
-        )
-      }
+    discretized <- .discretize_k_marginal_vars(
+      synth_data = synth_data,
+      conf_data = conf_data,
+      vars = shared_vars,
+      bins = bins,
+      discretize_method = discretize_method
+    )
 
-      synth_data[[var]] <- cut(x = synth_data[[var]], breaks = breaks)
-      conf_data[[var]] <- cut(x = conf_data[[var]], breaks = breaks)
-    }
+    synth_data <- discretized$synth_data
+    conf_data <- discretized$conf_data
   }
 
   if (length(shared_vars) < k) {
@@ -316,113 +262,21 @@
     }
   }
 
-  kmarginals_vars <- t(utils::combn(x = shared_vars, m = k))
-
-  # sample combinations down to n_marginals, always keeping combinations that
-  # contain a priority variable
-  if (nrow(kmarginals_vars) > n_marginals) {
-    is_priority <- apply(
-      X = kmarginals_vars,
-      MARGIN = 1,
-      FUN = \(vars) any(vars %in% priority_vars)
-    )
-
-    n_sampled <- min(max(n_marginals - sum(is_priority), 0), sum(!is_priority))
-
-    sampled_rows <- sample(x = which(!is_priority), size = n_sampled)
-
-    kmarginals_vars <- kmarginals_vars[
-      sort(c(which(is_priority), sampled_rows)), ,
-      drop = FALSE
-    ]
-  }
-
-  # cell proportions for one dataset over one set of variables; weighted
-  # proportions are weight shares instead of row shares
-  process_data <- function(data, vars, prop_name, allow_empty) {
-    if (na.rm) {
-      data <- dplyr::filter(
-        data,
-        !dplyr::if_any(.cols = dplyr::all_of(vars), .fns = is.na)
-      )
-
-      if (nrow(data) == 0 && !allow_empty) {
-        stop(
-          "no rows remain for the marginal over ",
-          paste(vars, collapse = ", "),
-          " after removing missing values"
-        )
-      }
-    }
-
-    if (is.null(weight_var)) {
-      counts <- dplyr::count(data, dplyr::across(dplyr::all_of(vars)))
-    } else {
-      counts <- dplyr::count(
-        data,
-        dplyr::across(dplyr::all_of(vars)),
-        wt = .data[[weight_var]]
-      )
-    }
-
-    props <- counts |>
-      dplyr::mutate("{prop_name}" := .data$n / sum(.data$n)) |>
-      dplyr::select(-"n")
-
-    return(props)
-  }
-
-  # per-cell differences for one set of variables; cells absent from one
-  # dataset count as 0
-  marginal_cells <- function(vars, synth_d, conf_d, allow_empty_synth) {
-    # only the synthetic side may be empty (a stratum the synthesis never
-    # produced); a confidential marginal with no rows has nothing to score
-    # against and errors inside process_data
-    cells <- dplyr::full_join(
-      process_data(
-        data = synth_d, vars = vars, prop_name = "prop_synth",
-        allow_empty = allow_empty_synth
-      ),
-      process_data(
-        data = conf_d, vars = vars, prop_name = "prop_conf",
-        allow_empty = FALSE
-      ),
-      by = vars
-    ) |>
-      tidyr::replace_na(replace = list(prop_synth = 0, prop_conf = 0)) |>
-      tidyr::unite(col = "cell", dplyr::all_of(vars), sep = ", ") |>
-      dplyr::mutate(
-        variables = paste(vars, collapse = ", "),
-        abs_diff = abs(.data$prop_synth - .data$prop_conf)
-      ) |>
-      dplyr::select(
-        "variables", "cell", "prop_synth", "prop_conf", "abs_diff"
-      )
-    # variables disambiguates cells across combinations and drives the
-    # per-combination summary; the prop columns show the direction of the
-    # discrepancy, not just its size
-    return(cells)
-  }
-
-  # all per-cell differences for one (synth, conf) pair of datasets
-  compute_cells <- function(synth_d, conf_d, allow_empty_synth) {
-    cells <- purrr::map(
-      .x = seq_len(nrow(kmarginals_vars)),
-      .f = \(i) marginal_cells(
-        vars = kmarginals_vars[i, ],
-        synth_d = synth_d,
-        conf_d = conf_d,
-        allow_empty_synth = allow_empty_synth
-      )
-    ) |>
-      purrr::list_rbind()
-
-    return(cells)
-  }
+  kmarginals_vars <- .select_k_marginal_combos(
+    shared_vars = shared_vars,
+    k = k,
+    n_marginals = n_marginals,
+    priority_vars = priority_vars
+  )
 
   if (is.null(group_by)) {
-    cells <- compute_cells(
-      synth_d = synth_data, conf_d = conf_data, allow_empty_synth = FALSE
+    cells <- .compute_marginal_cells(
+      synth_data = synth_data,
+      conf_data = conf_data,
+      combos = kmarginals_vars,
+      weight_var = weight_var,
+      na.rm = na.rm,
+      allow_empty_synth = FALSE
     ) |>
       dplyr::arrange(dplyr::desc(.data$abs_diff))
 
@@ -433,11 +287,9 @@
 
     # mean of the MabsDDs, rescaled to an ascending measure on [0, 1000];
     # computed from all marginals before any truncation
-    score <- (1 - mean(marginals$madd)) * 1000
-
     result <- structure(
       list(
-        score = score,
+        score = (1 - mean(marginals$madd)) * 1000,
         marginals = utils::head(marginals, n = keep_marginals),
         cells = utils::head(cells, n = keep_cells)
       ),
@@ -447,73 +299,21 @@
     return(result)
   }
 
-  # strata are defined by the confidential data; a stratum with no synthetic
-  # rows scores against all-zero synthetic proportions. Shares are the
-  # confidential row (or weight) share of each stratum, computed once
-  conf_totals <- if (is.null(weight_var)) {
-    rep(1, nrow(conf_data))
-  } else {
-    conf_data[[weight_var]]
-  }
-
-  strata <- conf_data |>
-    dplyr::mutate(.stratum_total = conf_totals) |>
-    dplyr::summarize(
-      .share = sum(.data$.stratum_total),
-      .by = dplyr::all_of(group_by)
-    ) |>
-    dplyr::mutate(.share = .data$.share / sum(.data$.share))
-
-  per_stratum <- purrr::map(
-    .x = seq_len(nrow(strata)),
-    .f = \(i) {
-      stratum <- strata[i, group_by, drop = FALSE]
-
-      synth_g <- dplyr::semi_join(synth_data, stratum, by = group_by)
-      conf_g <- dplyr::semi_join(conf_data, stratum, by = group_by)
-
-      cells_g <- compute_cells(
-        synth_d = synth_g, conf_d = conf_g, allow_empty_synth = TRUE
-      )
-
-      marginals_g <- cells_g |>
-        dplyr::summarize(madd = mean(.data$abs_diff), .by = "variables")
-
-      list(
-        cells = dplyr::bind_cols(stratum, cells_g),
-        marginals = dplyr::bind_cols(stratum, marginals_g),
-        group_scores = dplyr::bind_cols(
-          stratum,
-          tibble::tibble(
-            share = strata$.share[i],
-            score = (1 - mean(marginals_g$madd)) * 1000
-          )
-        )
-      )
-    }
+  stratified <- .stratify_k_marginals(
+    synth_data = synth_data,
+    conf_data = conf_data,
+    combos = kmarginals_vars,
+    group_by = group_by,
+    weight_var = weight_var,
+    na.rm = na.rm
   )
 
-  cells <- purrr::list_rbind(purrr::map(per_stratum, "cells")) |>
-    dplyr::arrange(dplyr::desc(.data$abs_diff))
-
-  marginals <- purrr::list_rbind(purrr::map(per_stratum, "marginals")) |>
-    dplyr::arrange(dplyr::desc(.data$madd))
-
-  group_scores <- purrr::list_rbind(
-    purrr::map(per_stratum, "group_scores")
-  ) |>
-    dplyr::arrange(.data$score)
-
-  # per-stratum scores roll up weighted by confidential shares, so small
-  # strata surface in group_scores without dominating the headline
-  score <- sum(group_scores$share * group_scores$score)
-
   result <- structure(
     list(
-      score = score,
-      marginals = utils::head(marginals, n = keep_marginals),
-      cells = utils::head(cells, n = keep_cells),
-      group_scores = group_scores
+      score = stratified$score,
+      marginals = utils::head(stratified$marginals, n = keep_marginals),
+      cells = utils::head(stratified$cells, n = keep_cells),
+      group_scores = stratified$group_scores
     ),
     class = "k_marginals"
   )
diff --git a/man/dot-compute_marginal_cells.Rd b/man/dot-compute_marginal_cells.Rd
new file mode 100644
index 0000000..4de3ee0
--- /dev/null
+++ b/man/dot-compute_marginal_cells.Rd
@@ -0,0 +1,43 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/compute_marginal_cells.R
+\name{.compute_marginal_cells}
+\alias{.compute_marginal_cells}
+\title{Compute per-cell proportion differences for the k-marginals metric}
+\usage{
+.compute_marginal_cells(
+  synth_data,
+  conf_data,
+  combos,
+  weight_var = NULL,
+  na.rm = FALSE,
+  allow_empty_synth = FALSE
+)
+}
+\arguments{
+\item{synth_data}{A tibble with synthetic data.}
+
+\item{conf_data}{A tibble with confidential data.}
+
+\item{combos}{A character matrix with one row per variable combination.}
+
+\item{weight_var}{Optional character name of a numeric sample-weight
+column; when set, cell proportions are weight shares instead of row
+shares. Defaults to \code{NULL} (unweighted).}
+
+\item{na.rm}{A logical for dropping rows with a missing value from each
+marginal that uses the affected variable.}
+
+\item{allow_empty_synth}{A logical for permitting the synthetic data to
+contribute no rows to a marginal (a stratum the synthesis never produced),
+in which case its proportions are zero. A confidential marginal with no
+rows always errors, since there is nothing to score against.}
+}
+\value{
+A tibble with one row per cell: \code{variables}, \code{cell},
+\code{prop_synth}, \code{prop_conf}, and \code{abs_diff}.
+}
+\description{
+For each supplied variable combination, computes the marginal
+cell proportions in the synthetic and confidential data and their absolute
+differences. Cells absent from one dataset count as proportion zero.
+}
diff --git a/man/dot-discretize_k_marginal_vars.Rd b/man/dot-discretize_k_marginal_vars.Rd
new file mode 100644
index 0000000..c853a41
--- /dev/null
+++ b/man/dot-discretize_k_marginal_vars.Rd
@@ -0,0 +1,39 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/discretize_k_marginal_vars.R
+\name{.discretize_k_marginal_vars}
+\alias{.discretize_k_marginal_vars}
+\title{Discretize numeric variables for the k-marginals metric}
+\usage{
+.discretize_k_marginal_vars(
+  synth_data,
+  conf_data,
+  vars,
+  bins,
+  discretize_method
+)
+}
+\arguments{
+\item{synth_data}{A tibble with synthetic data.}
+
+\item{conf_data}{A tibble with confidential data.}
+
+\item{vars}{Character vector of candidate variables; only numeric ones are
+discretized.}
+
+\item{bins}{Single integer >= 2 giving the number of bins. Fewer bins are
+produced, with a warning, if tied cut points collapse.}
+
+\item{discretize_method}{Method used to place bin breaks: "width" for
+fixed binwidths, "ntile" for quantile bins, or "cluster" for univariate
+k-means clustering (set a seed before calling for reproducible clusters).}
+}
+\value{
+A list with the discretized \code{synth_data} and \code{conf_data}.
+}
+\description{
+Discretizes every numeric variable among \code{vars} into \code{bins}
+bins with interior cut points derived from the observed (non-missing)
+confidential values and applied to both datasets. The outer bins extend to
++/-Inf so synthetic values outside the confidential range land in edge
+bins.
+}
diff --git a/man/dot-select_k_marginal_combos.Rd b/man/dot-select_k_marginal_combos.Rd
new file mode 100644
index 0000000..8f0aabe
--- /dev/null
+++ b/man/dot-select_k_marginal_combos.Rd
@@ -0,0 +1,30 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/select_k_marginal_combos.R
+\name{.select_k_marginal_combos}
+\alias{.select_k_marginal_combos}
+\title{Select the variable combinations for the k-marginals metric}
+\usage{
+.select_k_marginal_combos(shared_vars, k, n_marginals, priority_vars = NULL)
+}
+\arguments{
+\item{shared_vars}{Character vector of variables available for marginals.}
+
+\item{k}{Scalar order of the k-marginal.}
+
+\item{n_marginals}{Single integer target maximum for the number of
+variable combinations, or \code{Inf} for no cap. All priority combinations are
+retained even when they alone exceed the target.}
+
+\item{priority_vars}{Optional character vector of variable names whose
+combinations always survive sampling. Defaults to \code{NULL}.}
+}
+\value{
+A character matrix with one row per selected combination and \code{k}
+columns.
+}
+\description{
+Enumerates all unique k-combinations of the supplied
+variables and, when the total exceeds \code{n_marginals}, samples the
+combinations down to the cap while always retaining combinations that
+contain a priority variable.
+}
diff --git a/man/dot-stratify_k_marginals.Rd b/man/dot-stratify_k_marginals.Rd
new file mode 100644
index 0000000..60813f4
--- /dev/null
+++ b/man/dot-stratify_k_marginals.Rd
@@ -0,0 +1,43 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/stratify_k_marginals.R
+\name{.stratify_k_marginals}
+\alias{.stratify_k_marginals}
+\title{Compute stratified k-marginals results}
+\usage{
+.stratify_k_marginals(
+  synth_data,
+  conf_data,
+  combos,
+  group_by,
+  weight_var = NULL,
+  na.rm = FALSE
+)
+}
+\arguments{
+\item{synth_data}{A tibble with synthetic data.}
+
+\item{conf_data}{A tibble with confidential data.}
+
+\item{combos}{A character matrix with one row per variable combination.}
+
+\item{group_by}{Character vector of grouping variable names.}
+
+\item{weight_var}{Optional character name of a numeric sample-weight
+column; when set, cell proportions and stratum shares are weight shares
+instead of row shares. Defaults to \code{NULL} (unweighted).}
+
+\item{na.rm}{A logical for dropping rows with a missing value from each
+marginal that uses the affected variable.}
+}
+\value{
+A list with \code{score} (the share-weighted mean of per-stratum
+scores), \code{marginals} and \code{cells} (stacked per-stratum tables with the
+grouping columns, worst first), and \code{group_scores} (one row per stratum
+with its share and score, worst first).
+}
+\description{
+Splits the synthetic and confidential data into the strata
+observed in the confidential data, computes per-cell differences and
+per-combination MabsDDs within each stratum, and rolls the per-stratum
+scores up weighted by each stratum's confidential share.
+}
diff --git a/tests/testthat/test-compute_marginal_cells.R b/tests/testthat/test-compute_marginal_cells.R
new file mode 100644
index 0000000..d76a40b
--- /dev/null
+++ b/tests/testthat/test-compute_marginal_cells.R
@@ -0,0 +1,59 @@
+test_that("cells report proportions and absolute differences per combo", {
+  synth_cm <- tibble::tibble(
+    a = c("x", "y", "y", "y"),
+    b = c("p", "q", "q", "p")
+  )
+
+  conf_cm <- tibble::tibble(
+    a = c("x", "x", "y", "y"),
+    b = c("p", "p", "p", "q")
+  )
+
+  combos <- matrix(c("a", "b"), ncol = 1)
+
+  cells <- .compute_marginal_cells(
+    synth_data = synth_cm,
+    conf_data = conf_cm,
+    combos = combos
+  )
+
+  expect_named(
+    cells,
+    c("variables", "cell", "prop_synth", "prop_conf", "abs_diff")
+  )
+
+  # one row per observed level of each single-variable marginal
+  expect_equal(nrow(cells), 4)
+  expect_equal(unique(cells$variables), c("a", "b"))
+})
+
+test_that("only the synthetic side may be empty", {
+  synth_empty <- tibble::tibble(a = c(NA_character_, NA_character_))
+  conf_ok <- tibble::tibble(a = c("x", "y"))
+
+  combos <- matrix("a", ncol = 1)
+
+  # empty synthetic side scores against zero proportions
+  cells <- .compute_marginal_cells(
+    synth_data = synth_empty,
+    conf_data = conf_ok,
+    combos = combos,
+    na.rm = TRUE,
+    allow_empty_synth = TRUE
+  )
+
+  expect_equal(cells$prop_synth, c(0, 0))
+  expect_equal(cells$prop_conf, c(0.5, 0.5))
+
+  # the confidential side has nothing to score against and always errors
+  expect_error(
+    .compute_marginal_cells(
+      synth_data = conf_ok,
+      conf_data = synth_empty,
+      combos = combos,
+      na.rm = TRUE,
+      allow_empty_synth = TRUE
+    ),
+    regexp = "no rows remain"
+  )
+})
diff --git a/tests/testthat/test-discretize_k_marginal_vars.R b/tests/testthat/test-discretize_k_marginal_vars.R
new file mode 100644
index 0000000..69472e0
--- /dev/null
+++ b/tests/testthat/test-discretize_k_marginal_vars.R
@@ -0,0 +1,90 @@
+test_that("only numeric variables are discretized", {
+  conf_mix <- tibble::tibble(v = c(1, 2, 3, 4), a = c("x", "y", "x", "y"))
+  synth_mix <- tibble::tibble(v = c(1, 1, 4, 4), a = c("x", "x", "y", "y"))
+
+  result <- .discretize_k_marginal_vars(
+    synth_data = synth_mix,
+    conf_data = conf_mix,
+    vars = c("v", "a"),
+    bins = 2,
+    discretize_method = "width"
+  )
+
+  expect_s3_class(result$conf_data$v, "factor")
+  expect_s3_class(result$synth_data$v, "factor")
+  expect_type(result$conf_data$a, "character")
+
+  # breaks derive from the confidential data and apply to both datasets
+  expect_equal(levels(result$conf_data$v), levels(result$synth_data$v))
+})
+
+test_that("invalid bins values error", {
+  conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
+
+  for (bad in list(1, 2.5, "2", NA_real_, c(2, 3))) {
+    expect_error(
+      .discretize_k_marginal_vars(
+        synth_data = conf_num,
+        conf_data = conf_num,
+        vars = "v",
+        bins = bad,
+        discretize_method = "width"
+      ),
+      regexp = "`bins` must be a single integer >= 2"
+    )
+  }
+})
+
+test_that("too few distinct confidential values throw an error", {
+  conf_const <- tibble::tibble(v = c(2, 2, 2, 2))
+  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
+
+  for (method in c("width", "ntile", "cluster")) {
+    expect_error(
+      .discretize_k_marginal_vars(
+        synth_data = synth_num,
+        conf_data = conf_const,
+        vars = "v",
+        bins = 2,
+        discretize_method = method
+      ),
+      regexp = "fewer distinct confidential values than `bins`"
+    )
+  }
+})
+
+test_that("tied quantile cut points collapse bins with a warning", {
+  # heavily tied data passes the distinct-value pre-check, but the 25th and
+  # 50th percentiles coincide at 1, collapsing a quantile bin
+  conf_ties <- tibble::tibble(v = c(1, 1, 1, 1, 1, 1, 2, 3, 4, 5))
+  synth_num <- tibble::tibble(v = c(1, 2, 3, 4, 5))
+
+  expect_warning(
+    result <- .discretize_k_marginal_vars(
+      synth_data = synth_num,
+      conf_data = conf_ties,
+      vars = "v",
+      bins = 4,
+      discretize_method = "ntile"
+    ),
+    regexp = "bins instead of 4 because of tied cut points"
+  )
+
+  expect_lt(length(levels(result$conf_data$v)), 4)
+})
+
+test_that("infinite confidential values refuse to discretize", {
+  conf_inf <- tibble::tibble(v = c(1, 2, 3, Inf))
+  synth_num <- tibble::tibble(v = c(1, 2, 3, 3))
+
+  expect_error(
+    .discretize_k_marginal_vars(
+      synth_data = synth_num,
+      conf_data = conf_inf,
+      vars = "v",
+      bins = 2,
+      discretize_method = "width"
+    ),
+    regexp = "must be finite to discretize"
+  )
+})
diff --git a/tests/testthat/test-select_k_marginal_combos.R b/tests/testthat/test-select_k_marginal_combos.R
new file mode 100644
index 0000000..275afe2
--- /dev/null
+++ b/tests/testthat/test-select_k_marginal_combos.R
@@ -0,0 +1,55 @@
+test_that("all combinations are enumerated when under the cap", {
+  combos <- .select_k_marginal_combos(
+    shared_vars = c("a", "b", "c"),
+    k = 2,
+    n_marginals = Inf
+  )
+
+  expect_equal(nrow(combos), 3)
+  expect_equal(ncol(combos), 2)
+})
+
+test_that("priority combinations are always kept", {
+  # the two combinations containing a fill the cap exactly, so the
+  # selection is deterministic despite sampling
+  combos <- .select_k_marginal_combos(
+    shared_vars = c("a", "b", "c"),
+    k = 2,
+    n_marginals = 2,
+    priority_vars = "a"
+  )
+
+  expect_equal(nrow(combos), 2)
+  expect_true(all(apply(combos, 1, \(vars) "a" %in% vars)))
+})
+
+test_that("priority combinations exceeding n_marginals are all kept", {
+  combos <- .select_k_marginal_combos(
+    shared_vars = c("a", "b", "c"),
+    k = 2,
+    n_marginals = 1,
+    priority_vars = "a"
+  )
+
+  expect_equal(nrow(combos), 2)
+  expect_true(all(apply(combos, 1, \(vars) "a" %in% vars)))
+})
+
+test_that("sampling fills remaining slots after priority combinations", {
+  # 4 variables, k = 2: 6 combinations, 3 containing a
+  set.seed(20250813)
+
+  combos <- .select_k_marginal_combos(
+    shared_vars = c("a", "b", "c", "d"),
+    k = 2,
+    n_marginals = 4,
+    priority_vars = "a"
+  )
+
+  # cap respected exactly: all 3 priority combos plus 1 sampled non-priority
+  expect_equal(nrow(combos), 4)
+
+  has_a <- apply(combos, 1, \(vars) "a" %in% vars)
+  expect_equal(sum(has_a), 3)
+  expect_equal(sum(!has_a), 1)
+})
diff --git a/tests/testthat/test-stratify_k_marginals.R b/tests/testthat/test-stratify_k_marginals.R
new file mode 100644
index 0000000..7df0a14
--- /dev/null
+++ b/tests/testthat/test-stratify_k_marginals.R
@@ -0,0 +1,86 @@
+test_that("stratified results roll up by confidential shares", {
+  # stratum A diverges (score 500), stratum B matches (score 1000);
+  # equal confidential shares give an overall score of 750
+  conf_st <- tibble::tibble(
+    g = c("A", "A", "B", "B"),
+    a = c("x", "y", "x", "y")
+  )
+
+  synth_st <- tibble::tibble(
+    g = c("A", "A", "B", "B"),
+    a = c("x", "x", "x", "y")
+  )
+
+  combos <- matrix("a", ncol = 1)
+
+  result <- .stratify_k_marginals(
+    synth_data = synth_st,
+    conf_data = conf_st,
+    combos = combos,
+    group_by = "g"
+  )
+
+  expect_equal(result$score, 750)
+
+  # group_scores are worst first with the grouping column attached
+  expect_named(result$group_scores, c("g", "share", "score"))
+  expect_equal(result$group_scores$g, c("A", "B"))
+  expect_equal(result$group_scores$share, c(0.5, 0.5))
+
+  # stacked detail tables carry the grouping column, worst first
+  expect_equal(result$marginals$g[1], "A")
+  expect_true("g" %in% names(result$cells))
+})
+
+test_that("stratum shares use weights when weight_var is set", {
+  # conf weight shares: A = 0.5, B = 0.5 (row shares would be 2/3, 1/3)
+  conf_stw <- tibble::tibble(
+    g = c("A", "A", "B"),
+    a = c("x", "y", "x"),
+    w = c(1, 1, 2)
+  )
+
+  synth_stw <- tibble::tibble(
+    g = c("A", "A", "B"),
+    a = c("x", "x", "x"),
+    w = c(1, 1, 1)
+  )
+
+  result <- .stratify_k_marginals(
+    synth_data = synth_stw,
+    conf_data = conf_stw,
+    combos = matrix("a", ncol = 1),
+    group_by = "g",
+    weight_var = "w"
+  )
+
+  expect_equal(sort(result$group_scores$share), c(0.5, 0.5))
+  expect_equal(result$score, 750)
+})
+
+test_that("multi-column group_by builds joint strata directly", {
+  # only stratum (A, p) diverges: conf (x = 0.5, y = 0.5), synth (x = 1)
+  # -> 500; the other three strata match -> 1000
+  # overall = 0.25 * 500 + 0.75 * 1000 = 875
+  conf_st2 <- tibble::tibble(
+    g1 = c("A", "A", "A", "A", "B", "B", "B", "B"),
+    g2 = c("p", "p", "q", "q", "p", "p", "q", "q"),
+    a = c("x", "y", "x", "y", "x", "y", "x", "y")
+  )
+
+  synth_st2 <- dplyr::mutate(
+    conf_st2,
+    a = c("x", "x", "x", "y", "x", "y", "x", "y")
+  )
+
+  result <- .stratify_k_marginals(
+    synth_data = synth_st2,
+    conf_data = conf_st2,
+    combos = matrix("a", ncol = 1),
+    group_by = c("g1", "g2")
+  )
+
+  expect_equal(result$score, 875)
+  expect_named(result$group_scores, c("g1", "g2", "share", "score"))
+  expect_equal(nrow(result$group_scores), 4)
+})
diff --git a/tests/testthat/test-util_k_marginals.R b/tests/testthat/test-util_k_marginals.R
index 1bba2ea..5d7cd1d 100644
--- a/tests/testthat/test-util_k_marginals.R
+++ b/tests/testthat/test-util_k_marginals.R
@@ -32,38 +32,29 @@ synth <- tibble::tibble(
 )
 
 test_that("k = 1 score matches hand-computed value", {
-
   expect_equal(
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score,
     750
   )
-
 })
 
 test_that("k = 2 score matches hand-computed value", {
-
   expect_equal(
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 2)$score,
     5000 / 6
   )
-
 })
 
 test_that("identical data scores exactly 1000 for every k", {
-
   for (k in 1:2) {
-
     expect_equal(
       .util_k_marginals(synth_data = conf, conf_data = conf, k = k)$score,
       1000
     )
-
   }
-
 })
 
 test_that("cells absent from one dataset count as proportion zero", {
-
   # conf has level y that synth lacks; synth is all x
   # marginal a: conf (x = 0.5, y = 0.5), synth (x = 1, y = 0)
   # MabsDD = mean(0.5, 0.5) = 0.5 -> score 500
@@ -74,11 +65,9 @@ test_that("cells absent from one dataset count as proportion zero", {
     .util_k_marginals(synth_data = synth_gap, conf_data = conf_gap, k = 1)$score,
     500
   )
-
 })
 
 test_that("k outside 1:3 throws an error", {
-
   expect_error(
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 4),
     regexp = "`k` must be a single integer between 1 and 3"
@@ -104,11 +93,9 @@ test_that("k outside 1:3 throws an error", {
     .util_k_marginals(synth_data = synth, conf_data = conf, k = "1"),
     regexp = "`k` must be a single integer between 1 and 3"
   )
-
 })
 
 test_that("zero-row inputs throw an error instead of returning NaN", {
-
   empty <- conf[0, ]
 
   expect_error(
@@ -120,20 +107,16 @@ test_that("zero-row inputs throw an error instead of returning NaN", {
     .util_k_marginals(synth_data = synth, conf_data = empty, k = 1),
     regexp = "at least one row"
   )
-
 })
 
 test_that("k exceeding the number of shared variables throws an error", {
-
   expect_error(
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 3),
     regexp = "`k` cannot exceed the number of variables available"
   )
-
 })
 
 test_that("variables not shared by both datasets are ignored", {
-
   # extra synth-only column must not create combinations
   synth_extra <- dplyr::mutate(synth, c = c("m", "m", "n", "n"))
 
@@ -141,37 +124,29 @@ test_that("variables not shared by both datasets are ignored", {
     .util_k_marginals(synth_data = synth_extra, conf_data = conf, k = 1)$score,
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
   )
-
 })
 
 test_that("util_k_marginals accepts an eval_data object", {
-
   ed <- eval_data(conf_data = conf, synth_data = synth)
 
   expect_equal(util_k_marginals(eval_data = ed, k = 1)$score, 750)
 
   expect_equal(util_k_marginals(eval_data = ed, k = 2)$score, 5000 / 6)
-
 })
 
 test_that("util_k_marginals maps over replicates", {
-
   ed <- eval_data(conf_data = conf, synth_data = list(synth, conf))
 
   result <- util_k_marginals(eval_data = ed, k = 1)
 
   expect_equal(purrr::map_dbl(result, "score"), c(750, 1000))
-
 })
 
 test_that("util_k_marginals rejects non-eval_data input", {
-
   expect_error(util_k_marginals(eval_data = synth, k = 1))
-
 })
 
 test_that("marginals and cells report worst-first detail", {
-
   result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 2)
 
   expect_s3_class(result, "k_marginals")
@@ -187,11 +162,9 @@ test_that("marginals and cells report worst-first detail", {
   )
   expect_equal(result$cells$abs_diff, c(0.25, 0.25, 0))
   expect_equal(result$cells$cell[3], "y, p")
-
 })
 
 test_that("cells absent from the synthetic data appear with proportion zero", {
-
   conf_gap <- tibble::tibble(a = c("x", "y"))
   synth_gap <- tibble::tibble(a = c("x", "x"))
 
@@ -201,20 +174,16 @@ test_that("cells absent from the synthetic data appear with proportion zero", {
 
   expect_equal(y_cell$prop_synth, 0)
   expect_equal(y_cell$prop_conf, 0.5)
-
 })
 
 test_that("print method reports the score and worst marginals", {
-
   result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
 
   expect_output(print(result), regexp = "k-marginals score: 750")
   expect_output(print(result), regexp = "Worst marginals:")
-
 })
 
 test_that("keep_marginals and keep_cells truncate the detail tables", {
-
   full <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
 
   kept <- .util_k_marginals(
@@ -231,11 +200,9 @@ test_that("keep_marginals and keep_cells truncate the detail tables", {
   # retained rows are the worst ones from the full tables
   expect_equal(kept$marginals, full$marginals[1, ])
   expect_equal(kept$cells, full$cells[1:2, ])
-
 })
 
 test_that("retention keeps the highest abs_diff cells", {
-
   # three levels with a strict worst cell: y has the largest abs_diff
   # conf: x = 0.500, y = 0.250, z = 0.250
   # synth: x = 0.250, y = 0.625, z = 0.125
@@ -252,11 +219,9 @@ test_that("retention keeps the highest abs_diff cells", {
 
   expect_equal(kept$cells$cell, "y")
   expect_equal(kept$cells$abs_diff, 0.375)
-
 })
 
 test_that("score is computed from all marginals, not the retained subset", {
-
   full <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
 
   kept <- .util_k_marginals(
@@ -268,13 +233,10 @@ test_that("score is computed from all marginals, not the retained subset", {
   )
 
   expect_equal(kept$score, full$score)
-
 })
 
 test_that("invalid retention arguments throw an error", {
-
   for (bad_keep in list(0, "5", 1.5, NA_real_, NaN, c(1, 2))) {
-
     expect_error(
       .util_k_marginals(
         synth_data = synth, conf_data = conf, k = 1, keep_marginals = bad_keep
@@ -288,7 +250,6 @@ test_that("invalid retention arguments throw an error", {
       ),
       regexp = "must be single integers >= 1 or Inf"
     )
-
   }
 
   # Inf remains valid: it is the documented keep-everything default
@@ -299,11 +260,9 @@ test_that("invalid retention arguments throw an error", {
     )$score,
     750
   )
-
 })
 
 test_that("util_k_marginals passes retention arguments through", {
-
   ed <- eval_data(conf_data = conf, synth_data = synth)
 
   result <- util_k_marginals(
@@ -313,24 +272,18 @@ test_that("util_k_marginals passes retention arguments through", {
   expect_equal(nrow(result$marginals), 1)
   expect_equal(nrow(result$cells), 2)
   expect_equal(result$score, 750)
-
 })
 
 test_that("non-integer and non-finite k values throw an error", {
-
   for (bad_k in list(1.5, NA_real_, NaN, Inf)) {
-
     expect_error(
       .util_k_marginals(synth_data = synth, conf_data = conf, k = bad_k),
       regexp = "`k` must be a single integer between 1 and 3"
     )
-
   }
-
 })
 
 test_that("marginals are sorted by descending madd across combinations", {
-
   # third shared variable, identical in conf but perturbed in synth
   # pair madds: (a, c) = 1/4, (a, b) = 1/6, (b, c) = 1/6
   conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
@@ -346,50 +299,40 @@ test_that("marginals are sorted by descending madd across combinations", {
     result$cells$abs_diff,
     sort(result$cells$abs_diff, decreasing = TRUE)
   )
-
 })
 
 test_that("each replicate result is a complete k_marginals object", {
-
   ed <- eval_data(conf_data = conf, synth_data = list(synth, conf))
 
   result <- util_k_marginals(eval_data = ed, k = 1)
 
   for (rep in result) {
-
     expect_s3_class(rep, "k_marginals")
     expect_named(rep, c("score", "marginals", "cells"))
     expect_gt(nrow(rep$marginals), 0)
     expect_gt(nrow(rep$cells), 0)
-
   }
-
 })
 
 test_that("conf-only extra columns are ignored", {
-
   conf_extra <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
 
   expect_equal(
     .util_k_marginals(synth_data = synth, conf_data = conf_extra, k = 1)$score,
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
   )
-
 })
 
 test_that("one-row datasets produce a valid result", {
-
   one_row <- tibble::tibble(a = "x", b = "p")
 
   result <- .util_k_marginals(synth_data = one_row, conf_data = one_row, k = 1)
 
   expect_equal(result$score, 1000)
   expect_equal(nrow(result$cells), 2)
-
 })
 
 test_that("print truncates the marginals display to n rows", {
-
   result <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
 
   out <- utils::capture.output(print(result, n = 1))
@@ -397,11 +340,9 @@ test_that("print truncates the marginals display to n rows", {
   # tibble rows print with a leading row number: row 1 only, no row 2
   expect_true(any(grepl("^1 ", out)))
   expect_false(any(grepl("^2 ", out)))
-
 })
 
 test_that("n_marginals caps the number of evaluated combinations", {
-
   conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
   synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
 
@@ -413,11 +354,9 @@ test_that("n_marginals caps the number of evaluated combinations", {
 
   expect_equal(nrow(result$marginals), 2)
   expect_true(result$score >= 0 && result$score <= 1000)
-
 })
 
 test_that("sampling is reproducible given a seed", {
-
   conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
   synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
 
@@ -432,47 +371,9 @@ test_that("sampling is reproducible given a seed", {
   )
 
   expect_equal(first, second)
-
-})
-
-test_that("priority_vars combinations are always evaluated", {
-
-  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
-  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
-
-  # the two combinations containing a fill the cap exactly, so the
-  # selection is deterministic despite sampling
-  result <- .util_k_marginals(
-    synth_data = synth_3,
-    conf_data = conf_3,
-    k = 2,
-    n_marginals = 2,
-    priority_vars = "a"
-  )
-
-  expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
-
-})
-
-test_that("priority combinations exceeding n_marginals are all kept", {
-
-  conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
-  synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
-
-  result <- .util_k_marginals(
-    synth_data = synth_3,
-    conf_data = conf_3,
-    k = 2,
-    n_marginals = 1,
-    priority_vars = "a"
-  )
-
-  expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
-
 })
 
 test_that("n_marginals at or above the combination count changes nothing", {
-
   conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
   synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
 
@@ -483,11 +384,9 @@ test_that("n_marginals at or above the combination count changes nothing", {
   )
 
   expect_equal(capped, full)
-
 })
 
 test_that("invalid sampling arguments throw an error", {
-
   expect_error(
     .util_k_marginals(
       synth_data = synth, conf_data = conf, k = 1, n_marginals = 1.5
@@ -508,11 +407,9 @@ test_that("invalid sampling arguments throw an error", {
     ),
     regexp = "`priority_vars` must be a character vector of variables available"
   )
-
 })
 
 test_that("util_k_marginals passes sampling arguments through", {
-
   conf_3 <- dplyr::mutate(conf, c = c("m", "m", "n", "n"))
   synth_3 <- dplyr::mutate(synth, c = c("m", "n", "n", "n"))
 
@@ -523,49 +420,17 @@ test_that("util_k_marginals passes sampling arguments through", {
   )
 
   expect_equal(sort(result$marginals$variables), c("a, b", "a, c"))
-
-})
-
-test_that("sampling fills remaining slots after priority combinations", {
-
-  # 4 shared variables, k = 2: 6 combinations, 3 containing a
-  conf_4 <- dplyr::mutate(
-    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
-  )
-  synth_4 <- dplyr::mutate(
-    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
-  )
-
-  set.seed(20250813)
-
-  result <- .util_k_marginals(
-    synth_data = synth_4,
-    conf_data = conf_4,
-    k = 2,
-    n_marginals = 4,
-    priority_vars = "a"
-  )
-
-  # cap respected exactly: all 3 priority combos plus 1 sampled non-priority
-  expect_equal(nrow(result$marginals), 4)
-
-  has_a <- grepl("a", result$marginals$variables)
-  expect_equal(sum(has_a), 3)
-  expect_equal(sum(!has_a), 1)
-  expect_true(
-    all(result$marginals$variables[!has_a] %in% c("b, c", "b, d", "c, d"))
-  )
-
 })
 
 test_that("n_marginals caps k = 3 combinations", {
-
   # 4 shared variables, k = 3: 4 combinations
   conf_4 <- dplyr::mutate(
-    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
+    conf,
+    c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
   )
   synth_4 <- dplyr::mutate(
-    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
+    synth,
+    c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
   )
 
   set.seed(20250813)
@@ -576,17 +441,17 @@ test_that("n_marginals caps k = 3 combinations", {
 
   expect_equal(nrow(result$marginals), 2)
   expect_true(result$score >= 0 && result$score <= 1000)
-
 })
 
 test_that("priority_vars applies to k = 3 combinations", {
-
   # priority a appears in 3 of the 4 triples; cap of 3 keeps exactly those
   conf_4 <- dplyr::mutate(
-    conf, c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
+    conf,
+    c = c("m", "m", "n", "n"), d = c("u", "v", "u", "v")
   )
   synth_4 <- dplyr::mutate(
-    synth, c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
+    synth,
+    c = c("m", "n", "n", "n"), d = c("v", "v", "u", "u")
   )
 
   result <- .util_k_marginals(
@@ -601,11 +466,9 @@ test_that("priority_vars applies to k = 3 combinations", {
     sort(result$marginals$variables),
     c("a, b, c", "a, b, d", "a, c, d")
   )
-
 })
 
 test_that("weighted proportions match hand-computed values", {
-
   # conf: weight shares x = 3/4, y = 1/4; synth: x = 1/2, y = 1/2
   # MabsDD = mean(0.25, 0.25) = 0.25 -> score 750
   conf_w <- tibble::tibble(a = c("x", "y"), w = c(3, 1))
@@ -620,11 +483,9 @@ test_that("weighted proportions match hand-computed values", {
     dplyr::filter(result$cells, .data$cell == "x")$prop_conf,
     0.75
   )
-
 })
 
 test_that("unit weights reproduce the unweighted result", {
-
   conf_w <- dplyr::mutate(conf, w = 1)
   synth_w <- dplyr::mutate(synth, w = 1)
 
@@ -635,11 +496,9 @@ test_that("unit weights reproduce the unweighted result", {
   unweighted <- .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)
 
   expect_equal(weighted, unweighted)
-
 })
 
 test_that("the weight column is never a marginal", {
-
   conf_w <- dplyr::mutate(conf, w = 1)
   synth_w <- dplyr::mutate(synth, w = 1)
 
@@ -648,11 +507,9 @@ test_that("the weight column is never a marginal", {
   )
 
   expect_equal(sort(result$marginals$variables), c("a", "b"))
-
 })
 
 test_that("invalid weight_var throws an error", {
-
   conf_w <- dplyr::mutate(conf, w = 1)
   synth_w <- dplyr::mutate(synth, w = 1)
 
@@ -681,11 +538,9 @@ test_that("invalid weight_var throws an error", {
     ),
     regexp = "`weight_var` must be a numeric column in both datasets"
   )
-
 })
 
 test_that("util_k_marginals passes weight_var through", {
-
   conf_w <- tibble::tibble(a = c("x", "y"), w = c(3, 1))
   synth_w <- tibble::tibble(a = c("x", "y"), w = c(1, 1))
 
@@ -695,22 +550,19 @@ test_that("util_k_marginals passes weight_var through", {
     util_k_marginals(eval_data = ed, k = 1, weight_var = "w")$score,
     750
   )
-
 })
 
 test_that("invalid weight values throw an error", {
-
   synth_w <- dplyr::mutate(synth, w = 1)
 
   bad_weights <- list(
-    c(1, 1, 1, -1),   # negative
-    c(1, 1, 1, NA),   # missing
-    c(1, 1, 1, Inf),  # non-finite
-    c(0, 0, 0, 0)     # zero total
+    c(1, 1, 1, -1), # negative
+    c(1, 1, 1, NA), # missing
+    c(1, 1, 1, Inf), # non-finite
+    c(0, 0, 0, 0) # zero total
   )
 
   for (bad_w in bad_weights) {
-
     conf_bad <- dplyr::mutate(conf, w = bad_w)
 
     expect_error(
@@ -730,13 +582,10 @@ test_that("invalid weight values throw an error", {
       ),
       regexp = "finite and non-negative with a positive total"
     )
-
   }
-
 })
 
 test_that("zero weights are valid when the total is positive", {
-
   # zero-weight rows drop out: conf weight shares x = 1, y = 0
   conf_w <- tibble::tibble(a = c("x", "y"), w = c(2, 0))
   synth_w <- tibble::tibble(a = c("x", "y"), w = c(1, 1))
@@ -746,11 +595,9 @@ test_that("zero weights are valid when the total is positive", {
   )
 
   expect_equal(result$score, 500)
-
 })
 
 test_that("width discretization matches hand-computed values", {
-
   # conf 1:4 with 2 bins: interior cut at 2.5, so low = {1, 2}, high = {3, 4}
   # conf shares (0.5, 0.5); synth c(1, 1, 1, 4) shares (0.75, 0.25)
   # MabsDD = mean(0.25, 0.25) = 0.25 -> score 750
@@ -763,11 +610,9 @@ test_that("width discretization matches hand-computed values", {
 
   expect_equal(result$score, 750)
   expect_equal(nrow(result$cells), 2)
-
 })
 
 test_that("ntile discretization uses confidential quantiles", {
-
   # conf quartile cut points at 25/50/75th percentiles of 1:8
   # 4 bins of 2 values each: conf shares 0.25 apiece
   # synth all in the lowest bin: shares (1, 0, 0, 0)
@@ -784,11 +629,9 @@ test_that("ntile discretization uses confidential quantiles", {
   )
 
   expect_equal(result$score, 625)
-
 })
 
 test_that("cluster discretization separates well-separated groups", {
-
   # two tight clusters around 1 and 10: the midpoint break lands between
   # them, so conf shares (0.5, 0.5) and synth (1, 0) -> score 500
   conf_num <- tibble::tibble(v = c(1, 1.1, 10, 10.1))
@@ -805,11 +648,9 @@ test_that("cluster discretization separates well-separated groups", {
   )
 
   expect_equal(result$score, 500)
-
 })
 
 test_that("synthetic values outside the confidential range land in edge bins", {
-
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(-100, -100, 100, 100))
 
@@ -819,11 +660,9 @@ test_that("synthetic values outside the confidential range land in edge bins", {
 
   # extremes split evenly across the two edge bins, matching conf shares
   expect_equal(result$score, 1000)
-
 })
 
 test_that("non-numeric variables are untouched by discretization", {
-
   conf_mix <- dplyr::mutate(conf, v = c(1, 2, 3, 4))
   synth_mix <- dplyr::mutate(synth, v = c(1, 2, 3, 4))
 
@@ -834,11 +673,9 @@ test_that("non-numeric variables are untouched by discretization", {
   # categorical marginals a and b keep their original levels
   a_cells <- dplyr::filter(result$cells, .data$variables == "a")
   expect_equal(sort(a_cells$cell), c("x", "y"))
-
 })
 
 test_that("bins = NULL leaves numeric variables as-is", {
-
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
 
@@ -848,46 +685,30 @@ test_that("bins = NULL leaves numeric variables as-is", {
 
   # every distinct value is its own cell
   expect_equal(nrow(result$cells), 4)
-
 })
 
-test_that("invalid discretization arguments throw an error", {
-
+test_that("invalid discretization arguments reach the helper's errors", {
+  # bad-bins shapes are covered in test-discretize_k_marginal_vars.R; one
+  # case here confirms the worker wires bins through to that validation
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
 
-  for (bad_bins in list(1, 2.5, "2", NA_real_, c(2, 3))) {
-
-    expect_error(
-      .util_k_marginals(
-        synth_data = synth_num, conf_data = conf_num, k = 1, bins = bad_bins
-      ),
-      regexp = "`bins` must be a single integer >= 2"
-    )
-
-  }
-
   expect_error(
     .util_k_marginals(
-      synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2,
-      discretize_method = "magic"
-    )
+      synth_data = synth_num, conf_data = conf_num, k = 1, bins = 1
+    ),
+    regexp = "`bins` must be a single integer >= 2"
   )
 
-  # non-finite confidential values cannot be discretized
-  conf_inf <- tibble::tibble(v = c(1, 2, 3, Inf))
-
   expect_error(
     .util_k_marginals(
-      synth_data = synth_num, conf_data = conf_inf, k = 1, bins = 2
-    ),
-    regexp = "must be finite to discretize"
+      synth_data = synth_num, conf_data = conf_num, k = 1, bins = 2,
+      discretize_method = "magic"
+    )
   )
-
 })
 
 test_that("util_k_marginals passes discretization arguments through", {
-
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
 
@@ -897,11 +718,9 @@ test_that("util_k_marginals passes discretization arguments through", {
     suppressMessages(util_k_marginals(eval_data = ed, k = 1, bins = 2))$score,
     750
   )
-
 })
 
 test_that("util_k_marginals messages the resolved discretization method", {
-
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
 
@@ -921,49 +740,9 @@ test_that("util_k_marginals messages the resolved discretization method", {
 
   # no discretization, no message
   expect_no_message(util_k_marginals(eval_data = ed, k = 1))
-
-})
-
-test_that("too few distinct confidential values throw an error", {
-
-  conf_const <- tibble::tibble(v = c(2, 2, 2, 2))
-  synth_num <- tibble::tibble(v = c(1, 1, 1, 4))
-
-  for (method in c("width", "ntile", "cluster")) {
-
-    expect_error(
-      .util_k_marginals(
-        synth_data = synth_num, conf_data = conf_const, k = 1, bins = 2,
-        discretize_method = method
-      ),
-      regexp = "fewer distinct confidential values than `bins`"
-    )
-
-  }
-
-})
-
-test_that("tied quantile cut points collapse bins with a warning", {
-
-  # heavily tied data passes the distinct-value pre-check, but the 25th and
-  # 50th percentiles coincide at 1, collapsing a quantile bin
-  conf_ties <- tibble::tibble(v = c(1, 1, 1, 1, 1, 1, 2, 3, 4, 5))
-  synth_num <- tibble::tibble(v = c(1, 2, 3, 4, 5))
-
-  expect_warning(
-    result <- .util_k_marginals(
-      synth_data = synth_num, conf_data = conf_ties, k = 1, bins = 4,
-      discretize_method = "ntile"
-    ),
-    regexp = "bins instead of 4 because of tied cut points"
-  )
-
-  expect_lt(nrow(result$cells), 5)
-
 })
 
 test_that("synth_varnames restricts the worker's variable universe", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q"),
@@ -997,11 +776,9 @@ test_that("synth_varnames restricts the worker's variable universe", {
     result_all$marginals$variables,
     c("a, b", "a, c", "b, c")
   )
-
 })
 
 test_that("synth_vars = TRUE keeps only synthesized variables for postsynth", {
-
   ed <- eval_data(
     conf_data = penguins_conf,
     synth_data = penguins_postsynth
@@ -1011,11 +788,9 @@ test_that("synth_vars = TRUE keeps only synthesized variables for postsynth", {
 
   # species and island are carried over from start_data, not synthesized
   expect_setequal(result$marginals$variables, ed$synth_vars)
-
 })
 
 test_that("synth_vars = FALSE includes carried-over variables", {
-
   ed <- eval_data(
     conf_data = penguins_conf,
     synth_data = penguins_postsynth
@@ -1036,11 +811,9 @@ test_that("synth_vars = FALSE includes carried-over variables", {
   )
 
   expect_false(isTRUE(all.equal(result$score, result_synth_only$score)))
-
 })
 
 test_that("synth_vars = TRUE is a no-op for plain data frame eval_data", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1056,11 +829,9 @@ test_that("synth_vars = TRUE is a no-op for plain data frame eval_data", {
   result <- util_k_marginals(eval_data = ed, k = 1, synth_vars = TRUE)
 
   expect_setequal(result$marginals$variables, c("a", "b"))
-
 })
 
 test_that("priority_vars excluded by synth_varnames error informatively", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q"),
@@ -1083,11 +854,9 @@ test_that("priority_vars excluded by synth_varnames error informatively", {
     ),
     regexp = "`priority_vars` must be a character vector of variables available"
   )
-
 })
 
 test_that("k is validated against the restricted variable universe", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q"),
@@ -1109,11 +878,9 @@ test_that("k is validated against the restricted variable universe", {
     ),
     regexp = "`k` cannot exceed"
   )
-
 })
 
 test_that("invalid synth_vars values error", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1127,18 +894,14 @@ test_that("invalid synth_vars values error", {
   ed <- eval_data(conf_data = conf_sv, synth_data = synth_sv)
 
   for (bad in list("x", c(TRUE, FALSE), NA, 1)) {
-
     expect_error(
       util_k_marginals(eval_data = ed, k = 1, synth_vars = bad),
       regexp = "`synth_vars` must be a single TRUE or FALSE"
     )
-
   }
-
 })
 
 test_that("invalid synth_varnames values error", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1150,7 +913,6 @@ test_that("invalid synth_varnames values error", {
   )
 
   for (bad in list(character(0), NA_character_, c("a", NA), 1)) {
-
     expect_error(
       .util_k_marginals(
         synth_data = synth_sv,
@@ -1160,13 +922,10 @@ test_that("invalid synth_varnames values error", {
       ),
       regexp = "`synth_varnames` must be a non-empty character vector"
     )
-
   }
-
 })
 
 test_that("synth_varnames with no shared variables errors informatively", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1186,11 +945,9 @@ test_that("synth_varnames with no shared variables errors informatively", {
     ),
     regexp = "`synth_varnames` matches no variables available"
   )
-
 })
 
 test_that("empty synthesized-variable metadata errors informatively", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1214,11 +971,9 @@ test_that("empty synthesized-variable metadata errors informatively", {
 
   # synth_vars = FALSE ignores the empty metadata
   expect_no_error(util_k_marginals(eval_data = ed, k = 1, synth_vars = FALSE))
-
 })
 
 test_that("wrapper restriction bounds k by the synthesized-variable set", {
-
   conf_sv <- tibble::tibble(
     a = c("x", "x", "y", "y"),
     b = c("p", "p", "p", "q")
@@ -1242,7 +997,6 @@ test_that("wrapper restriction bounds k by the synthesized-variable set", {
 
   # the same call succeeds once the restriction is lifted
   expect_no_error(util_k_marginals(eval_data = ed, k = 2, synth_vars = FALSE))
-
 })
 
 # NA handling
@@ -1279,7 +1033,6 @@ synth_na <- tibble::tibble(
 )
 
 test_that("NA values form their own level by default", {
-
   result <- suppressMessages(
     .util_k_marginals(synth_data = synth_na, conf_data = conf_na, k = 1)
   )
@@ -1287,11 +1040,9 @@ test_that("NA values form their own level by default", {
   expect_equal(result$score, 19000 / 24)
 
   expect_true("NA" %in% result$cells$cell)
-
 })
 
 test_that("na.rm = TRUE drops missing values per marginal", {
-
   result <- .util_k_marginals(
     synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = TRUE
   )
@@ -1299,11 +1050,9 @@ test_that("na.rm = TRUE drops missing values per marginal", {
   expect_equal(result$score, 17000 / 24)
 
   expect_false("NA" %in% result$cells$cell)
-
 })
 
 test_that("missing data triggers a message when na.rm = FALSE", {
-
   expect_message(
     .util_k_marginals(synth_data = synth_na, conf_data = conf_na, k = 1),
     regexp = "contain missing data: a"
@@ -1314,26 +1063,20 @@ test_that("missing data triggers a message when na.rm = FALSE", {
       synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = TRUE
     )
   )
-
 })
 
 test_that("invalid na.rm values error", {
-
   for (bad in list("x", c(TRUE, FALSE), NA, 1)) {
-
     expect_error(
       .util_k_marginals(
         synth_data = synth_na, conf_data = conf_na, k = 1, na.rm = bad
       ),
       regexp = "`na.rm` must be a single TRUE or FALSE"
     )
-
   }
-
 })
 
 test_that("a literal 'NA' level alongside true NA values errors", {
-
   conf_lit <- tibble::tibble(a = c("NA", "x", NA))
   synth_lit <- tibble::tibble(a = c("x", "x", "x"))
 
@@ -1343,11 +1086,9 @@ test_that("a literal 'NA' level alongside true NA values errors", {
     ),
     regexp = "'NA' already exists"
   )
-
 })
 
 test_that("numeric NA values land in an NA bin or are dropped", {
-
   conf_num <- tibble::tibble(v = c(1, 2, 3, 4))
   synth_num <- tibble::tibble(v = c(1, 4, NA, NA))
 
@@ -1370,11 +1111,9 @@ test_that("numeric NA values land in an NA bin or are dropped", {
   # with the NAs dropped, the two synthetic values split evenly like the
   # confidential data, so the marginal matches exactly
   expect_equal(dropped$score, 1000)
-
 })
 
 test_that("a marginal with no complete rows errors under na.rm = TRUE", {
-
   conf_all_na <- tibble::tibble(a = c(NA_character_, NA_character_))
   synth_ok <- tibble::tibble(a = c("x", "y"))
 
@@ -1384,22 +1123,18 @@ test_that("a marginal with no complete rows errors under na.rm = TRUE", {
     ),
     regexp = "no rows remain"
   )
-
 })
 
 test_that("util_k_marginals passes na.rm through", {
-
   ed <- eval_data(conf_data = conf_na, synth_data = synth_na)
 
   expect_equal(
     util_k_marginals(eval_data = ed, k = 1, na.rm = TRUE)$score,
     17000 / 24
   )
-
 })
 
 test_that("na.rm = TRUE drops rows per combination, not globally", {
-
   # a's NA sits on a different row in each dataset; b and c are identical
   # across datasets
   #
@@ -1435,11 +1170,9 @@ test_that("na.rm = TRUE drops rows per combination, not globally", {
   bc <- dplyr::filter(result$marginals, .data$variables == "b, c")
 
   expect_equal(bc$madd, 0)
-
 })
 
 test_that("the missing-data message lists every affected variable", {
-
   conf_two <- tibble::tibble(
     a = c("x", NA),
     b = c(NA, "q"),
@@ -1456,11 +1189,9 @@ test_that("the missing-data message lists every affected variable", {
     .util_k_marginals(synth_data = synth_two, conf_data = conf_two, k = 1),
     regexp = "contain missing data: a, b"
   )
-
 })
 
 test_that("a literal 'NA' level in the synthetic data also errors", {
-
   conf_lit <- tibble::tibble(a = c("x", "x", "x"))
   synth_lit <- tibble::tibble(a = c("NA", "x", NA))
 
@@ -1478,11 +1209,9 @@ test_that("a literal 'NA' level in the synthetic data also errors", {
     ),
     regexp = "'NA' already exists"
   )
-
 })
 
 test_that("variables excluded by synth_varnames do not drive NA handling", {
-
   # a has missing values but is filtered out, so no message and no NA cells
   conf_excl <- tibble::tibble(
     a = c("x", NA),
@@ -1504,11 +1233,9 @@ test_that("variables excluded by synth_varnames do not drive NA handling", {
   )
 
   expect_false("NA" %in% result$cells$cell)
-
 })
 
 test_that("weighted proportions drop missing rows before computing shares", {
-
   # na.rm = TRUE drops each dataset's NA row, and weight shares are computed
   # from the surviving rows' weights:
   #   conf keeps weights 1, 1, 2 -> x = 2/4, y = 2/4
@@ -1535,11 +1262,9 @@ test_that("weighted proportions drop missing rows before computing shares", {
   )
 
   expect_equal(result$score, 750)
-
 })
 
 test_that("confidential numeric NA values discretize under both na.rm modes", {
-
   # breaks derive from the observed confidential values (1:4, cut at 2.5)
   conf_num_na <- tibble::tibble(v = c(1, 2, 3, 4, NA))
   synth_num_na <- tibble::tibble(v = c(1, 2, 4, 4, NA))
@@ -1562,21 +1287,6 @@ test_that("confidential numeric NA values discretize under both na.rm modes", {
 
   expect_false("NA" %in% dropped$cells$cell)
   expect_equal(dropped$score, 1000)
-
-})
-
-test_that("infinite confidential values still refuse to discretize", {
-
-  conf_inf <- tibble::tibble(v = c(1, 2, 3, Inf))
-  synth_num <- tibble::tibble(v = c(1, 2, 3, 3))
-
-  expect_error(
-    .util_k_marginals(
-      synth_data = synth_num, conf_data = conf_inf, k = 1, bins = 2
-    ),
-    regexp = "must be finite to discretize"
-  )
-
 })
 
 # group_by stratification
@@ -1605,17 +1315,14 @@ synth_g <- tibble::tibble(
 )
 
 test_that("group_by stratifies the score by confidential shares", {
-
   result <- .util_k_marginals(
     synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
   )
 
   expect_equal(result$score, 750)
-
 })
 
 test_that("grouped output gains group columns and group_scores", {
-
   result <- .util_k_marginals(
     synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
   )
@@ -1637,11 +1344,9 @@ test_that("grouped output gains group columns and group_scores", {
   ungrouped <- .util_k_marginals(synth_data = synth_g, conf_data = conf_g, k = 1)
 
   expect_null(ungrouped$group_scores)
-
 })
 
 test_that("an empty synthetic stratum scores against zero proportions", {
-
   # synth has no B rows: stratum B conf cells (x = 0.5, y = 0.5) face
   # synthetic proportions of 0 -> MabsDD = 0.5 -> score 500
   # stratum A: conf (x = 0.5, y = 0.5), synth (x = 0.5, y = 0.5) -> 1000
@@ -1656,11 +1361,9 @@ test_that("an empty synthetic stratum scores against zero proportions", {
   )
 
   expect_equal(result$score, 750)
-
 })
 
 test_that("group shares use weights when weight_var is set", {
-
   # conf weight shares: A = 2/4, B = 2/4 (row shares would be 2/3, 1/3)
   # stratum A: conf (x = 0.5, y = 0.5), synth (x = 1) -> score 500
   # stratum B: conf (x = 1), synth (x = 1) -> score 1000
@@ -1686,11 +1389,9 @@ test_that("group shares use weights when weight_var is set", {
   )
 
   expect_equal(result$score, 750)
-
 })
 
 test_that("invalid group_by values error", {
-
   expect_error(
     .util_k_marginals(
       synth_data = synth_g, conf_data = conf_g, k = 1, group_by = 1
@@ -1715,11 +1416,9 @@ test_that("invalid group_by values error", {
     ),
     regexp = "`group_by` cannot include `weight_var`"
   )
-
 })
 
 test_that("group variables are excluded from the marginal universe", {
-
   result <- .util_k_marginals(
     synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
   )
@@ -1733,32 +1432,26 @@ test_that("group variables are excluded from the marginal universe", {
     ),
     regexp = "`k` cannot exceed"
   )
-
 })
 
 test_that("grouped print shows group scores", {
-
   result <- .util_k_marginals(
     synth_data = synth_g, conf_data = conf_g, k = 1, group_by = "g"
   )
 
   expect_output(print(result), regexp = "Worst groups:")
-
 })
 
 test_that("util_k_marginals passes group_by through", {
-
   ed <- eval_data(conf_data = conf_g, synth_data = synth_g)
 
   expect_equal(
     util_k_marginals(eval_data = ed, k = 1, group_by = "g")$score,
     750
   )
-
 })
 
 test_that("missing group values follow na.rm", {
-
   conf_gna <- tibble::tibble(
     g = c("A", "A", NA, NA),
     a = c("x", "y", "x", "y")
@@ -1788,11 +1481,9 @@ test_that("missing group values follow na.rm", {
 
   expect_equal(dropped$score, 500)
   expect_equal(nrow(dropped$group_scores), 1)
-
 })
 
 test_that("multi-column group_by stratifies by joint combinations", {
-
   # four joint strata of two rows each; only stratum (A, p) diverges:
   #   conf (x = 0.5, y = 0.5), synth (x = 1) -> MabsDD = 0.5 -> score 500
   # the other three strata are identical -> 1000
@@ -1828,11 +1519,9 @@ test_that("multi-column group_by stratifies by joint combinations", {
 
   # neither grouping variable enters the marginal universe
   expect_false(any(result$marginals$variables %in% c("g1", "g2")))
-
 })
 
 test_that("partially missing joint strata follow na.rm", {
-
   # g2 is missing on rows 3-4; only the (A, NA) stratum diverges:
   #   conf (x = 0.5, y = 0.5), synth (x = 1) -> score 500
   conf_gpart <- tibble::tibble(
@@ -1863,11 +1552,9 @@ test_that("partially missing joint strata follow na.rm", {
 
   expect_equal(dropped$score, 1000)
   expect_equal(nrow(dropped$group_scores), 1)
-
 })
 
 test_that("empty, missing, and duplicate group_by values error", {
-
   expect_error(
     .util_k_marginals(
       synth_data = synth_g, conf_data = conf_g, k = 1,
@@ -1891,11 +1578,9 @@ test_that("empty, missing, and duplicate group_by values error", {
     ),
     regexp = "must not contain duplicate"
   )
-
 })
 
 test_that("group_by composes with synth_varnames", {
-
   # b is shared but unsynthesized; g stratifies; the universe is {a} only
   conf_gsv <- tibble::tibble(
     g = c("A", "A", "B", "B"),
@@ -1921,11 +1606,9 @@ test_that("group_by composes with synth_varnames", {
   )
 
   expect_equal(unique(result_gname$marginals$variables), "a")
-
 })
 
 test_that("literal 'NA' strings in grouping columns collide with true NA", {
-
   # single grouping column carrying both a literal "NA" level and a true NA
   conf_collide <- tibble::tibble(
     g = c("NA", "A", NA),
@@ -1994,11 +1677,9 @@ test_that("literal 'NA' strings in grouping columns collide with true NA", {
   )
 
   expect_equal(result$score, 750)
-
 })
 
 test_that("grouped results map over replicates with the same structure", {
-
   ed <- eval_data(
     conf_data = conf_g,
     synth_data = list(synth_g, conf_g)
@@ -2011,21 +1692,17 @@ test_that("grouped results map over replicates with the same structure", {
   expect_length(result, 2)
 
   for (rep in result) {
-
     expect_s3_class(rep, "k_marginals")
     expect_named(rep$group_scores, c("g", "share", "score"))
     expect_true("g" %in% names(rep$marginals))
     expect_true("g" %in% names(rep$cells))
-
   }
 
   # first replicate diverges in stratum A, second is identical data
   expect_equal(purrr::map_dbl(result, "score"), c(750, 1000))
-
 })
 
 test_that("a confidential stratum emptied by na.rm errors instead of NaN", {
-
   # stratum B's confidential rows are all missing on a, so per-marginal NA
   # removal leaves nothing to score against
   conf_gna2 <- tibble::tibble(
@@ -2057,11 +1734,9 @@ test_that("a confidential stratum emptied by na.rm errors instead of NaN", {
 
   # stratum B scores conf (x = 0.5, y = 0.5) against zero synth -> 500
   expect_equal(result$score, 750)
-
 })
 
 test_that("priority_vars with NA entries hits the intended error", {
-
   # %in% never propagates NA, so the membership test is FALSE, not NA, and
   # the package error fires rather than a base R condition failure
   expect_error(
@@ -2079,11 +1754,9 @@ test_that("priority_vars with NA entries hits the intended error", {
     ),
     regexp = "`priority_vars` must be a character vector of variables available"
   )
-
 })
 
 test_that("priority_vars = character(0) behaves like NULL", {
-
   # an empty priority set passes validation and simply guarantees nothing
   result <- .util_k_marginals(
     synth_data = synth, conf_data = conf, k = 1,
@@ -2094,5 +1767,4 @@ test_that("priority_vars = character(0) behaves like NULL", {
     result$score,
     .util_k_marginals(synth_data = synth, conf_data = conf, k = 1)$score
   )
-
 })

From fbb57cbc2b4b25747ca505e6f1642c6466581a08 Mon Sep 17 00:00:00 2001
From: Thiyaghessan <thiyaghessan@gmail.com>
Date: Fri, 14 Aug 2026 16:41:59 -0400
Subject: [PATCH 10/10] Fix Rd cross-reference warning from bracketed score
 range

---
 R/util_k_marginals.R        | 4 ++--
 man/dot-util_k_marginals.Rd | 2 +-
 man/util_k_marginals.Rd     | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/R/util_k_marginals.R b/R/util_k_marginals.R
index a6f1a90..bb28a46 100644
--- a/R/util_k_marginals.R
+++ b/R/util_k_marginals.R
@@ -60,7 +60,7 @@
 #' reproducible clusters). Defaults to "width".
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
-#' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
+#' range 0 to 1000 where a higher value denotes lower MabsDDs and consequently
 #' greater similarity between confidential and synthetic data; `marginals`, a
 #' tibble with the MabsDD for each combination of variables, worst first; and
 #' `cells`, a tibble with the synthetic and confidential proportions and their
@@ -402,7 +402,7 @@ print.k_marginals <- function(x, n = 5, ...) {
 #' reproducible clusters). Defaults to "width".
 #'
 #' @return A `k_marginals` object with three elements: `score`, a value in
-#' range [0, 1000] where a higher value denotes lower MabsDDs and consequently
+#' range 0 to 1000 where a higher value denotes lower MabsDDs and consequently
 #' greater similarity between confidential and synthetic data; `marginals`, a
 #' tibble with the MabsDD for each combination of variables, worst first; and
 #' `cells`, a tibble with the synthetic and confidential proportions and their
diff --git a/man/dot-util_k_marginals.Rd b/man/dot-util_k_marginals.Rd
index 651c67b..7185a48 100644
--- a/man/dot-util_k_marginals.Rd
+++ b/man/dot-util_k_marginals.Rd
@@ -85,7 +85,7 @@ reproducible clusters). Defaults to "width".}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
-range \link{0, 1000} where a higher value denotes lower MabsDDs and consequently
+range 0 to 1000 where a higher value denotes lower MabsDDs and consequently
 greater similarity between confidential and synthetic data; \code{marginals}, a
 tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their
diff --git a/man/util_k_marginals.Rd b/man/util_k_marginals.Rd
index 9e1f0d1..86ccc00 100644
--- a/man/util_k_marginals.Rd
+++ b/man/util_k_marginals.Rd
@@ -82,7 +82,7 @@ reproducible clusters). Defaults to "width".}
 }
 \value{
 A \code{k_marginals} object with three elements: \code{score}, a value in
-range \link{0, 1000} where a higher value denotes lower MabsDDs and consequently
+range 0 to 1000 where a higher value denotes lower MabsDDs and consequently
 greater similarity between confidential and synthetic data; \code{marginals}, a
 tibble with the MabsDD for each combination of variables, worst first; and
 \code{cells}, a tibble with the synthetic and confidential proportions and their