添加Numba @jit/@njit装饰器后，Brushfire算法向量场代码报错求助

You're running into classic Numba compatibility issues here—Numba's JIT compiler (especially njit, which runs in no-Python mode) has strict rules about the types of operations it can optimize. Let's unpack the two errors first:

• TypingError with @njit: Numba can't infer the type of the dynamic dictionary you're building in get_direction because you're deleting keys at runtime. It needs static, predictable structures to compile efficiently.
• LoweringError with @jit: The lambda function you're using to find the neighbor with the minimum cost isn't supported in Numba's lower-level code generation, especially when combined with dictionary operations.

Let's fix these issues step by step with Numba-friendly code refactors:

Step 1: Refactor get_direction for Numba Compatibility #

Instead of building a dynamic dictionary and deleting invalid keys, we'll use a fixed list of direction tuples, filter valid pairs, and return structured NumPy arrays. This gives Numba clear, static types to work with.

from numba import njit
import numpy as np

@njit
def get_direction(point, rows, cols):
    x, y = point
    # Define all possible direction pairs as a fixed static list
    all_directions = [
        ((x - 1, y + 1), 225),
        ((x, y + 1), 180),
        ((x + 1, y + 1), 135),
        ((x - 1, y), 270),
        ((x + 1, y), 90),
        ((x - 1, y - 1), 315),
        ((x, y - 1), 0),
        ((x + 1, y - 1), 45),
    ]
    valid_neighbors = []
    valid_angles = []
    # Filter out out-of-bounds neighbors
    for neighbor, angle in all_directions:
        nx, ny = neighbor
        if 0 <= nx < rows and 0 <= ny < cols:
            valid_neighbors.append(neighbor)
            valid_angles.append(angle)
    # Convert to NumPy arrays for efficient, Numba-friendly indexing
    return (np.array(valid_neighbors, dtype=np.int32),
            np.array(valid_angles, dtype=np.int32))

Step 2: Rewrite vectorfield to Avoid Lambda Functions #

We'll replace the lambda-based min() call with an explicit loop to find the neighbor with the lowest cost. This eliminates the lambda that was causing the LoweringError.

@njit
def vectorfield(inarray, goal, r, c):
    vectorarray = np.zeros((r, c), dtype=np.int32)
    goal_x, goal_y = goal
    for y in range(c):
        for x in range(r):
            # Skip the goal point to handle it separately later
            if x == goal_x and y == goal_y:
                continue
            # Get valid neighbors and their corresponding angles
            neighbors, angles = get_direction((x, y), r, c)
            # Explicitly find the neighbor with the minimum cost
            min_cost = np.inf
            best_angle = 0
            for idx in range(len(neighbors)):
                nx, ny = neighbors[idx]
                current_cost = inarray[nx, ny]
                if current_cost < min_cost:
                    min_cost = current_cost
                    best_angle = angles[idx]
            vectorarray[x, y] = best_angle
    # Set the goal point's value to -1 as required
    vectorarray[goal_x, goal_y] = -1
    return vectorarray

Key Changes That Fix the Errors #

1. No dynamic dictionaries: We replaced runtime dictionary modifications with a fixed list and explicit filtering, which Numba can type easily.
2. No lambda functions: The explicit loop for finding the minimum cost avoids Numba's limitations with lambda expressions in no-Python mode.
3. Structured NumPy arrays: Returning NumPy arrays instead of dictionaries ensures Numba can infer consistent types across function calls.

With these changes, both functions will compile successfully with @njit and deliver the performance boost you're looking for while maintaining the original functionality.

内容的提问来源于stack exchange，提问作者Epic Boss