augment method for #401

Author: topepo

NAMESPACE

@@ -1,5 +1,6 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method(augment,model_fit)
 S3method(fit,model_spec)
 S3method(fit_xy,model_spec)
 S3method(glance,model_fit)
@@ -103,6 +104,7 @@ export(.x)
 export(.y)
 export(C5.0_train)
 export(add_rowindex)
+export(augment)
 export(boost_tree)
 export(check_empty_ellipse)
 export(check_final_param)
@@ -208,6 +210,7 @@ importFrom(dplyr,starts_with)
 importFrom(dplyr,summarise)
 importFrom(dplyr,tally)
 importFrom(dplyr,vars)
+importFrom(generics,augment)
 importFrom(generics,fit)
 importFrom(generics,fit_xy)
 importFrom(generics,glance)

R/augment.R

@@ -0,0 +1,73 @@
+#' Augment data with predictions
+#'
+#' `augment()` will add column(s) for predictions to the given data.
+#'
+#' For regression models, a `.pred` column is added. If `x` was created using
+#' [fit()] and `new_data` contains the outcome column, a `.resid` column is
+#' also added.
+#'
+#' For classification models, the results include a column called `.pred_class`
+#' as well as class probability columns named `.pred_{level}`.
+#' @param x A `model_fit` object produced by [fit()] or [fit_xy()].
+#' @param new_data A data frame or matrix.
+#' @param ... Not currently used.
+#' @export
+#' @examples
+#' reg_form <-
+#'   linear_reg() %>%
+#'   set_engine("lm") %>%
+#'   fit(mpg ~ ., data = mtcars)
+#' reg_xy <-
+#'   linear_reg() %>%
+#'   set_engine("lm") %>%
+#'   fit_xy(mtcars[, -1], mtcars$mpg)
+#'
+#' augment(reg_form, head(mtcars))
+#' augment(reg_form, head(mtcars[, -1]))
+#'
+#' augment(reg_xy, head(mtcars))
+#' augment(reg_xy, head(mtcars[, -1]))
+#'
+#' # ------------------------------------------------------------------------------
+#'
+#' data(two_class_dat, package = "modeldata")
+#'
+#' cls_form <-
+#'   logistic_reg() %>%
+#'   set_engine("glm") %>%
+#'   fit(Class ~ ., data = two_class_dat)
+#' cls_xy <-
+#'   logistic_reg() %>%
+#'   set_engine("glm") %>%
+#'   fit_xy(two_class_dat[, -3],
+#'   two_class_dat$Class)
+#'
+#' augment(cls_form, head(two_class_dat))
+#' augment(cls_form, head(two_class_dat[, -3]))
+#'
+#' augment(cls_xy, head(two_class_dat))
+#' augment(cls_xy, head(two_class_dat[, -3]))
+#'
+augment.model_fit <- function(x, new_data, ...) {
+  if (x$spec$mode == "regression") {
+    new_data <-
+      new_data %>%
+      dplyr::bind_cols(
+        predict(x, new_data = new_data)
+      )
+    if (length(x$preproc$y_var) > 0) {
+      y_nm <- x$preproc$y_var
+      if (any(names(new_data) == y_nm)) {
+        new_data <- dplyr::mutate(new_data, .resid = !!rlang::sym(y_nm) - .pred)
+      }
+    }
+  } else {
+    new_data <-
+      new_data %>%
+      dplyr::bind_cols(
+        predict(x, new_data = new_data, type = "class"),
+        predict(x, new_data = new_data, type = "prob")
+      )
+  }
+  new_data
+}

_pkgdown.yml

@@ -49,8 +49,7 @@ reference:
       - set_mode
       - tidy.model_fit
       - translate
-      - varying
-      - varying_args
+      - augment.model.fit
 
   - title: Developer Tools
     contents:

man/augment.model_fit.Rd

@@ -0,0 +1,72 @@
+
+context("augment")
+
+# ------------------------------------------------------------------------------
+
+test_that('regression models', {
+  x <- linear_reg() %>% set_engine("lm")
+
+  reg_form <- x %>% fit(mpg ~ ., data = mtcars)
+  reg_xy <- x %>% fit_xy(mtcars[, -1], mtcars$mpg)
+
+  expect_equal(
+    colnames(augment(reg_form, head(mtcars))),
+    c("mpg", "cyl", "disp", "hp", "drat", "wt", "qsec", "vs", "am",
+      "gear", "carb", ".pred", ".resid")
+  )
+  expect_equal(nrow(augment(reg_form, head(mtcars))), 6)
+  expect_equal(
+    colnames(augment(reg_form, head(mtcars[, -1]))),
+    c("cyl", "disp", "hp", "drat", "wt", "qsec", "vs", "am",
+      "gear", "carb", ".pred")
+  )
+  expect_equal(nrow(augment(reg_form, head(mtcars[, -1]))), 6)
+
+  expect_equal(
+    colnames(augment(reg_xy, head(mtcars))),
+    c("mpg", "cyl", "disp", "hp", "drat", "wt", "qsec", "vs", "am",
+      "gear", "carb", ".pred")
+  )
+  expect_equal(nrow(augment(reg_xy, head(mtcars))), 6)
+  expect_equal(
+    colnames(augment(reg_xy, head(mtcars[, -1]))),
+    c("cyl", "disp", "hp", "drat", "wt", "qsec", "vs", "am",
+      "gear", "carb", ".pred")
+  )
+  expect_equal(nrow(augment(reg_xy, head(mtcars[, -1]))), 6)
+
+})
+
+
+
+test_that('classification models', {
+  data(two_class_dat, package = "modeldata")
+  x <- logistic_reg() %>% set_engine("glm")
+
+  cls_form <- x %>% fit(Class ~ ., data = two_class_dat)
+  cls_xy <- x %>% fit_xy(two_class_dat[, -3], two_class_dat$Class)
+
+  expect_equal(
+    colnames(augment(cls_form, head(two_class_dat))),
+    c("A", "B", "Class", ".pred_class", ".pred_Class1", ".pred_Class2")
+  )
+  expect_equal(nrow(augment(cls_form, head(two_class_dat))), 6)
+  expect_equal(
+    colnames(augment(cls_form, head(two_class_dat[, -3]))),
+    c("A", "B", ".pred_class", ".pred_Class1", ".pred_Class2")
+  )
+  expect_equal(nrow(augment(cls_form, head(two_class_dat[, -3]))), 6)
+
+  expect_equal(
+    colnames(augment(cls_xy, head(two_class_dat))),
+    c("A", "B", "Class", ".pred_class", ".pred_Class1", ".pred_Class2")
+  )
+  expect_equal(nrow(augment(cls_xy, head(two_class_dat))), 6)
+  expect_equal(
+    colnames(augment(cls_xy, head(two_class_dat[, -3]))),
+    c("A", "B", ".pred_class", ".pred_Class1", ".pred_Class2")
+  )
+  expect_equal(nrow(augment(cls_xy, head(two_class_dat[, -3]))), 6)
+
+})
+