tools: Add script to generate bandwidth graphs

This adds a script to generate bandwidth graphs from an exported
bandwidth csv. It will color the nodes and links depending on the load
and also adjust the thickness of the edges.

2 files changed, 118 insertions, 4 deletions

diff --git a/rumba/storyboard.py b/rumba/storyboard.py
index beadf91..5ff7a98 100644
--- a/rumba/storyboard.py
+++ b/rumba/storyboard.py
@@ -1046,12 +1046,9 @@ class StoryBoard(_SBEntity):
         if self.experiment is None:
             raise ValueError("An experiment is needed to schedule commands.")
 
-        if self._script is None:
-            self._script = _Script(self)
-
         action = functools.partial(self.experiment.export_dif_bandwidth,
                                    filename, dif)
-        self._script.add_event(Event(action, ev_time=t))
+        self.add_event(Event(action, ev_time=t))
 
     def schedule_link_state(self, t, dif, state):
         """
diff --git a/tools/bw_graph.py b/tools/bw_graph.py
new file mode 100755
index 0000000..ab772d4
--- /dev/null
+++ b/tools/bw_graph.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+
+import argparse
+import csv
+import math
+import sys
+
+try:
+    from colour import Color
+except:
+    print('Colour module not installed, please install it')
+    sys.exit(1)
+
+try:
+    import pydot
+except:
+    print('Pydot module not installed, please install it')
+    sys.exit(1)
+
+color_range = list(Color("DarkGreen").range_to(Color("DarkRed"), 20))
+multiplier = float(len(color_range))
+
+def get_color(value, max):
+    if value < max:
+        val = math.floor(multiplier * value / (max))
+    else:
+        val = math.floor(multiplier - 1)
+
+    color = color_range[val]
+
+    return color.get_hex()
+
+def get_penwidth(value, max):
+    if max == 0:
+        return "1.0"
+
+    r = (4.0 * float(value) / (0.5 * max)) + 1
+    return str(r)
+
+description = "CSV to bandwidth graph script"
+epilog = "2018 Sander Vrijders <sander.vrijders@ugent.be>"
+
+parser = argparse.ArgumentParser(description=description,
+                                 epilog=epilog)
+
+parser.add_argument('-o', '--output', type=str, required=True,
+                    help='Output PDF filename.')
+
+parser.add_argument('-i', '--input', type=str, required=True,
+                    help='Input CSV filename.')
+
+parser.add_argument('-m', '--max', type=int, default=10000, required=False,
+                    help='Maximum bandwidth in the network.')
+
+args = parser.parse_args()
+
+gvizg = pydot.Dot(graph_type='graph', layout='fdp')
+
+csvfile = open(args.input, 'r')
+reader = csv.reader(csvfile, delimiter=';', quotechar='|')
+
+bandwidth = {}
+nodes = set()
+for row in reader:
+    nodes.add(row[0])
+    if (row[1], row[0]) in bandwidth:
+        bandwidth[row[1], row[0]] += int(row[2])
+    else:
+        bandwidth[row[0], row[1]] = int(row[2])
+
+node_bandwidth = {}
+for key, value in bandwidth.items():
+    n0, n1 = key
+    if n0 in node_bandwidth:
+        node_bandwidth[n0] += value
+    else:
+        node_bandwidth[n0] = value
+    if n1 in node_bandwidth:
+        node_bandwidth[n1] += value
+    else:
+        node_bandwidth[n1] = value
+
+max_b = args.max
+for key, value in bandwidth.items():
+    max_b = max(value, max_b)
+
+for node in nodes:
+    try:
+        color = get_color(value, 2 * max_b)
+        g_node = pydot.Node(node,
+                            label=node,
+                            style="filled",
+                            fillcolor=color,
+                            color='white',
+                            fontcolor='white')
+        gvizg.add_node(g_node)
+    except Exception as e:
+        print('Failed to create pydot Node: ' + str(e))
+
+for key, value in bandwidth.items():
+    try:
+        color = get_color(value, max_b)
+        pw = get_penwidth(value, max_b)
+        edge = pydot.Edge(key[0], key[1],
+                          color = color,
+                          penwidth=pw)
+        gvizg.add_edge(edge)
+    except Exception as e:
+        print('Failed to create pydot Edge: ' + str(e))
+
+try:
+    gvizg.write_pdf(args.output)
+    print('Generated PDF of bandwidth graph')
+except Exception as e:
+    print('Failed to write PDF: ' + str(e))
+
+csvfile.close()