first release k means simulation

Author: AtsushiSakai

Mapping/kmean_clustering/kmean_clustering.py

@@ -1,6 +1,6 @@
 """
 
-Object clustering with k-mean algorithm
+Object clustering with k-means algorithm
 
 author: Atsushi Sakai (@Atsushi_twi)
 
@@ -25,9 +25,23 @@ def __init__(self, x, y, nlabel):
         self.cy = [0.0 for _ in range(nlabel)]
 
 
-def init_clusters(rx, ry, nc):
+def kmeans_clustering(rx, ry, nc):
 
     clusters = Clusters(rx, ry, nc)
+    clusters = calc_centroid(clusters)
+
+    MAX_LOOP = 10
+    DCOST_TH = 0.1
+    pcost = 100.0
+    for loop in range(MAX_LOOP):
+        #  print("Loop:", loop)
+        clusters, cost = update_clusters(clusters)
+        clusters = calc_centroid(clusters)
+
+        dcost = abs(cost - pcost)
+        if dcost < DCOST_TH:
+            break
+        pcost = cost
 
     return clusters
 
@@ -62,36 +76,23 @@ def update_clusters(clusters):
     return clusters, cost
 
 
-def kmean_clustering(rx, ry, nc):
-
-    clusters = init_clusters(rx, ry, nc)
-    clusters = calc_centroid(clusters)
+def calc_labeled_points(ic, clusters):
 
-    MAX_LOOP = 10
-    DCOST_TH = 1.0
-    pcost = 100.0
-    for loop in range(MAX_LOOP):
-        print("Loop:", loop)
-        clusters, cost = update_clusters(clusters)
-        clusters = calc_centroid(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)
 
-        dcost = abs(cost - pcost)
-        if dcost < DCOST_TH:
-            break
-        pcost = cost
+    x = tx[inds]
+    y = ty[inds]
 
-    return clusters
+    return x, y
 
 
-def calc_raw_data():
+def calc_raw_data(cx, cy, npoints, rand_d):
 
     rx, ry = [], []
 
-    cx = [0.0, 5.0]
-    cy = [0.0, 5.0]
-    npoints = 30
-    rand_d = 3.0
-
     for (icx, icy) in zip(cx, cy):
         for _ in range(npoints):
             rx.append(icx + rand_d * (random.random() - 0.5))
@@ -100,32 +101,75 @@ def calc_raw_data():
     return rx, ry
 
 
-def calc_labeled_points(ic, clusters):
+def update_positions(cx, cy):
 
-    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)
+    DX1 = 0.4
+    DY1 = 0.5
 
-    x = tx[inds]
-    y = ty[inds]
+    cx[0] += DX1
+    cy[0] += DY1
 
-    return x, y
+    DX2 = -0.3
+    DY2 = -0.5
+
+    cx[1] += DX2
+    cy[1] += DY2
+
+    return cx, cy
+
+
+def calc_association(cx, cy, clusters):
+
+    inds = []
+
+    for ic in range(len(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!!")
 
-    rx, ry = calc_raw_data()
-
+    cx = [0.0, 8.0]
+    cy = [0.0, 8.0]
+    npoints = 10
+    rand_d = 3.0
     ncluster = 2
-    clusters = kmean_clustering(rx, ry, ncluster)
-
-    for ic in range(clusters.nlabel):
-        x, y = calc_labeled_points(ic, clusters)
-        plt.plot(x, y, "x")
-    plt.plot(clusters.cx, clusters.cy, "o")
-    plt.show()
+    sim_time = 15.0
+    dt = 1.0
+    time = 0.0
+
+    while time <= sim_time:
+        print("Time:", time)
+        time += dt
+
+        # simulate objects
+        cx, cy = update_positions(cx, cy)
+        rx, ry = calc_raw_data(cx, cy, npoints, rand_d)
+
+        clusters = kmeans_clustering(rx, ry, ncluster)
+
+        # for animation
+        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")
+        plt.xlim(-2.0, 10.0)
+        plt.ylim(-2.0, 10.0)
+        plt.pause(dt)
+
+    print("Done")
 
 
 if __name__ == '__main__':