lib/grid: playing with A* planning.

src/nppilot/lib/grid.py

@@ -0,0 +1,155 @@
+import math
+import collections
+
+import numpy as np
+import matplotlib.patches
+import matplotlib.pyplot as plt
+
+INF = 1e3
+
+
+class Grid:
+    def __init__(self, w: float, h: float, s: float):
+        self._w = w
+        self._h = h
+        self._s = s
+
+        self._stride = int(math.ceil(w / s))
+        self._cells = np.zeros((self._stride, int(math.ceil(h / s))))
+
+    def get(self, p: tuple) -> float:
+        if not self.contains(p):
+            return 0
+        return self._cells[p[0]][p[1]]
+
+    def set(self, p: tuple, v: float) -> None:
+        self._cells[p[0]][p[1]] = v
+
+    def line(self, p: tuple, dp: tuple, c: int, v: float) -> None:
+        for i in range(c):
+            self.set((int(p[0]), int(p[1])), v)
+            p = (p[0] + dp[0], p[1] + dp[1])
+
+    def smear(self) -> None:
+        out = np.zeros(self._cells.shape)
+        for y in range(len(self._cells)):
+            for x in range(len(self._cells[0])):
+                acc = self.get((y, x))
+                if acc >= INF:
+                    out[y][x] = acc
+                else:
+                    for d in (-1, 0), (1, 0), (0, 1), (0, -1):
+                        v = self.get((y + d[0], x + d[1]))
+                        if v >= INF:
+                            v = 1
+                        acc += v
+                    out[y][x] = acc / 5
+        self._cells = out
+
+    def contains(self, p):
+        y, x = p
+        if y < 0 or x < 0:
+            return False
+        if y >= len(self._cells) or x >= len(self._cells[0]):
+            return False
+        return True
+
+
+def plot(g: Grid, keepout: Grid):
+    cells = g._cells / np.max(g._cells)
+    s = g._s
+    w = s * 0.9
+    cmap = plt.get_cmap('hot')
+    plt.figure()
+    ax = plt.gca()
+    for y, row in enumerate(cells):
+        yy = y * s - s / 2
+        for x, cell in enumerate(row):
+            xx = x * s - s / 2
+            if keepout is not None and keepout._cells[y][x] >= INF:
+                color = 'purple'
+            else:
+                color = cmap(cell)
+            r = matplotlib.patches.Rectangle((xx, yy), w, w, facecolor=color, linewidth=s)
+            ax.add_patch(r)
+
+
+def h(n, objective):
+    return abs(n[0] - objective[0]) + abs(n[1] - objective[1])
+
+
+def gf(g, p):
+    v = g.get(p)
+    if v <= 0:
+        return INF
+    return v
+
+
+def kf(p, keepout):
+    return keepout.get(p)
+
+
+DIRS = {
+    None: ((1, 0), (-1, 0), (0, 1), (0, -1)),
+    0: ((1, 0), (1, -1), (1, 1)),  # North
+    1: ((0, 1), (1, 1), (-1, 1)),  # East
+    2: ((-1, 0), (-1, -1), (-1, 1)),  # South
+    3: ((0, -1), (1, -1), (-1, -1)),  # West
+}
+
+
+def plan(g: Grid, keepout: Grid, start, goal):
+    closed = set()
+    open = [start]
+    came_from = {}
+    g.set(start, 0.001)
+
+    while open:
+        c = open[0]
+        if c == goal:
+            path = []
+            p = goal
+            while True:
+                path.append(p)
+                if p == start:
+                    path.reverse()
+                    return path
+                p = came_from[p]
+        open = open[1:]
+        closed.add(c)
+
+        for d in DIRS[None]:
+            n = (c[0] + d[0], c[1] + d[1])
+            if n in closed:
+                continue
+            if not g.contains(n):
+                continue
+            gn = g.get(c) + 1
+            if n not in open:
+                open.append(n)
+            elif gn >= g.get(n):
+                continue
+            came_from[n] = c
+            g.set(n, gn)
+        open.sort(key=lambda x: gf(g, x) + h(x, goal) + kf(x, keepout))
+
+
+def main():
+    g = Grid(30, 30, 1)
+    keepout = Grid(30, 30, 1)
+    keepout.line((10, 0), (0, 1), 5, INF)
+    keepout.line((12, 12), (1, 0), 15, INF)
+    keepout.smear()
+
+    path = plan(g, keepout, (1, 2), (29, 18))
+    path = np.array(path)
+    plot(g, keepout)
+    plt.plot(path[:, 1], path[:, 0])
+    plt.axis('equal')
+    #    plt.plot(X[:, 0, ], X[:, 1])
+    plt.savefig('fig.png')
+    plt.show()
+
+
+if __name__ == '__main__':
+    main()