Fixed kmeans clustering bug AtsushiSakai#251

Author: AtsushiSakai

Mapping/kmeans_clustering/kmeans_clustering.py

@@ -6,146 +6,116 @@
 
 """
 
-import numpy as np
 import math
 import matplotlib.pyplot as plt
 import random
 
+# k means parameters
+MAX_LOOP = 10
+DCOST_TH = 0.1
 show_animation = True
 
 
-class Clusters:
-
-    def __init__(self, x, y, nlabel):
-        self.x = x
-        self.y = y
-        self.ndata = len(self.x)
-        self.nlabel = nlabel
-        self.labels = [random.randint(0, nlabel - 1)
-                       for _ in range(self.ndata)]
-        self.cx = [0.0 for _ in range(nlabel)]
-        self.cy = [0.0 for _ in range(nlabel)]
-
-
 def kmeans_clustering(rx, ry, nc):
-
     clusters = Clusters(rx, ry, nc)
-    clusters = calc_centroid(clusters)
+    clusters.calc_centroid()
 
-    MAX_LOOP = 10
-    DCOST_TH = 0.1
-    pcost = 100.0
+    pre_cost = float("inf")
     for loop in range(MAX_LOOP):
-        #  print("Loop:", loop)
-        clusters, cost = update_clusters(clusters)
-        clusters = calc_centroid(clusters)
+        print("loop:", loop)
+        cost = clusters.update_clusters()
+        clusters.calc_centroid()
 
-        dcost = abs(cost - pcost)
-        if dcost < DCOST_TH:
+        d_cost = abs(cost - pre_cost)
+        if d_cost < DCOST_TH:
             break
-        pcost = cost
+        pre_cost = cost
 
     return clusters
 
 
-def calc_centroid(clusters):
-
-    for ic in range(clusters.nlabel):
-        x, y = calc_labeled_points(ic, clusters)
-        ndata = len(x)
-        clusters.cx[ic] = sum(x) / ndata
-        clusters.cy[ic] = sum(y) / ndata
-
-    return clusters
-
-
-def update_clusters(clusters):
-    cost = 0.0
-
-    for ip in range(clusters.ndata):
-        px = clusters.x[ip]
-        py = clusters.y[ip]
-
-        dx = [icx - px for icx in clusters.cx]
-        dy = [icy - py for icy in clusters.cy]
-
-        dlist = [math.sqrt(idx**2 + idy**2) for (idx, idy) in zip(dx, dy)]
-        mind = min(dlist)
-        min_id = dlist.index(mind)
-        clusters.labels[ip] = min_id
-        cost += mind
-
-    return clusters, cost
-
-
-def calc_labeled_points(ic, clusters):
-
-    inds = np.array([i for i in range(clusters.ndata)
-                     if clusters.labels[i] == ic])
-    tx = np.array(clusters.x)
-    ty = np.array(clusters.y)
-
-    x = tx[inds]
-    y = ty[inds]
-
-    return x, y
-
-
-def calc_raw_data(cx, cy, npoints, rand_d):
+class Clusters:
 
+    def __init__(self, x, y, n_label):
+        self.x = x
+        self.y = y
+        self.n_data = len(self.x)
+        self.n_label = n_label
+        self.labels = [random.randint(0, n_label - 1)
+                       for _ in range(self.n_data)]
+        self.center_x = [0.0 for _ in range(n_label)]
+        self.center_y = [0.0 for _ in range(n_label)]
+
+    def plot_cluster(self):
+        for label in set(self.labels):
+            x, y = self._get_labeled_x_y(label)
+            plt.plot(x, y, ".")
+
+    def calc_centroid(self):
+        for label in set(self.labels):
+            x, y = self._get_labeled_x_y(label)
+            n_data = len(x)
+            self.center_x[label] = sum(x) / n_data
+            self.center_y[label] = sum(y) / n_data
+
+    def update_clusters(self):
+        cost = 0.0
+
+        for ip in range(self.n_data):
+            px = self.x[ip]
+            py = self.y[ip]
+
+            dx = [icx - px for icx in self.center_x]
+            dy = [icy - py for icy in self.center_y]
+
+            dist_list = [math.sqrt(idx ** 2 + idy ** 2) for (idx, idy) in zip(dx, dy)]
+            min_dist = min(dist_list)
+            min_id = dist_list.index(min_dist)
+            self.labels[ip] = min_id
+            cost += min_dist
+
+        return cost
+
+    def _get_labeled_x_y(self, target_label):
+        x = [self.x[i] for i, label in enumerate(self.labels) if label == target_label]
+        y = [self.y[i] for i, label in enumerate(self.labels) if label == target_label]
+        return x, y
+
+
+def calc_raw_data(cx, cy, n_points, rand_d):
     rx, ry = [], []
 
     for (icx, icy) in zip(cx, cy):
-        for _ in range(npoints):
+        for _ in range(n_points):
             rx.append(icx + rand_d * (random.random() - 0.5))
             ry.append(icy + rand_d * (random.random() - 0.5))
 
     return rx, ry
 
 
 def update_positions(cx, cy):
-
+    # object moving parameters
     DX1 = 0.4
     DY1 = 0.5
-
-    cx[0] += DX1
-    cy[0] += DY1
-
     DX2 = -0.3
     DY2 = -0.5
 
+    cx[0] += DX1
+    cy[0] += DY1
     cx[1] += DX2
     cy[1] += DY2
 
     return cx, cy
 
 
-def calc_association(cx, cy, clusters):
-
-    inds = []
-
-    for ic, _ in enumerate(cx):
-        tcx = cx[ic]
-        tcy = cy[ic]
-
-        dx = [icx - tcx for icx in clusters.cx]
-        dy = [icy - tcy for icy in clusters.cy]
-
-        dlist = [math.sqrt(idx**2 + idy**2) for (idx, idy) in zip(dx, dy)]
-        min_id = dlist.index(min(dlist))
-        inds.append(min_id)
-
-    return inds
-
-
 def main():
     print(__file__ + " start!!")
 
     cx = [0.0, 8.0]
     cy = [0.0, 8.0]
-    npoints = 10
+    n_points = 10
     rand_d = 3.0
-    ncluster = 2
+    n_cluster = 2
     sim_time = 15.0
     dt = 1.0
     time = 0.0
@@ -154,20 +124,18 @@ def main():
         print("Time:", time)
         time += dt
 
-        # simulate objects
+        # objects moving simulation
         cx, cy = update_positions(cx, cy)
-        rx, ry = calc_raw_data(cx, cy, npoints, rand_d)
+        raw_x, raw_y = calc_raw_data(cx, cy, n_points, rand_d)
 
-        clusters = kmeans_clustering(rx, ry, ncluster)
+        clusters = kmeans_clustering(raw_x, raw_y, n_cluster)
 
         # for animation
         if show_animation:  # pragma: no cover
             plt.cla()
-            inds = calc_association(cx, cy, clusters)
-            for ic in inds:
-                x, y = calc_labeled_points(ic, clusters)
-                plt.plot(x, y, "x")
-            plt.plot(cx, cy, "o")
+            clusters.plot_cluster()
+
+            plt.plot(cx, cy, "or")
             plt.xlim(-2.0, 10.0)
             plt.ylim(-2.0, 10.0)
             plt.pause(dt)