Merge pull request #446 from maximlt/improve_viz_riverml

martindurant · web-flow · commit 8744c83a0fad · 2021-12-24T13:36:59.000-05:00
RiverML: update the plots less frequently
diff --git a/examples/river_kmeans.py b/examples/river_kmeans.py
@@ -1,4 +1,6 @@
+import functools
 import random
+import time
 
 import pandas as pd
 
@@ -8,6 +10,7 @@
 from river import cluster
 import holoviews as hv
 from panel.pane.holoviews import HoloViews
+import panel as pn
 hv.extension('bokeh')
 
 model = cluster.KMeans(n_clusters=3, sigma=0.1, mu=0.5)
@@ -32,14 +35,49 @@ def get_clusters(model):
 
 
 def main(viz=True):
+    cadance = 0.01
+
+    ex = pd.DataFrame({'x': [0.5], 'y': [0.5]})
+    pipe_in = hv.streams.Pipe(data=ex)
+    pipe_out = hv.streams.Pipe(data=ex)
+
     # setup pipes
-    cadance = 0.16 if viz else 0.01
     s = Stream.from_periodic(gen, cadance)
+
+    # Branch 0: Input/Observations
+    obs = s.map(lambda x: pd.DataFrame([x]))
+
+    # Branch 1: Output/River ML clusters
     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]})
-    ooo = s.map(lambda x: pd.DataFrame([x])).to_dataframe(example=ex)
-    out = km.map(get_clusters)
+    clusters = km.map(get_clusters)
+
+    concat = functools.partial(pd.concat, ignore_index=True)
+
+    def accumulate(previous, new, last_lines=50):
+        return concat([previous, new]).iloc[-last_lines:, :]
+
+    partition_obs = 10
+    particion_clusters = 10
+    backlog_obs = 100
+
+    # .partition is used to gather x number of points
+    # before sending them to the plots
+    # .accumulate allows to generate a backlog
+
+    (
+        obs
+        .partition(partition_obs)
+        .map(concat)
+        .accumulate(functools.partial(accumulate, last_lines=backlog_obs))
+        .sink(pipe_in.send)
+    )
+    (
+        clusters
+        .partition(particion_clusters)
+        .map(pd.concat)
+        .sink(pipe_out.send)
+    )
 
     # start things
     s.emit(gen())  # set initial model
@@ -48,23 +86,47 @@ def main(viz=True):
         model.centers[i]['y'] = y
 
     print("starting")
-    s.start()
 
     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))
+        button_start = pn.widgets.Button(name='Start')
+        button_stop = pn.widgets.Button(name='Stop')
+
+        t0 = 0
+
+        def start(event):
+            s.start()
+            global t0
+            t0 = time.time()
+
+        def stop(event):
+            print(count, "events")
+            global t0
+            t_spent = time.time() - t0
+            print("frequency", count[0] / t_spent, "Hz")
+            print("Current centres", centres)
+            print("Output centres", [list(c.values()) for c in model.centers.values()])
+            s.stop()
+
+        button_start.on_click(start)
+        button_stop.on_click(stop)
+
+        scatter_dmap_input = hv.DynamicMap(hv.Scatter, streams=[pipe_in]).opts(color="blue")
+        scatter_dmap_output = hv.DynamicMap(hv.Scatter, streams=[pipe_out]).opts(color="red")
+        pl = scatter_dmap_input * scatter_dmap_output
         pl.opts(xlim=(-0.2, 1.2), ylim=(-0.2, 1.2), height=600, width=600)
+
         pan = HoloViews(pl)
-        pan.show()
+        app = pn.Row(pn.Column(button_start, button_stop), pan)
+        app.show()
     else:
-        import time
+        s.start()
         time.sleep(5)
         print(count, "events")
+        print("frequency", count[0] / 5, "Hz")
         print("Current centres", centres)
         print("Output centres", [list(c.values()) for c in model.centers.values()])
-    s.stop()
+        s.stop()
 
 if __name__ == "__main__":
     main(viz=True)
