R/ei_summary_random.R

In eiballots: Ballot-Level Microdata and Summaries for Ecological Inference (Florida 2000)

# R/ei_summary_random.R
# Random aggregation of voters into artificial units for ecological inference.

# ----------------------------------------------------------------------------
# ei_summary_random()
# ----------------------------------------------------------------------------

#' Compute ecological inference summaries with random voter aggregation
#'
#' Randomly assigns voters to artificial units and computes ecological
#' inference summaries, ignoring the geographic boundaries of real
#' precincts. This allows assessing how sensitive the results are to
#' the specific way voters are grouped — a concern known in the
#' literature as the Modifiable Areal Unit Problem (MAUP).
#'
#' The total number of votes for each option is identical to that
#' returned by [ei_summary()] on the same data — only the distribution
#' of votes across units changes.
#'
#' The function returns a standard [ei_summary] object, fully compatible
#' with [to_lphom()], [to_eipack()], and [subset_elections()]. The
#' randomisation parameters are stored in `meta$random` for
#' reproducibility and auditing.
#'
#' To repeat the randomisation `n` times (e.g. for a sensitivity
#' analysis), call the function inside `lapply()`, varying `seed` each
#' time:
#' ```r
#' results <- lapply(1:100, function(i)
#'   ei_summary_random(c("PRE", "USS"), data = d, seed = i))
#' ```
#'
#' @param elections Character vector of race codes (e.g. `c("PRE", "USS")`).
#'   All codes must be keys in [election_catalog].See [ei_summary()].
#' @param data Optional `data.frame`. If provided, must contain columns
#'   `COUNTY`, `PRECINCT`, and all codes in `elections`. If `NULL`
#'   (default), data are loaded automatically via [get_election_data()].
#' @param source Character string or `NULL`. Data source override; see
#'   [get_county_data()]. Ignored when `data` is provided.
#' @param format Character. Output format for the `margins` element:
#'   \itemize{
#'     \item `"lphom"` *(default)*: a named list of `data.frame`s, one per
#'       race. Compatible with `lphom` and related packages.
#'     \item `"eipack"`: a single wide `data.frame` with all races side by
#'       side (columns named `<RACE>_<OPTION>`, e.g. `PRE_R`, `PRE_D`).
#'   }
#' @param n_units Integer or `NULL`. Number of random units to create.
#'   `NULL` (default) uses the same number of precincts as in the data.
#'   When `unit_size` or `size_range` is also supplied, see the
#'   precedence rules in the Details section.
#' @param unit_size Integer or `NULL`. Fixed number of voters per unit.
#'   If `N` (the number of eligible voters) is not an exact multiple of
#'   `unit_size`, a smaller residual unit containing the remaining voters
#'   is created. Takes precedence over `size_range` and `n_units`; both
#'   are ignored with a warning if `unit_size` is specified.
#' @param size_range Integer vector of length 2, `c(min, max)`, or
#'   `NULL`. When used alone (without `n_units`), units are created by
#'   drawing sizes uniformly from `[min, max]` until all `N` voters are
#'   assigned; the last unit may be smaller than `min`. When used with
#'   `n_units`, the effective interval is the largest symmetric range
#'   around the required mean (`N / n_units`) that fits within
#'   `[min, max]`; the last unit absorbs the residual. Takes precedence
#'   over `n_units`. Ignored when `unit_size` is specified.
#' @param seed Integer or `NULL`. Random seed passed to [set.seed()]
#'   before shuffling voters, enabling reproducible results. `NULL`
#'   (default) uses the current random state.
#'
#' @return An object of class `"ei_summary"` with the same structure as
#'   [ei_summary()], which contains the components: `margins`, `joint_precinct`,
#'   `joint_total` and `meta`; see [ei_summary()] for details.
#'   To document the randomisation the `meta` component includes,
#'   in addition to the fields `elections`, `counties`, `n_precincts`,
#'   `n_voters`, `opt_levels`, `format`, `precinct_info`, `created_at`
#'   the field `random`, with the following components:
#'   \describe{
#'     \item{`method`}{Character string describing the aggregation method
#'       used and the resulting unit structure.}
#'     \item{`n_units`}{Integer. Number of random units created.}
#'     \item{`unit_sizes`}{Integer vector of length `n_units`. Size
#'       (number of voters) of each unit. Always sums to
#'       `meta$n_voters`.}
#'     \item{`seed`}{The seed supplied to `seed`, or `NULL` if none.}
#'     \item{`n_units_arg`}{Value of the `n_units` argument as supplied
#'       by the user.}
#'     \item{`unit_size_arg`}{Value of the `unit_size` argument as
#'       supplied by the user.}
#'     \item{`size_range_arg`}{Value of the `size_range` argument as
#'       supplied by the user.}
#'   }
#'
#' @section Argument precedence:
#' When multiple sizing arguments are supplied simultaneously, they are
#' resolved in this order (highest to lowest priority):
#' \enumerate{
#'   \item **`unit_size`** — fixed size per unit; `n_units` and
#'     `size_range` are ignored with a warning.
#'   \item **`size_range`** — random sizes within bounds; if `n_units`
#'     is also supplied, the effective sampling interval is the largest
#'     symmetric interval around the required mean (`N / n_units`) that
#'     fits within `[min, max]`. The last unit absorbs the residual and
#'     may fall slightly outside the effective interval. A warning is
#'     issued if the required mean falls outside `[min, max]`, in which
#'     case all units are assigned approximately equal sizes.
#'   \item **`n_units`** — number of units; sizes are drawn
#'     proportionally from the real precinct size distribution.
#'   \item **Default** (no sizing argument) — same number of units as
#'     real precincts; sizes match the real precinct size distribution;
#'     voters are randomly shuffled across units.
#' }
#'
#' @section Universe:
#' As in [ei_summary()], only voters non-`NA` in **all** `elections`
#' simultaneously are included (the intersection universe). The
#' randomisation is applied to this universe, not to the full dataset.
#' The `meta$n_voters` field of the returned object always refers to
#' this intersection count.
#'
#' @section Relationship to ei_summary():
#' `ei_summary_random()` produces an object of the same class and
#' structure as [ei_summary()], and is fully compatible with
#' [to_lphom()], [to_eipack()], and [subset_elections()]. The only
#' difference is that voters are grouped into random units rather than
#' real (geographic) precincts. To compare random and real aggregations
#' directly, run both functions on the same `data` object:
#' ```r
#' d    <- get_election_data(c("PRE", "STS35"))
#' real <- ei_summary(c("PRE", "STS35"), data = d)
#' rnd  <- ei_summary_random(c("PRE", "STS35"), data = d, seed = 1)
#' ```
#'
#' @examples
#' # Using the built-in example dataset
#' obj <- ei_summary(c("PRE", "USS"), data = example_ballots)
#' print(obj)
#' summary(obj)
#'
#' \donttest{
#' # Default: same number and size distribution as real precincts
#' obj <- ei_summary_random(c("PRE", "STS35"), seed = 42)
#'
#' # 200 units with sizes proportional to the real precinct distribution
#' obj <- ei_summary_random(c("PRE", "STS35"), n_units = 200, seed = 42)
#'
#' # Fixed size of 500 voters per unit (residual unit if N not divisible)
#' obj <- ei_summary_random(c("PRE", "STS35"), unit_size = 500, seed = 42)
#'
#' # Random size between 200 and 800 voters per unit
#' obj <- ei_summary_random(c("PRE", "STS35"),
#'                          size_range = c(200, 800), seed = 42)
#'
#' # 200 units with random sizes between 100 and 600
#' obj <- ei_summary_random(c("PRE", "STS35"),
#'                          n_units = 200, size_range = c(100, 600),
#'                          seed = 42)
#'
#' # Sensitivity analysis: 100 random aggregations
#' results <- lapply(1:100, function(i)
#'   ei_summary_random(c("PRE", "USS"), seed = i))
#' }
#'
#' @seealso [ei_summary()], [subset_elections()], [to_lphom()],
#'   [to_eipack()], [get_election_data()]
#' @export
ei_summary_random <- function(elections,
                               data       = NULL,
                               source     = NULL,
                               format     = c("lphom", "eipack"),
                               n_units    = NULL,
                               unit_size  = NULL,
                               size_range = NULL,
                               seed       = NULL) {

  format <- match.arg(format)
  .validate_elections(elections)

  # --- Load data if not supplied ---
  if (is.null(data)) {
    data <- get_election_data(elections, source = source)
  } else {
    missing_cols <- setdiff(c("COUNTY", "PRECINCT", elections), names(data))
    if (length(missing_cols) > 0L)
      stop(
        sprintf("Missing columns in 'data': %s",
                paste(missing_cols, collapse = ", ")),
        call. = FALSE
      )
  }

  # --- Intersection universe (same rule as ei_summary) ---
  data$precinct_id <- paste(data$COUNTY, data$PRECINCT, sep = "_")
  is_complete      <- complete.cases(data[, elections, drop = FALSE])
  universe         <- data[is_complete, , drop = FALSE]

  if (nrow(universe) == 0L)
    stop(
      paste0(
        "No voters are eligible in all selected races simultaneously.\n",
        "Check that the chosen races have overlapping universes."
      ),
      call. = FALSE
    )

  N <- nrow(universe)

  # --- Resolve unit sizes (argument precedence logic) ---
  sizes_info <- .resolve_unit_sizes(
    N          = N,
    n_units    = n_units,
    unit_size  = unit_size,
    size_range = size_range,
    real_sizes = as.integer(table(universe$precinct_id))
  )

  # --- Shuffle voters and assign to random units ---
  if (!is.null(seed)) set.seed(seed)
  shuffled_idx <- sample.int(N)

  unit_ids <- .assign_units(shuffled_idx, sizes_info$sizes)

  # Replace PRECINCT with random unit label.
  # IMPORTANT: use only the unit number as precinct_id, NOT paste(COUNTY, unit)
  # — if multiple counties are present, the same unit would otherwise get
  # different precinct_ids per county and be split into separate units in
  # the joint array.
  universe$COUNTY      <- "rnd"
  universe$PRECINCT    <- sprintf("unit_%0*d", nchar(max(unit_ids)), unit_ids)
  universe$precinct_id <- universe$PRECINCT

  # --- Build ei_summary from the randomly aggregated data ---
  # Re-use the internal array builder from ei_summary()
  built         <- .build_joint_array(universe, elections)
  precinct_info <- built$precinct_info
  opt_levels    <- built$opt_levels

  if (length(elections) == 1L) {
    margins_list   <- list(.array_to_margin_df(t(built$joint_precinct),
                                               precinct_info))
    names(margins_list) <- elections
    joint_precinct <- NULL
    joint_total    <- NULL

  } else {
    joint_precinct <- built$joint_precinct
    margins_list   <- .derive_margins(joint_precinct, elections, precinct_info)
    names(margins_list) <- elections

    other_dims  <- seq_along(elections)
    joint_total <- apply(joint_precinct, other_dims, sum)
  }

  margins <- if (format == "eipack") {
    .to_eipack_wide(margins_list)
  } else {
    margins_list
  }

  structure(
    list(
      margins        = margins,
      joint_precinct = joint_precinct,
      joint_total    = joint_total,
      meta           = list(
        elections     = elections,
        counties      = unique(precinct_info$COUNTY),
        n_precincts   = nrow(precinct_info),
        n_voters      = built$n_universe,
        opt_levels    = opt_levels,
        format        = format,
        precinct_info = precinct_info,
        created_at    = Sys.time(),
        random        = list(
          method        = sizes_info$method,
          n_units       = length(sizes_info$sizes),
          unit_sizes    = sizes_info$sizes,
          seed          = seed,
          n_units_arg   = n_units,
          unit_size_arg = unit_size,
          size_range_arg = size_range
        )
      )
    ),
    class = c("ei_summary", "list")
  )
}


