R/plot.feasible.r

In CICI: Causal Inference with Continuous (Multiple Time Point) Interventions

plot.feasible <- function(x,
                          x.axis   = c("strategy", "time"),
                          which    = c("feasible", "nonoverlap"),
                          facet    = c("none", "time", "strategy"),
                          ...) {
  dat <- attr(x, "summary")
  if (is.null(dat)) stop("No summary available in object (attr 'summary' not found).")
  
  x0 <- x1 <- Abar <- y0 <- y1 <- NULL
  
  # args
  x.axis <- match.arg(x.axis)
  which  <- match.arg(which)
  facet  <- match.arg(facet)
  
  # mapping structure
  mapping_by_time <- split(dat[, c("Abar", "Strategy")], dat$time)
  .norm_map <- function(df) {
    out <- df[order(df$Abar), , drop = FALSE]
    rownames(out) <- NULL
    out
  }
  base_map <- .norm_map(mapping_by_time[[1]])
  all_identical <- all(vapply(mapping_by_time, function(df) identical(.norm_map(df), base_map), logical(1)))
  one_to_one <- all_identical && length(unique(dat$Abar)) == length(unique(dat$Strategy))
  
  # x / colour
  if (x.axis == "time") {
    x_vals       <- dat$time
    x_breaks     <- sort(unique(dat$time))
    x_label      <- "Time"
    if (one_to_one) {
      color_vals   <- factor(dat$Abar)
      legend_title <- "Target"
    } else {
      color_vals   <- factor(dat$Strategy)
      legend_title <- "Strategy"
    }
  } else {
    if (one_to_one) {
      x_vals       <- dat$Abar
      x_breaks     <- sort(unique(dat$Abar))
      x_label      <- "Target"
      color_vals   <- factor(dat$time)
      legend_title <- "Time"
    } else {
      x_vals       <- dat$Strategy
      x_breaks     <- sort(unique(dat$Strategy))
      x_label      <- "Strategy"
      color_vals   <- factor(dat$time)
      legend_title <- "Time"
    }
  }
  
  # y vars
  Feasible <- dat$Feasible
  Low      <- dat$Low
  Abar     <- dat$Abar
  
  # palette
  base_cols <- c("black", "orangered3", "dodgerblue4", "springgreen3",
                 "gold", "greenyellow", "purple")
  ext_cols  <- grDevices::rainbow(25)
  levels_color <- levels(color_vals)
  mycolors  <- rep_len(c(base_cols, ext_cols), length(levels_color))
  names(mycolors) <- levels_color
  
  # theme
  base_theme <- ggplot2::theme_bw() +
    ggplot2::theme(
      plot.title   = ggplot2::element_text(hjust = 0.5, size = 14, face = "bold"),
      axis.title.x = ggplot2::element_text(size = 12),
      axis.text.x  = ggplot2::element_text(size = 10, hjust = 1),
      axis.title.y = ggplot2::element_text(size = 12, angle = 90),
      axis.text.y  = ggplot2::element_text(size = 11),
      legend.text  = ggplot2::element_text(size = 11),
      legend.title = ggplot2::element_text(size = 12, face = "bold"),
      legend.position = "right"
    )
  
  # y-limits from Abar
  abar_unique <- sort(unique(dat$Abar))
  if (is.numeric(dat$Abar)) {
    a_range <- range(abar_unique, na.rm = TRUE)
    f_range <- range(Feasible, na.rm = TRUE)
    use_abar_limits <- length(abar_unique) > 1 &&
      all(is.finite(a_range)) &&
      f_range[1] >= a_range[1] &&
      f_range[2] <= a_range[2]
  } else {
    use_abar_limits <- FALSE
  }
  y_limits_abar <- if (use_abar_limits) a_range else NULL
  y_breaks_abar <- if (use_abar_limits) abar_unique else ggplot2::waiver()
  
  # flags
  x_is_time   <- (x.axis == "time")
  x_is_target <- (!x_is_time && identical(x_label, "Target"))
  xy_limits   <- if (x_is_target) range(x_breaks, na.rm = TRUE) else NULL
  
  # ticks for time
  if (x_is_time) {
    abar_ref_time <- data.frame(
      time       = dat$time,
      Strategy   = dat$Strategy,
      Abar       = dat$Abar,
      color_vals = color_vals
    )
    xb <- sort(unique(x_breaks))
    if (length(xb) > 1) {
      gap_time <- min(diff(xb))
    } else {
      gap_time <- 1
    }
    tick_w_time <- 0.15 * gap_time
    abar_ref_time$x0 <- abar_ref_time$time - tick_w_time / 2
    abar_ref_time$x1 <- abar_ref_time$time + tick_w_time / 2
  }
  
  # ticks for strategy
  if (!x_is_time && !x_is_target) {
    abar_ref_strat <- dat[, c("Strategy", "time", "Abar")]
    abar_ref_strat$color_vals <- factor(abar_ref_strat$time)
    xb2 <- sort(unique(x_breaks))
    if (length(xb2) > 1) {
      gap_strat <- min(diff(xb2))
    } else {
      gap_strat <- 1
    }
    tick_w_strat <- 0.15 * gap_strat
    abar_ref_strat$x0 <- abar_ref_strat$Strategy - tick_w_strat / 2
    abar_ref_strat$x1 <- abar_ref_strat$Strategy + tick_w_strat / 2
  }
  
  # legend labels in 1:1 case, x = time
  if (x_is_time && one_to_one) {
    nT <- length(unique(dat$time))
    target_labels <- vapply(
      levels_color,
      function(a) paste(rep(a, nT), collapse = ","),
      FUN.VALUE = character(1L)
    )
  } else {
    target_labels <- NULL
  }
  
  # diagonal segment data (for x = Target)
  ref_diag <- NULL
  if (x_is_target) {
    x_rng <- if (!is.null(xy_limits)) xy_limits else range(x_vals, na.rm = TRUE)
    y_rng <- if (!is.null(y_limits_abar)) y_limits_abar else range(Feasible, na.rm = TRUE)
    if (all(is.finite(c(x_rng, y_rng)))) {
      diag_min <- max(x_rng[1], y_rng[1])
      diag_max <- min(x_rng[2], y_rng[2])
      if (diag_min < diag_max) {
        ref_diag <- data.frame(
          x0 = diag_min,
          x1 = diag_max,
          y0 = diag_min,
          y1 = diag_max
        )
      }
    }
  }
  
  # subtitle
  subtitle_text <- if (x_is_target) {
    "Diagonal line: Mean feasible = target"
  } else {
    "Ticks: target Abar; points: mean feasible"
  }
  
  # base plot
  p1 <- ggplot2::ggplot(
    dat,
    ggplot2::aes(x = x_vals, y = Feasible,
                 color = color_vals, group = color_vals)
  ) + base_theme
  
  # background diag line (no legend)
  if (x_is_target && !is.null(ref_diag)) {
    p1 <- p1 +
      ggplot2::geom_segment(
        data = ref_diag,
        ggplot2::aes(x = x0, xend = x1,
                     y = y0, yend = y1),
        inherit.aes = FALSE,
        linetype     = "dashed",
        linewidth    = 0.7,
        show.legend  = FALSE
      )
  }
  
  # main layers and scales
  p1 <- p1 +
    ggplot2::geom_line(linewidth = 0.9, alpha = 0.6) +
    ggplot2::geom_point(size = 1.2, alpha = 0.7) +
    ggplot2::scale_color_manual(
      values = mycolors,
      labels = if (!is.null(target_labels)) target_labels else ggplot2::waiver()
    ) +
    ggplot2::scale_x_continuous(
      name   = x_label,
      breaks = x_breaks,
      limits = if (x_is_target) xy_limits else NULL
    ) +
    ggplot2::scale_y_continuous(
      name   = "Mean Feasible",
      breaks = if (use_abar_limits) y_breaks_abar else if (x_is_target) x_breaks else ggplot2::waiver(),
      limits = if (use_abar_limits) y_limits_abar else if (x_is_target) xy_limits else NULL
    ) +
    ggplot2::labs(
      title    = "Feasible vs Target",
      subtitle = subtitle_text
    )
  
  # ticks
  if (x_is_time) {
    p1 <- p1 +
      ggplot2::geom_segment(
        data = abar_ref_time,
        ggplot2::aes(x = x0, xend = x1,
                     y = Abar, yend = Abar,
                     color = color_vals),
        inherit.aes = FALSE,
        linewidth = 0.7
      )
  } else if (!x_is_target) {
    p1 <- p1 +
      ggplot2::geom_segment(
        data = abar_ref_strat,
        ggplot2::aes(x = x0, xend = x1,
                     y = Abar, yend = Abar,
                     color = color_vals,
                     group = color_vals),
        inherit.aes = FALSE,
        linewidth = 0.7
      )
  }
  
  # guides
  p1 <- p1 +
    ggplot2::guides(
      color    = ggplot2::guide_legend(title = legend_title),
      linetype = "none"
    )
  
  # non-overlap plot
  p2 <- ggplot2::ggplot(dat, ggplot2::aes(x = x_vals, y = Low,
                                          color = color_vals, group = color_vals)) +
    ggplot2::geom_line(linewidth = 0.9, alpha = 0.6) +
    ggplot2::geom_point(size = 1.2, alpha = 0.7) +
    ggplot2::scale_color_manual(values = mycolors) +
    ggplot2::scale_x_continuous(name = x_label, breaks = x_breaks) +
    ggplot2::scale_y_continuous("Non-overlap Ratio", limits = c(0, 1), breaks = seq(0, 1, by = 0.2)) +
    ggplot2::ggtitle("Non-overlap Ratio") +
    base_theme +
    ggplot2::guides(color = ggplot2::guide_legend(title = legend_title))
  
  # faceting
  if (facet == "time") {
    p1 <- p1 + ggplot2::facet_wrap(~ time, labeller = ggplot2::label_both)
    p2 <- p2 + ggplot2::facet_wrap(~ time, labeller = ggplot2::label_both)
  } else if (facet == "strategy") {
    p1 <- p1 + ggplot2::facet_wrap(~ Strategy, labeller = ggplot2::label_both)
    p2 <- p2 + ggplot2::facet_wrap(~ Strategy, labeller = ggplot2::label_both)
  }
  
  # output
  if (which == "feasible") {
    suppressWarnings(print(p1))
    invisible(p1)
  } else {
    suppressWarnings(print(p2))
    invisible(p2)
  }
}