R/plot_continuous.R

In jsdtools: Jensen-Shannon Divergence Estimation, Confidence Intervals, and Distribution Plots

#' Plot two continuous distributions
#'
#' @param x Numeric vector for group 1.
#' @param y Numeric vector for group 2.
#' @param group_names Group labels.
#' @param bins Approximate number of histogram bins.
#' @param style One of `"both"`, `"hist"`, or `"density"`.
#' @param main Plot title.
#' @param xlab X-axis label.
#' @param ylab Y-axis label.
#' @param col_x Fill color for group 1.
#' @param col_y Fill color for group 2.
#' @param line_col_x Line color for group 1 density.
#' @param line_col_y Line color for group 2 density.
#' @param lwd Line width for density curves.
#' @param na_rm Logical; remove missing values?
#' @param show_jsd Logical; whether to display JSD on the plot.
#' @param jsd_digits Number of digits for displayed JSD.
#'
#' @return Invisibly returns plotting data.
#' @export
#' @importFrom grDevices rgb adjustcolor
#' @importFrom graphics axis box hist legend lines par points rect segments text
#' @importFrom stats density
plot_continuous <- function(x, y,
                            group_names = c("Group 1", "Group 2"),
                            bins = 30,
                            style = c("both", "hist", "density"),
                            main = "Two-group raw distributions",
                            xlab = "Value",
                            ylab = "Density",
                            col_x = rgb(0.2, 0.4, 0.8, 0.4),
                            col_y = rgb(0.8, 0.2, 0.2, 0.4),
                            line_col_x = "#2F5FB3",
                            line_col_y = "#CC3333",
                            lwd = 2,
                            na_rm = TRUE,
                            show_jsd = TRUE,
                            jsd_digits = 3) {
  style <- match.arg(style)

  cleaned <- validate_xy(x, y, min_n = 2, na_rm = na_rm, finite_only = TRUE)
  x <- cleaned$x
  y <- cleaned$y

  all_vals <- c(x, y)
  breaks <- pretty(range(all_vals), n = bins)

  hx <- hist(x, breaks = breaks, plot = FALSE)
  hy <- hist(y, breaks = breaks, plot = FALSE)

  dx <- density(x)
  dy <- density(y)

  ymax <- max(
    if (style %in% c("both", "hist")) c(hx$density, hy$density) else 0,
    if (style %in% c("both", "density")) c(dx$y, dy$y) else 0
  )

  jsd_value <- NA_real_
  if (show_jsd) {
    jsd_value <- tryCatch(
      jsd(x, y)$estimate,
      error = function(e) NA_real_
    )
  }

  plot(range(breaks), c(0, ymax * 1.15), type = "n",
       xlab = xlab, ylab = ylab, main = main)

  if (style %in% c("both", "hist")) {
    hist(x, breaks = breaks, freq = FALSE, col = col_x, border = "white", add = TRUE)
    hist(y, breaks = breaks, freq = FALSE, col = col_y, border = "white", add = TRUE)
  }

  if (style %in% c("both", "density")) {
    lines(dx, lwd = lwd, col = line_col_x)
    lines(dy, lwd = lwd, col = line_col_y)
  }

  legend("topright",
         legend = group_names,
         fill = c(col_x, col_y),
         border = NA,
         bty = "n")

  if (show_jsd && is.finite(jsd_value)) {
    usr <- par("usr")
    x_pos <- usr[1] + 0.03 * (usr[2] - usr[1])
    y_pos <- usr[4] - 0.06 * (usr[4] - usr[3])

    text(x_pos, y_pos,
         labels = paste0("JSD = ", formatC(jsd_value, digits = jsd_digits, format = "f")),
         adj = c(0, 1),
         cex = 1)
  }

  invisible(list(
    x_density = dx,
    y_density = dy,
    x_hist = hx,
    y_hist = hy,
    jsd = jsd_value
  ))
}