swdata/draw_network.py
2018-11-17 07:13:44 +08:00

67 lines
2.1 KiB
Python

import networkx as nx
from matplotlib import pyplot as plt
from island.match import Match
from island.matches import Matches
class Draw:
def __init__(self, show):
self.matches = Matches.from_profile_expr(lambda r: True)
self.show = show
def generate_graph(self, match):
g = nx.empty_graph()
for r in match.query('player', 'join').raw_data:
g.add_node(r['pid'])
for r in match.query('action', 'done').raw_data:
if r['a'] > r['b']:
f, t = r['b'], r['a']
else:
f, t = r['a'], r['b']
if g.has_edge(f, t):
g.edges[f, t]['weight'] += 1
else:
g.add_edge(f, t, weight=1, t=1)
for r in match.query('neighbor', 'create').raw_data:
if r['a'] > r['b']:
f, t = r['b'], r['a']
else:
f, t = r['a'], r['b']
if g.has_edge(f, t):
g.edges[f, t]['t'] |= 2
else:
g.add_edge(f, t, weight=1, t=2)
return g
def draw_graph(self, g, m):
social_edges = []
game_edges = []
weight = []
for e in g.edges():
if g[e[0]][e[1]]['t'] & 1 == 1:
game_edges.append(e)
weight.append(g[e[0]][e[1]]['weight'])
else:
social_edges.append(e)
pos = nx.spring_layout(g, scale=0.3)
nx.draw_networkx_edges(g, edgelist=social_edges, pos=pos, edge_color='gray', alpha=0.8, style='dashed')
nx.draw_networkx_edges(g, edgelist=game_edges, pos=pos, width=weight, alpha=0.8, edge_color='#d63031')
nx.draw_networkx_nodes(g, pos=pos, node_color='#0984e3', node_size=300)
if self.show:
plt.show()
else:
plt.savefig("graph/%s-cluster.eps"%m.name)
def draw_a_graph(self, name):
for m in self.matches.data:
if m.name == name:
g = self.generate_graph(m)
self.draw_graph(g, m)
if __name__ == '__main__':
d = Draw(True)
d.draw_a_graph('G398')