Using dplyr and broom to compute kmeans on a training and test set

broom, dplyr, k-means, r

Solution

One approach would be to...

- extract the table specifying the centroids of your clusters (built on the training set) via broom.

- calculate the distance of each point in the test set from each of the cluster centroids built using the training set. Could do this via fuzzyjoin package.

- the cluster centroid that a test point has the shortest Euclidian distance from represents its assigned cluster.

- From there you can calculate any metrics of interest.

See below using a simpler dataset pulled from clustering example from tidymodels.

library(tidyverse)
library(rsample)
library(broom)
library(fuzzyjoin)

# data and train / test set-up
set.seed(27)
centers <- tibble(
  cluster = factor(1:3), 
  num_points = c(100, 150, 50),  # number points in each cluster
  x1 = c(5, 0, -3),              # x1 coordinate of cluster center
  x2 = c(-1, 1, -2)              # x2 coordinate of cluster center
)

labelled_points <- 
  centers %>%
  mutate(
    x1 = map2(num_points, x1, rnorm),
    x2 = map2(num_points, x2, rnorm)
  ) %>% 
  select(-num_points) %>% 
  unnest(cols = c(x1, x2))

points <- 
  labelled_points %>% 
  select(-cluster)

set.seed(1234)

split <- rsample::initial_split(points)
train <- rsample::training(split)
test <- rsample::testing(split)

# Fit kmeans on train then assign clusters to test
kclust <- kmeans(train, centers = 3)

clust_centers <- kclust %>% 
  tidy() %>% 
  select(-c(size, withinss))

test_clusts <- fuzzyjoin::distance_join(mutate(test, index = row_number()), 
                         clust_centers,
                         max_dist = Inf,
                         method = "euclidean",
                         distance_col = "dist") %>% 
  group_by(index) %>% 
  filter(dist == min(dist)) %>% 
  ungroup()
#> Joining by: c("x1", "x2")

# resulting table
test_clusts
#> # A tibble: 75 x 7
#>     x1.x    x2.x index  x1.y  x2.y cluster  dist
#>    <dbl>   <dbl> <int> <dbl> <dbl> <fct>   <dbl>
#>  1  4.24 -0.946      1  5.07 -1.10 3       0.847
#>  2  3.54  0.287      2  5.07 -1.10 3       2.06 
#>  3  3.71 -1.67       3  5.07 -1.10 3       1.47 
#>  4  5.03 -0.788      4  5.07 -1.10 3       0.317
#>  5  6.57 -2.49       5  5.07 -1.10 3       2.04 
#>  6  4.97  0.233      6  5.07 -1.10 3       1.34 
#>  7  4.43 -1.89       7  5.07 -1.10 3       1.01 
#>  8  5.34 -0.0705     8  5.07 -1.10 3       1.07 
#>  9  4.60  0.196      9  5.07 -1.10 3       1.38 
#> 10  5.68 -1.55      10  5.07 -1.10 3       0.758
#> # ... with 65 more rows

# calc within clusts SS on test
test_clusts %>% 
  group_by(cluster) %>% 
  summarise(size = n(),
            withinss = sum(dist^2),
            withinss_avg = withinss / size)
#> # A tibble: 3 x 4
#>   cluster  size withinss withinss_avg
#>   <fct>   <int>    <dbl>        <dbl>
#> 1 1          11     32.7         2.97
#> 2 2          35     78.9         2.26
#> 3 3          29     62.0         2.14

# compare to on train
tidy(kclust) %>% 
  mutate(withinss_avg = withinss / size)
#> # A tibble: 3 x 6
#>        x1    x2  size withinss cluster withinss_avg
#>     <dbl> <dbl> <int>    <dbl> <fct>          <dbl>
#> 1 -3.22   -1.91    40     76.8 1               1.92
#> 2  0.0993  1.06   113    220.  2               1.95
#> 3  5.07   -1.10    72    182.  3               2.53

# plot of test and train points
test_clusts %>% 
  select(x1 = x1.x, x2 = x2.x, cluster) %>% 
  mutate(type = "test") %>% 
  bind_rows(
    augment(kclust, train) %>% 
      mutate(type = "train") %>% 
      rename(cluster = .cluster)
    ) %>% 
  ggplot(aes(x = x1, 
             y = x2, 
             color = as.factor(cluster)))+
  geom_point()+
  facet_wrap(~fct_rev(as.factor(type)))+
  coord_fixed()+
  labs(title = "Cluster Assignment on Training and Holdout Datasets",
       color = "Cluster")+
  theme_bw()

Created on 2021-08-19 by the reprex package (v2.0.0)

(See comment on OP for link to conversations on making this easier within tidymodels.)

Problem

Original source