# ----------------------------------------------------------------------------
# Internal: resolve unit sizes
# ----------------------------------------------------------------------------

# Resolve the vector of unit sizes from the user's arguments.
#
# Returns a list with:
#   - sizes:  integer vector of unit sizes (sums to N)
#   - method: character description for meta$random$method

.resolve_unit_sizes <- function(N, n_units, unit_size, size_range,
                                real_sizes) {

  # ---- Priority 1: unit_size (fixed size per unit) -------------------------
  if (!is.null(unit_size)) {

    if (!is.null(size_range))
      warning(
        "'size_range' ignored because 'unit_size' was specified.",
        call. = FALSE
      )
    if (!is.null(n_units))
      warning(
        "'n_units' ignored because 'unit_size' was specified.",
        call. = FALSE
      )

    unit_size <- as.integer(unit_size)
    if (unit_size < 1L || unit_size > N)
      stop(
        sprintf("'unit_size' must be between 1 and N (%d).", N),
        call. = FALSE
      )

    k        <- N %/% unit_size
    residual <- N %% unit_size
    sizes    <- rep(unit_size, k)
    if (residual > 0L) {
      sizes  <- c(sizes, residual)
      method <- sprintf(
        "Fixed unit size of %d voters; %d full units + 1 residual unit of %d voters.",
        unit_size, k, residual
      )
    } else {
      method <- sprintf(
        "Fixed unit size of %d voters; %d units.",
        unit_size, k
      )
    }
    return(list(sizes = sizes, method = method))
  }

  # ---- Priority 2: size_range (random sizes within bounds) -----------------
  if (!is.null(size_range)) {

    if (length(size_range) != 2L || !is.numeric(size_range) ||
        size_range[1L] > size_range[2L] || size_range[1L] < 1L)
      stop(
        "'size_range' must be a numeric vector c(min, max) with 0 < min <= max.",
        call. = FALSE
      )

    size_min <- as.integer(size_range[1L])
    size_max <- as.integer(size_range[2L])

    if (!is.null(n_units)) {
      k         <- as.integer(n_units)
      mean_size <- N / k

      if (mean_size < size_min || mean_size > size_max) {
        # Mean falls outside the requested range: warn and use fixed size
        warning(sprintf(
          paste0("The required mean unit size (N/n_units = %.1f) falls ",
                 "outside size_range = c(%d, %d). ",
                 "All units will have approximately equal size."),
          mean_size, size_min, size_max
        ), call. = FALSE)
        sizes  <- .scale_sizes_to_N(rep(1L, k), N)
        method <- sprintf(
          paste0("n_units = %d with size_range = c(%d, %d): mean size %.1f ",
                 "outside range; units assigned approximately equal sizes ",
                 "(N = %d voters)."),
          k, size_min, size_max, mean_size, N
        )
      } else {
        # Compute the largest symmetric interval around mean_size within
        # [size_min, size_max]: half-width = min(mean - min, max - mean)
        half     <- floor(min(mean_size - size_min, size_max - mean_size))
        eff_min  <- as.integer(mean_size - half)
        eff_max  <- as.integer(mean_size + half)

        # Draw k-1 sizes uniformly from [eff_min, eff_max]; last = residual
        if (eff_min == eff_max) {
          # Degenerate interval: all units the same size
          sizes <- .scale_sizes_to_N(rep(1L, k), N)
        } else {
          sizes       <- integer(k)
          sizes[1:(k - 1L)] <- sample(eff_min:eff_max, k - 1L, replace = TRUE)
          sizes[k]    <- N - sum(sizes[1:(k - 1L)])
        }

        method <- sprintf(
          paste0("n_units = %d, size_range = c(%d, %d): effective symmetric ",
                 "interval [%d, %d] (centered on mean = %.1f); ",
                 "last unit absorbs residual (N = %d voters)."),
          k, size_min, size_max, eff_min, eff_max, mean_size, N
        )
      }

    } else {
      # No n_units: draw sizes until all N voters are assigned
      sizes <- integer(0L)
      while (sum(sizes) < N) {
        s         <- sample(size_min:size_max, 1L)
        remaining <- N - sum(sizes)
        sizes     <- c(sizes, min(s, remaining))
      }
      method <- sprintf(
        paste0("Random unit sizes drawn uniformly from [%d, %d] until ",
               "N = %d voters assigned; %d units created."),
        size_min, size_max, N, length(sizes)
      )
    }
    return(list(sizes = sizes, method = method))
  }

  # ---- Priority 3: n_units (number of units, sizes from real distribution) -
  if (!is.null(n_units)) {

    k <- as.integer(n_units)
    if (k < 1L || k > N)
      stop(
        sprintf("'n_units' must be between 1 and N (%d).", N),
        call. = FALSE
      )

    raw_sizes <- .resample_sizes(real_sizes, k, N)
    method <- sprintf(
      paste0("Number of units fixed at %d; sizes drawn proportionally ",
             "from the real precinct size distribution (N = %d voters)."),
      k, N
    )
    return(list(sizes = raw_sizes, method = method))
  }

  # ---- Default: same number of units as real precincts, same sizes ---------
  sizes  <- sort(real_sizes, decreasing = TRUE)
  method <- sprintf(
    paste0("Default: %d units with the same size distribution as the ",
           "real precincts; voters randomly shuffled (N = %d)."),
    length(sizes), N
  )
  list(sizes = sizes, method = method)
}

