R/prep.R

In mtfairbanks/gts: Parallel 'fable' tools

#' Convert tsibble to forecast ready object
#'
#' @param .data A data.frame or data.table
#'
#' @export
#'
#' @examples
#' your_df %>%
#'   ts_prep(key = c(level_1, level_2), index = index, target = revenue)
ts_prep <- function(.data,
                    key = NULL,
                    index = index,
                    target = NULL) {

  if (!is.data.frame(.data)) stop(".data must be a data.frame or data.table")
  if (!data.table::is.data.table(.data)) .data <- data.table::as.data.table(.data)

  index <- enexpr(index)
  key <- enexpr(key)
  target <- enexpr(target)

  index_class <- class(eval_tidy(expr('$'(.data, !!index))))

  if(is.null(target)) abort("target must be given")

  if (is.null(key)) {

    .data <- .data %>%
      nest(data = everything(), .key = "time_series")

  } else {
    groups <- tidy_vec_selector(.data, !!key)

    .data <- .data %>%
      group_nest(!!!groups, .key = "time_series")

  }

  if (any(index_class %in% c("integer", "numeric"))) {
    print("NOTE: Index column is detected as an integer")

    .data <- .data %>%
      mutate(
        time_series = map(time_series, as_tsibble, index = !!index)
      )
  } else if (index_class == "Date") {
    print("NOTE: Index column is detected as Date")

    .data <- .data %>%
      mutate(time_series = map(time_series,
                               ~mutate(.x, !!index := yearmonth(!!index)) %>%
                                 as_tsibble(index = !!index))
      )
  } else {
    abort("Index format not recognizable. Make sure its either in integer (period) or date format.")
  }
  .data
}