machine_learning: pin numeric output of k_means_clust with doctests (#15285)

Author: priya-sundaram-dev

machine_learning/k_means_clust.py

@@ -81,6 +81,13 @@ def centroid_pairwise_dist(x, centroids):
 
 
 def assign_clusters(data, centroids):
+    """Assign each data point to the index of its nearest centroid.
+
+    >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]])
+    >>> centroids = np.array([[0.0, 0.0], [10.0, 10.0]])
+    >>> assign_clusters(data, centroids).tolist()
+    [0, 0, 1, 1]
+    """
     # Compute distances between each data point and the set of centroids:
     # Fill in the blank (RHS only)
     distances_from_centroids = centroid_pairwise_dist(data, centroids)
@@ -93,6 +100,13 @@ def assign_clusters(data, centroids):
 
 
 def revise_centroids(data, k, cluster_assignment):
+    """Recompute each centroid as the mean of the points assigned to it.
+
+    >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]])
+    >>> assignment = np.array([0, 0, 1, 1])
+    >>> revise_centroids(data, 2, assignment).tolist()
+    [[0.0, 0.5], [10.0, 10.5]]
+    """
     new_centroids = []
     for i in range(k):
         # Select all data points that belong to cluster i. Fill in the blank (RHS only)
@@ -106,6 +120,16 @@ def revise_centroids(data, k, cluster_assignment):
 
 
 def compute_heterogeneity(data, k, centroids, cluster_assignment):
+    """Sum of squared distances from each point to its assigned centroid.
+
+    This is the objective k-means minimises; lower is a tighter clustering.
+
+    >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]])
+    >>> centroids = np.array([[0.0, 0.5], [10.0, 10.5]])
+    >>> assignment = np.array([0, 0, 1, 1])
+    >>> float(compute_heterogeneity(data, 2, centroids, assignment))
+    1.0
+    """
     heterogeneity = 0.0
     for i in range(k):
         # Select all data points that belong to cluster i. Fill in the blank (RHS only)
@@ -154,7 +178,16 @@ def kmeans(
                           as function of iterations
                           if None, do not store the history.
     verbose: if True, print how many data points changed their cluster labels in
-                          each iteration"""
+                          each iteration
+
+    >>> data = np.array([[0.0, 0.0], [0.0, 1.0], [10.0, 10.0], [10.0, 11.0]])
+    >>> initial_centroids = np.array([[0.0, 0.0], [10.0, 10.0]])
+    >>> centroids, assignment = kmeans(data, 2, initial_centroids, maxiter=10)
+    >>> centroids.tolist()
+    [[0.0, 0.5], [10.0, 10.5]]
+    >>> assignment.tolist()
+    [0, 0, 1, 1]
+    """
     centroids = initial_centroids[:]
     prev_cluster_assignment = None