Hi, i am using numba to accelerate a dijkstra graph search on 2d grid.
The code is pretty slow (~1s for 1000x1000 image). Is there something obvious i could do to speed this up?
import heapq
import numpy as np
from numba import jit, njit
from numba.typed import List
@jit
def dijkstra_run(matrix, costs, priority_queue, visited, transitions, dxy):
x_max, y_max = matrix.shape
side = np.sqrt(x_max**2 + y_max**2)
while priority_queue:
cur_cost, (cur_x, cur_y) = heapq.heappop(priority_queue)
if visited[cur_x, cur_y]:
continue
for dx, dy in dxy:
x, y = cur_x + dx, cur_y + dy
if x < 0 or x >= x_max or y < 0 or y >= y_max:
continue
if ~visited[x,y]:
if matrix[x,y] + costs[cur_x, cur_y] < costs[x,y]:
costs[x,y] = matrix[x,y] + costs[cur_x,cur_y]
#distance = np.sqrt( (x_max-x)**2 + (y_max-y)**2 ) # activate to get Astar
heuristic = costs[x,y]
heapq.heappush(priority_queue, (heuristic, (x, y)))
transitions[x,y,0] = cur_x
transitions[x,y,1] = cur_y
visited[cur_x,cur_y] = 1
# retrieve the path
cur_x, cur_y = x_max - 1, y_max - 1
on_path = List()
on_path.append((cur_x,cur_y))
while (cur_x, cur_y) != (0, 0):
cur_x, cur_y = transitions[(cur_x, cur_y)]
on_path.append((cur_x,cur_y))
return on_path
def dijkstra(matrix, diag=False):
x_max, y_max = matrix.shape
dxy = List()
dxy.append((0,1))
dxy.append((1,0))
if diag:
dxy.append((1,1))
costs = np.full_like(matrix, 1.0e10)
costs[0,0] = matrix[0,0]
visited = np.zeros(matrix.shape, dtype=bool)
transitions = np.zeros((x_max, y_max, 2), dtype=np.int32)
priority_queue = List()
priority_queue.append((matrix[0,0], (0,0)))
return dijkstra_run(matrix, costs, priority_queue, visited, transitions, dxy)
# TEST
dijkstra(np.random.randn(1000,1000))