better working river example

Author: Martin Durant

examples/river_kmeans.py

@@ -14,15 +14,18 @@
 import tornado.ioloop
 hv.extension('bokeh')
 
-model = cluster.STREAMKMeans(n_clusters=3)
+model = cluster.KMeans(n_clusters=3, sigma=0.1, mu=0.5)
 centres = [[random.random(), random.random()] for _ in range(3)]
+count = [0]
 
-def gen(split_chance=0.01):
+def gen(move_chance=0.05):
     centre = int(random.random() * 3)  # 3x faster than random.randint(0, 2)
-    if random.random() < split_chance:
-        centres[centre] = [random.random(), random.random()]
+    if random.random() < move_chance:
+        centres[centre][0] += random.random() / 5 - 0.1
+        centres[centre][1] += random.random() / 5 - 0.1
     value = {'x': random.random() / 20 + centres[centre][0],
              'y': random.random() / 20 + centres[centre][1]}
+    count[0] += 1
     return value
 
 
@@ -32,9 +35,10 @@ def get_clusters(model):
     return pd.DataFrame(data, index=range(3))
 
 
-def main():
+def main(viz=True):
     # setup pipes
-    s = Stream.from_periodic(gen, 0.05)
+    cadance = 0.16 if viz else 0.01
+    s = Stream.from_periodic(gen, cadance)
     km = RiverTrain(model, pass_model=True)
     s.map(lambda x: (x,)).connect(km)  # learn takes a tuple of (x,[ y[, w]])
     ex = pd.DataFrame({'x': [0.5], 'y': [0.5]})
@@ -47,16 +51,23 @@ def main():
         model.centers[i]['x'] = x
         model.centers[i]['y'] = y
 
+    print("starting")
     s.start()
 
-    # plot
-    pout = out.to_dataframe(example=ex)
-    pl = (ooo.hvplot.scatter('x', 'y', color="blue", backlog=100) *
-          pout.hvplot.scatter('x', 'y', color="red", backlog=3))
-    pl.opts(xlim=(-0.5, 1.5), ylim=(-0.5, 1.5), height=600, width=600)
-    pan = panel.pane.holoviews.HoloViews(pl)
-    pan.show()
-
+    if viz:
+        # plot
+        pout = out.to_dataframe(example=ex)
+        pl = (ooo.hvplot.scatter('x', 'y', color="blue", backlog=50) *
+              pout.hvplot.scatter('x', 'y', color="red", backlog=3))
+        pl.opts(xlim=(-0.2, 1.2), ylim=(-0.2, 1.2), height=600, width=600)
+        pan = panel.pane.holoviews.HoloViews(pl)
+        pan.show(threaded=True)
+    else:
+        import time
+        time.sleep(5)
+        print(count, "events")
+        print("Current centres", centres)
+        print("Output centres", [list(c.values()) for c in model.centers.values()])
 
 if __name__ == "__main__":
-    main()
+    main(viz=True)