R/smote_impl.R

In themis: Extra Recipes Steps for Dealing with Unbalanced Data

#' SMOTE Algorithm
#'
#' SMOTE generates new examples of the minority class using nearest neighbors
#' of these cases.
#'
#' @inheritParams step_smote
#' @param df data.frame or tibble. Must have 1 factor variable and remaining
#'  numeric variables.
#' @param var Character, name of variable containing factor variable.
#' @param k An integer. Number of nearest neighbor that are used
#'  to generate the new examples of the minority class.
#'
#' @return A data.frame or tibble, depending on type of `df`.
#' @export
#'
#' @details
#' The parameter `neighbors` controls the way the new examples are created.
#' For each currently existing minority class example X new examples will be
#' created (this is controlled by the parameter `over_ratio` as mentioned
#' above). These examples will be generated by using the information from the
#' `neighbors` nearest neighbor of each example of the minority class.
#' The parameter `neighbors` controls how many of these neighbor are used.
#
#' All columns used in this function must be numeric with no missing data.
#'
#' @references Chawla, N. V., Bowyer, K. W., Hall, L. O., and Kegelmeyer,
#'  W. P. (2002). Smote: Synthetic minority over-sampling technique.
#'  Journal of Artificial Intelligence Research, 16:321-357.
#'
#' @seealso [step_smote()] for step function of this method
#' @family Direct Implementations
#'
#' @examples
#' circle_numeric <- circle_example[, c("x", "y", "class")]
#'
#' res <- smote(circle_numeric, var = "class")
#'
#' res <- smote(circle_numeric, var = "class", k = 10)
#'
#' res <- smote(circle_numeric, var = "class", over_ratio = 0.8)
smote <- function(df, var, k = 5, over_ratio = 1) {
  if (length(var) != 1) {
    rlang::abort("Please select a single factor variable for `var`.")
  }

  var <- rlang::arg_match(var, colnames(df))

  if (!(is.factor(df[[var]]) | is.character(df[[var]]))) {
    rlang::abort(glue("`{var}` should be a factor or character variable."))
  }

  if (length(k) != 1) {
    rlang::abort("`k` must be length 1.")
  }

  if (k < 1) {
    rlang::abort("`k` must be non-negative.")
  }

  predictors <- setdiff(colnames(df), var)

  check_numeric(df[, predictors])
  check_na(select(df, -all_of(var)))

  smote_impl(df, var, k, over_ratio)
}

smote_impl <- function(df, var, k, over_ratio, call = caller_env()) {
  data <- split(df, df[[var]])
  majority_count <- max(table(df[[var]]))
  ratio_target <- majority_count * over_ratio
  which_upsample <- which(table(df[[var]]) < ratio_target)
  samples_needed <- ratio_target - table(df[[var]])[which_upsample]
  min_names <- names(samples_needed)

  out_dfs <- list()

  for (i in seq_along(samples_needed)) {
    minority_df <- data[[min_names[i]]]
    minority <- as.matrix(minority_df[names(minority_df) != var])

    if (nrow(minority) <= k) {
      rlang::abort(
        glue(
          "Not enough observations of '{min_names[i]}' to perform SMOTE."
        ),
        call = call
      )
    }

    synthetic <- smote_data(minority, k = k, n_samples = samples_needed[i])
    out_df <- as.data.frame(synthetic)
    names(out_df) <- setdiff(names(df), var)
    out_df_nrow <- min(nrow(out_df), 1)
    out_df[var] <- data[[names(samples_needed)[i]]][[var]][out_df_nrow]
    out_df <- out_df[names(df)]
    out_dfs[[i]] <- out_df
  }

  final <- rbind(df, do.call(rbind, out_dfs))
  final[[var]] <- factor(final[[var]], levels = levels(df[[var]]))
  rownames(final) <- NULL
  final
}

smote_data <- function(data, k, n_samples, smote_ids = seq_len(nrow(data))) {
  ids <- RANN::nn2(data, k = k + 1, searchtype = "priority")$nn.idx
  indexes <- rep(sample(smote_ids), length.out = n_samples)
  index_len <- tabulate(indexes, NROW(data))
  out <- matrix(0, nrow = n_samples, ncol = ncol(data))
  sampleids <- sample.int(k, n_samples, TRUE)
  runif_ids <- stats::runif(n_samples)

  iii <- 0
  for (row_num in smote_ids) {
    index_selection <- iii + seq_len(index_len[row_num])
    # removes itself as nearest neighbour
    id_knn <- ids[row_num, ids[row_num, ] != row_num]
    dif <- data[id_knn[sampleids[index_selection]], ] -
      data[rep(row_num, index_len[row_num]), ]
    gap <- dif * runif_ids[index_selection]
    out[index_selection, ] <- data[rep(row_num, index_len[row_num]), ] + gap
    iii <- iii + index_len[row_num]
  }

  out
}