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…
Reference in new issue