dijkstra.py 3.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. import queue
  2. import env
  3. import plotting
  4. class Dijkstra:
  5. def __init__(self, x_start, x_goal):
  6. self.xI, self.xG = x_start, x_goal
  7. self.Env = env.Env()
  8. self.plotting = plotting.Plotting(self.xI, self.xG)
  9. self.u_set = self.Env.motions # feasible input set
  10. self.obs = self.Env.obs # position of obstacles
  11. [self.path, self.policy, self.visited] = self.searching(self.xI, self.xG)
  12. self.fig_name = "Dijkstra's Algorithm"
  13. self.plotting.animation(self.path, self.visited, self.fig_name) # animation generate
  14. def searching(self, xI, xG):
  15. """
  16. Searching using Dijkstra.
  17. :return: planning path, action in each node, visited nodes in the planning process
  18. """
  19. q_dijk = queue.QueuePrior() # priority queue
  20. q_dijk.put(xI, 0)
  21. parent = {xI: xI} # record parents of nodes
  22. action = {xI: (0, 0)} # record actions of nodes
  23. visited = [] # record visited nodes
  24. cost = {xI: 0}
  25. while not q_dijk.empty():
  26. x_current = q_dijk.get()
  27. if x_current == xG: # stop condition
  28. break
  29. visited.append(x_current)
  30. for u_next in self.u_set: # explore neighborhoods of current node
  31. x_next = tuple([x_current[i] + u_next[i] for i in range(len(x_current))])
  32. if x_next not in self.obs: # node not visited and not in obstacles
  33. new_cost = cost[x_current] + self.get_cost(x_current, u_next)
  34. if x_next not in cost or new_cost < cost[x_next]:
  35. cost[x_next] = new_cost
  36. priority = new_cost
  37. q_dijk.put(x_next, priority) # put node into queue using cost to come as priority
  38. parent[x_next], action[x_next] = x_current, u_next
  39. [path, policy] = self.extract_path(xI, xG, parent, action)
  40. return path, policy, visited
  41. def get_cost(self, x, u):
  42. """
  43. Calculate cost for this motion
  44. :param x: current node
  45. :param u: input
  46. :return: cost for this motion
  47. :note: cost function could be more complicate!
  48. """
  49. return 1
  50. def extract_path(self, xI, xG, parent, policy):
  51. """
  52. Extract the path based on the relationship of nodes.
  53. :param xI: Starting node
  54. :param xG: Goal node
  55. :param parent: Relationship between nodes
  56. :param policy: Action needed for transfer between two nodes
  57. :return: The planning path
  58. """
  59. path_back = [xG]
  60. acts_back = [policy[xG]]
  61. x_current = xG
  62. while True:
  63. x_current = parent[x_current]
  64. path_back.append(x_current)
  65. acts_back.append(policy[x_current])
  66. if x_current == xI: break
  67. return list(path_back), list(acts_back)
  68. if __name__ == '__main__':
  69. x_Start = (5, 5) # Starting node
  70. x_Goal = (49, 5) # Goal node
  71. dijkstra = Dijkstra(x_Start, x_Goal)