R/Chi_test.R

In autotestR: Automated Functions for Basic Statistical Tests

#' Chi-square test
#'
#' Applies the Pearson chi-square test or Fisher's exact test to assess association
#' between two categorical variables.
#'
#' @param x Categorical vector or data frame with two columns (group 1 and group 2).
#' @param y Categorical vector (group 2). Required if x is a vector.
#' @param title Plot title (string). Default: "Chi-square Test".
#' @param xlab X-axis label in the plot (string). Default: NULL (uses variable name).
#' @param ylab Y-axis label in the plot (string). Default: "Proportion".
#' @param style Plot style generated by the function.
#' @param show_table Logical. If TRUE, prints the contingency table to the console.
#'        Default: TRUE.
#' @param help Logical. If TRUE, displays a detailed explanation of the function.
#'        Default: FALSE.
#' @param verbose Logical. If TRUE, prints messages about the test and expected
#'        frequencies. Default: TRUE.
#' @return Test result and contingency table.
#' @export
#'
#' @examples
#'
#' data <- data.frame(
#'   control = c(rep("healthy", 50), rep("sick", 150)),
#'   treatment = c(rep("healthy", 100), rep("sick", 100))
#' )
#' test.chi(data)

test.chi <- function(x, y = NULL,
                     title = "Chi-square Test",
                     xlab = NULL,
                     ylab = "Proportion",
                     style = c("stacked", "barplot", "mosaic", "pie"),
                     show_table = TRUE,
                     help = FALSE,
                     verbose = TRUE) {

  style <- match.arg(style)

  # Help
  if (help || missing(x)) {
    if (verbose) {
      message(
        "Function test.chi()

Description:
  Performs Pearson's Chi-square test to evaluate association between
  two categorical variables.

When to use:
  - Two categorical variables with large sample sizes
  - Contingency tables may have more than 2 categories per variable
  - Expected cell frequencies should be >= 5

Limitations:
  - Not recommended for small tables; use 'fisher_test()' instead.

Example:
  data <- data.frame(
    control = c(rep('healthy', 50), rep('sick', 150)),
    treatment = c(rep('healthy', 100), rep('sick', 100))
  )
  test.chi(data)
"
      )
    }
    return(invisible(NULL))
  }

  # Required packages
  required_packages <- c("ggplot2", "dplyr", "tidyr", "vcd")
  for (pkg in required_packages) {
    if (!requireNamespace(pkg, quietly = TRUE)) {
      stop(
        paste0(
          "Package '", pkg,
          "' is not installed. Install it with install.packages('", pkg, "')"
        ),
        call. = FALSE
      )
    }
  }

  # Case 1: x is a data frame with two columns
  if (is.data.frame(x)) {
    if (ncol(x) != 2) {
      stop(
        "The data frame must contain exactly two categorical columns.",
        call. = FALSE
      )
    }

    # Store original column names
    column_names <- colnames(x)

    # Convert to long format
    data_long <- tidyr::pivot_longer(
      x,
      cols = tidyselect::everything(),
      names_to = "group",
      values_to = "category"
    )

    # Preserve original column names as factor levels
    data_long$group <- factor(
      data_long$group,
      levels = column_names,
      labels = column_names
    )

    group <- data_long$group
    category <- data_long$category

    name_x <- column_names[1]
    name_y <- column_names[2]

  } else {
    # Case 2: x and y are vectors
    if (is.null(y)) {
      stop("Argument 'y' must be provided if 'x' is a vector.", call. = FALSE)
    }
    if (length(x) != length(y)) {
      stop("Both variables must have the same length.", call. = FALSE)
    }

    group <- x
    category <- y

    name_x <- deparse(substitute(x))
    name_y <- deparse(substitute(y))
    name_x <- sub(".*\\$", "", name_x)
    name_y <- sub(".*\\$", "", name_y)
  }

  if (is.null(xlab)) xlab <- name_x
  if (is.null(ylab)) ylab <- "Proportion"

  # Contingency table
  contingency_table <- table(group, category)

  if (verbose && show_table) {
    message("Observed contingency table:")
    print(contingency_table)
  }

  # Chi-square test
  test <- suppressWarnings(stats::chisq.test(contingency_table))
  expected_freq <- test$expected

  small_exp <- any(expected_freq < 5)

  if (verbose && small_exp) {
    warning(
      "Some expected frequencies < 5. Chi-square approximation may be unreliable. ",
      "Consider Fisher's exact test."
    )
  }

  # -----------------------------
  # Effect size: Cramér's V
  # -----------------------------

  v <- .cramers_v(contingency_table)
  boot_v <- .boot_cramers_v(contingency_table)


  # -----------------------------
  # Data preparation for plotting
  # -----------------------------
  df_plot <- data.frame(group = group, category = category)

  df_prop <- df_plot |>
    dplyr::group_by(group, category) |>
    dplyr::summarise(n = dplyr::n(), .groups = "drop") |>
    dplyr::group_by(group) |>
    dplyr::mutate(prop = n / sum(n))

  # Subtitle
  subtitle_text <- .make_subtitle_chi(
    cramers_v = v,
    p_value = test$p.value,
    small_expected = small_exp
  )

  # Labels and colors
  vivid_colors <- scales::hue_pal()(length(unique(df_prop$group)))

  # --------------------------
  # STYLE 1 (Stacked bar plot)
  # --------------------------
  if (style == "stacked") {
    g <- ggplot2::ggplot(
      df_prop,
      ggplot2::aes(x = group, y = prop, fill = category)
    ) +
      ggplot2::geom_bar(stat = "identity", color = NA) +
      ggplot2::scale_fill_manual(values = vivid_colors) +
      ggplot2::labs(
        title = title,
        subtitle = subtitle_text,
        x = "",
        y = ylab,
        fill = name_y
      ) +
      ggplot2::theme_minimal(base_size = 12) +
      ggplot2::theme(
        legend.position = "right",
        axis.text.x = ggplot2::element_text(
          angle = 45, hjust = 1, size = 12
        )
      )
  }

  # --------------------------
  # STYLE 2 (Side-by-side bars)
  # --------------------------
  if (style == "barplot") {
    g <- ggplot2::ggplot(
      df_prop,
      ggplot2::aes(x = group, y = prop, fill = category)
    ) +
      ggplot2::geom_bar(
        stat = "identity",
        position = ggplot2::position_dodge(width = 0.8)
      ) +
      ggplot2::scale_fill_manual(values = vivid_colors) +
      ggplot2::labs(
        title = title,
        subtitle = subtitle_text,
        x = "",
        y = ylab,
        fill = name_y
      ) +
      ggplot2::theme_minimal(base_size = 12) +
      ggplot2::theme(
        legend.position = "right",
        axis.text.x = ggplot2::element_text(
          angle = 45, hjust = 1, size = 12
        )
      )
  }

  # --------------------------
  # STYLE 3 (Mosaic plot)
  # --------------------------
  if (style == "mosaic") {
    if (!requireNamespace("vcd", quietly = TRUE)) {
      stop(
        "Style = 'mosaic' requires the 'vcd' package. Install it with install.packages('vcd')"
      )
    }

    vcd::mosaic(
      contingency_table,
      shade = TRUE,
      legend = TRUE,
      main = paste0(title, "\n", subtitle_text)
    )
  }

  # --------------------------
  # STYLE 4 (Pie chart)
  # --------------------------
  if (style == "pie") {

    g <- ggplot2::ggplot(
      df_prop,
      ggplot2::aes(x = "", y = prop, fill = category)
    ) +
      ggplot2::geom_bar(stat = "identity", width = 1) +
      ggplot2::coord_polar("y") +
      ggplot2::facet_wrap(~ group) +
      ggplot2::scale_fill_manual(values = vivid_colors) +
      ggplot2::theme_void(base_size = 12) +
      ggplot2::labs(
        title = title,
        subtitle = subtitle_text,
        fill = name_y
      ) +
      ggplot2::theme(
        plot.title = ggplot2::element_text(
          hjust = 0.5,
          size = 14
        ),
        plot.subtitle = ggplot2::element_text(
          hjust = 0.5,
          size = 11,
          margin = ggplot2::margin(b = 10)
        ),
        strip.text = ggplot2::element_text(
          size = 12
        )
      )
  }

  if (style != "mosaic") print(g)

  # --------------------------
  # Output
  # --------------------------
  obj <- list(
    type = "Chi-square",
    statistic = as.numeric(test$statistic),
    df = as.numeric(test$parameter),
    p = test$p.value,
    cramers_v = v,
    cramers_ci = c(
      boot_v$ci_low,
      boot_v$ci_high
    ),
    expected = expected_freq,
    table = contingency_table,
    small_expected = small_exp,
    data = df_plot
  )

  # -----------------------------
  # Return
  # -----------------------------

  if (verbose) {

    .print_header("Chi-square Test")

    .print_block("Statistics", function() {

      cat(
        "Chi-square = ",
        round(test$statistic, 3),
        " | df = ",
        test$parameter,
        " | p = ",
        .format_pval(test$p.value),
        "\n",
        sep = ""
      )

      cat(
        "Cramer's V = ",
        round(v, 3),
        " [",
        round(boot_v$ci_low, 3), ", ",
        round(boot_v$ci_high, 3),
        "]\n",
        sep = ""
      )

      if (small_exp) {
        cat(
          "Note: Some expected counts < 5 (interpret with caution)\n"
        )
      }

    })
  }

  return(invisible(list(result = obj)))

  }