knn multi_predict

Author: topepo

R/nearest_neighbor.R

@@ -178,3 +178,43 @@ translate.nearest_neighbor <- function(x, engine = x$engine, ...) {
   }
   x
 }
+
+
+# ------------------------------------------------------------------------------
+
+#' @importFrom purrr map_df
+#' @importFrom dplyr starts_with
+#' @rdname multi_predict
+#' @param neighbors An integer vector for the number of nearest neighbors.
+#' @export
+multi_predict._train.kknn <-
+  function(object, new_data, type = NULL, neighbors = NULL, ...) {
+    if (any(names(enquos(...)) == "newdata"))
+      stop("Did you mean to use `new_data` instead of `newdata`?", call. = FALSE)
+
+    if (is.null(neighbors))
+      neighbors <- rlang::eval_tidy(tt$fit$call$ks)
+    neighbors <- sort(neighbors)
+
+    if (is.null(type)) {
+      if (object$spec$mode == "classification")
+        type <- "class"
+      else
+        type <- "numeric"
+    }
+
+    res <-
+      purrr::map_df(neighbors, knn_by_k, object = object,
+                    new_data = new_data, type = type, ...)
+    res <- dplyr::arrange(res, .row, neighbors)
+    res <- split(res[, -1], res$.row)
+    names(res) <- NULL
+    dplyr::tibble(.pred = res)
+  }
+
+knn_by_k <- function(k, object, new_data, type, ...) {
+  object$fit$call$ks <- k
+  predict(object, new_data = new_data, type = type, ...) %>%
+    dplyr::mutate(neighbors = k, .row = dplyr::row_number()) %>%
+    dplyr::select(.row, neighbors, dplyr::starts_with(".pred"))
+}