Package nltk_lite :: Package contrib :: Package fst :: Module draw_graph
[hide private]
[frames] | no frames]

Source Code for Module nltk_lite.contrib.fst.draw_graph

  1  # Natural Language Toolkit: Graph Visualization 
  2  # 
  3  # Copyright (C) 2001 University of Pennsylvania 
  4  # Author: Edward Loper <edloper@gradient.cis.upenn.edu> 
  5  # 
  6  # URL: <http://nltk.sf.net> 
  7  # For license information, see LICENSE.TXT 
  8  # 
  9  # $Id: graph.py,v 1.2 2004/03/18 21:02:36 edloper Exp $ 
 10   
 11  """ 
 12  Graphically display a graph.  This module defines two new canvas 
 13  widgets: L{GraphEdgeWidget}, and L{GraphWidget}.  Together, these two 
 14  widgets can be used to display directed graphs. 
 15   
 16  C{GraphEdgeWidget} is an arrow, optionally annotated with a 'label', 
 17  which can be any canvas widget.  In addition to a source location and 
 18  a destination location, it has a 'curve' attribute, which can be used 
 19  to define how curved it is (positive values curve one way, and 
 20  negative values the other).  This is useful, e.g., if you want to have 
 21  two separate graph edges with the same source and the same 
 22  destination.  It is also useful for drawing arrows that have the same 
 23  source and destination (i.e., loops). 
 24   
 25  The C{GraphWidget} widget is used to display a single directed graph. 
 26  It is a container widget, containing zero or more I{node widgets}, 
 27  which are connected by zero or more I{edge widgets}.  Any canvas 
 28  widget can be used as a node widget.  E.g., a StackWidget containing 
 29  an OvalWidget and a LabelWidget could be used to draw a circle with a 
 30  label below it.  Edge widgets must be C{GraphEdgeWidgets}.  The 
 31  C{GraphWidget} is responsible for adjusting the start and end 
 32  positions of edge widgets whenever node widgets move.  Thus, you can 
 33  make a node widget draggable, and when the user drags it, the edges 
 34  will update automatically.  The C{GraphWidget} also defines a method 
 35  C{arrange}, which will automatically choose a layout for the nodes, 
 36  attempting to minimize crossing edges. 
 37  """ 
 38   
 39  import math 
 40  from nltk_lite.draw import * 
 41   
