From 6e7083411ef710b214027f0452b9bb438478afe3 Mon Sep 17 00:00:00 2001
From: David Holstius <david.holstius@gmail.com>
Date: Mon, 29 Sep 2014 12:01:38 -0700
Subject: [PATCH] Add pred_grid.rq

---
 R/compute_model_prediction.R | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/R/compute_model_prediction.R b/R/compute_model_prediction.R
index bf368eff..69a2d493 100644
--- a/R/compute_model_prediction.R
+++ b/R/compute_model_prediction.R
@@ -15,7 +15,7 @@
 #'   returning predictions with \code{\link{predict}}. If not supplied, will use
 #'   \code{\link{loess}} for <= 1000 points, otherwise it will use
 #'   \code{\link[mgcv]{gam}}. Other modelling functions that will work include
-#'   \code{\link{lm}}, \code{\link{glm}} and \code{\link[MASS]{rlm}}.
+#'   \code{\link{lm}}, \code{\link{glm}}, \code{\link[quantreg]{rq}}, and \code{\link[MASS]{rlm}}.
 #' @param formula Formula passed to modelling function. Can use any variables
 #'   from data.
 #' @param se include standard errors in output? Requires appropriate method of
@@ -237,6 +237,36 @@ pred_grid.lm <- function(model, data, domain = NULL, n = 80, se = FALSE,
   }
 }
 
+#' @export
+pred_grid.rq <- function(model, data, domain = NULL, n = 80, se = FALSE,
+                         level = 0.95) {
+  x_var <- get_predict_vars(terms(model))
+  if (length(x_var) > 1) {
+    stop("Only know how to make grid for one variable", call. = FALSE)
+  }
+
+  x_rng <- domain %||% range(data[[x_var]], na.rm = TRUE)
+  x_grid <- seq(x_rng[1], x_rng[2], length = n)
+  grid <- setNames(data.frame(x_grid), x_var)
+
+  # Much like `pred_grid.lm` but *do not* pass `se` to `predict.rq` here
+  resp <- predict(model, newdata = grid,
+                  level = level, interval = if(se) "confidence" else "none")
+
+  if (!se) {
+    data.frame(
+      pred_ = x_grid,
+      resp_ = as.vector(resp)
+    )
+  } else {
+    data.frame(
+      pred_ = x_grid,
+      resp_ = resp[, "fit"],
+      resp_lwr_ = resp[, "lower"],
+      resp_upr_ = resp[, "higher"]
+    )
+  }
+}
 
 # Given a formula object, return a character vector of predictor variables
 get_predict_vars <- function(f) {