Move viz code + changes to search

.travis.yml

@@ -12,6 +12,8 @@ install:
   - pip install flake8
   - pip install ipython
   - pip install matplotlib
+  - pip install networkx
+  - pip install ipywidgets
 
 script:
   - py.test

notebook.py

@@ -886,3 +886,159 @@ def draw_table(self):
         self.fill(0, 0, 0)
         self.text_n(self.table[self.context[0]][self.context[1]] if self.context else "Click for text", 0.025, 0.975)
         self.update()
+
+############################################################################################################
+
+#####################           Functions to assist plotting in search.ipynb            ####################
+
+############################################################################################################
+import networkx as nx
+import matplotlib.pyplot as plt
+from matplotlib import lines
+
+from ipywidgets import interact
+import ipywidgets as widgets
+from IPython.display import display
+import time
+from search import GraphProblem, romania_map
+
+def show_map(graph_data, node_colors = None):
+    G = nx.Graph(graph_data['graph_dict'])
+    node_colors = node_colors or graph_data['node_colors']
+    node_positions = graph_data['node_positions']
+    node_label_pos = graph_data['node_label_positions']
+    edge_weights= graph_data['edge_weights']
+
+    # set the size of the plot
+    plt.figure(figsize=(18,13))
+    # draw the graph (both nodes and edges) with locations from romania_locations
+    nx.draw(G, pos = {k : node_positions[k] for k in G.nodes()},
+                node_color = [node_colors[node] for node in G.nodes()], linewidths = 0.3, edgecolors = 'k')
+
+    # draw labels for nodes
+    node_label_handles = nx.draw_networkx_labels(G, pos = node_label_pos, font_size = 14)
+
+    # add a white bounding box behind the node labels
+    [label.set_bbox(dict(facecolor='white', edgecolor='none')) for label in node_label_handles.values()]
+
+    # add edge lables to the graph
+    nx.draw_networkx_edge_labels(G, pos = node_positions, edge_labels = edge_weights, font_size = 14)
+
+    # add a legend
+    white_circle = lines.Line2D([], [], color="white", marker='o', markersize=15, markerfacecolor="white")
+    orange_circle = lines.Line2D([], [], color="orange", marker='o', markersize=15, markerfacecolor="orange")
+    red_circle = lines.Line2D([], [], color="red", marker='o', markersize=15, markerfacecolor="red")
+    gray_circle = lines.Line2D([], [], color="gray", marker='o', markersize=15, markerfacecolor="gray")
+    green_circle = lines.Line2D([], [], color="green", marker='o', markersize=15, markerfacecolor="green")
+    plt.legend((white_circle, orange_circle, red_circle, gray_circle, green_circle),
+               ('Un-explored', 'Frontier', 'Currently Exploring', 'Explored', 'Final Solution'),
+               numpoints=1,prop={'size':16}, loc=(.8,.75))
+
+    # show the plot. No need to use in notebooks. nx.draw will show the graph itself.
+    plt.show()
+
+## helper functions for visualisations
+
+def final_path_colors(initial_node_colors, problem, solution):
+    "returns a node_colors dict of the final path provided the problem and solution"
+
+    # get initial node colors
+    final_colors = dict(initial_node_colors)
+    # color all the nodes in solution and starting node to green
+    final_colors[problem.initial] = "green"
+    for node in solution:
+        final_colors[node] = "green"  
+    return final_colors
+
+def display_visual(graph_data, user_input, algorithm=None, problem=None):
+    initial_node_colors = graph_data['node_colors']
+    if user_input == False:
+        def slider_callback(iteration):
+            # don't show graph for the first time running the cell calling this function
+            try:
+                show_map(graph_data, node_colors = all_node_colors[iteration])
+            except:
+                pass
+        def visualize_callback(Visualize):
+            if Visualize is True:
+                button.value = False
+
+                global all_node_colors
+
+                iterations, all_node_colors, node = algorithm(problem)
+                solution = node.solution()
+                all_node_colors.append(final_path_colors(all_node_colors[0], problem, solution))
+
+                slider.max = len(all_node_colors) - 1
+
+                for i in range(slider.max + 1):
+                    slider.value = i
+                     #time.sleep(.5)
+
+        slider = widgets.IntSlider(min=0, max=1, step=1, value=0)
+        slider_visual = widgets.interactive(slider_callback, iteration = slider)
+        display(slider_visual)
+
+        button = widgets.ToggleButton(value = False)
+        button_visual = widgets.interactive(visualize_callback, Visualize = button)
+        display(button_visual)
+
+    if user_input == True:
+        node_colors = dict(initial_node_colors)
+        if isinstance(algorithm, dict):
+            assert set(algorithm.keys()).issubset(set(["Breadth First Tree Search", 
+                                                       "Depth First Tree Search", 
+                                                       "Breadth First Search", 
+                                                       "Depth First Graph Search", 
+                                                       "Uniform Cost Search", 
+                                                       "A-star Search"]))
+
+            algo_dropdown = widgets.Dropdown(description = "Search algorithm: ",
+                                             options = sorted(list(algorithm.keys())),
+                                             value = "Breadth First Tree Search")
+            display(algo_dropdown)
+        elif algorithm is None:
+            print("No algorithm to run.")
+            return 0
+
+        def slider_callback(iteration):
+            # don't show graph for the first time running the cell calling this function
+            try:
+                show_map(graph_data, node_colors = all_node_colors[iteration])
+            except:
+                pass
+
+        def visualize_callback(Visualize):
+            if Visualize is True:
+                button.value = False
+
+                problem = GraphProblem(start_dropdown.value, end_dropdown.value, romania_map)
+                global all_node_colors
+
+                user_algorithm = algorithm[algo_dropdown.value]
+
+                iterations, all_node_colors, node = user_algorithm(problem)
+                solution = node.solution()
+                all_node_colors.append(final_path_colors(all_node_colors[0], problem, solution))
+
+                slider.max = len(all_node_colors) - 1
+
+                for i in range(slider.max + 1):
+                    slider.value = i
+                    #time.sleep(.5)
+
+        start_dropdown = widgets.Dropdown(description = "Start city: ",
+                                          options = sorted(list(node_colors.keys())), value = "Arad")
+        display(start_dropdown)
+
+        end_dropdown = widgets.Dropdown(description = "Goal city: ",
+                                        options = sorted(list(node_colors.keys())), value = "Fagaras")
+        display(end_dropdown)
+
+        button = widgets.ToggleButton(value = False)
+        button_visual = widgets.interactive(visualize_callback, Visualize = button)
+        display(button_visual)
+
+        slider = widgets.IntSlider(min=0, max=1, step=1, value=0)
+        slider_visual = widgets.interactive(slider_callback, iteration = slider)
+        display(slider_visual)