-
Notifications
You must be signed in to change notification settings - Fork 16
/
multimode.py
44 lines (38 loc) · 1.57 KB
/
multimode.py
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
import networkx as net
import matplotlib.pyplot as plot
from collections import defaultdict
def plot_multimode(m,layout=net.spring_layout, type_string='type', with_labels=True, filename_prefix='',output_type='pdf'):
## create a default color order and an empty color-map
colors=['r','g','b','c','m','y','k']
colormap={}
d=net.degree(m) #we use degree for sizing nodes
pos=layout(m) #compute layout
#Now we need to find groups of nodes that need to be colored differently
nodesets=defaultdict(list)
for n in m.nodes():
try:
t=m.node[n][type_string]
except KeyError:
##this happens if a node doesn't have a type_string -- give it a None value
t='None'
nodesets[t].append(n)
## Draw each group of nodes separately, using its own color settings
print "drawing nodes..."
i=0
for key in nodesets.keys():
ns=[d[n]*100 for n in nodesets[key]]
net.draw_networkx_nodes(m,pos,nodelist=nodesets[key], node_size=ns, node_color=colors[i], alpha=0.6)
colormap[key]=colors[i]
i+=1
if i==len(colors):
i=0 ### wrap around the colormap if we run out of colors
print colormap
## Draw edges using a default drawing mechanism
print "drawing edges..."
net.draw_networkx_edges(m,pos,width=0.5,alpha=0.5)
print "drawing labels..."
if with_labels:
net.draw_networkx_labels(m,pos,font_size=12)
plot.axis('off')
if filename_prefix is not '':
plot.savefig(filename_prefix+'.'+output_type)