42 -class GraphEdgeWidget(CanvasWidget):
43 """ 44 A canvas widget used to display graph edges. 45 46 @todo: Add an 'arrow' attribute, which can be used to control the 47 direction and/or shape of the arrow. 48 """
49 - def __init__(self, canvas, x1, y1, x2, y2, label=None, **attribs):
50 self._curve = 0 51 coords = self._line_coords((x1, y1), (x2, y2)) 52 self._line = canvas.create_line(arrow='last', smooth=1, *coords) 53 canvas.lower(self._line) 54 self._label = label 55 if label is not None: 56 self._add_child_widget(label) 57 58 CanvasWidget.__init__(self, canvas, **attribs)
59
60 - def __setitem__(self, attr, value):
61 if attr == 'curve': 62 self._curve = value 63 coords = self._line_coords(self.start(), self.end()) 64 self.canvas().coords(self._line, *coords) 65 elif attr == 'color': 66 self.canvas().itemconfig(self._line, fill=value) 67 elif attr == 'width': 68 self.canvas().itemconfig(self._line, width=value) 69 else: 70 CanvasWidget.__setitem__(self, attr, value)
71
72 - def __getitem__(self, attr):
73 if attr == 'curve': 74 return self._curve 75 elif attr == 'color': 76 return self.canvas().itemcget(self._line, fill) 77 elif attr == 'width': 78 return self.canvas().itemcget(self._line, width) 79 else: 80 return CanvasWidget.__getitem__(self, attr)
81
82 - def _tags(self): return [self._line]
83
84 - def __repr__(self):
85 return '[GraphEdge: %r %r->%r]' % (self._label, self.start(), 86 self.end())
87
88 - def start(self):
89 return self.canvas().coords(self._line)[:2]
90
91 - def end(self):
92 return self.canvas().coords(self._line)[-2:]
93
94 - def set_start(self, x, y):
95 coords = self._line_coords((x, y), self.end()) 96 self.canvas().coords(self._line, *coords) 97 self.update(self._label)
98
99 - def set_end(self, x, y):
100 coords = self._line_coords(self.start(), (x, y)) 101 self.canvas().coords(self._line, *coords) 102 self.update(self._label)
103
104 - def _update(self, child):
105 # The label moved? 106 (x1, y1, x2, y2) = child.bbox() 107 (x, y) = self._label_coords() 108 child.move(x-(x1+x2)/2, y-(y1+y2)/2)
109
110 - def _manage(self):
111 if self._label is not None: 112 self._update(self._label)
113
114 - def _label_coords(self):
115 line_coords = self.canvas().coords(self._line) 116 (x1, y1) = line_coords[:2] 117 (x2, y2) = line_coords[-2:] 118 if (x1, y1) == (x2, y2): 119 # Self-loops. Emperically, this formula seems about 120 # right, but it wasn't derived mathmatically. 121 radius = 0 122 return (x1, y1 + 0.81*(150*self._curve+radius) + 10) 123 elif self._curve >= 0: 124 # Normal edges. 125 r = max(math.sqrt((x1-x2)**2 + (y1-y2)**2), 1) 126 labelx = (x1+x2)*0.5 + (y2-y1)*(self._curve*.6) 127 labely = (y1+y2)*0.5 - (x2-x1)*(self._curve*.6) 128 return (int(labelx), int(labely)) 129 else: 130 # Normal edges. 131 r = max(math.sqrt((x1-x2)**2 + (y1-y2)**2), 1) 132 labelx = (x1+x2)*0.5 + (y2-y1)*(self._curve/2 + 8/r) 133 labely = (y1+y2)*0.5 - (x2-x1)*(self._curve/2 + 8/r) 134 return (int(labelx), int(labely))
135
136 - def _line_coords(self, (startx, starty), (endx, endy)):
137 (x1, y1) = int(startx), int(starty) 138 (x2, y2) = int(endx), int(endy) 139 radius1 = 0 140 radius2 = 0 141 142 if abs(x1-x2)+abs(y1-y2) < 5: 143 # Self-loops 144 x3 = x1 - 70*self._curve - radius1 145 y3 = y1 + 70*self._curve + radius1 146 x4 = x1 147 y4 = y1 + 140*self._curve + radius1 148 x5 = x1 + 70*self._curve + radius1 149 y5 = y1 + 70*self._curve + radius1 150 return (int(x1), int(y1), int(x3), int(y3), int(x4), 151 int(y4), int(x5), int(y5), int(x1), int(y1)) 152 else: 153 # Normal edges. 154 x3 = (x1+x2)*0.5 + (y2-y1)*self._curve 155 y3 = (y1+y2)*0.5 - (x2-x1)*self._curve 156 # Adjust endpoints so they end at the node parimeter. 157 r = max(math.sqrt((x1-x2)**2 + (y1-y2)**2), 0.001) 158 (dx, dy) = (x2-x1, y2-y1) 159 x1 += dx/r * radius1 160 y1 += dy/r * radius1 161 x2 -= dx/r * radius2 162 y2 -= dy/r * radius2 163 return (int(x1), int(y1), int(x3), int(y3), int(x2), int(y2))
164
165 -class GraphWidget(CanvasWidget):
166 """ 167 A canvas widget used to display directed graphs. This container 168 widget contains zero or more 'node widgets', which are connected by 169 zero or more C{GraphEdgeWidget}s. The C{GraphWidget} is responsible 170 for updating the edge widgets when nodes move; and for initially 171 arranging the nodes. 172 """
173 - def __init__(self, canvas, nodes, edges, **attrs):
174 """ 175 @param edges: A list of tuples (n1, n2, e), where n1 is a 176 CanvasWidget in C{nodes}; n2 is a CanvasWidget in 177 C{nodes}; and e is a GraphEdgeWidget. 178 """ 179 if len(nodes) == 0: 180 # dummy node, since we need to have a bbox. 181 nodes = [SpaceWidget(canvas,0,0)] 182 183 self._nodes = set(nodes) 184 185 # Management parameters. I should add attributes for these. 186 self._arrange = 'dfs' 187 self._xspace = attrs.pop('xspace', 50) 188 189 self._yspace = attrs.pop('yspace', 50) 190 self._orientation = attrs.pop('orientation', 'horizontal') 191 192 # Attributes for edges. 193 194 # Out & in edges for a given node 195 self._outedges = {} 196 self._inedges = {} 197 198 # Start & end nodes for a given edge. 199 self._startnode = {} 200 self._endnode = {} 201 202 # Keep track of edge widgets. 203 self._edgewidgets = {} 204 205 self._initialized = False 206 for node in self._nodes: 207 self.add_node(node) 208 for (start, end, edgewidget) in edges: 209 self.add_edge(start, end, edgewidget) 210 self._initialized = True 211 212 CanvasWidget.__init__(self, canvas, **attrs)
213
214 - def add_node(self, node):
215 """ 216 Add a new node to the graph. 217 """ 218 self._add_child_widget(node) 219 self._nodes.add(node)
220
221 - def add_edge(self, start, end, edgewidget):
222 """ 223 Add a new edge to the graph. 224 @param start: The start node 225 @type start: C{CanvasWidget} 226 @param end: The end node 227 @type end: C{CanvasWidget} 228 @param edgewidget: The edge 229 @type edgewidget: C{GraphEdgeWidget} 230 """ 231 num_edges = (len(self._edgewidgets.get( (start, end), [])) + 232 len(self._edgewidgets.get( (end, start), []))) 233 if start is end: 234 num_edges = num_edges/2+1 235 curve = 0.3 * ((num_edges+1)/2) * (num_edges%2*2-1) 236 else: 237 curve = 0.4 * ((num_edges+1)/2) * (num_edges%2*2-1) 238 edgewidget['curve'] = curve 239 240 # Add the edge to the outedge & inedge dictionaries 241 self._outedges.setdefault(start, []).append(edgewidget) 242 self._inedges.setdefault(end, []).append(edgewidget) 243 244 self._startnode[edgewidget] = start 245 self._endnode[edgewidget] = end 246 247 self._edgewidgets.setdefault((start, end),[]).append(edgewidget) 248 self._add_child_widget(edgewidget) 249 if self._initialized: self._update_edge(edgewidget)
250
251 - def remove_edge(self, edge):
252 """ 253 Remove an edge from the graph (but don't destroy it). 254 @type edge: L{GraphEdgeWidget} 255 """ 256 print 'remove', edge 257 # Get the edge's start & end nodes. 258 start, end = self._startnode[edge], self._endnode[edge] 259 260 # Remove the edge from the node->edge maps 261 self._outedges[start].remove(edge) 262 self._inedges[end].remove(edge) 263 264 # Remove the edge from the edge->node maps. 265 del self._startnode[edge] 266 del self._endnode[edge] 267 268 # Remove the edge from the list of edge widgets that connect 2 269 # nodes. (Recompute curves?) 270 self._edgewidgets[start, end].remove(edge) 271 272 # Remove the edge from our list of child widgets. 273 self._remove_child_widget(edge)
274
275 - def remove_node(self, node):
276 """ 277 Remove a node from the graph (but don't destroy it). 278 @type node: L{CanvasWidget} 279 @return: A list of widgets that were removed from the 280 graph. Note that this will include any edges that 281 connected to C{node}. 282 """ 283 # Remove all edges that connect to this node. 284 removed_edges = [] 285 for edge in self._outedges.get(node, [])[:]: 286 self.remove_edge(edge) 287 removed_edges.append(edge) 288 for edge in self._inedges.get(node, [])[:]: 289 self.remove_edge(edge) 290 removed_edges.append(edge) 291 292 # Remove the node from the node->edges map 293 try: del self._outedges[node] 294 except KeyError: pass 295 try: del self._inedges[node] 296 except KeyError: pass 297 298 # Remove the node from our list of nodes 299 self._nodes.remove(node) 300 301 # Remove the node from our list of child widgets. 302 self._remove_child_widget(node) 303 304 # Return the list of removed widgets 305 return removed_edges + [node]
306
307 - def destroy_edge(self, edge):
308 """ 309 Remove an edge from the graph, and destroy the edge. 310 """ 311 self.remove_edge(edge) 312 edge.destroy()
313
314 - def destroy_node(self, node):
315 """ 316 Remove a node from the graph, and destroy the node. 317 """ 318 print 'removing', node 319 for widget in self.remove_node(node): 320 print 'destroying', widget 321 widget.destroy()
322
323 - def _tags(self): return []
324
325 - def _update(self, child):
326 """ 327 Make sure all edges/nodes are connected correctly. 328 """ 329 if isinstance(child, GraphEdgeWidget): 330 # Moved an edge. 331 pass 332 else: 333 # Moved a node. 334 for outedge in self._outedges.get(child, []): 335 self._update_edge(outedge) 336 for inedge in self._inedges.get(child, []): 337 self._update_edge(inedge)
338
339 - def _update_edge(self, edge):
340 curve = edge['curve'] 341 # Set the start. 342 src_x, src_y = self._node_center(self._endnode[edge]) 343 x, y = self._node_port(self._startnode[edge], src_x, src_y, curve) 344 edge.set_start(x, y) 345 # Set the end. 346 src_x, src_y = x, y#self._node_center(self._startnode[edge]) 347 x, y = self._node_port(self._endnode[edge], src_x, src_y, curve) 348 edge.set_end(x, y)
349
350 - def _node_port(self, node, src_x, src_y, curve):
351 x1, y1, x2, y2 = node.bbox() 352 x, y = (x1+x2)/2, (y1+y2)/2 353 w, h = abs(x2-x1), abs(y2-y1) 354 dx, dy = x-src_x, y-src_y 355 356 if dx > abs(dy)/5: return x-w/2, y 357 elif dx < -abs(dy)/5: return x+w/2, y 358 #if dx > abs(dy): return x-w/2, y 359 #elif dx < -abs(dy): return x+w/2, y 360 elif dy > 0: return x, y-h/2 361 elif dy < 0: return x, y+h/2 362 elif curve > 0: 363 return x, y+h/2 364 else: 365 return x, y-h/2
366
367 - def _node_center(self, node):
368 (x1, y1, x2, y2) = node.bbox() 369 return (x1+x2)/2, (y1+y2)/2
370
371 - def _manage(self):
372 self.arrange()
373 374 ##//////////////////////// 375 ## Graph Layout 376 ##////////////////////////
377 - def arrange(self, arrange_algorithm=None, toplevel=None):
378 """ 379 Set the node positions. This routine should attempt to 380 minimize the number of crossing edges, in order to make the 381 graph easier to read. 382 """ 383 if arrange_algorithm is not None: 384 self._arrange = arrange_algorithm 385 386 self._arrange_into_levels(toplevel) 387 self._arrange_levels() 388 389 (old_left, old_top) = self.bbox()[:2] 390 for node in self._nodes: 391 (x1, y1) = node.bbox()[:2] 392 node.move(-x1, -y1) 393 394 # Now we want to minimize crossing edges.. how, again? :) 395 for i in range(len(self._levels)): 396 for j in range(len(self._levels[i])): 397 if self._levels[i][j] is not None: 398 node = self._levels[i][j] 399 if self._orientation == 'horizontal': 400 node.move(i*self._xspace, j*self._yspace) 401 else: 402 node.move(j*self._xspace, i*self._yspace) 403 404 # If there is an edge from a node at the same 405 # position within its level, but whose level is at 406 # least 2 levels prior, then it's likely that that 407 # edge goes through an intervening node; so if its 408 # curve is zero, then increment it. 409 for edge in self._inedges.get(node,[]): 410 from_node = self._startnode[edge] 411 from_levelnum = self._nodelevel[from_node] 412 from_level = self._levels[from_levelnum] 413 if (abs(i-from_levelnum)>1 and 414 len(from_level) > j and 415 from_node == from_level[j] and 416 edge['curve'] == 0): 417 edge['curve'] = -0.25 418 419 (left, top) = self.bbox()[:2] 420 self.move(int(old_left-left), int(old_top-top))
421
422 - def _arrange_levels(self):
423 """ 424 Re-arrange each level to (locally) minimize the number of 425 crossing edges. 426 """ 427 # For now, put nodes with more incoming edges further down. 428 for levelnum in range(len(self._levels)): 429 self._arrange_level(levelnum)
430
431 - def _arrange_level(self, levelnum):
432 """ 433 Arrange a given level.. This algorithm is simple and pretty 434 heuristic.. 435 """ 436 if levelnum == 0: return 437 438 # For each position where we might want to put a node, create 439 # a scores dictionary, mapping candidate nodes to scores. We 440 # will then use these scores to distribute nodes to level positions. 441 scores = [{} for i in range(max(len(self._levels[levelnum]), 442 len(self._levels[levelnum-1])))] 443 for node in self._levels[levelnum]: 444 # All else being equal, put nodes with more incoming 445 # connections towards the end (=bottom) of the level. 446 for pos in range(len(scores)): 447 scores[pos][node] = 1.0/len(self._inedges.get(node, [])) 448 449 # Try to put a node at level position x if nodes 450 # in previous levels at position x point to it. 451 for edge in self._inedges.get(node, []): 452 from_node = self._startnode[edge] 453 from_levelnum = self._nodelevel[from_node] 454 if from_levelnum < levelnum: 455 from_pos = self._levels[from_levelnum].index(from_node) 456 score = (scores[from_pos].get(node, 0) + 1.0 / 457 (levelnum - from_levelnum)) 458 scores[from_pos][node] = score 459 460 # Get the list of nodes that we need to fill in, and empty the 461 # level. 462 nodes = self._levels[levelnum] 463 self._levels[levelnum] = [None] * len(scores) 464 level = self._levels[levelnum] 465 466 # Fill in nodes, picking the best first.. 467 while len(nodes) > 0: 468 best = (None, None, -1) # node, position, score. 469 for pos in range(len(scores)): 470 for (node, score) in scores[pos].items(): 471 if (score > best[2] and level[pos] is None and 472 node in nodes): 473 best = (node, pos, score) 474 elif (score == best[2] and pos<best[1] and 475 level[pos] is None and node in nodes): 476 # Put higher scores at lower level positions 477 best = (node, pos, score) 478 nodes.remove(best[0]) 479 level[best[1]] = best[0]
480
481 - def _arrange_into_levels(self, toplevel):
482 """ 483 Assign a level to each node. 484 """ 485 # Mapping from node to level. 486 self._nodelevel = {} 487 self._levels = [] 488 489 # Find any nodes that have no incoming edges; put all of these 490 # in level 0. 491 if toplevel is None: 492 toplevel = [] 493 for node in self._nodes: 494 if len(self._inedges.get(node, [])) == 0: 495 toplevel.append(node) 496 self._nodelevel[node] = 0 497 else: 498 for node in toplevel: 499 self._nodelevel[node] = 0 500 501 # Expand all of their children. 502 self._levels = [toplevel] 503 self._add_descendants(toplevel, 1) 504 505 # If we didn't get all the nodes, we'll have to start picking 506 # nodes that do have incoming transitions. Pick the ones that 507 # have the most reachable nodes. (n.b., this implementation 508 # isn't terribly efficient, but we dont' expect to be 509 # displaying huge graphs, so it should be ok) 510 while len(self._nodelevel) < len(self._nodes): 511 expand_node = None 512 max_reachable = -1 513 514 for node in self._nodes: 515 reachable = self._reachable(node) 516 if reachable >= max_reachable: 517 max_reachable = reachable 518 expand_node = node 519 520 # Expand the new node's children. 521 self._levels[0].append(expand_node) 522 self._nodelevel[expand_node] = 0 523 self._add_descendants(toplevel, 1)
524
525 - def _reachable(self, node, reached=None):
526 """ 527 How many *unexpanded* nodes can be reached from the given node? 528 """ 529 if self._nodelevel.has_key(node): return 0 530 if reached is None: reached = {} 531 if not reached.has_key(node): 532 reached[node] = 1 533 for edge in self._outedges.get(node, []): 534 self._reachable(self._endnode[edge], reached) 535 return len(reached)
536
537 - def _add_descendants(self, parent_level, levelnum):
538 """ 539 Add all the descendants of the nodes in the list parent_level 540 to the structures self._level and self._nodelevel. 541 """ 542 if self._arrange == 'bfs': 543 self._add_descendants_bfs(parent_level, levelnum) 544 elif self._arrange == 'dfs': 545 self._add_descendants_dfs(parent_level, levelnum) 546 else: 547 # Default to dfs 548 self._add_descendants_dfs(parent_level, levelnum)
549
550 - def _add_descendants_dfs(self, parent_level, levelnum):
551 if levelnum >= len(self._levels): self._levels.append([]) 552 for parent_node in parent_level: 553 # Add the parent node 554 if not self._nodelevel.has_key(parent_node): 555 self._levels[levelnum-1].append(parent_node) 556 self._nodelevel[parent_node] = levelnum-1 557 558 # Recurse to its children 559 child_nodes = [self._endnode[edge] 560 for edge in self._outedges.get(parent_node, []) 561 if not self._nodelevel.has_key(self._endnode[edge])] 562 if len(child_nodes) > 0: 563 self._add_descendants_dfs(child_nodes, levelnum+1)
564
565 - def _add_descendants_bfs(self, parent_level, levelnum):
566 frontier_nodes = [] 567 if levelnum >= len(self._levels): self._levels.append([]) 568 for parent_node in parent_level: 569 child_nodes = [self._endnode[edge] 570 for edge in self._outedges.get(parent_node, [])] 571 for node in child_nodes: 572 if not self._nodelevel.has_key(node): 573 self._levels[levelnum].append(node) 574 self._nodelevel[node] = levelnum 575 frontier_nodes.append(node) 576 if len(frontier_nodes) > 0: 577 self._add_descendants_bfs(frontier_nodes, levelnum+1)
578
579 - def _add_descendants_bfs2(self, parent_level, levelnum):
580 frontier_nodes = [] 581 if levelnum >= len(self._levels): self._levels.append([]) 582 for parent_node in parent_level: 583 child_nodes = [self._endnode[edge] 584 for edge in self._outedges.get(parent_node, [])] 585 child_nodes += [self._startnode[edge] 586 for edge in self._inedges.get(parent_node, [])] 587 for node in child_nodes: 588 if not self._nodelevel.has_key(node): 589 self._levels[levelnum].append(node) 590 self._nodelevel[node] = levelnum 591 frontier_nodes.append(node) 592 if len(frontier_nodes) > 0: 593 self._add_descendants_bfs2(frontier_nodes, levelnum+1)
594