Package networkx :: Package drawing :: Module nx_pylab
[hide private]
[frames] | no frames]

Source Code for Module networkx.drawing.nx_pylab

  1  """ 
  2  Draw networks with matplotlib (pylab). 
  3   
  4  Provides: 
  5   
  6   - draw() 
  7   - draw_networkx() 
  8   - draw_networkx_nodes() 
  9   - draw_networkx_edges() 
 10   - draw_networkx_labels() 
 11   - draw_circular 
 12   - draw_random 
 13   - draw_spectral 
 14   - draw_spring 
 15   - draw_shell 
 16   - draw_graphviz 
 17   
 18  References: 
 19   - matplotlib:     http://matplotlib.sourceforge.net/ 
 20   - pygraphviz:     http://networkx.lanl.gov/pygraphviz/ 
 21   
 22  """ 
 23  __author__ = """Aric Hagberg (hagberg@lanl.gov)""" 
 24  __date__ = "$Date: 2005-06-15 11:29:39 -0600 (Wed, 15 Jun 2005) $" 
 25  __credits__ = """""" 
 26  __revision__ = "$Id" 
 27  #    Copyright (C) 2004,2005 by  
 28  #    Aric Hagberg <hagberg@lanl.gov> 
 29  #    Dan Schult <dschult@colgate.edu> 
 30  #    Pieter Swart <swart@lanl.gov> 
 31  #    Distributed under the terms of the GNU Lesser General Public License 
 32  #    http://www.gnu.org/copyleft/lesser.html 
 33   
 34  import networkx 
 35  import sys 
 36   
 37  try: 
 38      import matplotlib 
 39      import matplotlib.cbook as cb 
 40      from matplotlib.colors import colorConverter,normalize,Colormap 
 41      from matplotlib.collections import LineCollection 
 42      from matplotlib.numerix import sin, cos, pi, sqrt, arctan2, asarray 
 43      from matplotlib.numerix.mlab import amin, amax, ravel 
 44      import matplotlib.pylab 
 45  except ImportError: 
 46      raise ImportError, "Import Error: not able to import matplotlib." 
 47   
