Merge pull request #166 from tidymodels/one-row-prediction-fix

Author: EmilHvitfeldt

NEWS.md

@@ -22,6 +22,8 @@
 
 * Engine specific documentation has been added for all models and engines. (#159)
 
+* Fixed bug where levels didn't match number of clusters if prediction on fewer number of observations. (#158)
+
 # tidyclust 0.1.2
 
 * The cluster specification methods for `generics::tune_args()` and `generics::tunable()` are now registered unconditionally (#115).

R/predict_helpers.R

@@ -150,5 +150,5 @@ make_predictions <- function(x, prefix, n_clusters) {
   }
   pred_clusts <- unique(clusters$.cluster)[pred_clusts_num]
 
-  return(factor(pred_clusts))
+  pred_clusts
 }

tests/testthat/test-hier_clust-stats.R

@@ -27,6 +27,21 @@ test_that("predicting", {
   )
 })
 
+test_that("all levels are preserved with 1 row predictions", {
+  set.seed(1234)
+  spec <- hier_clust(num_clusters = 3) %>%
+    set_engine("stats")
+
+  res <- fit(spec, ~., mtcars)
+
+  preds <- predict(res, mtcars[1, ])
+
+  expect_identical(
+    levels(preds$.pred_cluster),
+    paste0("Cluster_", 1:3)
+  )
+})
+
 test_that("extract_centroids() works", {
   set.seed(1234)
   spec <- hier_clust(num_clusters = 3) %>%

tests/testthat/test-k_means-clustMixType.R

@@ -40,6 +40,21 @@ test_that("predicting", {
   )
 })
 
+test_that("all levels are preserved with 1 row predictions", {
+  set.seed(1234)
+  spec <- k_means(num_clusters = 3) %>%
+    set_engine("clustMixType")
+
+  res <- fit(spec, ~., iris)
+
+  preds <- predict(res, iris[1, ])
+
+  expect_identical(
+    levels(preds$.pred_cluster),
+    paste0("Cluster_", 1:3)
+  )
+})
+
 test_that("extract_centroids() works", {
   skip_if_not_installed("clustMixType")
 

tests/testthat/test-k_means-clusterR.R

@@ -32,6 +32,21 @@ test_that("predicting", {
   )
 })
 
+test_that("all levels are preserved with 1 row predictions", {
+  set.seed(1234)
+  spec <- k_means(num_clusters = 3) %>%
+    set_engine("ClusterR")
+
+  res <- fit(spec, ~., mtcars)
+
+  preds <- predict(res, mtcars[1, ])
+
+  expect_identical(
+    levels(preds$.pred_cluster),
+    paste0("Cluster_", 1:3)
+  )
+})
+
 test_that("extract_centroids() works", {
   skip_if_not_installed("ClusterR")
 

tests/testthat/test-k_means-klaR.R

@@ -51,6 +51,28 @@ test_that("predicting", {
   )
 })
 
+test_that("all levels are preserved with 1 row predictions", {
+  skip_if_not_installed("klaR")
+  skip_if_not_installed("modeldata")
+
+  data("ames", package = "modeldata")
+
+  ames_cat <- dplyr::select(ames, dplyr::where(is.factor))
+
+  set.seed(1234)
+  spec <- k_means(num_clusters = 3) %>%
+    set_engine("klaR")
+
+  res <- fit(spec, ~., ames_cat)
+
+  preds <- predict(res, ames_cat[1, ])
+
+  expect_identical(
+    levels(preds$.pred_cluster),
+    paste0("Cluster_", 1:3)
+  )
+})
+
 test_that("predicting ties argument works", {
   skip_if_not_installed("klaR")
 

tests/testthat/test-k_means-stats.R

@@ -27,6 +27,21 @@ test_that("predicting", {
   )
 })
 
+test_that("all levels are preserved with 1 row predictions", {
+  set.seed(1234)
+  spec <- k_means(num_clusters = 3) %>%
+    set_engine("stats")
+
+  res <- fit(spec, ~., mtcars)
+
+  preds <- predict(res, mtcars[1, ])
+
+  expect_identical(
+    levels(preds$.pred_cluster),
+    paste0("Cluster_", 1:3)
+  )
+})
+
 test_that("extract_centroids() works", {
   set.seed(1234)
   spec <- k_means(num_clusters = 3) %>%