Astar

Astar.py

import math
import numpy as np
import bisect
from heapq import *

def a_star(B, E, D, H, N):
  """
  B : begin Point
  E : end Point
  D : distance (Point, Point) -> Real
  H : heuristic distance (Point, Point) -> Real
  N : neighbors (Point) -> [Point]
  """

  class node_astar:
    def __init__(self, db, de, parent=None):
      self.db = db
      self.de = de
      self.parent = parent
      self.d = self.db + self.de

    def __lt__(self, other):
      return self.d < other.d

  # Init
  open_set, close_set = [], {}
  heappush(open_set, (node_astar(0, H(E, B)), E))

  while open_set:
    # Chose the nearly point
    node, pos = heappop(open_set)
    close_set[pos] = node

    if pos == B:
      # Success
      path = []
      while node.parent:
        path.append(pos)
        pos, node = node.parent, close_set[node.parent]
        yield pos
      return
      # yield from reversed(path)

    # Add the neighbors
    for neighbor in (n for n in N(pos) if n not in close_set):
      db = node.db + D(pos, neighbor)
      de = H(neighbor, B)
      heappush(open_set, (node_astar(db, de, pos), neighbor))

  # Fail
  return

class Map:
  def __init__(self, size, layers=None):
    if layers is None:
      layers = list()

    assert len(size) == 2
    assert all(l.shape == size for l in layers)

    self.size = size
    self.layers = layers
    self.cached_potential = {}
    self.cached_distance = {}

  def neighbors(self, pos):
    assert len(pos) == 2
    assert pos[0] >= 0
    assert pos[1] >= 0
    assert pos[0] < self.size[0]
    assert pos[1] < self.size[1]

    candidate = list()
    if pos[0] >= 1:
      candidate.append((pos[0] - 1, pos[1]))
    if pos[1] >= 1:
      candidate.append((pos[0], pos[1] - 1))
    if pos[0] <= self.size[0] - 2:
      candidate.append((pos[0] + 1, pos[1]))
    if pos[1] <= self.size[1] - 2:
      candidate.append((pos[0], pos[1] + 1))

    def f(p):
      pot = self.potential(p)
      return not np.isinf(pot) and not np.isnan(pot)
    yield from filter(f, candidate)

  def potential(self, pos):
    if pos in self.cached_potential:
      return self.cached_potential[pos]
    re = sum(l[pos] for l in self.layers)
    self.cached_potential[pos] = re
    return re

  def distance(self, p1, p2):
    assert len(p1) == 2
    assert len(p2) == 2
    assert p1[0] >= 0
    assert p1[1] >= 0
    assert p1[0] < self.size[0]
    assert p1[1] < self.size[1]
    assert p2[0] >= 0
    assert p2[1] >= 0
    assert p2[0] < self.size[0]
    assert p2[1] < self.size[1]

    if (p1, p2) in self.cached_distance:
      return self.cached_distance[(p1, p2)]
    d = (p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2
    re = d + sum(l[p1] - l[p2] for l in self.layers)
    self.cached_distance[(p1, p2)] = re
    return re

  def path(self, p1, p2):
    func_D = lambda p1, p2: self.distance(p1, p2)
    func_H = func_D
    func_N = lambda pos: self.neighbors(pos)
    yield from a_star(p1, p2, func_D, func_H, func_N)

if __name__ == '__main__':
  from scipy.misc import imsave, imread


  wall = imread('./output.png').astype(np.double)
  wall[wall > 0] = np.inf
  walls = [wall[:,:] for _ in range(20)]

  size = (1500, 1000)
  m = Map(size, walls)

  import time
  b = time.time()
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  p = list(m.path((0, 0), (size[0] - 1, size[1] - 1)))
  print(time.time() - b)

  visited = np.zeros((1500, 1000))
  for a in p:
    visited[a] = 42
  imsave('test.png', visited)