R/cits_plot.R

In citsr: Controlled Interrupted Time Series Analysis and Visualization

#' CITS Automatic Analysis and Visualization
#'
#' Visualizes the results of a controlled interrupted time series (CITS) model
#' fitted using `cits()`. The function plots observed data points, fitted values,
#' and 95% confidence intervals. An optional vertical intervention line can also be added.
#'
#' @param res A list returned by `cits()`, containing model output and fitted values.
#' @param y_col Name of the outcome variable (string). Corresponds to `y` in `cits()`.
#' @param T_col Name of the time index variable (string). Corresponds to `T` in `cits()`.
#' @param E_col Name of the group indicator variable (string). Corresponds to `E` in `cits()`.
#' @param intervention_time Optional numeric. If provided, a vertical dashed line is drawn
#'        at this time to mark the intervention point.
#'
#' @return A ggplot object showing observed points, fitted lines, confidence ribbons,
#'         and (optionally) the intervention line.
#'
#' @importFrom stats qnorm
#' @importFrom ggplot2 ggplot aes geom_point geom_line geom_ribbon geom_vline annotate
#'         labs scale_color_manual scale_fill_manual theme_minimal
#'
#' @examples
#' # Synthetic example
#' df <- data.frame(
#'   T = 1:100,
#'   E = rep(c(0,1), each = 100),
#'   I = c(rep(0,50), rep(1,50), rep(0,50), rep(1,50)),
#'   y = rnorm(200)
#' )
#'
#' # Use lightweight ARMA search for examples (CRAN speed requirement)
#' res <- cits(
#'   df,
#'   y_col = "y",
#'   T_col = "T",
#'   I_col = "I",
#'   E_col = "E",
#'   p_range = 0:1,
#'   q_range = 0:0
#' )
#'
#' plot_cits_result(res)
#'
#' @export
plot_cits_result <- function(res, y_col = "y", T_col = "T", E_col = "E", intervention_time = NULL) {

  df <- as.data.frame(res$data)
  df$fitted <- as.numeric(df$fitted)
  df$se     <- as.numeric(df$se)

  # Compute 95% confidence intervals
  alpha <- 0.05
  z <- qnorm(1 - alpha / 2)
  df$lwr <- df$fitted - z * df$se
  df$upr <- df$fitted + z * df$se

  # Ensure group variable is a factor
  df[[E_col]] <- factor(df[[E_col]])

  # Base plot
  plt <- ggplot(df, aes(x = .data[[T_col]], y = .data[[y_col]], color = .data[[E_col]])) +
    geom_point(alpha = 0.5) +
    geom_line(aes(y = fitted), linewidth = 0.8) +
    geom_ribbon(aes(ymin = lwr, ymax = upr, fill = .data[[E_col]]),
                alpha = 0.2, color = NA) +
    labs(x = "Time", y = "Outcome", color = "Group", fill = "Group") +
    scale_color_manual(values = c("0" = "darkred", "1" = "blue")) +
    scale_fill_manual(values = c("0" = "pink", "1" = "skyblue")) +
    theme_minimal()

  # Optional intervention marker
  if (!is.null(intervention_time)) {
    plt <- plt +
      geom_vline(xintercept = intervention_time, linetype = "dashed", color = "black") +
      annotate("text",
               x = intervention_time,
               y = max(df[[y_col]]) * 1.02,
               label = "Intervention",
               vjust = 0, hjust = 0.5, fontface = "bold")
  }

  return(plt)
}