Networkxで2次元格子上のパス計算を行う

簡単なパス計算を行うのに、PythonライブラリのNetworkxを使うと楽そうだったので、使ってみました。

まず、以下のような2次元格子の環境を考えます。格子サイズは20x20です。

import networkx as nx
import matplotlib.pyplot as plt

ngrid = 20
gp = nx.grid_graph(dim=[ngrid, ngrid])
nx.draw(gp,
        pos=dict((n, n) for n in gp.nodes()),
        node_color=[gp.node[n].get('color', 'white') for n in gp.nodes_iter()],
        node_size=200)
plt.axis('equal')
plt.show()

格子状の斜めパスもOKにしたかったので、各ノードに斜めのエッジを追加します。

import networkx as nx
import matplotlib.pyplot as plt

def add_cross_edge(gp, shape):
    """
    2DGridのグラフに斜め方向のエッジを追加する
    """
    for node in gp.nodes_iter():
        nx_node = (node[0] + 1, node[1] + 1)
        if nx_node[0] < shape[0] and nx_node[1] < shape[1]:
            gp.add_edge(node, nx_node)
        nx_node = (node[0] + 1, node[1] - 1)
        if nx_node[0] < shape[0] and nx_node[1] >= 0:
            gp.add_edge(node, nx_node)

ngrid = 20
gp = nx.grid_graph(dim=[ngrid, ngrid])
add_cross_edge(gp, [ngrid, ngrid])
nx.draw(gp,
        pos=dict((n, n) for n in gp.nodes()),
        node_color=[gp.node[n].get('color', 'white') for n in gp.nodes_iter()],
        node_size=200)
plt.axis('equal')
plt.show()

最後に、障害物とスタート・ゴールを設定して、環境自体の構築の完了です。
ここでは障害物を黒、スタートを緑、ゴールを赤で色分けしています。
この環境を用いて、障害物を避けてスタートとゴールを繋ぐパスの計算を行います。

import networkx as nx
import matplotlib.pyplot as plt

def add_cross_edge(gp, shape):
    """
    2DGridのグラフに斜め方向のエッジを追加する
    """
    for node in gp.nodes_iter():
        nx_node = (node[0] + 1, node[1] + 1)
        if nx_node[0] < shape[0] and nx_node[1] < shape[1]:
            gp.add_edge(node, nx_node)
        nx_node = (node[0] + 1, node[1] - 1)
        if nx_node[0] < shape[0] and nx_node[1] >= 0:
            gp.add_edge(node, nx_node)

ngrid = 20
gp = nx.grid_graph(dim=[ngrid, ngrid])
add_cross_edge(gp, [ngrid, ngrid])
idcs = np.random.choice(len(gp.nodes()), int(ngrid * ngrid * 0.2), replace=False)
# スタート・ゴール・障害物を設定する
st, gl, obs = gp.nodes()[idcs[0]], gp.nodes()[idcs[1]], [gp.nodes()[i] for i in idcs[2:]]
gp.node[st]['color'] = 'green'
gp.node[gl]['color'] = 'red'
for o in obs:
    gp.node[o]['color'] = 'black'
nx.draw(gp,
        pos=dict((n, n) for n in gp.nodes()),
        node_color=[gp.node[n].get('color', 'white') for n in gp.nodes_iter()],
        node_size=200)
plt.axis('equal')
plt.show()

次に障害物を避けるために各エッジ間にWeightを設定します。
ノードが障害物とかぶっている際にペナルティを課すようにしています。

import numpy as np

def cost(a, b):
    """
    コスト関数
    """
    x1 = np.array(a, dtype=np.float32)
    x2 = np.array(b, dtype=np.float32)
    dist = np.linalg.norm(x1 - x2)
    if any([(a == o or b == o) for o in obs]):
        penalty = 1.0e6
    else:
        penalty = 0.0
    return dist + penalty

for u, v, d in gp.edges_iter(data=True):
    d['weight'] = cost(u, v)

Networkxの中でパス探索関数はいくつかあるのですが、今回はA*アルゴリズムを使います。A*アルゴリズムではヒューリスティック関数と呼ばれるものが必要で、以下のように設定します。

def dist(a, b):
    """
    ヒューリスティック関数
    """
    x1 = np.array(a, dtype=np.float32)
    x2 = np.array(b, dtype=np.float32)
    return np.linalg.norm(x1 - x2)

path = nx.astar_path(gp, st, gl, dist)
length = nx.astar_path_length(gp, st, gl, dist)
print(path)
print(length)

最後にパス探索結果を描画してみます。
パス上のノードを青で示しています。

for p in path[1:-1]:
    if gp.node[p].get('color', '') == 'black':
        print('Invalid path')
        continue
    gp.node[p]['color'] = 'blue'

nx.draw(gp,
        pos=dict((n, n) for n in gp.nodes()),
        node_color=[gp.node[n].get('color', 'white') for n in gp.nodes_iter()],
        node_size=200)
plt.axis('equal')
plt.show()

コード全体は以下のようになります。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

def add_cross_edge(gp, shape):
    """
    2DGridのグラフに斜め方向のエッジを追加する
    """
    for node in gp.nodes_iter():
        nx_node = (node[0] + 1, node[1] + 1)
        if nx_node[0] < shape[0] and nx_node[1] < shape[1]:
            gp.add_edge(node, nx_node)
        nx_node = (node[0] + 1, node[1] - 1)
        if nx_node[0] < shape[0] and nx_node[1] >= 0:
            gp.add_edge(node, nx_node)

ngrid = 20
gp = nx.grid_graph(dim=[ngrid, ngrid])
add_cross_edge(gp, [ngrid, ngrid])
idcs = np.random.choice(len(gp.nodes()), int(ngrid * ngrid * 0.2), replace=False)
# スタート・ゴール・障害物を設定する
st, gl, obs = gp.nodes()[idcs[0]], gp.nodes()[idcs[1]], [gp.nodes()[i] for i in idcs[2:]]
gp.node[st]['color'] = 'green'
gp.node[gl]['color'] = 'red'
for o in obs:
    gp.node[o]['color'] = 'black'

def dist(a, b):
    """
    ヒューリスティック関数
    """
    x1 = np.array(a, dtype=np.float32)
    x2 = np.array(b, dtype=np.float32)
    return np.linalg.norm(x1 - x2)

def cost(a, b, k1=1.0, k2=10.0, kind='intsct'):
    """
    コスト関数
    """
    x1 = np.array(a, dtype=np.float32)
    x2 = np.array(b, dtype=np.float32)
    dist = np.linalg.norm(x1 - x2)
    if any([(a == o or b == o) for o in obs]):
        penalty = 1.0e6
    else:
        penalty = 0.0
    return dist + penalty

for u, v, d in gp.edges_iter(data=True):
    d['weight'] = cost(u, v)

path = nx.astar_path(gp, st, gl, dist)
length = nx.astar_path_length(gp, st, gl, dist)
print(path)
print(length)
for p in path[1:-1]:
    if gp.node[p].get('color', '') == 'black':
        print('Invalid path')
        continue
    gp.node[p]['color'] = 'blue'

nx.draw(gp,
        pos=dict((n, n) for n in gp.nodes()),
        node_color=[gp.node[n].get('color', 'white') for n in gp.nodes_iter()],
        node_size=200)
plt.axis('equal')
plt.show()

Networkx便利ですね。
今回設定した2次元格子の障害物あり環境は、強化学習とかにも使えそうです。