lib/grid: playing with A* planning.

master
Michael Hope 4 years ago
parent dc53ab3fa7
commit 898cf63df0

@ -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()
Loading…
Cancel
Save