yggdrasil-map/web/graphPlotter.py

import pygraphviz as pgv
import time
import json
import collections
import math
import networkx as nx
from networkx.algorithms import centrality

def position_nodes(nodes, edges):
    G = pgv.AGraph(strict=True, directed=False, size='10!')

    for n in nodes.values():
        G.add_node(n.ip, label=n.label, version=n.version)

    for e in edges:
        G.add_edge(e.a.ip, e.b.ip, len=1.0)

    G.layout(prog='neato', args='-Gepsilon=0.0001 -Gmaxiter=100000')

    return G

def compute_betweenness(G):
    ng = nx.Graph()
    for start in G.iternodes():
        others = G.neighbors(start)
        for other in others:
            ng.add_edge(start, other)

    c = centrality.betweenness_centrality(ng)

    for k, v in c.items():
        c[k] = v

    return c

def canonalize_ip(ip):
    return ':'.join( i.rjust(4, '0') for i in ip.split(':') )

def load_db():
    with open('nodedb/nodes') as f:
        return dict([ (canonalize_ip(v[0]), v[1]) for v in [ l.split(None)[:2] for l in f.readlines() ] if len(v) > 1 ])

def get_graph_json(G):
    max_neighbors = 1
    for n in G.iternodes():
        neighbors = len(G.neighbors(n))
        if neighbors > max_neighbors:
            max_neighbors = neighbors
    print 'Max neighbors: %d' % max_neighbors

    out_data = {
        'created': int(time.time()),
        'nodes': [],
        'edges': []
    }

    centralities = compute_betweenness(G)
    db = load_db()

    for n in G.iternodes():
        neighbor_ratio = len(G.neighbors(n)) / float(max_neighbors)
        pos = n.attr['pos'].split(',', 1)
        centrality = centralities.get(n.name, 0)
        pcentrality = (centrality + 0.0001) * 500
        size = (pcentrality ** 0.3 / 500) * 1000 + 1
        name = db.get(n.name)

        out_data['nodes'].append({
            'id': n.name,
            'label': name if name else n.attr['label'],
            'name': name,
            'version': n.attr['version'],
            'x': float(pos[0]),
            'y': float(pos[1]),
            'color': _gradient_color(neighbor_ratio, [(100, 100, 100), (0, 0, 0)]),
            'size': size,
            'centrality': '%.4f' % centrality
        })

    for e in G.iteredges():
        out_data['edges'].append({
            'sourceID': e[0],
            'targetID': e[1]
        })

    return json.dumps(out_data)


def _gradient_color(ratio, colors):
    jump = 1.0 / (len(colors) - 1)
    gap_num = int(ratio / (jump + 0.0000001))

    a = colors[gap_num]
    b = colors[gap_num + 1]

    ratio = (ratio - gap_num * jump) * (len(colors) - 1)

    r = a[0] + (b[0] - a[0]) * ratio
    g = a[1] + (b[1] - a[1]) * ratio
    b = a[2] + (b[2] - a[2]) * ratio

    return '#%02x%02x%02x' % (r, g, b)