aboutsummaryrefslogtreecommitdiff
path: root/tools/bw_graph.py
blob: ab772d4237ae6d0ec3803f7f02caee1e96885a68 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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()