R/ordinalscore.R

In tidymodels/recipes: Preprocessing and Feature Engineering Steps for Modeling

#' Convert ordinal factors to numeric scores
#'
#' `step_ordinalscore()` creates a *specification* of a recipe step that will
#' convert ordinal factor variables into numeric scores.
#'
#' @inheritParams step_center
#' @inheritParams step_pca
#' @param convert A function that takes an ordinal factor vector
#'  as an input and outputs a single numeric variable.
#' @template step-return
#' @family dummy variable and encoding steps
#' @export
#' @details Dummy variables from ordered factors with `C`
#'  levels will create polynomial basis functions with `C-1`
#'  terms. As an alternative, this step can be used to translate the
#'  ordered levels into a single numeric vector of values that
#'  represent (subjective) scores. By default, the translation uses
#'  a linear scale (1, 2, 3, ... `C`) but custom score
#'  functions can also be used (see the example below).
#'
#' # Tidying
#'
#' When you [`tidy()`][tidy.recipe()] this step, a tibble is returned with
#' columns `terms` and `id`:
#'
#' \describe{
#'   \item{terms}{character, the selectors or variables selected}
#'   \item{id}{character, id of this step}
#' }
#'
#' @template case-weights-not-supported
#'
#' @examples
#' fail_lvls <- c("meh", "annoying", "really_bad")
#'
#' ord_data <-
#'   data.frame(
#'     item = c("paperclip", "twitter", "airbag"),
#'     fail_severity = factor(fail_lvls,
#'       levels = fail_lvls,
#'       ordered = TRUE
#'     )
#'   )
#'
#' model.matrix(~fail_severity, data = ord_data)
#'
#' linear_values <- recipe(~ item + fail_severity, data = ord_data) %>%
#'   step_dummy(item) %>%
#'   step_ordinalscore(fail_severity)
#'
#' linear_values <- prep(linear_values, training = ord_data)
#'
#' bake(linear_values, new_data = NULL)
#'
#' custom <- function(x) {
#'   new_values <- c(1, 3, 7)
#'   new_values[as.numeric(x)]
#' }
#'
#' nonlin_scores <- recipe(~ item + fail_severity, data = ord_data) %>%
#'   step_dummy(item) %>%
#'   step_ordinalscore(fail_severity, convert = custom)
#'
#' tidy(nonlin_scores, number = 2)
#'
#' nonlin_scores <- prep(nonlin_scores, training = ord_data)
#'
#' bake(nonlin_scores, new_data = NULL)
#'
#' tidy(nonlin_scores, number = 2)
step_ordinalscore <-
  function(recipe,
           ...,
           role = NA,
           trained = FALSE,
           columns = NULL,
           convert = as.numeric,
           skip = FALSE,
           id = rand_id("ordinalscore")) {
    add_step(
      recipe,
      step_ordinalscore_new(
        terms = enquos(...),
        role = role,
        trained = trained,
        columns = columns,
        convert = convert,
        skip = skip,
        id = id
      )
    )
  }

step_ordinalscore_new <-
  function(terms, role, trained, columns, convert, skip, id) {
    step(
      subclass = "ordinalscore",
      terms = terms,
      role = role,
      trained = trained,
      columns = columns,
      convert = convert,
      skip = skip,
      id = id
    )
  }

#' @export
prep.step_ordinalscore <- function(x, training, info = NULL, ...) {
    col_names <- recipes_eval_select(x$terms, training, info)
    check_type(training[, col_names], types = "ordered")
    check_function(x$convert, arg = "convert")

    step_ordinalscore_new(
      terms = x$terms,
      role = x$role,
      trained = TRUE,
      columns = col_names,
      convert = x$convert,
      skip = x$skip,
      id = x$id
    )
  }

#' @export
bake.step_ordinalscore <- function(object, new_data, ...) {
  col_names <- object$columns
  check_new_data(col_names, object, new_data)

  for (col_name in col_names) {
    score <- object$convert(new_data[[col_name]])
    score <- vctrs::vec_cast(score, integer())
    new_data[[col_name]] <- score
  }

  new_data
}

#' @export
print.step_ordinalscore <-
  function(x, width = max(20, options()$width - 30), ...) {
    title <- "Scoring for "
    print_step(x$columns, x$terms, x$trained, title, width)
    invisible(x)
  }


#' @rdname tidy.recipe
#' @export
tidy.step_ordinalscore <- function(x, ...) {
  res <- simple_terms(x, ...)
  res$id <- x$id
  res
}