使用Numba循环Numpy数组时出现TypingError与TypeError报错求助

第一版代码报错原因 #

你对已经是二维numpy数组的raw_array_调用np.array()做无意义转换，Numba的nopython模式下没有对应np.array(二维数组)的重载实现，触发类型匹配错误。同时你直接把np.round_返回的浮点值作为数组索引，也会触发后续类型报错。

第二版代码报错原因 #

错误使用了@nb.generated_jit装饰器，该装饰器用于根据输入参数类型动态生成JIT实现，运行时传入函数的是Numba的类型对象而非实际数组，自然无法调用len()方法获取长度。

可正常运行的修复代码 #

import numba as nb
import numpy as np

@nb.njit
def _complicated(heat_list_, raw_array_, x_min_, x_step_, y_min_, y_step_):
    # 直接对输入的raw_array切片，不需要重复用np.array转换
    z_val = raw_array_[2, :]
    x_val = raw_array_[0, :]
    y_val = raw_array_[1, :]

    for i in range(len(z_val)):
        # 把np.round_的结果转成整数才能作为数组索引
        y_idx = int(np.round_((y_val[i] - y_min_) / y_step_))
        x_idx = int(np.round_((x_val[i] - x_min_) / x_step_))
        heat_list_[y_idx, x_idx] = z_val[i]
    return heat_list_

额外注意事项 #

• 运行前请确认计算出的y_idx、x_idx不会超出heat_list_的索引边界，否则会触发越界错误
• 如果需要拆分x/y/z数组单独传参，直接把函数入参改成对应三个一维数组，其余逻辑不变即可，不要使用@nb.generated_jit装饰器

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