# ----------------------------------------------------------------------------
# Internal helpers
# ----------------------------------------------------------------------------

# Assign shuffled voter indices to units of given sizes.
# Returns an integer vector of unit labels (1, 2, ..., k), one per voter,
# in the order of `shuffled_idx`.
.assign_units <- function(shuffled_idx, sizes) {
  unit_labels <- rep(seq_along(sizes), times = sizes)
  # unit_labels[i] = unit assigned to the i-th voter in shuffled order
  order_back  <- order(shuffled_idx)   # undo shuffle to align with original rows
  unit_labels[order_back]
}


# Scale a vector of raw sizes so they sum exactly to N,
# preserving relative proportions as closely as possible.
.scale_sizes_to_N <- function(raw, N) {
  scaled   <- round(raw * N / sum(raw))
  diff     <- N - sum(scaled)
  # Distribute rounding residual to the largest units
  if (diff != 0L) {
    idx     <- order(scaled, decreasing = diff > 0L)[seq_len(abs(diff))]
    scaled[idx] <- scaled[idx] + sign(diff)
  }
  pmax(scaled, 1L)   # ensure no unit has 0 voters
}


# Draw k sizes that sum to N, using the real precinct sizes as a template.
# If k == length(real_sizes), returns the real sizes shuffled.
# If k < length(real_sizes), merges adjacent sizes.
# If k > length(real_sizes), splits sizes.
.resample_sizes <- function(real_sizes, k, N) {
  if (k == length(real_sizes)) {
    return(sample(real_sizes))   # shuffle order only
  }
  # General case: draw k sizes with replacement from real_sizes,
  # then scale to sum exactly N.
  raw <- sample(real_sizes, k, replace = TRUE)
  .scale_sizes_to_N(raw, N)
}