Invisibly return data in plot (#94)

r-package/maq/DESCRIPTION

@@ -22,6 +22,7 @@ Depends:
     R (>= 3.5.0)
 Suggests:
     grf (>= 2.3.0),
+    ggplot2,
     testthat (>= 3.0.0)
 SystemRequirements: GNU make
 URL: https://github.com/grf-labs/maq

r-package/maq/R/plot.R

@@ -17,6 +17,39 @@
 #'  `max(floor(length(path.length) / 1000), 1)` where path.length is the size of the
 #'  grid underlying the estimated Qini curve.
 #'
+#' @return A data.frame with the data making up the plot (point estimates and lower/upper 95% CIs)
+#'
+#' @examples
+#' \donttest{
+#' if (require("ggplot2", quietly = TRUE)) {
+#' # Generate toy data and customize plots.
+#' n = 500
+#' K = 1
+#' reward = matrix(1 + rnorm(n * K), n, K)
+#' scores = reward + matrix(rnorm(n * K), n, K)
+#' cost = 1
+#'
+#' # Fit Qini curves.
+#' qini.avg <- maq(reward, cost, scores, R = 200, target.with.covariates = FALSE)
+#' qini <- maq(reward, cost, scores, R = 200)
+#'
+#' # In some settings we may want to plot using one of R's many plot libraries.
+#' # The plot method invisibly returns the plot data we can use for this purpose.
+#' df.qini.baseline <- plot(qini.avg)
+#' df.qini <- plot(qini, add = TRUE, col = 2)
+#'
+#' # Make an alternate plot style, using, for example, ggplot.
+#' ggplot(df.qini, aes(x = spend, y = gain)) +
+#'   geom_ribbon(aes(ymin = gain - 1.96 * std.err,
+#'                   ymax = gain + 1.96 * std.err),
+#'               fill = "lightgray") +
+#'   geom_line(linewidth = 2) +
+#'   ylab("Policy value") +
+#'   xlab("Fraction treated") +
+#'   geom_line(data = df.qini.baseline, aes(x = spend, y = gain), lty = 2)
+#' }
+#' }
+#'
 #' @method plot maq
 #' @export
 plot.maq <- function(x,
@@ -83,4 +116,6 @@ plot.maq <- function(x,
     do.call(graphics::lines, c(list(x = spend, y = lb), lines.args))
     do.call(graphics::lines, c(list(x = spend, y = ub), lines.args))
   }
+
+  invisible(data.frame(spend, gain, std.err, lb, ub))
 }