dsa.draw

  1""" Module to access graphic drawing functions for Trees, Heaps, Tries and Graphs. """
  2
  3from dsa.tree import Tree, TreeNode
  4from dsa.heap import Heap
  5from dsa.trie import Trie
  6
  7import networkx as nx
  8import matplotlib.pyplot as plt
  9
 10class Draw:
 11    """
 12    A base class for drawing various data structures.
 13    
 14    This class provides basic functionalities for rendering, saving, and displaying
 15    visual representations of data structures.
 16    """
 17    def __init__(self):
 18        """
 19        Initialize the Draw class with default figure size.
 20        """
 21        self.figsize = (5, 3)
 22        
 23    def render(self, **kwargs):
 24        """
 25        Render the visual representation of the data structure.
 26        
 27        This method should be overridden by subclasses to provide specific rendering logic.
 28        Args:
 29            **kwargs: Additional keyword arguments.
 30        """
 31        pass
 32
 33    def set_figsize(self, figsize):
 34        """
 35        Set the figure size for the plot.
 36        
 37        Args:
 38            figsize (tuple): A tuple representing the figure size (width, height).
 39        """
 40        self.figsize = figsize
 41
 42    def save(self, filename, **kwargs):
 43        """
 44        Save the rendered plot to a file.
 45        
 46        Args:
 47            filename (str): The name of the file to save the plot to.
 48            **kwargs: Additional keyword arguments.
 49        """
 50        plt = self.render(**kwargs)
 51        plt.axis('off')
 52        plt.savefig(filename)
 53
 54    def draw(self, **kwargs):
 55        """
 56        Display the rendered plot.
 57
 58        Args:
 59            **kwargs: Additional keyword arguments.
 60        """
 61        plt = self.render(**kwargs)
 62        plt.axis('off')
 63        plt.show()
 64
 65class TreeDraw(Draw):
 66    """
 67    A class for drawing a tree structure using the NetworkX library.
 68
 69    This class extends the `Draw` class to visualize tree structures. It organizes the nodes
 70    in a hierarchical tree layout and provides options for customization through the `render` method.
 71
 72    Attributes:
 73        tree (Tree): The tree structure to be drawn.
 74
 75    Usage Example:
 76        t = Tree(root_node)  # Define your tree structure with a root node
 77        td = TreeDraw(t)
 78        td.draw()
 79    """
 80
 81    def __init__(self, tree: Tree):
 82        """ 
 83        Initializes the TreeDraw object.
 84
 85        Args:
 86            tree (Tree): The tree structure to be drawn.
 87        """
 88        super().__init__()
 89        self.tree = tree
 90        
 91    def add_edges(self, graph, node, pos=None, x: float=0, y: float=0, layer=1):
 92        """
 93        Recursively adds edges to the graph and positions the nodes in a tree layout.
 94    
 95        Args:
 96            graph (networkx.DiGraph): The graph object where edges are added.
 97            node (TreeNode): The current node in the tree.
 98            pos (dict, optional): A dictionary to store the positions of the nodes. Defaults to None.
 99            x (float): The x-coordinate for the current node. Defaults to 0.
100            y (float): The y-coordinate for the current node. Defaults to 0.
101            layer (int): The current layer/level of the node in the tree. Defaults to 1.
102
103        Returns:
104            dict: A dictionary containing the positions of all nodes in the tree.
105        """
106        if pos is None:
107            pos = {}
108        if node is not None:
109            pos[node.value] = (x, y)
110            if node.left:
111                graph.add_edge(node.value, node.left.value)
112                pos = self.add_edges(graph, node.left, pos=pos, x=x-1/layer, y=y-1, layer=layer+1)
113            if node.right:
114                graph.add_edge(node.value, node.right.value)
115                pos = self.add_edges(graph, node.right, pos=pos, x=x+1/layer, y=y-1, layer=layer+1)
116        return pos
117    
118    def render(self, **kwargs):
119        """ 
120        Renders the tree using Matplotlib. Not to be called directly. Call draw() instead.
121
122        This method generates a graphical representation of the tree with nodes positioned
123        in a hierarchical layout. Customization options can be provided via keyword arguments.
124
125        Args:
126            **kwargs: Additional keyword arguments for customization.
127
128        Returns:
129            matplotlib.pyplot: The Matplotlib plot object for further customization or display.
130        """
131        super().render(**kwargs)
132        graph = nx.DiGraph()
133        pos = self.add_edges(graph, self.tree.root)
134        plt.figure(figsize=self.figsize)
135        nx.draw(graph, pos, with_labels=True, arrows=False, node_color="tab:blue", node_size=800, font_size=12, font_color="white") 
136        return plt
137
138class HeapDraw(Draw):
139    """
140    A class for drawing a heap structure using the NetworkX library.
141
142    This class extends the `Draw` class to visualize heap structures, such as binary heaps or min-heaps.
143
144    Attributes:
145        heap (Heap): The heap structure to be drawn.
146
147    Usage Example:
148        h = MinHeap()  # Define your heap, e.g., MinHeap or Heap
149        hd = HeapDraw(h)
150        hd.draw()
151    """
152
153    def __init__(self, heap: Heap, **kwargs):
154        """
155        Initializes the HeapDraw object.
156
157        Args:
158            heap (Heap): The heap structure to be drawn.
159            **kwargs: Additional keyword arguments passed to the parent `Draw` class.
160        """
161        super().__init__(**kwargs)
162        self.heap = heap
163
164    def array_to_node(self, index: int, array):
165        """
166        Converts an array-based heap into a tree node structure.
167
168        This helper function recursively constructs a tree from the array representation
169        of the heap, reflecting the binary tree structure of the heap.
170
171        Args:
172            index (int): The current index in the array representing the node.
173            array (list): The array containing heap values, organized as a complete binary tree.
174
175        Returns:
176            TreeNode: The root node of the constructed subtree.
177        """
178        if index >= len(array):
179            return None
180        else:
181            value = array[index]
182            left_index = index * 2 + 1
183            right_index = index * 2 + 2
184            node = TreeNode(value)
185            node.left = self.array_to_node(left_index, array)
186            node.right = self.array_to_node(right_index, array)
187            return node
188
189    def render(self, **kwargs):
190        """
191        Renders the heap as a tree using Matplotlib. Not to be called directly. Call draw() instead.
192
193        This method converts the heap into a tree structure and then uses the `TreeDraw` class
194        to render it visually. Customization options can be provided via keyword arguments.
195
196        Returns:
197            matplotlib.pyplot: The Matplotlib plot object for further customization or display.
198        """
199        node = self.array_to_node(0, [node[1] for node in self.heap.enumerate()])
200        tree = Tree(node)
201
202        tree_draw = TreeDraw(tree)
203        tree_draw.set_figsize(self.figsize)
204        return tree_draw.render(**kwargs)
205    
206class TrieDraw(Draw):
207    """
208    A class for drawing a Trie structure using the NetworkX library.
209
210    This class extends the `Draw` class to visualize Trie structures, commonly used for storing strings
211    or prefix trees. It provides methods to convert the Trie into a networkx graph, arrange nodes
212    hierarchically, and render the visualization using Matplotlib.
213
214    Attributes:
215        trie (Trie): The Trie structure to be drawn.
216
217    Usage Example:
218        trie = Trie()  # Initialize your Trie and populate it with words
219        trd = TrieDraw(trie)
220        trd.draw()
221    """
222
223    def __init__(self, trie: Trie, **kwargs):
224        """
225        Initializes the TrieDraw object.
226
227        Args:
228            trie (Trie): The Trie structure to be drawn.
229            **kwargs: Additional keyword arguments passed to the parent `Draw` class.
230        """
231        super().__init__(**kwargs)
232        self.figsize = (10, 6)
233        self.trie = trie
234        
235    def _add_edges(self, graph, node, current_path):
236        """
237        Recursively adds edges to the networkx graph based on the Trie structure.
238
239        This helper function traverses the Trie, adding edges between nodes in the
240        networkx graph and associating them with characters from the Trie.
241
242        Args:
243            graph (networkx.DiGraph): The directed graph to which edges are added.
244            node (TrieNode): The current node in the Trie.
245            current_path (str): The path representing the current prefix in the Trie.
246        """
247        if node is None:
248            return
249        for char, child in node.children.items():
250            new_path = current_path + char
251            graph.add_edge(current_path, new_path, label=char)
252            self._add_edges(graph, child, new_path)
253    
254    def to_networkx(self) -> nx.DiGraph:
255        """
256        Converts the Trie into a NetworkX directed graph (DiGraph).
257
258        This method creates a networkx graph representation of the Trie, where each node
259        represents a prefix and each edge represents a character transition.
260
261        Returns:
262            networkx.DiGraph: The networkx graph representation of the Trie.
263        """
264        graph = nx.DiGraph()
265        self._add_edges(graph, self.trie.root, "")
266        return graph
267    
268    def hierarchical_pos(self, G, root=None, width=1., vert_gap=0.2, vert_loc=0, xcenter=0.5):
269        """
270        Computes the hierarchical position of nodes in the graph for visualization.
271
272        This method arranges the nodes of the Trie in a hierarchical layout, which is
273        particularly useful for visualizing tree-like structures such as Tries.
274
275        Args:
276            G (networkx.Graph): The graph for which to compute positions.
277            root (str, optional): The root node of the graph. Defaults to None, which means
278                                  the root will be determined automatically.
279            width (float): The width of the entire drawing. Defaults to 1.
280            vert_gap (float): The gap between levels in the hierarchy. Defaults to 0.2.
281            vert_loc (float): The vertical location of the root node. Defaults to 0.
282            xcenter (float): The horizontal center of the root node. Defaults to 0.5.
283
284        Returns:
285            dict: A dictionary mapping each node to its (x, y) position in the layout.
286        """
287        if root is None:
288            root = next(iter(nx.topological_sort(G)))
289    
290        def _hierarchical_pos(G, node, width: float=1., vert_gap: float=0.2, vert_loc:float=0, xcenter:float=0.5, pos=None, parent=None, parsed=[]):
291            if pos is None:
292                pos = {node: (xcenter, vert_loc)}
293            else:
294                pos[node] = (xcenter, vert_loc)
295    
296            children = list(G.neighbors(node))
297            if not isinstance(G, nx.DiGraph) and parent is not None:
298                children.remove(parent)
299    
300            if len(children) != 0:
301                dx = width / len(children)
302                nextx = xcenter - width / 2 - dx / 2
303                for child in children:
304                    nextx += dx
305                    pos = _hierarchical_pos(G, child, width=dx, vert_gap=vert_gap, vert_loc=vert_loc - vert_gap, xcenter=nextx, pos=pos, parent=node, parsed=parsed)
306            return pos
307    
308        return _hierarchical_pos(G, root, width, vert_gap, vert_loc, xcenter)
309    
310    def render_rectangle(self, **kwargs):
311        """
312        Renders the Trie as a hierarchical graph using Matplotlib. Not to be called directly. Call draw() instead.
313
314        This method uses the hierarchical positions of the nodes to render the Trie
315        as a directed graph. Nodes are drawn as rectangles, and edges represent the transitions
316        between prefixes.
317
318        Returns:
319            matplotlib.pyplot: The Matplotlib plot object for further customization or display.
320        """
321        super().render(**kwargs)
322        trie_graph = self.to_networkx()
323        
324        pos = self.hierarchical_pos(trie_graph)
325        
326        plt.figure(figsize=self.figsize)
327        nx.draw_networkx_edges(trie_graph, pos, arrows=True)
328        
329        ax = plt.gca()
330        rect_width = 0.05
331        rect_height = 0.15
332        for node in pos:
333            x, y = pos[node]
334            rectangle = plt.Rectangle((x - (rect_width / 2), y - (rect_height / 2)), rect_width, rect_height, color="tab:blue")
335            ax.add_patch(rectangle)
336            plt.text(x, y, node[-1] if node else "", verticalalignment='center', horizontalalignment='center', fontsize=12, color="white")
337        return plt
338
339    def render(self, **kwargs):
340        """
341        Renders the Trie as a hierarchical graph using Matplotlib. Not to be called directly. Call draw() instead.
342
343        This version uses NetworkX's default drawing tools with circular nodes for simplicity and clarity.
344        """
345        super().render(**kwargs)
346        trie_graph = self.to_networkx()
347        pos = self.hierarchical_pos(trie_graph)
348
349        plt.figure(figsize=self.figsize)
350
351        # Draw the graph using built-in node and edge drawing
352        nx.draw(
353            trie_graph,
354            pos,
355            with_labels=True,
356            labels={node: node[-1] if node else "" for node in trie_graph.nodes},
357            node_shape='o',
358            node_color='tab:blue',
359            font_color='white',
360            font_size=10,
361            edgecolors='none',
362            arrows=False,
363            **kwargs
364        )
365
366        return plt
367
368class GraphDraw(Draw):
369    """
370    A class for drawing graphs using the NetworkX library.
371
372    This class extends the `Draw` class to visualize graphs. It supports both directed 
373    and undirected graphs, as well as weighted and unweighted graphs. Additionally, 
374    it provides an option to display the Minimum Spanning Tree (MST) of the graph.
375
376    Attributes:
377        graph: The graph to be drawn
378        directed (bool): Specifies if the graph is directed. Defaults to False
379        weighted (bool): Specifies if the graph has weighted edges. Defaults to False
380
381    Usage Example:
382        gd = GraphDraw(g) # g is a Graph type (AdjacencyMatrixGraph, AdjacencyMatrixWeightedGraph, AdjacencyListWeightedGraph, AdjacencyListGraph)
383        gd.draw()
384    """
385    def __init__(self, graph, directed=False, weighted=False):
386        """
387        Initializes the GraphDraw object.
388        """
389        super().__init__()
390        self.graph = graph
391        self.directed = directed
392        self.weighted = weighted
393            
394    def render(self, pos=None, show_mst=False, mst_only=False, **kwargs):
395        """
396        Renders the graph using Matplotlib. Not to be called directly. Call draw() instead.
397        """
398        super().render(**kwargs)
399        edges = self.graph.edges()
400
401        if self.directed:
402            g = nx.DiGraph()
403        else:
404            g = nx.Graph()
405
406        if self.weighted:
407            g.add_weighted_edges_from(edges)
408        else:
409            g.add_edges_from(edges)
410    
411        if pos is None:
412            pos = nx.shell_layout(g) 
413
414        plt.figure(figsize=self.figsize)
415
416        if mst_only:
417            show_mst = True
418            nx.draw_networkx_nodes(g, pos, node_color="tab:blue", node_size=800)
419            nx.draw_networkx_labels(g, pos, font_size=10, font_color="white")
420        else:
421            nx.draw(g, pos, with_labels=True, node_color="tab:blue", node_size=800, font_size=10, font_color="white")
422
423        if self.weighted:
424            edge_labels = {(u, v): d["weight"] for u, v, d in g.edges(data=True)}
425            nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels, font_size=14, label_pos=0.5)
426
427        if show_mst:
428            T = nx.minimum_spanning_tree(g)
429            nx.draw_networkx_edges(T, pos, edge_color="tab:red", width=2)
430        return plt