R/ei_summary.R

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

# eiballots/R/ei_summary.R
# Main function: ei_summary()

# ----------------------------------------------------------------------------
# ei_summary()
# ----------------------------------------------------------------------------

#' Compute ecological inference summaries from ballot microdata
#'
#' Builds the objects needed for ecological inference over a set of races:
#' marginal distributions per precinct for each race, a joint contingency
#' array across all races and precincts, and its precinct-aggregated versions.
#'
#' All summaries share a common **intersection universe**: voters who were
#' eligible to vote in **every** race in `elections` (i.e. non-`NA` in all
#' selected columns). To obtain the full universe of a single race, pass
#' only that race.
#'
#' @param elections Character vector of race codes (e.g. `c("PRE", "USS")`).
#'   All codes must be keys in [election_catalog].
#' @param data Optional `data.frame`. If provided, it 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`).
#'   }
#'
#' @return An object of class `"ei_summary"` (a named list) with components:
#' \describe{
#'   \item{`margins`}{Marginal distributions per precinct. A named list of
#'     `data.frame`s (format `"lphom"`) or a single wide `data.frame`
#'     (format `"eipack"`). Each `data.frame` contains columns
#'     `precinct_id`, `county_id`, `n_voters`, and one column per vote
#'     option. `n_voters` is identical across all races because the
#'     universe is shared.}
#'   \item{`joint_precinct`}{Integer array of dimensions
#'     `[opt_1 x ... x opt_k x precincts]`. `NULL` when only one race is
#'     requested. Summing over any race dimension reproduces the
#'     corresponding margin.}
#'   \item{`joint_total`}{Same array collapsed over precincts:
#'     `[opt_1 x ... x opt_k]`. `NULL` when only one race is requested.}
#'   \item{`meta`}{Named list of metadata: `elections`, `counties`,
#'     `n_precincts`, `n_voters`, `opt_levels`, `format`,
#'     `precinct_info`, `created_at`.}
#' }
#'
#' @details
#' ### Vote codes
#' In the microdata, `"A"` denotes a blank vote and `"I"` an invalid vote;
#' both are included as active categories. `NA` means the voter was not
#' eligible in that race and defines the universe boundary.
#'
#' ### Multi-member races
#' For races where voters could select more than one candidate (e.g.
#' `"HOS3"`), each observed combination (e.g. `"BauDuf"`) is treated as an
#' atomic category. See [election_catalog] for the `candidates` field
#' listing the individual candidates on the ballot.
#'
#' ### Subsetting
#' Use the `[` operator to extract a subset of races from an existing
#' `ei_summary` object without reloading data:
#' ```r
#' obj3 <- ei_summary(c("PRE", "USS", "HOS3"))
#' obj2 <- obj3[c("PRE", "USS")]   # subset to 2 races
#' ```
#'
#' @examples
#' # Using the built-in example dataset
#' obj <- ei_summary(c("PRE", "USS"), data = example_ballots)
#' print(obj)
#' summary(obj)
#'
#' \donttest{
#' # Single race: margins only, no joint array
#' pres <- ei_summary("PRE")
#' summary(pres)
#'
#' # Two races: margins + 3-D joint array [PRE x USS x precinct]
#' obj <- ei_summary(c("PRE", "USS"))
#' print(obj)
#' summary(obj)
#' dim(obj$joint_precinct)
#'
#' # Three races
#' obj3 <- ei_summary(c("PRE", "USS", "HOS3"))
#' dim(obj3$joint_precinct)   # [opts_PRE x opts_USS x opts_HOS3 x n_precincts]
#'
#' # Wide format for eipack-style packages
#' obj_wide <- ei_summary(c("PRE", "USS"), format = "eipack")
#' }
#'
#' @seealso [ei_summary_random()]] [get_election_data()], [get_county_data()],
#'   [to_lphom()], [to_eipack()], [list_elections()]
#' @export
ei_summary <- function(elections,
                       data   = NULL,
                       source = NULL,
                       format = c("lphom", "eipack")) {

  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
      )
  }

  # --- Build joint array (intersection universe) ---
  built         <- .build_joint_array(data, elections)
  precinct_info <- built$precinct_info
  opt_levels    <- built$opt_levels

  # --- Single race: no joint objects ---
  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

  # --- Multiple races: derive margins from joint array ---
  } else {
    joint_precinct <- built$joint_precinct
    margins_list   <- .derive_margins(joint_precinct, elections, precinct_info)
    names(margins_list) <- elections

    # Collapse over precincts (now the LAST dimension)
    other_dims  <- seq_along(elections)
    joint_total <- apply(joint_precinct, other_dims, sum)
  }

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

  # --- Assemble output ---
  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,   # retained for subsetting / to_lphom()
        created_at    = Sys.time()
      )
    ),
    class = c("ei_summary", "list")
  )
}


# ----------------------------------------------------------------------------
# Format converters
# ----------------------------------------------------------------------------

#' Convert an ei_summary to lphom list format
#'
#' Converts the `margins` element to a named list of `data.frame`s (one per
#' race). This is the default format of [ei_summary()] and the input format
#' expected by `lphom` and related ecological inference functions.
#'
#' @param x An `"ei_summary"` object.
#' @return An `"ei_summary"` object with `margins` as a named list of
#'   `data.frame`s and `meta$format == "lphom"`.
#'
#' @seealso [to_eipack()], [ei_summary()]
#' @export
to_lphom <- function(x) {
  stopifnot(inherits(x, "ei_summary"))
  if (x$meta$format == "lphom") return(x)

  elections <- x$meta$elections

  # Re-derive margins from joint_precinct when available
  margins_list <- if (is.null(x$joint_precinct)) {
    # Single race: reconstruct from the wide data.frame
    lst <- list(.wide_to_single_df(x$margins, elections))
    names(lst) <- elections
    lst
  } else {
    m <- .derive_margins(x$joint_precinct, elections, x$meta$precinct_info)
    names(m) <- elections
    m
  }

  x$margins     <- margins_list
  x$meta$format <- "lphom"
  x
}


#' Convert an ei_summary to eipack wide format
#'
#' Converts the `margins` element to a single wide `data.frame` where
#' columns are named `<RACE>_<OPTION>` (e.g. `PRE_R`, `PRE_D`).
#'
#' @param x An `"ei_summary"` object.
#' @return An `"ei_summary"` object with `margins` as a wide `data.frame`
#'   and `meta$format == "eipack"`.
#'
#' @seealso [to_lphom()], [ei_summary()]
#' @export
to_eipack <- function(x) {
  stopifnot(inherits(x, "ei_summary"))
  if (x$meta$format == "eipack") return(x)

  # Ensure list form before converting
  src <- if (x$meta$format == "lphom") x else to_lphom(x)

  x$margins     <- .to_eipack_wide(src$margins)
  x$meta$format <- "eipack"
  x
}


# Internal helper: pull a single-race margin df from a wide df
.wide_to_single_df <- function(wide_df, election) {
  prefix   <- paste0(election, "_")
  opt_cols <- grep(paste0("^", prefix), names(wide_df), value = TRUE)
  df       <- wide_df[, c("precinct_id", "county_id", "n_voters", opt_cols),
                       drop = FALSE]
  names(df) <- sub(prefix, "", names(df))
  df
}