R/outcome-names.R

In tune: Tidy Tuning Tools

#' Determine names of the outcome data in a workflow
#'
#' @param x An object.
#' @param ... Not used.
#' @return A character string of variable names
#' @keywords internal
#' @examples
#' library(dplyr)
#' lm(cbind(mpg, wt) ~ ., data = mtcars) |>
#'   purrr::pluck(terms) |>
#'   outcome_names()
#' @export
outcome_names <- function(x, ...) {
  UseMethod("outcome_names")
}

#' @export
#' @rdname outcome_names
outcome_names.terms <- function(x, ...) {
  if (length(x) == 2) {
    res <- character(0)
  } else {
    res <- all.vars(x[[2]])
  }
  res
}

#' @export
#' @rdname outcome_names
outcome_names.formula <- outcome_names.terms


#' @export
#' @rdname outcome_names
outcome_names.recipe <- function(x, ...) {
  y <- summary(x)
  y$variable[y$role == "outcome" & !is.na(y$role)]
}

#' @export
#' @rdname outcome_names
outcome_names.workflow <- function(x, ...) {
  if (!is.null(x$pre$mold)) {
    y_vals <- extract_mold(x)$outcomes
    res <- colnames(y_vals)
  } else {
    preprocessor <- extract_preprocessor(x)
    res <- outcome_names(preprocessor, ...)
  }
  res
}

#' @export
#' @rdname outcome_names
outcome_names.tune_results <- function(x, ...) {
  att <- attributes(x)
  if (any(names(att) == "outcomes")) {
    res <- att$outcomes
  } else {
    res <- NA_character_
  }
  res
}

#' @export
#' @rdname outcome_names
outcome_names.workflow_variables <- function(x, data = NULL, ...) {
  if (!inherits(data, "data.frame")) {
    cli::cli_abort(
      "When using {.fn outcome_names} on an object for class {.cls {class(x)}}, the
      argument {.arg data} should be a data frame, not {.obj_type_friendly {data}}."
    )
  }
  names(dplyr::select(data, !!x$outcomes))
}