Merge pull request #108 from tidymodels/fix-94

make sure ClusterR predict method also rearranges

Author: EmilHvitfeldt

NEWS.md

@@ -1,5 +1,7 @@
 # tidyclust (development version)
 
+* Fixed bug where `extract_cluster_assignment()` and `predict()` sometimes didn't have agreement of clusters. (#94)
+
 * `silhouette()` and `silhouette_avg()` now return NAs instead of erroring when applied to a clustering object with 1 cluster. (#104)
 
 * Fixed bug where `extract_cluster_assignment()` doesn't work for `hier_clust()` models in workflows where `num_clusters` is specified in `extract_cluster_assignment()`.

R/extract_assignment.R

@@ -62,7 +62,7 @@ cluster_assignment_tibble <- function(clusters,
                                       n_clusters,
                                       ...,
                                       prefix = "Cluster_") {
-  reorder_clusts <- order(unique(clusters))
+  reorder_clusts <- order(union(unique(clusters), seq_len(n_clusters)))
   names <- paste0(prefix, seq_len(n_clusters))
   res <- names[reorder_clusts][clusters]
 

R/predict_helpers.R

@@ -7,8 +7,13 @@ stats_kmeans_predict <- function(object, new_data, prefix = "Cluster_") {
 }
 
 clusterR_kmeans_predict <- function(object, new_data, prefix = "Cluster_") {
-  res <- predict(object, new_data)
-  res <- paste0(prefix, res)
+  clusters <- predict(object, new_data)
+  n_clusters <- length(object$obs_per_cluster)
+
+  reorder_clusts <- order(union(unique(clusters), seq_len(n_clusters)))
+  names <- paste0(prefix, seq_len(n_clusters))
+  res <- names[reorder_clusts][clusters]
+
   factor(res)
 }