48 -def draw(G, pos=None, ax=None, hold=None, **kwds):
49 """Draw the graph G with matplotlib (pylab). 50 51 This is a pylab friendly function that will use the 52 current pylab figure axes (e.g. subplot). 53 54 pos is a dictionary keyed by vertex with a two-tuple 55 of x-y positions as the value. 56 See networkx.layout for functions that compute node positions. 57 58 Usage: 59 60 >>> from networkx import * 61 >>> G=dodecahedral_graph() 62 >>> draw(G) 63 >>> pos=graphviz_layout(G) 64 >>> draw(G,pos) 65 >>> draw(G,pos=spring_layout(G)) 66 67 Also see doc/examples/draw_* 68 69 :Parameters: 70 71 - `nodelist`: list of nodes to be drawn (default=G.nodes()) 72 - `edgelist`: list of edges to be drawn (default=G.edges()) 73 - `node_size`: scalar or array of the same length as nodelist (default=300) 74 - `node_color`: single color string or numeric/numarray array of floats (default='r') 75 - `node_shape`: node shape (default='o'), or 'so^>v<dph8' see pylab.scatter 76 - `alpha`: transparency (default=1.0) 77 - `cmap`: colormap for mapping intensities (default=None) 78 - `vmin,vmax`: min and max for colormap scaling (default=None) 79 - `width`: line width of edges (default =1.0) 80 - `edge_color`: scalar or array (default='k') 81 - `edge_cmap`: colormap for edge intensities (default=None) 82 - `edge_vmin,edge_vmax`: min and max for colormap edge scaling (default=None) 83 - `style`: edge linestyle (default='solid') (solid|dashed|dotted,dashdot) 84 - `labels`: dictionary keyed by node of text labels (default=None) 85 - `font_size`: size for text labels (default=12) 86 - `font_color`: (default='k') 87 - `font_weight`: (default='normal') 88 - `font_family`: (default='sans-serif') 89 - `ax`: matplotlib axes instance 90 91 for more see pylab.scatter 92 93 NB: this has the same name as pylab.draw so beware when using 94 95 >>> from networkx import * 96 97 since you will overwrite the pylab.draw function. 98 99 A good alternative is to use 100 101 >>> import pylab as P 102 >>> import networkx as NX 103 >>> G=NX.dodecahedral_graph() 104 105 and then use 106 107 >>> NX.draw(G) # networkx draw() 108 109 and 110 >>> P.draw() # pylab draw() 111 112 """ 113 if pos is None: 114 pos=networkx.drawing.spring_layout(G) # default to spring layout 115 116 if ax is None: 117 ax=matplotlib.pylab.gca() 118 # allow callers to override the hold state by passing hold=True|False 119 b = ax.ishold() 120 if hold is not None: 121 matplotlib.pylab.hold(hold) 122 try: 123 # turn of axes ticks and labels 124 ax.set_xticks([]) 125 ax.set_yticks([]) 126 draw_networkx(G, pos, ax=ax, **kwds) 127 except: 128 matplotlib.pylab.hold(b) 129 raise 130 matplotlib.pylab.hold(b)
131
132 -def draw_networkx(G, pos, with_labels=True, **kwds):
133 """Draw the graph G with given node positions pos 134 135 Usage: 136 137 >>> from networkx import * 138 >>> import pylab as P 139 >>> ax=P.subplot(111) 140 >>> G=dodecahedral_graph() 141 >>> pos=spring_layout(G) 142 >>> draw_networkx(G,pos,ax=ax) 143 144 This is same as 'draw' but the node positions *must* be 145 specified in the variable pos. 146 pos is a dictionary keyed by vertex with a two-tuple 147 of x-y positions as the value. 148 See networkx.layout for functions that compute node positions. 149 150 An optional matplotlib axis can be provided through the 151 optional keyword ax. 152 153 with_labels contols text labeling of the nodes 154 155 Also see: 156 157 draw_networkx_nodes() 158 draw_networkx_edges() 159 draw_networkx_labels() 160 """ 161 from matplotlib.pylab import draw_if_interactive 162 node_collection=draw_networkx_nodes(G, pos, **kwds) 163 edge_collection=draw_networkx_edges(G, pos, **kwds) 164 if with_labels: 165 draw_networkx_labels(G, pos, **kwds) 166 draw_if_interactive()
167
168 -def draw_networkx_nodes(G, pos, 169 nodelist=None, 170 node_size=300, 171 node_color='r', 172 node_shape='o', 173 alpha=1.0, 174 cmap=None, 175 vmin=None, 176 vmax=None, 177 ax=None, 178 linewidths=None, 179 **kwds):
180 """Draw nodes of graph G 181 182 This draws only the nodes of the graph G. 183 184 pos is a dictionary keyed by vertex with a two-tuple 185 of x-y positions as the value. 186 See networkx.layout for functions that compute node positions. 187 188 nodelist is an optional list of nodes in G to be drawn. 189 If provided only the nodes in nodelist will be drawn. 190 191 see draw_networkx for the list of other optional parameters. 192 193 """ 194 if ax is None: 195 ax=matplotlib.pylab.gca() 196 197 if nodelist is None: 198 nodelist=G.nodes() 199 200 if not nodelist or len(nodelist)==0: # empty nodelist, no drawing 201 return None 202 203 try: 204 xy=asarray([pos[v] for v in nodelist]) 205 except KeyError,e: 206 raise networkx.NetworkXError('Node %s has no position.'%e) 207 except ValueError: 208 raise networkx.NetworkXError('Bad value in node positions.') 209 210 211 node_collection=ax.scatter(xy[:,0], xy[:,1], 212 s=node_size, 213 c=node_color, 214 marker=node_shape, 215 cmap=cmap, 216 vmin=vmin, 217 vmax=vmax, 218 alpha=alpha, 219 linewidths=linewidths) 220 221 node_collection.set_zorder(2) 222 return node_collection
223 224
225 -def draw_networkx_edges(G, pos, 226 edgelist=None, 227 width=1.0, 228 edge_color='k', 229 style='solid', 230 alpha=1.0, 231 edge_cmap=None, 232 edge_vmin=None, 233 edge_vmax=None, 234 ax=None, 235 arrows=True, 236 **kwds):
237 """Draw the edges of the graph G 238 239 This draws only the edges of the graph G. 240 241 pos is a dictionary keyed by vertex with a two-tuple 242 of x-y positions as the value. 243 See networkx.layout for functions that compute node positions. 244 245 edgelist is an optional list of the edges in G to be drawn. 246 If provided, only the edges in edgelist will be drawn. 247 248 edgecolor can be a list of matplotlib color letters such as 'k' or 249 'b' that lists the color of each edge; the list must be ordered in 250 the same way as the edge list. Alternatively, this list can contain 251 numbers and those number are mapped to a color scale using the color 252 map edge_cmap. 253 254 For directed graphs, "arrows" (actually just thicker stubs) are drawn 255 at the head end. Arrows can be turned off with keyword arrows=False. 256 257 See draw_networkx for the list of other optional parameters. 258 259 """ 260 if ax is None: 261 ax=matplotlib.pylab.gca() 262 263 if edgelist is None: 264 edgelist=G.edges() 265 266 if not edgelist or len(edgelist)==0: # no edges! 267 return None 268 269 # set edge positions 270 edge_pos=asarray([(pos[e[0]],pos[e[1]]) for e in edgelist]) 271 272 if not cb.iterable(width): 273 lw = (width,) 274 else: 275 lw = width 276 277 if not cb.is_string_like(edge_color) \ 278 and cb.iterable(edge_color) \ 279 and len(edge_color)==len(edge_pos): 280 if matplotlib.numerix.alltrue([cb.is_string_like(c) 281 for c in edge_color]): 282 # (should check ALL elements) 283 # list of color letters such as ['k','r','k',...] 284 edge_colors = tuple([colorConverter.to_rgba(c,alpha) 285 for c in edge_color]) 286 elif matplotlib.numerix.alltrue([not cb.is_string_like(c) 287 for c in edge_color]): 288 # numbers (which are going to be mapped with a colormap) 289 edge_colors = None 290 else: 291 raise ValueError('edge_color must consist of either color names or numbers') 292 else: 293 if len(edge_color)==1: 294 edge_colors = ( colorConverter.to_rgba(edge_color, alpha), ) 295 else: 296 raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges') 297 298 edge_collection = LineCollection(edge_pos, 299 colors = edge_colors, 300 linewidths = lw, 301 antialiaseds = (1,), 302 linestyle = style, 303 transOffset = ax.transData, 304 ) 305 edge_collection.set_alpha(alpha) 306 307 # need 0.87.7 or greater for edge colormaps 308 mpl_version=matplotlib.__version__ 309 if mpl_version.endswith('svn'): 310 mpl_version=matplotlib.__version__[0:-3] 311 if mpl_version.endswith('pre'): 312 mpl_version=matplotlib.__version__[0:-3] 313 if map(int,mpl_version.split('.'))>=[0,87,7]: 314 if edge_colors is None: 315 if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap)) 316 edge_collection.set_array(asarray(edge_color)) 317 edge_collection.set_cmap(edge_cmap) 318 if edge_vmin is not None or edge_vmax is not None: 319 edge_collection.set_clim(edge_vmin, edge_vmax) 320 else: 321 edge_collection.autoscale() 322 # else: 323 # sys.stderr.write(\ 324 # """matplotlib version >= 0.87.7 required for colormapped edges. 325 # (version %s detected)."""%matplotlib.__version__) 326 # raise UserWarning(\ 327 # """matplotlib version >= 0.87.7 required for colormapped edges. 328 # (version %s detected)."""%matplotlib.__version__) 329 330 arrow_collection=None 331 332 if G.is_directed() and arrows: 333 334 # a directed graph hack 335 # draw thick line segments at head end of edge 336 # waiting for someone else to implement arrows that will work 337 arrow_colors = ( colorConverter.to_rgba('k', alpha), ) 338 a_pos=[] 339 p=1.0-0.25 # make head segment 25 percent of edge length 340 for src,dst in edge_pos: 341 x1,y1=src 342 x2,y2=dst 343 dx=x2-x1 # x offset 344 dy=y2-y1 # y offset 345 d=sqrt(float(dx**2+dy**2)) # length of edge 346 if d==0: # source and target at same position 347 continue 348 if dx==0: # vertical edge 349 xa=x2 350 ya=dy*p+y1 351 if dy==0: # horizontal edge 352 ya=y2 353 xa=dx*p+x1 354 else: 355 theta=arctan2(dy,dx) 356 xa=p*d*cos(theta)+x1 357 ya=p*d*sin(theta)+y1 358 359 a_pos.append(((xa,ya),(x2,y2))) 360 361 arrow_collection = LineCollection(a_pos, 362 colors = arrow_colors, 363 linewidths = [4*ww for ww in lw], 364 antialiaseds = (1,), 365 transOffset = ax.transData, 366 ) 367 368 # update view 369 minx = amin(ravel(edge_pos[:,:,0])) 370 maxx = amax(ravel(edge_pos[:,:,0])) 371 miny = amin(ravel(edge_pos[:,:,1])) 372 maxy = amax(ravel(edge_pos[:,:,1])) 373 374 375 376 w = maxx-minx 377 h = maxy-miny 378 padx, pady = 0.05*w, 0.05*h 379 corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady) 380 ax.update_datalim( corners) 381 ax.autoscale_view() 382 383 edge_collection.set_zorder(1) # edges go behind nodes 384 ax.add_collection(edge_collection) 385 if arrow_collection: 386 arrow_collection.set_zorder(1) # edges go behind nodes 387 ax.add_collection(arrow_collection) 388 389 390 return edge_collection
391 392
393 -def draw_networkx_labels(G, pos, 394 labels=None, 395 font_size=12, 396 font_color='k', 397 font_family='sans-serif', 398 font_weight='normal', 399 alpha=1.0, 400 ax=None, 401 **kwds):
402 """Draw node labels on the graph G 403 404 pos is a dictionary keyed by vertex with a two-tuple 405 of x-y positions as the value. 406 See networkx.layout for functions that compute node positions. 407 408 labels is an optional dictionary keyed by vertex with node labels 409 as the values. If provided only labels for the keys in the dictionary 410 are drawn. 411 412 See draw_networkx for the list of other optional parameters. 413 414 """ 415 if ax is None: 416 ax=matplotlib.pylab.gca() 417 418 if labels is None: 419 labels=dict(zip(G.nodes(),G.nodes())) 420 421 text_items={} # there is no text collection so we'll fake one 422 for (n,label) in labels.items(): 423 (x,y)=pos[n] 424 if not cb.is_string_like(label): 425 label=str(label) # this will cause "1" and 1 to be labeled the same 426 t=ax.text(x, y, 427 label, 428 size=font_size, 429 color=font_color, 430 family=font_family, 431 weight=font_weight, 432 horizontalalignment='center', 433 verticalalignment='center', 434 transform = ax.transData, 435 ) 436 text_items[n]=t 437 438 return text_items
439
440 -def draw_circular(G, **kwargs):
441 """Draw the graph G with a circular layout""" 442 from networkx.drawing.layout import circular_layout 443 draw(G,circular_layout(G),**kwargs)
444
445 -def draw_random(G, **kwargs):
446 """Draw the graph G with a random layout.""" 447 from networkx.drawing.layout import random_layout 448 draw(G,random_layout(G),**kwargs)
449
450 -def draw_spectral(G, **kwargs):
451 """Draw the graph G with a spectral layout.""" 452 from networkx.drawing.layout import spectral_layout 453 draw(G,spectral_layout(G),**kwargs)
454
455 -def draw_spring(G, **kwargs):
456 """Draw the graph G with a spring layout""" 457 from networkx.drawing.layout import spring_layout 458 draw(G,spring_layout(G),**kwargs)
459
460 -def draw_shell(G, **kwargs):
461 """Draw networkx graph with shell layout""" 462 from networkx.drawing.layout import shell_layout 463 nlist = kwargs.get('nlist', None) 464 if nlist != None: 465 del(kwargs['nlist']) 466 draw(G,shell_layout(G,nlist=nlist),**kwargs)
467
468 -def draw_graphviz(G, prog="neato", **kwargs):
469 """Draw networkx graph with graphviz layout""" 470 pos=networkx.drawing.graphviz_layout(G,prog) 471 draw(G,pos,**kwargs)
472
473 -def draw_nx(G,pos,**kwds):
474 """For backward compatibility; use draw or draw_networkx""" 475 draw(G,pos,**kwds)
476
477 -def _test_suite():
478 import doctest 479 suite = doctest.DocFileSuite('tests/drawing/nx_pylab.txt',\ 480 package='networkx') 481 return suite
482 483 if __name__ == "__main__": 484 import os 485 import sys 486 import unittest 487 488 if sys.version_info[:2] < (2, 4): 489 print "Python version 2.4 or later required (%d.%d detected)." \ 490 % sys.version_info[:2] 491 sys.exit(-1) 492 # directory of networkx package (relative to this) 493 nxbase=sys.path[0]+os.sep+os.pardir 494 sys.path.insert(0,nxbase) # prepend to search path 495 unittest.TextTestRunner().run(_test_suite()) 496