Add grid search visualization (#1155)

Issue #1155 asked for a visualization of how a search algorithm expands and finds
the path (à la redblobgames), on a 2-D grid. Adds:

- search.py: GridProblem — shortest-path finding on a grid with obstacles (4- or
  8-connected, unit step cost, Manhattan h), a clean reusable Problem.
- notebook_utils.py: grid_search_steps(problem, strategy) records the cells in the
  order they are expanded (strategies: bfs/ucs/greedy/astar) and returns the path;
  plot_grid_search() shades the expanded cells by order and draws the path,
  obstacles and start/goal (matplotlib, works headless).
- search.ipynb: a 'GRID SEARCH VISUALIZATION' section comparing BFS vs greedy vs A*
  on the same grid (with the rendered figure) — A* expands far fewer cells than the
  uninformed search while still returning an optimal path.

Tests: GridProblem (optimal path around a wall, unreachable goal) and the grid
visualization (A* expands <= BFS; plot renders). 34 passed in test_search.

Author: dmeoli

aima/notebook_utils.py

@@ -1,3 +1,4 @@
+import heapq
 import os
 import time
 from collections import defaultdict
@@ -16,7 +17,7 @@
 from aima.games import TicTacToe, alpha_beta_player, random_player, Fig52Extended
 from aima.learning import DataSet
 from aima.logic import parse_definite_clause, standardize_variables, unify_mm, subst
-from aima.search import GraphProblem, romania_map
+from aima.search import GraphProblem, romania_map, Node
 
 # repo root (the directory containing the `aima` package), so cwd-relative
 # data/image paths resolve regardless of where a notebook is launched from
@@ -1236,3 +1237,75 @@ def plot_model_boundary(dataset, attr1, attr2, model=None):
     plt.xlim(xx.min(), xx.max())
     plt.ylim(yy.min(), yy.max())
     plt.show()
+
+
+# ______________________________________________________________________________
+# Visualizing search on a 2-D grid
+
+
+def grid_search_steps(problem, strategy='astar'):
+    """Run a search on a :class:`~aima.search.GridProblem`, recording how it explores.
+
+    Returns ``(explored, path)``: ``explored`` is the list of cells in the order
+    they are expanded (popped from the frontier), and ``path`` is the list of cells
+    from start to goal (``None`` if the goal is unreachable). ``strategy`` selects
+    the evaluation function -- ``'astar'`` (g + h), ``'ucs'`` (g), ``'bfs'`` (depth)
+    or ``'greedy'`` (h) -- which is what makes the different exploration patterns.
+    """
+    evaluators = {
+        'astar': lambda n: n.path_cost + problem.h(n),
+        'ucs': lambda n: n.path_cost,
+        'bfs': lambda n: n.depth,
+        'greedy': lambda n: problem.h(n),
+    }
+    f = evaluators[strategy]
+    start = Node(problem.initial)
+    frontier = [(f(start), 0, start)]   # (priority, tie-breaker, node)
+    best_cost = {problem.initial: start.path_cost}
+    explored, expanded = [], set()
+    counter = 1
+    while frontier:
+        _, _, node = heapq.heappop(frontier)
+        if node.state in expanded:
+            continue
+        expanded.add(node.state)
+        explored.append(node.state)
+        if problem.goal_test(node.state):
+            return explored, [n.state for n in node.path()]
+        for child in node.expand(problem):
+            if child.state not in best_cost or child.path_cost < best_cost[child.state]:
+                best_cost[child.state] = child.path_cost
+                heapq.heappush(frontier, (f(child), counter, child))
+                counter += 1
+    return explored, None
+
+
+def plot_grid_search(problem, explored, path=None, ax=None, title=None):
+    """Visualise a grid search: obstacles, the cells expanded (shaded by *when* they
+    were expanded), the solution path, and the start/goal. ``explored`` and ``path``
+    are as returned by :func:`grid_search_steps`. Returns the matplotlib Axes.
+    """
+    width, height = problem.width, problem.height
+    shade = np.full((height, width), np.nan)
+    for order, (x, y) in enumerate(explored):
+        shade[y, x] = order
+
+    if ax is None:
+        _, ax = plt.subplots(figsize=(width * 0.45 + 1, height * 0.45 + 1))
+    ax.imshow(shade, origin='lower', cmap='YlGnBu',
+              extent=(-0.5, width - 0.5, -0.5, height - 0.5))
+    for (x, y) in problem.obstacles:
+        ax.add_patch(plt.Rectangle((x - 0.5, y - 0.5), 1, 1, color='0.2'))
+    if path:
+        ax.plot([c[0] for c in path], [c[1] for c in path],
+                color='orange', linewidth=2, marker='.', zorder=2)
+    (sx, sy), (gx, gy) = problem.initial, problem.goal
+    ax.scatter([sx], [sy], c='lime', s=130, marker='o', edgecolors='k', zorder=3)
+    ax.scatter([gx], [gy], c='red', s=160, marker='*', edgecolors='k', zorder=3)
+    ax.set_xticks(range(width))
+    ax.set_yticks(range(height))
+    ax.grid(True, color='0.85', linewidth=0.5)
+    ax.set_xlim(-0.5, width - 0.5)
+    ax.set_ylim(-0.5, height - 0.5)
+    ax.set_title(title or '{} cells expanded'.format(len(explored)))
+    return ax

aima/search.py

@@ -1103,6 +1103,43 @@ def value(self, state):
         return self.grid[x][y]
 
 
+class GridProblem(Problem):
+    """Shortest-path finding on a 2-D grid.
+
+    States are ``(x, y)`` cells. The agent steps between in-bounds, obstacle-free
+    cells -- 4-connected via ``directions4`` by default, or pass ``directions8``
+    for 8-connected movement. Every step costs 1; ``h`` is the Manhattan distance
+    to the goal (admissible for the 4-connected, unit-cost case).
+    """
+
+    def __init__(self, initial, goal, width, height, obstacles=(), directions=directions4):
+        super().__init__(initial, goal)
+        self.width = width
+        self.height = height
+        self.obstacles = set(map(tuple, obstacles))
+        self.directions = directions
+
+    def passable(self, cell):
+        """True if ``cell`` is on the grid and not an obstacle."""
+        x, y = cell
+        return 0 <= x < self.width and 0 <= y < self.height and tuple(cell) not in self.obstacles
+
+    def actions(self, state):
+        return [a for a, d in self.directions.items() if self.passable(vector_add(state, d))]
+
+    def result(self, state, action):
+        return vector_add(state, self.directions[action])
+
+    def path_cost(self, c, state1, action, state2):
+        return c + 1
+
+    def h(self, node):
+        """Manhattan distance from ``node`` (or a state) to the goal."""
+        x, y = node.state if isinstance(node, Node) else node
+        gx, gy = self.goal
+        return abs(x - gx) + abs(y - gy)
+
+
 class OnlineDFSAgent:
     """
     [Figure 4.21]

notebooks/search.ipynb

@@ -991,6 +991,33 @@ def color_city_and_update_map(node, color):
 assert puzzle.check_solvability((2, 4, 3, 1, 5, 6, 7, 8, 0))
 recursive_best_first_search(puzzle).solution()
 
+# %% [markdown]
+# ## GRID SEARCH VISUALIZATION
+#
+# The visualizations above color the nodes of a *graph* (the Romania road map) as the search expands. On a 2-D **grid** the same idea is even more intuitive: we can watch the frontier spread out across the cells and see how an informed search is pulled towards the goal.
+#
+# `GridProblem` (in `search.py`) is shortest-path finding on a grid with obstacles, and `grid_search_steps` / `plot_grid_search` (in `notebook_utils.py`) run a strategy and draw the cells in the order they were expanded, shaded light-to-dark, with the solution path in orange and the start (green) and goal (red star).
+#
+# Below, the same grid (with a wall that has a gap near the top) is solved by an **uninformed** breadth-first search, a **greedy** best-first search, and **A\***:
+
+# %%
+# %matplotlib inline
+from aima.notebook_utils import grid_search_steps, plot_grid_search
+import matplotlib.pyplot as plt
+
+walls = [(4, y) for y in range(8)]          # vertical wall x=4, with a gap at y=8,9
+grid = GridProblem((0, 0), (8, 2), 12, 10, obstacles=walls)
+
+fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+for ax, strategy in zip(axes, ['bfs', 'greedy', 'astar']):
+    explored, path = grid_search_steps(grid, strategy)
+    plot_grid_search(grid, explored, path, ax=ax,
+                     title='{}: {} cells expanded'.format(strategy.upper(), len(explored)))
+plt.show()
+
+# %% [markdown]
+# Breadth-first search fans out in all directions and expands many cells; greedy best-first heads straight for the goal (expanding the fewest, but it can be misled by obstacles and is not optimal in general); **A\*** balances the two -- it is pulled toward the goal yet still returns an optimal path, expanding far fewer cells than the uninformed search.
+
 # %% [markdown]
 # ## A* HEURISTICS
 #

notebooks/search.py

@@ -483,5 +483,34 @@ def test_plan_route():
     assert problem.result(WumpusPosition(1, 1, 'RIGHT'), 'Forward').get_location() == (1, 1)
 
 
+def test_grid_problem():
+    # 5x5 grid with a wall blocking the direct route (gap at y=4)
+    walls = [(2, y) for y in range(4)]
+    problem = GridProblem((0, 0), (4, 0), 5, 5, obstacles=walls)
+    astar = astar_search(problem)
+    bfs = breadth_first_graph_search(problem)
+    assert astar is not None
+    assert astar.path_cost == bfs.path_cost                            # both optimal
+    assert all(problem.passable(node.state) for node in astar.path())  # avoids walls
+    # a goal walled off on every side is unreachable
+    boxed = GridProblem((0, 0), (4, 4), 5, 5, obstacles=[(3, 4), (4, 3)])
+    assert astar_search(boxed) is None
+
+
+def test_grid_search_visualization():
+    import matplotlib
+    matplotlib.use('Agg')
+    from aima.notebook_utils import grid_search_steps, plot_grid_search
+    walls = [(3, y) for y in range(7)]
+    problem = GridProblem((0, 0), (6, 0), 10, 10, obstacles=walls)
+    expl_bfs, path_bfs = grid_search_steps(problem, 'bfs')
+    expl_astar, path_astar = grid_search_steps(problem, 'astar')
+    assert path_bfs[0] == path_astar[0] == (0, 0)
+    assert path_bfs[-1] == path_astar[-1] == (6, 0)
+    assert len(path_astar) == len(path_bfs)            # same optimal path length
+    assert len(expl_astar) <= len(expl_bfs)            # informed search is more focused
+    assert plot_grid_search(problem, expl_astar, path_astar) is not None
+
+
 if __name__ == '__main__':
     pytest.main()