Merge pull request #415 from tidymodels/augment

augment method for model_fit objects

Author: topepo

DESCRIPTION

@@ -22,7 +22,7 @@ Imports:
     purrr,
     utils,
     tibble (>= 2.1.1),
-    generics,
+    generics (>= 0.1.0),
     glue,
     magrittr,
     stats,

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,80 @@
+#' 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
+#' car_trn <- mtcars[11:32,]
+#' car_tst <- mtcars[ 1:10,]
+#'
+#' reg_form <-
+#'   linear_reg() %>%
+#'   set_engine("lm") %>%
+#'   fit(mpg ~ ., data = car_trn)
+#' reg_xy <-
+#'   linear_reg() %>%
+#'   set_engine("lm") %>%
+#'   fit_xy(car_trn[, -1], car_trn$mpg)
+#'
+#' augment(reg_form, car_tst)
+#' augment(reg_form, car_tst[, -1])
+#'
+#' augment(reg_xy, car_tst)
+#' augment(reg_xy, car_tst[, -1])
+#'
+#' # ------------------------------------------------------------------------------
+#'
+#' data(two_class_dat, package = "modeldata")
+#' cls_trn <- two_class_dat[-(1:10), ]
+#' cls_tst <- two_class_dat[  1:10 , ]
+#'
+#' cls_form <-
+#'   logistic_reg() %>%
+#'   set_engine("glm") %>%
+#'   fit(Class ~ ., data = cls_trn)
+#' cls_xy <-
+#'   logistic_reg() %>%
+#'   set_engine("glm") %>%
+#'   fit_xy(cls_trn[, -3],
+#'   cls_trn$Class)
+#'
+#' augment(cls_form, cls_tst)
+#' augment(cls_form, cls_tst[, -3])
+#'
+#' augment(cls_xy, cls_tst)
+#' augment(cls_xy, cls_tst[, -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 if (x$spec$mode == "classification") {
+    new_data <-
+      new_data %>%
+      dplyr::bind_cols(
+        predict(x, new_data = new_data, type = "class"),
+        predict(x, new_data = new_data, type = "prob")
+      )
+  } else {
+    rlang::abort(paste("Unknown mode:", x$spec$mode))
+  }
+  new_data
+}

R/reexports.R

@@ -15,7 +15,10 @@ generics::fit_xy
 #' @export
 generics::tidy
 
-
 #' @importFrom generics glance
 #' @export
 generics::glance
+
+#' @importFrom generics augment
+#' @export
+generics::augment

R/zzz.R

@@ -1,13 +1,14 @@
 # nocov start
 
 .onLoad <- function(libname, pkgname) {
-  s3_register("broom::tidy", "model_fit")
-  s3_register("broom::tidy", "nullmodel")
-  s3_register("broom::tidy", "_elnet")
-  s3_register("broom::tidy", "_lognet")
-  s3_register("broom::tidy", "_multnet")
-  s3_register("broom::tidy", "_fishnet")
-  s3_register("broom::glance", "model_fit")
+  s3_register("generics::tidy", "model_fit")
+  s3_register("generics::tidy", "nullmodel")
+  s3_register("generics::tidy", "_elnet")
+  s3_register("generics::tidy", "_lognet")
+  s3_register("generics::tidy", "_multnet")
+  s3_register("generics::tidy", "_fishnet")
+  s3_register("generics::glance", "model_fit")
+  s3_register("generics::augment", "model_fit")
 }
 
 

_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,76 @@
+
+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)
+
+  reg_form$spec$mode <- "depeche"
+
+  expect_error(augment(reg_form, head(mtcars[, -1])), "Unknown mode: depeche")
+
+})
+
+
+
+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)
+
+})
+