Numba @jitclass：实践、踩雷注意事项

@jitclass

允許定義類並進行 JIT 編譯。然而，在使用 jitclass 時，以下是一些重要的注意事項。
如果你也在numba jit debug地狱，希望这个对你有所帮助

注意事项

设定明确的spec

每个class中的var都需要设定，简言之就是使用self.xxx的都要

from numba import jit, int32, float64, boolean, int64
spec = [
    ('a', int32),
    ('b', int32),
    ('c', int32[:,:,:]),
    ('d', float64[:,:])
]
@jitclass(spec)
class foo:
	def __init():
		self.a = ...
		self.b = ...
		self.c = ...
		self.d = ...

避免重新分配 Numpy array

在处理 Numpy 数组时，避免在初始创建后重新分配数组。 Numba 的优化基于数据结构不改变形状的假设。

错误范例

# class开头声明
np_arr = np.array((2,3), )
# 后来操作
np_arr*= 0
np_arr= np.array(temp_HitsReport.shape)

尽量用for loop取代如找值、清零

明确声明循环变量

在 Numba 中，必须在循环内明确声明循环变量。与标准 Python 不同，Numba 无法从其他作用域推断循环内声明的变量。

错误范例
以下范例，在python中可以成立的状况，但numba无法识别

# class开头
for i in range (10)
	print(i)
self.a = i 

在初始化之外不要运行 Numpy 数组操作

在 Numba 中，最好避免在初始化之外运行 Numpy 数组操作，这样可以避免一些潜在的问题。

错误范例

# 非初始化处执行这个操作
now_array = np.array(C2_lst[self.arr_trigger_condition[0][2] - 6].copy(), dtype=np.int32)

可行做法

# 使用复制方法
now_array = xxx.copy()

直接加载数据而非Call其他函数来加载

当加载数据时，应该在调用函数的地方直接加载，而不是通过其他类或函数进行转发。

总结

「直接一点」：哪边call function就哪边放变数，不要再丢到其他class function执行，减少Class的深度，这样读起来才不会报错。