pathdemo.py

"""
Demo pathfinding, with obstacles
"""

from matplotlib import pyplot as plt
import networkx as nx


GRID_SIZE = [100, 100]
G = nx.grid_graph(dim=GRID_SIZE)


def obstacle(lower_left, width, height):
    """
    Given the (x,y) coord of the lower left corner of a box, and its width and
    height, return the coords covered by this box
    """
    x, y = lower_left
    for i in range(x, x + width):
        for j in range(y, y + height):
            yield (i, j)


def dist(a, b):
    """
    euclidean distance metric
    """
    (x1, y1) = a
    (x2, y2) = b
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5


def gridplot(graph):
    """
    Plots the grid graph as black dots
    """
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111)
    x, y = zip(*G.nodes())
    ax.plot(x, y, 'k,')
    ax.axis('equal')


def pathplot(graph, start, stop, **kwargs):
    """
    Plots the shortest path via A* and returns the list of x and y coords.

    Extra kwargs are passed to the line plot.
    """
    path = nx.astar_path(G, start, stop, dist)
    x, y = zip(*path)
    plt.plot(x, y, **kwargs)
    return path

# Create some obstacles and remove those points from consideration by deleting
# them from the graph
obstacles = [
    ((5, 25), 70, 5),    # a long wall
    ((30, 50), 30, 30),  # a box
]
for lower_left, width, height in obstacles:
    for coord in obstacle(lower_left, width, height):
        G.remove_node(coord)

gridplot(G)
start = (47, 10)
stop = (45, 89)
pathplot(G, start, stop, color='r', linewidth=2)

plt.show()