diff --git a/DESCRIPTION b/DESCRIPTION index e11acf4..f8c3ded 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -23,6 +23,7 @@ Imports: propensity (>= 0.0.0.9000), purrr, rlang, + scales, smd, tibble, tidyr, diff --git a/NAMESPACE b/NAMESPACE index 68b6e38..068dc03 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -19,6 +19,7 @@ export(bind_matches) export(check_auc) export(check_balance) export(check_calibration) +export(check_ess) export(contains) export(ends_with) export(ess) @@ -38,6 +39,7 @@ export(one_of) export(peek_vars) export(plot_balance) export(plot_calibration) +export(plot_ess) export(plot_mirror_distributions) export(plot_qq) export(plot_roc_auc) diff --git a/R/check_auc.R b/R/check_auc.R index e9669ef..8300024 100644 --- a/R/check_auc.R +++ b/R/check_auc.R @@ -226,7 +226,7 @@ roc_curve <- function( } weights <- extract_weight_data(weights) - + # Handle zero and negative weights if (any(weights <= 0, na.rm = TRUE)) { n_zero_neg <- sum(weights <= 0, na.rm = TRUE) diff --git a/R/check_ess.R b/R/check_ess.R new file mode 100644 index 0000000..244c558 --- /dev/null +++ b/R/check_ess.R @@ -0,0 +1,182 @@ +#' Check Effective Sample Size +#' +#' Computes the effective sample size (ESS) for one or more weighting schemes, +#' optionally stratified by treatment groups. ESS reflects how many observations +#' you would have if all were equally weighted. +#' +#' @details +#' The effective sample size (ESS) is calculated using the classical formula: +#' \eqn{ESS = (\sum w)^2 / \sum(w^2)}. +#' +#' When weights vary substantially, the ESS can be much smaller than the actual +#' number of observations, indicating that a few observations carry +#' disproportionately large weights. +#' +#' When `.group` is provided, ESS is calculated separately for each group level: +#' - For binary/categorical exposures: ESS is computed within each treatment level +#' - For continuous exposures: The variable is divided into quantiles (using +#' `dplyr::ntile()`) and ESS is computed within each quantile +#' +#' The function returns results in a tidy format suitable for plotting or +#' further analysis. +#' +#' @inheritParams check_params +#' @param .group Optional grouping variable. When provided, ESS is calculated +#' separately for each group level. For continuous variables, groups are +#' created using quantiles. +#' @param n_tiles For continuous `.group` variables, the number of quantile +#' groups to create. Default is 4 (quartiles). +#' @param tile_labels Optional character vector of labels for the quantile groups +#' when `.group` is continuous. If NULL, uses "Q1", "Q2", etc. +#' +#' @return A tibble with columns: +#' \item{method}{Character. The weighting method ("observed" or weight variable name).} +#' \item{group}{Character. The group level (if `.group` is provided).} +#' \item{n}{Integer. The number of observations in the group.} +#' \item{ess}{Numeric. The effective sample size.} +#' \item{ess_pct}{Numeric. ESS as a percentage of the actual sample size.} +#' +#' @family balance functions +#' @seealso [ess()] for the underlying ESS calculation, [plot_ess()] for visualization +#' +#' @examples +#' # Overall ESS for different weighting schemes +#' check_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atm)) +#' +#' # ESS by treatment group (binary exposure) +#' check_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) +#' +#' # ESS by treatment group (categorical exposure) +#' check_ess(nhefs_weights, .wts = w_cat_ate, .group = alcoholfreq_cat) +#' +#' # ESS by quartiles of a continuous variable +#' check_ess(nhefs_weights, .wts = w_ate, .group = age, n_tiles = 4) +#' +#' # Custom labels for continuous groups +#' check_ess(nhefs_weights, .wts = w_ate, .group = age, +#' n_tiles = 3, tile_labels = c("Young", "Middle", "Older")) +#' +#' # Without unweighted comparison +#' check_ess(nhefs_weights, .wts = w_ate, .group = qsmk, +#' include_observed = FALSE) +#' +#' @export +check_ess <- function( + .data, + .wts = NULL, + .group = NULL, + include_observed = TRUE, + n_tiles = 4, + tile_labels = NULL +) { + # Validate inputs + validate_data_frame(.data) + + # Handle group variable + group_quo <- rlang::enquo(.group) + has_group <- !rlang::quo_is_null(group_quo) + + if (has_group) { + group_name <- get_column_name(group_quo, ".group") + validate_column_exists(.data, group_name, ".group") + group_var <- .data[[group_name]] + + # Check if continuous (numeric and more than 10 unique values) + is_continuous <- is.numeric(group_var) && + length(unique(stats::na.omit(group_var))) > 10 + + if (is_continuous) { + # Create quantile groups + if (!is.null(tile_labels) && length(tile_labels) != n_tiles) { + abort( + "Length of {.arg tile_labels} must equal {.arg n_tiles}", + error_class = "halfmoon_length_error" + ) + } + + # Create tile groups + .data$.ess_group <- dplyr::ntile(group_var, n_tiles) + + # Apply labels + if (is.null(tile_labels)) { + tile_labels <- paste0("Q", seq_len(n_tiles)) + } + .data$.ess_group <- factor( + .data$.ess_group, + levels = seq_len(n_tiles), + labels = tile_labels + ) + group_col <- ".ess_group" + } else { + group_col <- group_name + } + } + + # Handle weights + wts_quo <- rlang::enquo(.wts) + + if (rlang::quo_is_null(wts_quo)) { + # No weights provided, just use observed + wts_names <- character() + } else { + wts_cols <- tidyselect::eval_select(wts_quo, .data) + wts_names <- names(wts_cols) + + # Convert psw weight columns to numeric + for (wts_name in wts_names) { + .data[[wts_name]] <- extract_weight_data(.data[[wts_name]]) + } + } + + # Add observed if requested + if (include_observed || length(wts_names) == 0) { + .data$.observed <- 1 + wts_names <- c(".observed", wts_names) + } + + # Reshape to long format + plot_data <- tidyr::pivot_longer( + .data, + cols = dplyr::all_of(wts_names), + names_to = "method", + values_to = "weight" + ) + + # Clean up method names + plot_data$method <- ifelse( + plot_data$method == ".observed", + "observed", + plot_data$method + ) + + # Calculate ESS + if (has_group) { + # Group-wise ESS + ess_data <- plot_data |> + dplyr::group_by(method, .data[[group_col]]) |> + dplyr::summarise( + n = dplyr::n(), + ess = ess(weight, na.rm = TRUE), + ess_pct = ess / n * 100, + .groups = "drop" + ) |> + dplyr::rename(group = !!group_col) + } else { + # Overall ESS + ess_data <- plot_data |> + dplyr::group_by(method) |> + dplyr::summarise( + n = dplyr::n(), + ess = ess(weight, na.rm = TRUE), + ess_pct = ess / n * 100, + .groups = "drop" + ) + } + + # Clean up temporary columns + if (has_group && is_continuous && ".ess_group" %in% names(ess_data)) { + ess_data <- dplyr::select(ess_data, -.ess_group) + } + + ess_data +} diff --git a/R/compute_balance.R b/R/compute_balance.R index 2c5982d..8edd41d 100644 --- a/R/compute_balance.R +++ b/R/compute_balance.R @@ -370,12 +370,14 @@ bal_ks <- function( p_ref <- if (is.null(weights)) { mean(covariate[idx_ref]) } else { - sum(extract_weight_data(weights)[idx_ref] * covariate[idx_ref]) / sum(extract_weight_data(weights)[idx_ref]) + sum(extract_weight_data(weights)[idx_ref] * covariate[idx_ref]) / + sum(extract_weight_data(weights)[idx_ref]) } p_other <- if (is.null(weights)) { mean(covariate[idx_other]) } else { - sum(extract_weight_data(weights)[idx_other] * covariate[idx_other]) / sum(extract_weight_data(weights)[idx_other]) + sum(extract_weight_data(weights)[idx_other] * covariate[idx_other]) / + sum(extract_weight_data(weights)[idx_other]) } return(abs(p_other - p_ref)) } @@ -384,7 +386,8 @@ bal_ks <- function( # Extract and weight x_ref <- covariate[idx_ref] x_other <- covariate[idx_other] - w_ref <- if (is.null(weights)) rep(1, length(x_ref)) else extract_weight_data(weights)[idx_ref] + w_ref <- if (is.null(weights)) rep(1, length(x_ref)) else + extract_weight_data(weights)[idx_ref] w_other <- if (is.null(weights)) rep(1, length(x_other)) else extract_weight_data(weights)[idx_other] w_ref <- w_ref / sum(w_ref) diff --git a/R/compute_qq.R b/R/compute_qq.R index 1efa20a..24ea4aa 100644 --- a/R/compute_qq.R +++ b/R/compute_qq.R @@ -259,7 +259,7 @@ compute_method_quantiles <- function( weighted_quantile <- function(values, quantiles, .wts) { # Extract numeric data from weights (handles both numeric and psw objects) .wts <- extract_weight_data(.wts) - + # Remove NA values if present na_idx <- is.na(values) | is.na(.wts) if (any(na_idx)) { diff --git a/R/ess.R b/R/ess.R index afa2d4c..0ac54c0 100644 --- a/R/ess.R +++ b/R/ess.R @@ -6,6 +6,7 @@ #' #' @param wts A numeric vector of weights (e.g., from survey or #' inverse-probability weighting). +#' @param na.rm Logical. Should missing values be removed? Default is FALSE. #' #' @return A single numeric value representing the effective sample size. #' @@ -42,8 +43,8 @@ #' ess(wts2) #' #' @export -ess <- function(wts) { +ess <- function(wts, na.rm = FALSE) { # Extract numeric data from psw weights if present wts <- extract_weight_data(wts) - sum(wts)^2 / sum(wts^2) + sum(wts, na.rm = na.rm)^2 / sum(wts^2, na.rm = na.rm) } diff --git a/R/geom_calibration.R b/R/geom_calibration.R index 3c2fbb0..b72c9d6 100644 --- a/R/geom_calibration.R +++ b/R/geom_calibration.R @@ -152,7 +152,13 @@ check_treatment_level <- function(group_var, treatment_level) { create_treatment_indicator(group_var, treatment_level) } -check_columns <- function(data, fitted_name, group_name, treatment_level, call = rlang::caller_env()) { +check_columns <- function( + data, + fitted_name, + group_name, + treatment_level, + call = rlang::caller_env() +) { if (is.null(treatment_level)) { if (!fitted_name %in% names(data)) { abort( @@ -591,7 +597,13 @@ compute_calibration_for_group <- function( # Compute calibration based on method calibration_result <- if (method == "breaks") { - compute_calibration_breaks_imp(df, bins, binning_method, conf_level, call = call) + compute_calibration_breaks_imp( + df, + bins, + binning_method, + conf_level, + call = call + ) } else if (method == "logistic") { compute_calibration_logistic_imp(df, smooth, conf_level, k = k, call = call) } else if (method == "windowed") { diff --git a/R/geom_mirrored_density.R b/R/geom_mirrored_density.R index e4e9601..da42829 100644 --- a/R/geom_mirrored_density.R +++ b/R/geom_mirrored_density.R @@ -62,7 +62,7 @@ StatMirrorDensity <- ggplot2::ggproto( .n_groups = length(unique(group)), .groups = "drop" ) - + # Check for panels with more than 2 groups if (any(panel_groups$.n_groups > 2)) { abort( @@ -70,19 +70,21 @@ StatMirrorDensity <- ggplot2::ggproto( error_class = "halfmoon_group_error" ) } - + # Join back to get panel group info for each row data <- dplyr::left_join(data, panel_groups, by = "PANEL") - + # Mark which groups should be mirrored (first group in each panel) - data$.should_mirror <- purrr::map2_lgl(data$group, data$.panel_groups, + data$.should_mirror <- purrr::map2_lgl( + data$group, + data$.panel_groups, ~ length(.y) == 2 && .x == .y[1] ) - + # Clean up temporary columns data$.panel_groups <- NULL data$.n_groups <- NULL - + data }, compute_group = function( @@ -109,15 +111,15 @@ StatMirrorDensity <- ggplot2::ggproto( error_class = "halfmoon_aes_error" ) } - + # Store mirroring flag should_mirror <- unique(data$.should_mirror) - + # Extract numeric data from psw weights if present if ("weight" %in% names(data)) { data$weight <- extract_weight_data(data$weight) } - + data <- ggplot2::StatDensity$compute_group( data = data, scales = scales, @@ -130,7 +132,7 @@ StatMirrorDensity <- ggplot2::ggproto( bounds = bounds, flipped_aes = flipped_aes ) - + # Apply mirroring if needed if (length(should_mirror) == 1 && should_mirror) { data$density <- -data$density @@ -138,7 +140,7 @@ StatMirrorDensity <- ggplot2::ggproto( data$scaled <- -data$scaled data$ndensity <- -data$ndensity } - + data } ) diff --git a/R/geom_mirrored_histogram.R b/R/geom_mirrored_histogram.R index f636de9..c7ac3cf 100644 --- a/R/geom_mirrored_histogram.R +++ b/R/geom_mirrored_histogram.R @@ -58,7 +58,7 @@ StatMirrorCount <- ggplot2::ggproto( .n_groups = length(unique(group)), .groups = "drop" ) - + # Check for panels with more than 2 groups if (any(panel_groups$.n_groups > 2)) { abort( @@ -66,19 +66,21 @@ StatMirrorCount <- ggplot2::ggproto( error_class = "halfmoon_group_error" ) } - + # Join back to get panel group info for each row data <- dplyr::left_join(data, panel_groups, by = "PANEL") - + # Mark which groups should be mirrored (first group in each panel) - data$.should_mirror <- purrr::map2_lgl(data$group, data$.panel_groups, + data$.should_mirror <- purrr::map2_lgl( + data$group, + data$.panel_groups, ~ length(.y) == 2 && .x == .y[1] ) - + # Clean up temporary columns data$.panel_groups <- NULL data$.n_groups <- NULL - + data }, compute_group = function( @@ -108,15 +110,15 @@ StatMirrorCount <- ggplot2::ggproto( error_class = "halfmoon_aes_error" ) } - + # Store mirroring flag should_mirror <- unique(data$.should_mirror) - + # Extract numeric data from psw weights if present if ("weight" %in% names(data)) { data$weight <- extract_weight_data(data$weight) } - + data <- ggplot2::StatBin$compute_group( data = data, scales = scales, @@ -132,12 +134,12 @@ StatMirrorCount <- ggplot2::ggproto( right = right, drop = drop ) - + # Apply mirroring if needed if (length(should_mirror) == 1 && should_mirror) { data$count <- -data$count } - + data } ) diff --git a/R/plot_ess.R b/R/plot_ess.R new file mode 100644 index 0000000..4c21d3a --- /dev/null +++ b/R/plot_ess.R @@ -0,0 +1,203 @@ +#' Plot Effective Sample Size +#' +#' Creates a bar plot visualization of effective sample sizes (ESS) for different +#' weighting schemes. ESS values are shown as percentages of the actual sample size, +#' with a reference line at 100% indicating no loss of effective sample size. +#' +#' @details +#' This function visualizes the output of `check_ess()` or computes ESS directly +#' from the provided data. The plot shows how much "effective" sample size remains +#' after weighting, which is a key diagnostic for assessing weight variability. +#' +#' When `.group` is not provided, the function displays overall ESS for each +#' weighting method. When `.group` is provided, ESS is shown separately for each +#' group level using dodged bars. +#' +#' For continuous grouping variables, the function automatically creates quantile +#' groups (quartiles by default) to show how ESS varies across the distribution +#' of the continuous variable. +#' +#' Lower ESS percentages indicate: +#' - Greater weight variability +#' - More extreme weights +#' - Potentially unstable weighted estimates +#' - Need for weight trimming or alternative methods +#' +#' @param .data A data frame, either: +#' - Output from `check_ess()` containing ESS calculations +#' - Raw data to compute ESS from (requires `.wts` to be specified) +#' @inheritParams check_ess +#' @param fill_color Color for the bars when `.group` is not provided. +#' Default is "#0172B1". +#' @param alpha Transparency level for the bars. Default is 0.8. +#' @param show_labels Logical. Show ESS percentage values as text labels on bars? +#' Default is TRUE. +#' @param label_size Size of text labels. Default is 3. +#' @param percent_scale Logical. Display ESS as percentage of sample size (TRUE) +#' or on original scale (FALSE)? Default is TRUE. +#' @param reference_line_color Color for the 100% reference line. Default is "gray50". +#' @param reference_line_type Line type for the reference line. Default is "dashed". +#' +#' @return A ggplot2 object. +#' +#' @seealso [check_ess()] for computing ESS values, [ess()] for the underlying calculation +#' +#' @examples +#' # Overall ESS for different weighting schemes +#' plot_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atm)) +#' +#' # ESS by treatment group (binary exposure) +#' plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) +#' +#' # ESS by treatment group (categorical exposure) +#' plot_ess(nhefs_weights, .wts = w_cat_ate, .group = alcoholfreq_cat) +#' +#' # ESS by age quartiles +#' plot_ess(nhefs_weights, .wts = w_ate, .group = age) +#' +#' # Customize quantiles for continuous variable +#' plot_ess(nhefs_weights, .wts = w_ate, .group = age, +#' n_tiles = 5, tile_labels = c("Youngest", "Young", "Middle", "Older", "Oldest")) +#' +#' # Without percentage labels +#' plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk, +#' show_labels = FALSE) +#' +#' # Custom styling +#' plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk, +#' alpha = 0.6, fill_color = "steelblue", +#' reference_line_color = "red") +#' +#' # Using pre-computed ESS data +#' ess_data <- check_ess(nhefs_weights, .wts = c(w_ate, w_att)) +#' plot_ess(ess_data) +#' +#' # Show ESS on original scale instead of percentage +#' plot_ess(nhefs_weights, .wts = c(w_ate, w_att), percent_scale = FALSE) +#' +#' @export +plot_ess <- function( + .data, + .wts = NULL, + .group = NULL, + include_observed = TRUE, + n_tiles = 4, + tile_labels = NULL, + fill_color = "#0172B1", + alpha = 0.8, + show_labels = TRUE, + label_size = 3, + percent_scale = TRUE, + reference_line_color = "gray50", + reference_line_type = "dashed" +) { + # Check if .data is already ESS output + is_ess_output <- all(c("method", "ess", "ess_pct", "n") %in% names(.data)) + + if (!is_ess_output) { + # Compute ESS from raw data + .data <- check_ess( + .data = .data, + .wts = {{ .wts }}, + .group = {{ .group }}, + include_observed = include_observed, + n_tiles = n_tiles, + tile_labels = tile_labels + ) + } + + # Determine if we have groups + has_groups <- "group" %in% names(.data) + + # Choose which variable to plot + y_var <- if (percent_scale) "ess_pct" else "ess" + + # Create the base plot + if (has_groups) { + # Plot with groups + p <- ggplot2::ggplot( + .data, + ggplot2::aes(x = method, y = .data[[y_var]], fill = group) + ) + + ggplot2::geom_col( + position = ggplot2::position_dodge(width = 0.9), + alpha = alpha + ) + + if (show_labels) { + p <- p + + ggplot2::geom_text( + ggplot2::aes( + label = if (percent_scale) { + scales::percent(ess_pct / 100, accuracy = 0.1) + } else { + scales::number(ess, accuracy = 0.1) + } + ), + position = ggplot2::position_dodge(width = 0.9), + vjust = -0.5, + size = label_size + ) + } + } else { + # Plot without groups + p <- ggplot2::ggplot( + .data, + ggplot2::aes(x = method, y = .data[[y_var]]) + ) + + ggplot2::geom_col( + fill = fill_color, + alpha = alpha + ) + + if (show_labels) { + p <- p + + ggplot2::geom_text( + ggplot2::aes( + label = if (percent_scale) { + scales::percent(ess_pct / 100, accuracy = 0.1) + } else { + scales::number(ess, accuracy = 0.1) + } + ), + vjust = -0.5, + size = label_size + ) + } + } + + # Add reference line + if (percent_scale) { + p <- p + + ggplot2::geom_hline( + yintercept = 100, + color = reference_line_color, + linetype = reference_line_type, + alpha = 0.7 + ) + } + + # Determine y-axis limits + if (percent_scale) { + max_ess <- max(.data$ess_pct, na.rm = TRUE) + y_upper <- ifelse(show_labels, max_ess * 1.1, max_ess * 1.05) + y_upper <- max(y_upper, 105) # Ensure we show at least to 105% + } else { + max_ess <- max(.data$ess, na.rm = TRUE) + y_upper <- ifelse(show_labels, max_ess * 1.1, max_ess * 1.05) + } + + # Add labels and formatting + p <- p + + ggplot2::labs( + x = "method", + y = if (percent_scale) "effective sample size (%)" else "effective sample size" + ) + + ggplot2::scale_y_continuous( + limits = c(0, y_upper), + expand = c(0, 0), + labels = if (percent_scale) \(x) scales::percent(x / 100) else scales::number + ) + + p +} diff --git a/R/plot_mirror_distributions.R b/R/plot_mirror_distributions.R index 66807a9..1968b65 100644 --- a/R/plot_mirror_distributions.R +++ b/R/plot_mirror_distributions.R @@ -82,15 +82,15 @@ #' .wts = w_ate, #' include_unweighted = FALSE #' ) -#' +#' #' # Categorical exposure - creates grid of comparisons #' plot_mirror_distributions( -#' nhefs_weights, -#' age, +#' nhefs_weights, +#' age, #' alcoholfreq_cat, #' type = "density" #' ) -#' +#' #' # Categorical with weights #' plot_mirror_distributions( #' nhefs_weights, @@ -140,7 +140,7 @@ plot_mirror_distributions <- function( } group_var <- .data[[group_name]] - + # Check if we have a categorical exposure (>2 levels) # Always use actual unique values in the data, not factor levels group_levels <- sort(unique(group_var[!is.na(group_var)])) @@ -148,23 +148,23 @@ plot_mirror_distributions <- function( # Convert to character to ensure proper comparison group_levels <- as.character(group_levels) } - + is_categorical <- length(group_levels) > 2 - + if (is_categorical) { # Categorical exposure reference_group <- determine_reference_group(group_var, reference_group) - + # Create binary comparisons using purrr comparison_levels <- setdiff(group_levels, reference_group) - + .data <- purrr::map_dfr(comparison_levels, \(level) { comparison_df <- .data |> dplyr::filter(.data[[group_name]] %in% c(reference_group, level)) |> dplyr::mutate(comparison = paste0(level, " vs ", reference_group)) # Update the group factor to only have the two levels comparison_df[[group_name]] <- factor( - comparison_df[[group_name]], + comparison_df[[group_name]], levels = c(reference_group, level) ) comparison_df @@ -269,7 +269,7 @@ plot_mirror_distributions <- function( na.rm = na.rm ) } - + # Add faceting for categorical exposures without weights if (is_categorical) { p <- p + ggplot2::facet_wrap(~comparison, scales = "free_y") diff --git a/R/plot_qq.R b/R/plot_qq.R index b7b044c..3f139e7 100644 --- a/R/plot_qq.R +++ b/R/plot_qq.R @@ -145,7 +145,7 @@ plot_qq <- function( for (wts_name in wts_names) { .data[[wts_name]] <- extract_weight_data(.data[[wts_name]]) } - + if (include_observed) { # Add observed as a weight column with value 1 .data$.observed <- 1 diff --git a/R/plot_stratified_residuals.R b/R/plot_stratified_residuals.R index 245d50d..065291a 100644 --- a/R/plot_stratified_residuals.R +++ b/R/plot_stratified_residuals.R @@ -120,7 +120,7 @@ plot_stratified_residuals.lm <- function( ... ) { plot_type <- rlang::arg_match(plot_type) - + # Validate required arguments if (missing(treatment)) { abort( diff --git a/R/utils-group.R b/R/utils-group.R index e720f08..4e11ac3 100644 --- a/R/utils-group.R +++ b/R/utils-group.R @@ -1,7 +1,11 @@ # Group handling helper functions for the halfmoon package # Extract and validate group levels -extract_group_levels <- function(group, require_binary = TRUE, call = rlang::caller_env()) { +extract_group_levels <- function( + group, + require_binary = TRUE, + call = rlang::caller_env() +) { levels <- if (is.factor(group)) { levels(group) } else { @@ -23,7 +27,11 @@ extract_group_levels <- function(group, require_binary = TRUE, call = rlang::cal } # Determine reference group with consistent logic -determine_reference_group <- function(group, reference_group = NULL, call = rlang::caller_env()) { +determine_reference_group <- function( + group, + reference_group = NULL, + call = rlang::caller_env() +) { levels <- extract_group_levels(group, require_binary = FALSE) if (is.null(reference_group)) { diff --git a/R/utils-validation.R b/R/utils-validation.R index 0117702..c6b7a66 100644 --- a/R/utils-validation.R +++ b/R/utils-validation.R @@ -1,7 +1,11 @@ # Validation helper functions for the halfmoon package # Numeric validation -validate_numeric <- function(x, arg_name = deparse(substitute(x)), call = rlang::caller_env()) { +validate_numeric <- function( + x, + arg_name = deparse(substitute(x)), + call = rlang::caller_env() +) { if (!is.numeric(x)) { abort( "{.arg {arg_name}} must be numeric, got {.cls {class(x)[1]}}", @@ -13,12 +17,17 @@ validate_numeric <- function(x, arg_name = deparse(substitute(x)), call = rlang: } # Weight validation -validate_weights <- function(weights, n, arg_name = "weights", call = rlang::caller_env()) { +validate_weights <- function( + weights, + n, + arg_name = "weights", + call = rlang::caller_env() +) { if (is.null(weights)) return(invisible(weights)) # Accept both numeric vectors and psw objects from propensity package is_valid_weights <- is.numeric(weights) || propensity::is_psw(weights) - + if (!is_valid_weights) { abort( "{.arg {arg_name}} must be numeric, a psw object, or {.code NULL}", @@ -44,7 +53,13 @@ validate_weights <- function(weights, n, arg_name = "weights", call = rlang::cal } # Length validation -validate_equal_length <- function(x, y, x_name = NULL, y_name = NULL, call = rlang::caller_env()) { +validate_equal_length <- function( + x, + y, + x_name = NULL, + y_name = NULL, + call = rlang::caller_env() +) { x_name <- x_name %||% deparse(substitute(x)) y_name <- y_name %||% deparse(substitute(y)) @@ -62,7 +77,11 @@ validate_equal_length <- function(x, y, x_name = NULL, y_name = NULL, call = rla } # Non-empty validation -validate_not_empty <- function(x, arg_name = deparse(substitute(x)), call = rlang::caller_env()) { +validate_not_empty <- function( + x, + arg_name = deparse(substitute(x)), + call = rlang::caller_env() +) { if (length(x) == 0) { abort( "{.arg {arg_name}} cannot be empty", @@ -74,7 +93,11 @@ validate_not_empty <- function(x, arg_name = deparse(substitute(x)), call = rlan } # Binary group validation -validate_binary_group <- function(group, arg_name = "group", call = rlang::caller_env()) { +validate_binary_group <- function( + group, + arg_name = "group", + call = rlang::caller_env() +) { levels <- unique(stats::na.omit(group)) if (length(levels) != 2) { abort( @@ -104,7 +127,11 @@ validate_reference_group <- function( } # Data frame validation -validate_data_frame <- function(data, arg_name = ".data", call = rlang::caller_env()) { +validate_data_frame <- function( + data, + arg_name = ".data", + call = rlang::caller_env() +) { if (!is.data.frame(data)) { abort( "{.arg {arg_name}} must be a data frame", @@ -116,7 +143,12 @@ validate_data_frame <- function(data, arg_name = ".data", call = rlang::caller_e } # Column existence validation -validate_column_exists <- function(data, column_name, arg_name = NULL, call = rlang::caller_env()) { +validate_column_exists <- function( + data, + column_name, + arg_name = NULL, + call = rlang::caller_env() +) { arg_name <- arg_name %||% column_name if (!column_name %in% names(data)) { abort( @@ -178,7 +210,11 @@ get_exposure_type <- function(group, call = rlang::caller_env()) { } # Validate exposure type -validate_exposure_type <- function(group, arg_name = "group", call = rlang::caller_env()) { +validate_exposure_type <- function( + group, + arg_name = "group", + call = rlang::caller_env() +) { exposure_type <- get_exposure_type(group, call = call) # For now, just return the type - validation happens in get_exposure_type diff --git a/R/utils.R b/R/utils.R index 6061a9e..7d64bed 100644 --- a/R/utils.R +++ b/R/utils.R @@ -53,7 +53,7 @@ get_column_name <- function(quo, arg_name, call = rlang::caller_env()) { call = call ) } - + # First try as_name (works for symbols and strings) tryCatch( { @@ -67,9 +67,9 @@ get_column_name <- function(quo, arg_name, call = rlang::caller_env()) { }, error = function(e2) { abort( - "{.code {arg_name}} must be a column name (quoted or unquoted)", - error_class = "halfmoon_type_error" - ) + "{.code {arg_name}} must be a column name (quoted or unquoted)", + error_class = "halfmoon_type_error" + ) } ) @@ -112,7 +112,12 @@ utils::globalVariables(c( "weight", ".var", ".weights", - "group_level" + "group_level", + "n", + ".ess_group", + "group", + "ess_pct", + "ess" )) # Error Classes for halfmoon @@ -121,7 +126,7 @@ utils::globalVariables(c( # Type errors # - halfmoon_type_error: Wrong input type (e.g., non-numeric when numeric expected) # -# Column/variable errors +# Column/variable errors # - halfmoon_column_error: Column not found in data frame # # Length/size errors diff --git a/R/zzz.R b/R/zzz.R index 290b0e3..d72440c 100644 --- a/R/zzz.R +++ b/R/zzz.R @@ -3,6 +3,6 @@ # This is needed for psw objects in nhefs_weights to work properly without # users having to explicitly load propensity loadNamespace("propensity") - + invisible() -} \ No newline at end of file +} diff --git a/man/bal_corr.Rd b/man/bal_corr.Rd index 3d75556..0b98a0f 100644 --- a/man/bal_corr.Rd +++ b/man/bal_corr.Rd @@ -41,6 +41,7 @@ Other balance functions: \code{\link{bal_vr}()}, \code{\link{check_auc}()}, \code{\link{check_balance}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/bal_ks.Rd b/man/bal_ks.Rd index 0e8366e..8378ad1 100644 --- a/man/bal_ks.Rd +++ b/man/bal_ks.Rd @@ -76,6 +76,7 @@ Other balance functions: \code{\link{bal_vr}()}, \code{\link{check_auc}()}, \code{\link{check_balance}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/bal_smd.Rd b/man/bal_smd.Rd index d4c2720..de7e845 100644 --- a/man/bal_smd.Rd +++ b/man/bal_smd.Rd @@ -76,6 +76,7 @@ Other balance functions: \code{\link{bal_vr}()}, \code{\link{check_auc}()}, \code{\link{check_balance}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/bal_vr.Rd b/man/bal_vr.Rd index 436edce..4012d2d 100644 --- a/man/bal_vr.Rd +++ b/man/bal_vr.Rd @@ -73,6 +73,7 @@ Other balance functions: \code{\link{bal_smd}()}, \code{\link{check_auc}()}, \code{\link{check_balance}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/check_auc.Rd b/man/check_auc.Rd index 1cefcf3..8602608 100644 --- a/man/check_auc.Rd +++ b/man/check_auc.Rd @@ -75,6 +75,7 @@ Other balance functions: \code{\link{bal_smd}()}, \code{\link{bal_vr}()}, \code{\link{check_balance}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/check_balance.Rd b/man/check_balance.Rd index 3b4bae0..8ac5cc4 100644 --- a/man/check_balance.Rd +++ b/man/check_balance.Rd @@ -136,6 +136,7 @@ Other balance functions: \code{\link{bal_smd}()}, \code{\link{bal_vr}()}, \code{\link{check_auc}()}, +\code{\link{check_ess}()}, \code{\link{plot_balance}()} } \concept{balance functions} diff --git a/man/check_ess.Rd b/man/check_ess.Rd new file mode 100644 index 0000000..0069ebf --- /dev/null +++ b/man/check_ess.Rd @@ -0,0 +1,101 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/check_ess.R +\name{check_ess} +\alias{check_ess} +\title{Check Effective Sample Size} +\usage{ +check_ess( + .data, + .wts = NULL, + .group = NULL, + include_observed = TRUE, + n_tiles = 4, + tile_labels = NULL +) +} +\arguments{ +\item{.data}{A data frame containing the variables to analyze.} + +\item{.wts}{Optional weighting variables. Can be unquoted variable names, +a character vector, or NULL. Multiple weights can be provided to compare +different weighting schemes.} + +\item{.group}{Optional grouping variable. When provided, ESS is calculated +separately for each group level. For continuous variables, groups are +created using quantiles.} + +\item{include_observed}{Logical. If using \code{.wts}, also calculate observed +(unweighted) metrics? Defaults to TRUE.} + +\item{n_tiles}{For continuous \code{.group} variables, the number of quantile +groups to create. Default is 4 (quartiles).} + +\item{tile_labels}{Optional character vector of labels for the quantile groups +when \code{.group} is continuous. If NULL, uses "Q1", "Q2", etc.} +} +\value{ +A tibble with columns: +\item{method}{Character. The weighting method ("observed" or weight variable name).} +\item{group}{Character. The group level (if \code{.group} is provided).} +\item{n}{Integer. The number of observations in the group.} +\item{ess}{Numeric. The effective sample size.} +\item{ess_pct}{Numeric. ESS as a percentage of the actual sample size.} +} +\description{ +Computes the effective sample size (ESS) for one or more weighting schemes, +optionally stratified by treatment groups. ESS reflects how many observations +you would have if all were equally weighted. +} +\details{ +The effective sample size (ESS) is calculated using the classical formula: +\eqn{ESS = (\sum w)^2 / \sum(w^2)}. + +When weights vary substantially, the ESS can be much smaller than the actual +number of observations, indicating that a few observations carry +disproportionately large weights. + +When \code{.group} is provided, ESS is calculated separately for each group level: +\itemize{ +\item For binary/categorical exposures: ESS is computed within each treatment level +\item For continuous exposures: The variable is divided into quantiles (using +\code{dplyr::ntile()}) and ESS is computed within each quantile +} + +The function returns results in a tidy format suitable for plotting or +further analysis. +} +\examples{ +# Overall ESS for different weighting schemes +check_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atm)) + +# ESS by treatment group (binary exposure) +check_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) + +# ESS by treatment group (categorical exposure) +check_ess(nhefs_weights, .wts = w_cat_ate, .group = alcoholfreq_cat) + +# ESS by quartiles of a continuous variable +check_ess(nhefs_weights, .wts = w_ate, .group = age, n_tiles = 4) + +# Custom labels for continuous groups +check_ess(nhefs_weights, .wts = w_ate, .group = age, + n_tiles = 3, tile_labels = c("Young", "Middle", "Older")) + +# Without unweighted comparison +check_ess(nhefs_weights, .wts = w_ate, .group = qsmk, + include_observed = FALSE) + +} +\seealso{ +\code{\link[=ess]{ess()}} for the underlying ESS calculation, \code{\link[=plot_ess]{plot_ess()}} for visualization + +Other balance functions: +\code{\link{bal_corr}()}, +\code{\link{bal_ks}()}, +\code{\link{bal_smd}()}, +\code{\link{bal_vr}()}, +\code{\link{check_auc}()}, +\code{\link{check_balance}()}, +\code{\link{plot_balance}()} +} +\concept{balance functions} diff --git a/man/ess.Rd b/man/ess.Rd index eb6adb0..b85e43a 100644 --- a/man/ess.Rd +++ b/man/ess.Rd @@ -4,11 +4,13 @@ \alias{ess} \title{Calculate the Effective Sample Size (ESS)} \usage{ -ess(wts) +ess(wts, na.rm = FALSE) } \arguments{ \item{wts}{A numeric vector of weights (e.g., from survey or inverse-probability weighting).} + +\item{na.rm}{Logical. Should missing values be removed? Default is FALSE.} } \value{ A single numeric value representing the effective sample size. diff --git a/man/plot_balance.Rd b/man/plot_balance.Rd index 9ea7497..eebd0f9 100644 --- a/man/plot_balance.Rd +++ b/man/plot_balance.Rd @@ -112,6 +112,7 @@ Other balance functions: \code{\link{bal_smd}()}, \code{\link{bal_vr}()}, \code{\link{check_auc}()}, -\code{\link{check_balance}()} +\code{\link{check_balance}()}, +\code{\link{check_ess}()} } \concept{balance functions} diff --git a/man/plot_ess.Rd b/man/plot_ess.Rd new file mode 100644 index 0000000..8834d22 --- /dev/null +++ b/man/plot_ess.Rd @@ -0,0 +1,129 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/plot_ess.R +\name{plot_ess} +\alias{plot_ess} +\title{Plot Effective Sample Size} +\usage{ +plot_ess( + .data, + .wts = NULL, + .group = NULL, + include_observed = TRUE, + n_tiles = 4, + tile_labels = NULL, + fill_color = "#0172B1", + alpha = 0.8, + show_labels = TRUE, + label_size = 3, + percent_scale = TRUE, + reference_line_color = "gray50", + reference_line_type = "dashed" +) +} +\arguments{ +\item{.data}{A data frame, either: +\itemize{ +\item Output from \code{check_ess()} containing ESS calculations +\item Raw data to compute ESS from (requires \code{.wts} to be specified) +}} + +\item{.wts}{Optional weighting variables. Can be unquoted variable names, +a character vector, or NULL. Multiple weights can be provided to compare +different weighting schemes.} + +\item{.group}{Optional grouping variable. When provided, ESS is calculated +separately for each group level. For continuous variables, groups are +created using quantiles.} + +\item{include_observed}{Logical. If using \code{.wts}, also calculate observed +(unweighted) metrics? Defaults to TRUE.} + +\item{n_tiles}{For continuous \code{.group} variables, the number of quantile +groups to create. Default is 4 (quartiles).} + +\item{tile_labels}{Optional character vector of labels for the quantile groups +when \code{.group} is continuous. If NULL, uses "Q1", "Q2", etc.} + +\item{fill_color}{Color for the bars when \code{.group} is not provided. +Default is "#0172B1".} + +\item{alpha}{Transparency level for the bars. Default is 0.8.} + +\item{show_labels}{Logical. Show ESS percentage values as text labels on bars? +Default is TRUE.} + +\item{label_size}{Size of text labels. Default is 3.} + +\item{percent_scale}{Logical. Display ESS as percentage of sample size (TRUE) +or on original scale (FALSE)? Default is TRUE.} + +\item{reference_line_color}{Color for the 100\% reference line. Default is "gray50".} + +\item{reference_line_type}{Line type for the reference line. Default is "dashed".} +} +\value{ +A ggplot2 object. +} +\description{ +Creates a bar plot visualization of effective sample sizes (ESS) for different +weighting schemes. ESS values are shown as percentages of the actual sample size, +with a reference line at 100\% indicating no loss of effective sample size. +} +\details{ +This function visualizes the output of \code{check_ess()} or computes ESS directly +from the provided data. The plot shows how much "effective" sample size remains +after weighting, which is a key diagnostic for assessing weight variability. + +When \code{.group} is not provided, the function displays overall ESS for each +weighting method. When \code{.group} is provided, ESS is shown separately for each +group level using dodged bars. + +For continuous grouping variables, the function automatically creates quantile +groups (quartiles by default) to show how ESS varies across the distribution +of the continuous variable. + +Lower ESS percentages indicate: +\itemize{ +\item Greater weight variability +\item More extreme weights +\item Potentially unstable weighted estimates +\item Need for weight trimming or alternative methods +} +} +\examples{ +# Overall ESS for different weighting schemes +plot_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atm)) + +# ESS by treatment group (binary exposure) +plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) + +# ESS by treatment group (categorical exposure) +plot_ess(nhefs_weights, .wts = w_cat_ate, .group = alcoholfreq_cat) + +# ESS by age quartiles +plot_ess(nhefs_weights, .wts = w_ate, .group = age) + +# Customize quantiles for continuous variable +plot_ess(nhefs_weights, .wts = w_ate, .group = age, + n_tiles = 5, tile_labels = c("Youngest", "Young", "Middle", "Older", "Oldest")) + +# Without percentage labels +plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk, + show_labels = FALSE) + +# Custom styling +plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk, + alpha = 0.6, fill_color = "steelblue", + reference_line_color = "red") + +# Using pre-computed ESS data +ess_data <- check_ess(nhefs_weights, .wts = c(w_ate, w_att)) +plot_ess(ess_data) + +# Show ESS on original scale instead of percentage +plot_ess(nhefs_weights, .wts = c(w_ate, w_att), percent_scale = FALSE) + +} +\seealso{ +\code{\link[=check_ess]{check_ess()}} for computing ESS values, \code{\link[=ess]{ess()}} for the underlying calculation +} diff --git a/man/plot_mirror_distributions.Rd b/man/plot_mirror_distributions.Rd index d9977b3..7565410 100644 --- a/man/plot_mirror_distributions.Rd +++ b/man/plot_mirror_distributions.Rd @@ -116,8 +116,8 @@ plot_mirror_distributions( # Categorical exposure - creates grid of comparisons plot_mirror_distributions( - nhefs_weights, - age, + nhefs_weights, + age, alcoholfreq_cat, type = "density" ) diff --git a/tests/testthat/Rplots.pdf b/tests/testthat/Rplots.pdf index 6885885..9776871 100644 Binary files a/tests/testthat/Rplots.pdf and b/tests/testthat/Rplots.pdf differ diff --git a/tests/testthat/_snaps/check_ess.md b/tests/testthat/_snaps/check_ess.md new file mode 100644 index 0000000..f1d9446 --- /dev/null +++ b/tests/testthat/_snaps/check_ess.md @@ -0,0 +1,24 @@ +# check_ess validates inputs + + Code + expr + Condition + Error in `check_ess()`: + ! `.data` must be a data frame + +--- + + Code + expr + Condition + Error in `check_ess()`: + ! Column `not_a_column` not found in `.group` + +--- + + Code + expr + Condition + Error in `check_ess()`: + ! Length of `tile_labels` must equal `n_tiles` + diff --git a/tests/testthat/_snaps/plot_ess/plot-ess-basic.svg b/tests/testthat/_snaps/plot_ess/plot-ess-basic.svg new file mode 100644 index 0000000..00ed250 --- /dev/null +++ b/tests/testthat/_snaps/plot_ess/plot-ess-basic.svg @@ -0,0 +1,60 @@ + + + + + + + + + + + + + + + + + + + + + + + + +100.0% +64.7% +68.3% + + + + +0% +25% +50% +75% +100% + + + + + + + + +observed +w_ate +w_att +method +effective sample size (%) +plot_ess_basic + + diff --git a/tests/testthat/_snaps/plot_ess/plot-ess-continuous.svg b/tests/testthat/_snaps/plot_ess/plot-ess-continuous.svg new file mode 100644 index 0000000..64f4b48 --- /dev/null +++ b/tests/testthat/_snaps/plot_ess/plot-ess-continuous.svg @@ -0,0 +1,75 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + +100.0% +100.0% +100.0% +56.1% +62.4% +80.0% + + + + +0% +25% +50% +75% +100% + + + + + + + +observed +w_ate +method +effective sample size (%) + +group + + + + + + +Q1 +Q2 +Q3 +plot_ess_continuous + + diff --git a/tests/testthat/_snaps/plot_ess/plot-ess-groups.svg b/tests/testthat/_snaps/plot_ess/plot-ess-groups.svg new file mode 100644 index 0000000..6fb2684 --- /dev/null +++ b/tests/testthat/_snaps/plot_ess/plot-ess-groups.svg @@ -0,0 +1,74 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + +100.0% +100.0% +97.0% +80.9% +68.4% +100.0% + + + + +0% +25% +50% +75% +100% + + + + + + + + +observed +w_ate +w_att +method +effective sample size (%) + +group + + + + +0 +1 +plot_ess_groups + + diff --git a/tests/testthat/_snaps/plot_ess/plot-ess-no-labels.svg b/tests/testthat/_snaps/plot_ess/plot-ess-no-labels.svg new file mode 100644 index 0000000..76b3055 --- /dev/null +++ b/tests/testthat/_snaps/plot_ess/plot-ess-no-labels.svg @@ -0,0 +1,57 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +0% +25% +50% +75% +100% + + + + + + + + +observed +w_ate +w_att +method +effective sample size (%) +plot_ess_no_labels + + diff --git a/tests/testthat/helper-snapshot.R b/tests/testthat/helper-snapshot.R index ddbfe7d..8802cd1 100644 --- a/tests/testthat/helper-snapshot.R +++ b/tests/testthat/helper-snapshot.R @@ -2,7 +2,7 @@ # These ensure we always capture both the message and the condition class # For testing errors - captures snapshot AND verifies error class if provided -# Usage: +# Usage: # expect_halfmoon_error(plot_balance(invalid_data)) # expect_halfmoon_error(plot_balance(invalid_data), "halfmoon_type_error") expect_halfmoon_error <- function(expr, class = NULL) { @@ -16,7 +16,7 @@ expect_halfmoon_error <- function(expr, class = NULL) { } # For testing warnings - captures snapshot AND verifies warning class if provided -# Usage: +# Usage: # expect_halfmoon_warning(check_calibration(small_data)) # expect_halfmoon_warning(check_calibration(small_data), "halfmoon_data_warning") expect_halfmoon_warning <- function(expr, class = NULL) { @@ -26,4 +26,4 @@ expect_halfmoon_warning <- function(expr, class = NULL) { cnd_class = TRUE, expr ) -} \ No newline at end of file +} diff --git a/tests/testthat/test-check_balance.R b/tests/testthat/test-check_balance.R index 765704b..e1f6e6f 100644 --- a/tests/testthat/test-check_balance.R +++ b/tests/testthat/test-check_balance.R @@ -416,7 +416,10 @@ test_that("check_balance validates inputs correctly", { ) # Test no variables selected - expect_halfmoon_error(check_balance_basic(data, c(), qsmk), "halfmoon_empty_error") + expect_halfmoon_error( + check_balance_basic(data, c(), qsmk), + "halfmoon_empty_error" + ) # Test group with wrong number of levels - should error data_bad_group <- data diff --git a/tests/testthat/test-check_ess.R b/tests/testthat/test-check_ess.R new file mode 100644 index 0000000..050babb --- /dev/null +++ b/tests/testthat/test-check_ess.R @@ -0,0 +1,131 @@ +test_that("check_ess works with no weights", { + result <- check_ess(nhefs_weights) + + expect_s3_class(result, "tbl_df") + expect_equal(nrow(result), 1) + expect_equal(result$method, "observed") + expect_equal(result$ess_pct, 100) + expect_equal(result$ess, result$n) +}) + +test_that("check_ess works with single weight", { + result <- check_ess(nhefs_weights, .wts = w_ate) + + expect_s3_class(result, "tbl_df") + expect_equal(nrow(result), 2) + expect_equal(result$method, c("observed", "w_ate")) + expect_true(all(result$ess_pct <= 100)) + expect_true(all(result$ess > 0)) +}) + +test_that("check_ess works with multiple weights", { + result <- check_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atm)) + + expect_s3_class(result, "tbl_df") + expect_equal(nrow(result), 4) + # Methods may be returned in different order due to tidyselect + expect_setequal(result$method, c("observed", "w_ate", "w_att", "w_atm")) +}) + +test_that("check_ess works without observed", { + result <- check_ess(nhefs_weights, .wts = w_ate, include_observed = FALSE) + + expect_equal(nrow(result), 1) + expect_equal(result$method, "w_ate") +}) + +test_that("check_ess works with binary groups", { + result <- check_ess(nhefs_weights, .wts = w_ate, .group = qsmk) + + expect_s3_class(result, "tbl_df") + expect_true("group" %in% names(result)) + expect_equal(length(unique(result$group)), 2) + expect_equal(nrow(result), 4) # 2 methods x 2 groups +}) + +test_that("check_ess works with categorical groups", { + result <- check_ess(nhefs_weights, .wts = w_cat_ate, .group = alcoholfreq_cat) + + expect_s3_class(result, "tbl_df") + expect_true("group" %in% names(result)) + expect_true(length(unique(result$group)) > 2) +}) + +test_that("check_ess works with continuous groups", { + result <- check_ess(nhefs_weights, .wts = w_ate, .group = age, n_tiles = 4) + + expect_s3_class(result, "tbl_df") + expect_true("group" %in% names(result)) + expect_equal(length(unique(result$group)), 4) + expect_true(all(grepl("^Q[1-4]$", unique(result$group)))) +}) + +test_that("check_ess works with custom tile labels", { + labels <- c("Young", "Middle", "Older") + result <- check_ess(nhefs_weights, .wts = w_ate, .group = age, + n_tiles = 3, tile_labels = labels) + + expect_setequal(as.character(unique(result$group)), labels) +}) + +test_that("check_ess handles tidyselect syntax", { + result1 <- check_ess(nhefs_weights, .wts = starts_with("w_a")) + result2 <- check_ess(nhefs_weights, .wts = c(w_ate, w_att, w_atc, w_atm, w_ato)) + + # Should have same number of methods (plus observed) + expect_equal(nrow(result1), nrow(result2)) +}) + +test_that("check_ess handles psw weight objects", { + # Assuming psw weights are numeric vectors with special class + # The extract_weight_data function should handle conversion + result <- check_ess(nhefs_weights, .wts = w_ate) + + expect_true(is.numeric(result$ess)) + expect_true(all(result$ess > 0)) +}) + +test_that("check_ess validates inputs", { + expect_halfmoon_error( + check_ess("not a data frame"), + "halfmoon_type_error" + ) + + expect_halfmoon_error( + check_ess(nhefs_weights, .group = "not_a_column"), + "halfmoon_column_error" + ) + + expect_halfmoon_error( + check_ess(nhefs_weights, .wts = w_ate, .group = age, + n_tiles = 3, tile_labels = c("Too", "Few")), + "halfmoon_length_error" + ) +}) + +test_that("ESS calculation is correct", { + # Create simple test case with known ESS + test_df <- data.frame( + wts = c(1, 1, 1, 1), # Equal weights -> ESS = n + wts2 = c(4, 0, 0, 0) # All weight on one obs -> ESS = 1 + ) + + result <- check_ess(test_df, .wts = c(wts, wts2), include_observed = FALSE) + + expect_equal(result$ess[result$method == "wts"], 4) + expect_equal(result$ess[result$method == "wts2"], 1) + expect_equal(result$ess_pct[result$method == "wts"], 100) + expect_equal(result$ess_pct[result$method == "wts2"], 25) +}) + +test_that("check_ess handles NA values", { + # Create data with NAs + test_df <- nhefs_weights + test_df$w_ate[1:10] <- NA + + result <- check_ess(test_df, .wts = w_ate) + + # Should still compute ESS on non-NA values + expect_true(all(!is.na(result$ess))) + expect_true(all(result$ess > 0)) +}) \ No newline at end of file diff --git a/tests/testthat/test-compute_balance.R b/tests/testthat/test-compute_balance.R index 2e67e46..7f7ee25 100644 --- a/tests/testthat/test-compute_balance.R +++ b/tests/testthat/test-compute_balance.R @@ -129,7 +129,7 @@ test_that("bal_smd error handling", { bal_smd(covariate = data$x_cont, group = rep(1, 100)), "halfmoon_group_error" ) - + # Now supports 3+ groups (categorical) expect_no_error(bal_smd( covariate = data$x_cont, @@ -144,7 +144,7 @@ test_that("bal_smd error handling", { ), "halfmoon_length_error" ) - + expect_halfmoon_error( bal_smd( covariate = data$x_cont, @@ -269,7 +269,7 @@ test_that("bal_vr error handling", { ), "halfmoon_group_error" ) - + # Now supports 3+ groups (categorical) expect_no_error(bal_vr( covariate = data$x_cont, @@ -284,7 +284,7 @@ test_that("bal_vr error handling", { ), "halfmoon_length_error" ) - + expect_halfmoon_error( bal_vr( covariate = data$x_cont, @@ -398,7 +398,7 @@ test_that("bal_ks error handling", { ), "halfmoon_length_error" ) - + expect_halfmoon_error( bal_ks( covariate = data$x_cont, @@ -508,7 +508,7 @@ test_that("bal_corr error handling", { bal_corr(data$x_cont[1:50], data$x_skewed), "halfmoon_length_error" ) - + expect_halfmoon_error( bal_corr( data$x_cont, @@ -1504,46 +1504,46 @@ test_that("balance functions work seamlessly with psw objects from propensity pa # This test ensures psw objects from the propensity package work throughout # the halfmoon package without requiring explicit conversion data(nhefs_weights) - + # Verify we have psw objects in the dataset expect_true(propensity::is_psw(nhefs_weights$w_cat_ate)) expect_true(propensity::is_psw(nhefs_weights$w_cat_att_none)) - + # Test that balance functions work directly with psw weights result_smd <- bal_smd( - nhefs_weights$age, - nhefs_weights$alcoholfreq_cat, + nhefs_weights$age, + nhefs_weights$alcoholfreq_cat, weights = nhefs_weights$w_cat_ate ) expect_true(all(is.finite(result_smd))) - + result_vr <- bal_vr( - nhefs_weights$wt71, - nhefs_weights$alcoholfreq_cat, + nhefs_weights$wt71, + nhefs_weights$alcoholfreq_cat, weights = nhefs_weights$w_cat_att_none ) expect_true(all(is.finite(result_vr) & result_vr > 0)) - + result_ks <- bal_ks( - nhefs_weights$age, - nhefs_weights$alcoholfreq_cat, + nhefs_weights$age, + nhefs_weights$alcoholfreq_cat, weights = nhefs_weights$w_cat_ato ) expect_true(all(is.finite(result_ks) & result_ks >= 0 & result_ks <= 1)) - + # Test check_balance works with psw weights balance_results <- check_balance( - nhefs_weights, - c(age, wt71), - alcoholfreq_cat, - .wts = w_cat_ate, + nhefs_weights, + c(age, wt71), + alcoholfreq_cat, + .wts = w_cat_ate, .metrics = "smd", include_observed = FALSE ) expect_s3_class(balance_results, "data.frame") expect_true(nrow(balance_results) > 0) expect_true(all(is.finite(balance_results$estimate))) - + # Test weighted_quantile works with psw weights quantiles <- weighted_quantile( nhefs_weights$age, diff --git a/tests/testthat/test-compute_balance_categorical.R b/tests/testthat/test-compute_balance_categorical.R index afb668a..a34b173 100644 --- a/tests/testthat/test-compute_balance_categorical.R +++ b/tests/testthat/test-compute_balance_categorical.R @@ -267,7 +267,10 @@ test_that("categorical exposure validation works correctly", { # Test get_exposure_type expect_equal(get_exposure_type(data$exposure), "categorical") expect_equal(get_exposure_type(c(0, 1, 1, 0)), "binary") - expect_halfmoon_error(get_exposure_type(c(1, 1, 1, 1)), "halfmoon_group_error") + expect_halfmoon_error( + get_exposure_type(c(1, 1, 1, 1)), + "halfmoon_group_error" + ) }) test_that("categorical balance handles missing values correctly", { diff --git a/tests/testthat/test-error-messages.R b/tests/testthat/test-error-messages.R index d8be7a6..f616d1a 100644 --- a/tests/testthat/test-error-messages.R +++ b/tests/testthat/test-error-messages.R @@ -1,4 +1,3 @@ - test_that("error messages show user-facing function names", { # Test plot_mirror_distributions with invalid reference group expect_snapshot( @@ -11,7 +10,7 @@ test_that("error messages show user-facing function names", { reference_group = "invalid" ) ) - + # Test with numeric out of bounds reference group expect_snapshot( error = TRUE, @@ -23,14 +22,14 @@ test_that("error messages show user-facing function names", { reference_group = 10 ) ) - + # Test missing required arguments expect_snapshot( error = TRUE, cnd_class = TRUE, plot_mirror_distributions(nhefs_weights) ) - + # Test with missing column expect_snapshot( error = TRUE, @@ -41,7 +40,7 @@ test_that("error messages show user-facing function names", { qsmk ) ) - + # Test plot_qq with invalid treatment level expect_snapshot( error = TRUE, @@ -53,7 +52,7 @@ test_that("error messages show user-facing function names", { treatment_level = "invalid" ) ) - + # Test plot_stratified_residuals with missing treatment model <- lm(mpg ~ wt, data = mtcars) expect_snapshot( @@ -61,7 +60,7 @@ test_that("error messages show user-facing function names", { cnd_class = TRUE, plot_stratified_residuals(model) ) - + # Test check_balance with wrong group levels expect_snapshot( error = TRUE, @@ -72,7 +71,7 @@ test_that("error messages show user-facing function names", { .group = rep(1, nrow(nhefs_weights)) ) ) - + # Test bal_prognostic_score with treatment in formula expect_snapshot( error = TRUE, @@ -96,7 +95,7 @@ test_that("validation errors show correct function context", { .group = "not_numeric" ) ) - + # Test empty data frame expect_snapshot( error = TRUE, @@ -120,7 +119,7 @@ test_that("errors have correct custom classes", { ), class = "halfmoon_reference_error" ) - + # Test range error class expect_halfmoon_error( plot_mirror_distributions( @@ -131,13 +130,13 @@ test_that("errors have correct custom classes", { ), class = "halfmoon_range_error" ) - + # Test arg error class expect_halfmoon_error( plot_mirror_distributions(nhefs_weights), class = "halfmoon_arg_error" ) - + # Test column error class expect_halfmoon_error( plot_mirror_distributions( @@ -147,7 +146,7 @@ test_that("errors have correct custom classes", { ), class = "halfmoon_column_error" ) - + # Test formula error class expect_halfmoon_error( bal_prognostic_score( @@ -158,4 +157,3 @@ test_that("errors have correct custom classes", { class = "halfmoon_formula_error" ) }) - diff --git a/tests/testthat/test-ess.R b/tests/testthat/test-ess.R index 330413f..8e8b1e7 100644 --- a/tests/testthat/test-ess.R +++ b/tests/testthat/test-ess.R @@ -25,3 +25,12 @@ test_that("ess gives `NaN` if all weights are 0", { # sum(wts) = 0, sum(wts^2) = 0 -> 0/0 is NaN expect_true(is.nan(ess(wts_zero))) }) + +test_that("ess handles NA values", { + # With na.rm = FALSE (default), should return NA + expect_true(is.na(ess(c(1, 2, NA, 3)))) + + # With na.rm = TRUE, should calculate ESS on non-NA values + expect_equal(ess(c(1, 1, NA, 1), na.rm = TRUE), 3) + expect_true(ess(c(0.5, 2, NA, 0.1), na.rm = TRUE) < 3) +}) diff --git a/tests/testthat/test-extract_weight_data.R b/tests/testthat/test-extract_weight_data.R index 2d8db65..0b6b0e2 100644 --- a/tests/testthat/test-extract_weight_data.R +++ b/tests/testthat/test-extract_weight_data.R @@ -11,12 +11,12 @@ test_that("extract_weight_data passes through numeric vectors unchanged", { result <- extract_weight_data(numeric_weights) expect_identical(result, numeric_weights) expect_type(result, "double") - + # Test with different numeric values uniform_weights <- runif(10, 0.5, 1.5) result2 <- extract_weight_data(uniform_weights) expect_identical(result2, uniform_weights) - + # Test with extreme values extreme_weights <- c(0.001, 100, 0.5, 2.5) result3 <- extract_weight_data(extreme_weights) @@ -25,38 +25,38 @@ test_that("extract_weight_data passes through numeric vectors unchanged", { test_that("extract_weight_data extracts data from psw objects", { skip_if_not_installed("propensity") - + # Create psw object from numeric data numeric_data <- c(1.2, 0.8, 1.5, 0.9, 1.1) psw_obj <- propensity::psw(numeric_data, estimand = "ate") - + # Extract should return the underlying numeric data result <- extract_weight_data(psw_obj) expect_equal(result, numeric_data) expect_type(result, "double") - + # Should not have psw class anymore expect_false(propensity::is_psw(result)) }) test_that("extract_weight_data preserves values from different psw estimands", { skip_if_not_installed("propensity") - + numeric_data <- c(2.1, 1.3, 0.7, 1.8, 1.0) - + # Test different estimands produce same numeric extraction psw_ate <- propensity::psw(numeric_data, estimand = "ate") psw_att <- propensity::psw(numeric_data, estimand = "att") psw_ato <- propensity::psw(numeric_data, estimand = "ato") - + result_ate <- extract_weight_data(psw_ate) result_att <- extract_weight_data(psw_att) result_ato <- extract_weight_data(psw_ato) - + expect_equal(result_ate, numeric_data) expect_equal(result_att, numeric_data) expect_equal(result_ato, numeric_data) - + # All should be identical expect_identical(result_ate, result_att) expect_identical(result_ate, result_ato) @@ -83,9 +83,9 @@ test_that("extract_weight_data works with edge case values", { edge_weights <- c(0, 1e-10, 1e10, Inf, -Inf) result <- extract_weight_data(edge_weights) expect_identical(result, edge_weights) - + # Test with single value single_weight <- 1.5 result_single <- extract_weight_data(single_weight) expect_identical(result_single, single_weight) -}) \ No newline at end of file +}) diff --git a/tests/testthat/test-geom_mirrored_histogram.R b/tests/testthat/test-geom_mirrored_histogram.R index fbf555c..38cdbc9 100644 --- a/tests/testthat/test-geom_mirrored_histogram.R +++ b/tests/testthat/test-geom_mirrored_histogram.R @@ -22,7 +22,7 @@ test_that("geom_mirrored_histogram errors correctly", { aes(group = education), bins = 50 ) - + expect_error( ggplot_build(edu_group), class = "halfmoon_group_error" diff --git a/tests/testthat/test-geom_qq2.R b/tests/testthat/test-geom_qq2.R index 9e4990c..8ad6803 100644 --- a/tests/testthat/test-geom_qq2.R +++ b/tests/testthat/test-geom_qq2.R @@ -143,7 +143,7 @@ test_that("geom_qq2 visual regression tests", { nhefs_for_pivot2 <- nhefs_weights nhefs_for_pivot2$w_ate <- vctrs::vec_data(nhefs_weights$w_ate) nhefs_for_pivot2$w_att <- vctrs::vec_data(nhefs_weights$w_att) - + long_data <- tidyr::pivot_longer( nhefs_for_pivot2, cols = c(w_ate, w_att), diff --git a/tests/testthat/test-numeric-psw-equivalence.R b/tests/testthat/test-numeric-psw-equivalence.R index 2c9c0b1..c736e9f 100644 --- a/tests/testthat/test-numeric-psw-equivalence.R +++ b/tests/testthat/test-numeric-psw-equivalence.R @@ -7,27 +7,29 @@ test_that("backward compatibility: all functions work with numeric weights", { x <- rnorm(n) g <- sample(c(0, 1), n, replace = TRUE) numeric_weights <- runif(n, 0.5, 2.0) - + # Test all balance functions with numeric weights expect_no_error(smd_result <- bal_smd(x, g, weights = numeric_weights)) expect_true(is.finite(smd_result)) - + expect_no_error(vr_result <- bal_vr(x, g, weights = numeric_weights)) expect_true(is.finite(vr_result) && vr_result > 0) - + expect_no_error(ks_result <- bal_ks(x, g, weights = numeric_weights)) expect_true(is.finite(ks_result) && ks_result >= 0 && ks_result <= 1) - + # Test correlation with numeric weights y <- 2 * x + rnorm(n, sd = 0.5) expect_no_error(corr_result <- bal_corr(x, y, weights = numeric_weights)) expect_true(is.finite(corr_result) && corr_result >= -1 && corr_result <= 1) - + # Test energy balance with numeric weights covariates <- data.frame(x1 = x, x2 = rnorm(n)) - expect_no_error(energy_result <- bal_energy(covariates, g, weights = numeric_weights)) + expect_no_error( + energy_result <- bal_energy(covariates, g, weights = numeric_weights) + ) expect_true(is.finite(energy_result) && energy_result >= 0) - + # Test check_balance with numeric weights test_data <- data.frame( x = x, @@ -36,7 +38,13 @@ test_that("backward compatibility: all functions work with numeric weights", { w = numeric_weights ) expect_no_error( - balance_results <- check_balance(test_data, c(x, y), g, .wts = w, .metrics = "smd") + balance_results <- check_balance( + test_data, + c(x, y), + g, + .wts = w, + .metrics = "smd" + ) ) expect_s3_class(balance_results, "data.frame") expect_true(all(is.finite(balance_results$estimate))) @@ -44,47 +52,47 @@ test_that("backward compatibility: all functions work with numeric weights", { test_that("numeric and psw weights produce identical results", { skip_if_not_installed("propensity") - + # Create test data set.seed(42) n <- 100 x <- rnorm(n) g <- sample(c(0, 1), n, replace = TRUE) - + # Create numeric weights numeric_weights <- runif(n, 0.5, 2.0) - + # Create equivalent psw objects psw_ate <- propensity::psw(numeric_weights, estimand = "ate") psw_att <- propensity::psw(numeric_weights, estimand = "att") psw_ato <- propensity::psw(numeric_weights, estimand = "ato") - + # Test bal_smd equivalence smd_numeric <- bal_smd(x, g, weights = numeric_weights) smd_psw_ate <- bal_smd(x, g, weights = psw_ate) - smd_psw_att <- bal_smd(x, g, weights = psw_att) + smd_psw_att <- bal_smd(x, g, weights = psw_att) smd_psw_ato <- bal_smd(x, g, weights = psw_ato) - + expect_identical(smd_numeric, smd_psw_ate) expect_identical(smd_numeric, smd_psw_att) expect_identical(smd_numeric, smd_psw_ato) - + # Test bal_vr equivalence vr_numeric <- bal_vr(x, g, weights = numeric_weights) vr_psw_ate <- bal_vr(x, g, weights = psw_ate) vr_psw_att <- bal_vr(x, g, weights = psw_att) vr_psw_ato <- bal_vr(x, g, weights = psw_ato) - + expect_identical(vr_numeric, vr_psw_ate) expect_identical(vr_numeric, vr_psw_att) expect_identical(vr_numeric, vr_psw_ato) - + # Test bal_ks equivalence ks_numeric <- bal_ks(x, g, weights = numeric_weights) ks_psw_ate <- bal_ks(x, g, weights = psw_ate) ks_psw_att <- bal_ks(x, g, weights = psw_att) ks_psw_ato <- bal_ks(x, g, weights = psw_ato) - + expect_identical(ks_numeric, ks_psw_ate) expect_identical(ks_numeric, ks_psw_att) expect_identical(ks_numeric, ks_psw_ato) @@ -92,29 +100,29 @@ test_that("numeric and psw weights produce identical results", { test_that("numeric and psw weights produce identical results for correlation", { skip_if_not_installed("propensity") - + # Create test data set.seed(42) n <- 100 x <- rnorm(n) y <- 2 * x + rnorm(n, sd = 0.5) - + # Create numeric weights numeric_weights <- runif(n, 0.5, 2.0) - + # Create equivalent psw object psw_weights <- propensity::psw(numeric_weights, estimand = "ate") - + # Test bal_corr equivalence corr_numeric <- bal_corr(x, y, weights = numeric_weights) corr_psw <- bal_corr(x, y, weights = psw_weights) - + expect_identical(corr_numeric, corr_psw) }) test_that("numeric and psw weights produce identical results for energy balance", { skip_if_not_installed("propensity") - + # Create test data set.seed(42) n <- 100 @@ -123,60 +131,60 @@ test_that("numeric and psw weights produce identical results for energy balance" x2 = rnorm(n) ) g <- sample(c(0, 1), n, replace = TRUE) - + # Create numeric weights numeric_weights <- runif(n, 0.5, 2.0) - + # Create equivalent psw object psw_weights <- propensity::psw(numeric_weights, estimand = "ate") - + # Test bal_energy equivalence energy_numeric <- bal_energy(covariates, g, weights = numeric_weights) energy_psw <- bal_energy(covariates, g, weights = psw_weights) - + expect_identical(energy_numeric, energy_psw) }) test_that("check_balance works equivalently with numeric and psw weights", { skip_if_not_installed("propensity") - + # Use nhefs_weights dataset data(nhefs_weights) - + # Extract numeric data from existing psw weights for comparison w_ate_numeric <- vctrs::vec_data(nhefs_weights$w_ate) w_att_numeric <- vctrs::vec_data(nhefs_weights$w_att) - + # Create new psw objects from the same numeric data w_ate_psw <- propensity::psw(w_ate_numeric, estimand = "ate") w_att_psw <- propensity::psw(w_att_numeric, estimand = "att") - + # Add numeric weights to test data test_data <- nhefs_weights[1:100, ] test_data$w_ate_numeric <- w_ate_numeric[1:100] test_data$w_att_numeric <- w_att_numeric[1:100] test_data$w_ate_psw <- w_ate_psw[1:100] test_data$w_att_psw <- w_att_psw[1:100] - + # Test single weight comparison result_numeric <- check_balance( - test_data, - c(age, wt71), - qsmk, + test_data, + c(age, wt71), + qsmk, .wts = w_ate_numeric, .metrics = "smd", include_observed = FALSE ) - + result_psw <- check_balance( test_data, c(age, wt71), qsmk, .wts = w_ate_psw, - .metrics = "smd", + .metrics = "smd", include_observed = FALSE ) - + # Results should be identical except for method names expect_equal(result_numeric$estimate, result_psw$estimate) expect_equal(result_numeric$variable, result_psw$variable) @@ -185,14 +193,14 @@ test_that("check_balance works equivalently with numeric and psw weights", { test_that("mixed numeric and psw weights work in same function call", { skip_if_not_installed("propensity") - + # Use nhefs_weights dataset with mixed weight types data(nhefs_weights) test_data <- nhefs_weights[1:100, ] - + # Create a numeric version of one weight test_data$w_ate_numeric <- vctrs::vec_data(test_data$w_ate) - + # This should work without error - one numeric, one psw result <- check_balance( test_data, @@ -202,7 +210,7 @@ test_that("mixed numeric and psw weights work in same function call", { .metrics = "smd", include_observed = FALSE ) - + expect_s3_class(result, "data.frame") expect_true(nrow(result) > 0) expect_true(all(is.finite(result$estimate))) @@ -212,82 +220,82 @@ test_that("mixed numeric and psw weights work in same function call", { test_that("edge cases work identically for numeric and psw weights", { skip_if_not_installed("propensity") - + # Test with extreme weights set.seed(42) n <- 50 x <- rnorm(n) g <- sample(c(0, 1), n, replace = TRUE) - + # Extreme numeric weights - extreme_numeric <- c(rep(0.001, n/2), rep(100, n/2)) + extreme_numeric <- c(rep(0.001, n / 2), rep(100, n / 2)) extreme_psw <- propensity::psw(extreme_numeric, estimand = "ate") - + # Both should handle extreme weights the same way smd_numeric <- bal_smd(x, g, weights = extreme_numeric) smd_psw <- bal_smd(x, g, weights = extreme_psw) expect_identical(smd_numeric, smd_psw) - + # Test with uniform weights (should be close to unweighted) uniform_numeric <- rep(1, n) uniform_psw <- propensity::psw(uniform_numeric, estimand = "ate") - + smd_unweighted <- bal_smd(x, g) smd_uniform_numeric <- bal_smd(x, g, weights = uniform_numeric) smd_uniform_psw <- bal_smd(x, g, weights = uniform_psw) - + expect_identical(smd_uniform_numeric, smd_uniform_psw) expect_equal(smd_unweighted, smd_uniform_numeric, tolerance = 1e-10) }) test_that("NA handling is identical for numeric and psw weights", { skip_if_not_installed("propensity") - + # Create test data with NAs set.seed(42) n <- 100 x <- rnorm(n) - x[1:5] <- NA # Add some NAs - g <- sample(c(0, 1), n, replace = TRUE) - + x[1:5] <- NA # Add some NAs + g <- sample(c(0, 1), n, replace = TRUE) + # Create weights numeric_weights <- runif(n, 0.5, 2.0) psw_weights <- propensity::psw(numeric_weights, estimand = "ate") - + # Test na.rm = FALSE (should return NA) smd_numeric_na <- bal_smd(x, g, weights = numeric_weights, na.rm = FALSE) smd_psw_na <- bal_smd(x, g, weights = psw_weights, na.rm = FALSE) - + expect_identical(smd_numeric_na, smd_psw_na) expect_true(is.na(smd_numeric_na)) - + # Test na.rm = TRUE smd_numeric_narm <- bal_smd(x, g, weights = numeric_weights, na.rm = TRUE) smd_psw_narm <- bal_smd(x, g, weights = psw_weights, na.rm = TRUE) - + expect_identical(smd_numeric_narm, smd_psw_narm) expect_true(is.finite(smd_numeric_narm)) }) test_that("all helper functions work with both numeric and psw weights", { skip_if_not_installed("propensity") - + # Test weighted_quantile set.seed(42) n <- 100 x <- rnorm(n) - + numeric_weights <- runif(n, 0.5, 2.0) psw_weights <- propensity::psw(numeric_weights, estimand = "ate") - + q_numeric <- weighted_quantile(x, c(0.25, 0.5, 0.75), numeric_weights) q_psw <- weighted_quantile(x, c(0.25, 0.5, 0.75), psw_weights) - + expect_identical(q_numeric, q_psw) - + # Test ess function ess_numeric <- ess(numeric_weights) ess_psw <- ess(psw_weights) - + expect_identical(ess_numeric, ess_psw) -}) \ No newline at end of file +}) diff --git a/tests/testthat/test-plot_balance.R b/tests/testthat/test-plot_balance.R index 5bfed0e..813672d 100644 --- a/tests/testthat/test-plot_balance.R +++ b/tests/testthat/test-plot_balance.R @@ -189,13 +189,13 @@ test_that("plot_balance handles categorical exposures with facet_grid", { .wts = w_cat_ate, .metrics = c("smd", "vr") ) - + p <- plot_balance(balance_cat) expect_s3_class(p, "ggplot") - + # Check that facet_grid was applied for categorical exposure expect_s3_class(p$facet, "FacetGrid") - + # Check that both group_level and metric are in the facet facet_vars <- names(p$facet$params$rows) expect_true("group_level" %in% facet_vars) @@ -212,10 +212,10 @@ test_that("plot_balance handles categorical exposures with single metric", { .wts = w_cat_ate, .metrics = "smd" ) - + p <- plot_balance(balance_cat_single) expect_s3_class(p, "ggplot") - + # Should use facet_wrap for single metric expect_s3_class(p$facet, "FacetWrap") }) @@ -229,12 +229,12 @@ test_that("plot_balance correctly identifies categorical vs binary exposures", { .wts = w_ate, .metrics = c("smd", "vr") ) - + p_binary <- plot_balance(balance_binary) - + # Should use facet_wrap for binary exposure expect_s3_class(p_binary$facet, "FacetWrap") - + # Categorical exposure balance_cat <- check_balance( nhefs_weights, @@ -242,9 +242,9 @@ test_that("plot_balance correctly identifies categorical vs binary exposures", { alcoholfreq_cat, .metrics = c("smd", "vr") ) - + p_cat <- plot_balance(balance_cat) - + # Should use facet_grid for categorical exposure with multiple metrics expect_s3_class(p_cat$facet, "FacetGrid") }) @@ -331,7 +331,7 @@ test_that("plot_balance visual tests", { vlinewidth = 1 ) ) - + # Categorical exposure with multiple metrics balance_cat_multi <- check_balance( nhefs_weights, @@ -340,12 +340,12 @@ test_that("plot_balance visual tests", { .wts = w_cat_ate, .metrics = c("smd", "vr", "ks") ) - + expect_doppelganger( "balance-plot-categorical-multi-metric", plot_balance(balance_cat_multi) ) - + # Categorical exposure with single metric balance_cat_single <- check_balance( nhefs_weights, @@ -354,12 +354,12 @@ test_that("plot_balance visual tests", { .wts = c(w_cat_ate, w_cat_att_2_3wk), .metrics = "smd" ) - + expect_doppelganger( "balance-plot-categorical-single-metric", plot_balance(balance_cat_single) ) - + # Categorical exposure without weights balance_cat_observed <- check_balance( nhefs_weights, @@ -367,7 +367,7 @@ test_that("plot_balance visual tests", { alcoholfreq_cat, .metrics = c("smd", "vr") ) - + expect_doppelganger( "balance-plot-categorical-observed", plot_balance(balance_cat_observed) @@ -385,12 +385,12 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = c("smd", "vr") ) - + expect_doppelganger( "balance-plot-categorical-ref-daily", plot_balance(balance_cat_ref) ) - + # Categorical with energy metric included balance_cat_energy <- check_balance( nhefs_weights, @@ -399,7 +399,7 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = c("smd", "energy") ) - + # Suppress ggplot2's "Removed 2 rows containing missing values" warning # This happens when energy distance produces NA values for small sample sizes suppressWarnings( @@ -408,7 +408,7 @@ test_that("plot_balance visual tests - more categorical scenarios", { plot_balance(balance_cat_energy) ) ) - + # Categorical with fixed scales balance_cat_fixed <- check_balance( nhefs_weights, @@ -417,12 +417,12 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = c("smd", "vr", "ks") ) - + expect_doppelganger( "balance-plot-categorical-fixed-scales", plot_balance(balance_cat_fixed, facet_scales = "fixed") ) - + # Categorical without absolute SMD balance_cat_multi_test <- check_balance( nhefs_weights, @@ -431,12 +431,12 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = c("smd", "vr", "ks") ) - + expect_doppelganger( "balance-plot-categorical-no-abs-smd", plot_balance(balance_cat_multi_test, abs_smd = FALSE) ) - + # Categorical with multiple ATT weights balance_cat_multi_att <- check_balance( nhefs_weights, @@ -445,12 +445,12 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = c(w_cat_att_none, w_cat_att_2_3wk, w_cat_att_daily), .metrics = "smd" ) - + expect_doppelganger( "balance-plot-categorical-multi-att", plot_balance(balance_cat_multi_att) ) - + # Categorical with only KS metric balance_cat_ks <- check_balance( nhefs_weights, @@ -459,12 +459,12 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = "ks" ) - + expect_doppelganger( "balance-plot-categorical-ks-only", plot_balance(balance_cat_ks) ) - + # Categorical with custom vline for SMD only balance_cat_vline <- check_balance( nhefs_weights, @@ -473,26 +473,40 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = w_cat_ate, .metrics = "smd" ) - + expect_doppelganger( "balance-plot-categorical-custom-vline", - plot_balance(balance_cat_vline, vline_xintercept = 0.05, vline_color = "red") + plot_balance( + balance_cat_vline, + vline_xintercept = 0.05, + vline_color = "red" + ) ) - + # Categorical with many variables to test scrolling/layout balance_cat_many_vars <- check_balance( nhefs_weights, - c(age, wt71, sex, race, education, smokeintensity, smokeyrs, exercise, active), + c( + age, + wt71, + sex, + race, + education, + smokeintensity, + smokeyrs, + exercise, + active + ), alcoholfreq_cat, .wts = w_cat_ate, .metrics = c("smd", "vr") ) - + expect_doppelganger( "balance-plot-categorical-many-vars", plot_balance(balance_cat_many_vars) ) - + # Categorical comparing observed vs multiple weights balance_cat_compare <- check_balance( nhefs_weights, @@ -501,7 +515,7 @@ test_that("plot_balance visual tests - more categorical scenarios", { .wts = c(w_cat_ate, w_cat_ato, w_cat_atm), .metrics = c("smd", "vr") ) - + expect_doppelganger( "balance-plot-categorical-compare-weights", plot_balance(balance_cat_compare) diff --git a/tests/testthat/test-plot_ess.R b/tests/testthat/test-plot_ess.R new file mode 100644 index 0000000..609a4d6 --- /dev/null +++ b/tests/testthat/test-plot_ess.R @@ -0,0 +1,152 @@ +test_that("plot_ess creates a ggplot object", { + p <- plot_ess(nhefs_weights, .wts = w_ate) + + expect_s3_class(p, "ggplot") + expect_s3_class(p, "gg") +}) + +test_that("plot_ess works with raw data", { + p <- plot_ess(nhefs_weights, .wts = c(w_ate, w_att)) + + expect_s3_class(p, "ggplot") + # Check that data was computed + expect_true("ess_pct" %in% names(p$data)) +}) + +test_that("plot_ess works with pre-computed ESS data", { + ess_data <- check_ess(nhefs_weights, .wts = c(w_ate, w_att)) + p <- plot_ess(ess_data) + + expect_s3_class(p, "ggplot") + expect_equal(nrow(p$data), nrow(ess_data)) +}) + +test_that("plot_ess works without groups", { + p <- plot_ess(nhefs_weights, .wts = c(w_ate, w_att)) + + # Should have single fill color (no group aesthetic) + expect_true(!"fill" %in% names(p$mapping)) + + # Check layers + expect_true(any(sapply(p$layers, function(l) inherits(l$geom, "GeomCol")))) +}) + +test_that("plot_ess works with groups", { + p <- plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) + + # Should have fill aesthetic for groups + expect_true("fill" %in% names(p$mapping)) + + # Check for dodged position + col_layer <- p$layers[[which(sapply(p$layers, function(l) inherits(l$geom, "GeomCol")))]] + expect_s3_class(col_layer$position, "PositionDodge") +}) + +test_that("plot_ess includes reference line", { + p <- plot_ess(nhefs_weights, .wts = w_ate) + + # Check for horizontal line at 100 + hline_layer <- sapply(p$layers, function(l) inherits(l$geom, "GeomHline")) + expect_true(any(hline_layer)) + + # Check that line is at 100 + hline_idx <- which(hline_layer)[1] + expect_equal(p$layers[[hline_idx]]$data$yintercept, 100) +}) + +test_that("plot_ess labels work correctly", { + p_with_labels <- plot_ess(nhefs_weights, .wts = w_ate, show_labels = TRUE) + p_without_labels <- plot_ess(nhefs_weights, .wts = w_ate, show_labels = FALSE) + + # Check for text layer + text_layer_with <- sapply(p_with_labels$layers, function(l) inherits(l$geom, "GeomText")) + text_layer_without <- sapply(p_without_labels$layers, function(l) inherits(l$geom, "GeomText")) + + expect_true(any(text_layer_with)) + expect_false(any(text_layer_without)) +}) + +test_that("plot_ess y-axis uses percent scale", { + p <- plot_ess(nhefs_weights, .wts = w_ate) + + # Check that y-axis has percent labels + y_scale <- p$scales$get_scales("y") + expect_true(is.function(y_scale$labels)) +}) + +test_that("plot_ess handles continuous groups", { + p <- plot_ess(nhefs_weights, .wts = w_ate, .group = age, n_tiles = 4) + + expect_s3_class(p, "ggplot") + expect_true("fill" %in% names(p$mapping)) +}) + +test_that("plot_ess customization works", { + p <- plot_ess( + nhefs_weights, + .wts = w_ate, + fill_color = "red", + alpha = 0.5, + reference_line_color = "blue", + reference_line_type = "solid" + ) + + expect_s3_class(p, "ggplot") + + # Check fill color for non-grouped plot + col_layer <- p$layers[[which(sapply(p$layers, function(l) inherits(l$geom, "GeomCol")))]] + expect_equal(col_layer$aes_params$fill, "red") + expect_equal(col_layer$aes_params$alpha, 0.5) + + # Check reference line + hline_layer <- p$layers[[which(sapply(p$layers, function(l) inherits(l$geom, "GeomHline")))]] + expect_equal(hline_layer$aes_params$colour, "blue") + expect_equal(hline_layer$aes_params$linetype, "solid") +}) + +test_that("plot_ess has correct labels", { + p <- plot_ess(nhefs_weights, .wts = w_ate) + + expect_equal(p$labels$x, "method") + expect_equal(p$labels$y, "effective sample size (%)") +}) + +test_that("plot_ess y-limits are appropriate", { + # Create data with low ESS + test_df <- data.frame( + wts = c(10, 0.1, 0.1, 0.1) # Very unequal weights + ) + + p <- plot_ess(test_df, .wts = wts, include_observed = FALSE) + + # Y-axis should start at 0 and go above the max value + y_scale <- p$scales$get_scales("y") + expect_equal(y_scale$limits[1], 0) + expect_true(y_scale$limits[2] >= 105) # At least to 105% +}) + +test_that("plot_ess snapshot tests", { + # Basic plot + expect_doppelganger( + "plot_ess_basic", + plot_ess(nhefs_weights, .wts = c(w_ate, w_att)) + ) + + # Plot with groups + expect_doppelganger( + "plot_ess_groups", + plot_ess(nhefs_weights, .wts = c(w_ate, w_att), .group = qsmk) + ) + + # Plot without labels + expect_doppelganger( + "plot_ess_no_labels", + plot_ess(nhefs_weights, .wts = c(w_ate, w_att), show_labels = FALSE) + ) + + # Plot with continuous groups + expect_doppelganger( + "plot_ess_continuous", + plot_ess(nhefs_weights, .wts = w_ate, .group = age, n_tiles = 3) + ) +}) \ No newline at end of file diff --git a/tests/testthat/test-plot_mirror_distributions.R b/tests/testthat/test-plot_mirror_distributions.R index 1261e1a..3507fe5 100644 --- a/tests/testthat/test-plot_mirror_distributions.R +++ b/tests/testthat/test-plot_mirror_distributions.R @@ -159,12 +159,12 @@ test_that("plot_mirror_distributions validates inputs", { plot_mirror_distributions(nhefs_weights), "halfmoon_arg_error" ) - + expect_halfmoon_error( plot_mirror_distributions(nhefs_weights, age), "halfmoon_arg_error" ) - + # Non-existent column expect_halfmoon_error( plot_mirror_distributions(nhefs_weights, nonexistent, qsmk), @@ -251,10 +251,10 @@ test_that("plot_mirror_distributions works with categorical exposures", { alcoholfreq_cat, type = "density" ) - + expect_s3_class(p_cat, "ggplot") expect_doppelganger("categorical exposure density", p_cat) - + # Categorical with histogram p_cat_hist <- plot_mirror_distributions( nhefs_weights, @@ -263,9 +263,9 @@ test_that("plot_mirror_distributions works with categorical exposures", { type = "histogram", bins = 20 ) - + expect_doppelganger("categorical exposure histogram", p_cat_hist) - + # Categorical with custom reference p_cat_ref <- plot_mirror_distributions( nhefs_weights, @@ -274,7 +274,7 @@ test_that("plot_mirror_distributions works with categorical exposures", { reference_group = "daily", type = "density" ) - + expect_doppelganger("categorical custom reference", p_cat_ref) }) @@ -288,10 +288,10 @@ test_that("plot_mirror_distributions works with categorical exposures and weight reference_group = "none", bins = 20 ) - + expect_s3_class(p_cat_wt, "ggplot") expect_doppelganger("categorical with weights", p_cat_wt) - + # Categorical with multiple weights p_cat_multi <- plot_mirror_distributions( nhefs_weights, @@ -301,9 +301,9 @@ test_that("plot_mirror_distributions works with categorical exposures and weight reference_group = "none", type = "density" ) - + expect_doppelganger("categorical multiple weights", p_cat_multi) - + # Without unweighted p_cat_no_obs <- plot_mirror_distributions( nhefs_weights, @@ -313,7 +313,7 @@ test_that("plot_mirror_distributions works with categorical exposures and weight include_unweighted = FALSE, type = "density" ) - + expect_doppelganger("categorical no unweighted", p_cat_no_obs) }) @@ -328,7 +328,7 @@ test_that("plot_mirror_distributions validates categorical reference group", { ), "halfmoon_reference_error" ) - + # Numeric reference out of range expect_halfmoon_error( plot_mirror_distributions( diff --git a/tests/testthat/test-plot_qq.R b/tests/testthat/test-plot_qq.R index 4c728ff..91cd351 100644 --- a/tests/testthat/test-plot_qq.R +++ b/tests/testthat/test-plot_qq.R @@ -52,7 +52,7 @@ test_that("plot_qq validates missing arguments", { plot_qq(nhefs_weights), "halfmoon_arg_error" ) - + expect_halfmoon_error( plot_qq(nhefs_weights, age), "halfmoon_arg_error"