R/cits_plot_counterfactual.R

In citsr: Controlled Interrupted Time Series Analysis and Visualization

#' CITS Automatic Analysis and Counterfactual Visualization
#'
#' Visualizes the results of a controlled interrupted time series (CITS) model
#' fitted using `cits()`, and generates a counterfactual trajectory for the
#' treatment group (E = 1) by setting the intervention indicator (`I`) to 0
#' after the intervention time. The function displays observed points, fitted
#' values, 95% confidence intervals, and a dashed counterfactual line, along
#' with a vertical marker for the intervention.
#'
#' @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 I_col Name of the intervention indicator variable (string). Corresponds to `I` in `cits()`.
#' @param E_col Name of the group indicator variable (string). Corresponds to `E` in `cits()`.
#' @param intervention_time Numeric. Time point at which the intervention occurs.
#'
#' @return A ggplot object showing observed points, fitted lines, confidence ribbons,
#'         a counterfactual trajectory for the treatment group, and an intervention marker.
#'
#' @importFrom stats qnorm predict formula
#' @importFrom ggplot2 ggplot aes geom_point geom_line geom_ribbon geom_vline annotate
#'         labs scale_color_manual scale_fill_manual theme_minimal
#'
#' @examples
#' 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_cf(res, intervention_time = 10)
#'
#' @export
plot_cits_result_cf <- function(res, y_col = "y", T_col = "T",
                                I_col = "I", E_col = "E", intervention_time) {

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

  z <- qnorm(0.975)
  df$lwr <- df$fitted - z * df$se
  df$upr <- df$fitted + z * df$se

  df[[E_col]] <- factor(df[[E_col]])

  # Counterfactual dataset
  df_cf <- df
  mask <- df_cf[[E_col]] == 1 & df_cf[[T_col]] >= intervention_time
  df_cf[[I_col]][mask] <- 0

  df_cf$TI  <- df_cf[[T_col]] * df_cf[[I_col]]
  df_cf$ET  <- as.numeric(df_cf[[E_col]]) * df_cf[[T_col]]
  df_cf$EI  <- as.numeric(df_cf[[E_col]]) * df_cf[[I_col]]
  df_cf$ETI <- as.numeric(df_cf[[E_col]]) * df_cf[[T_col]] * df_cf[[I_col]]

  model_vars <- all.vars(formula(res$model))[-1]
  df_cf_pred <- df_cf[, model_vars, drop = FALSE]
  for (v in model_vars) df_cf_pred[[v]] <- as.numeric(df_cf_pred[[v]])

  df_cf$fitted_cf <- predict(res$model, newdata = df_cf_pred)

  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_line(data = df_cf[mask, ], aes(y = fitted_cf),
              color = "blue", linetype = "dashed", linewidth = 0.8) +
    geom_ribbon(aes(ymin = lwr, ymax = upr, fill = .data[[E_col]]),
                alpha = 0.2, color = NA) +
    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") +
    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()

  return(plt)
}