尝试在 python 上加速 DP 算法,numba 似乎是一个合适的候选者。
我正在用提供 3D 数组的 1D 数组减去 2D 数组。然后我使用.argmin()第三维来获得一个二维数组。这适用于 numpy,但不适用于 numba。
重现问题的玩具代码:
from numba import jit
import numpy as np
inflow = np.arange(1,0,-0.01) # Dim [T]
actions = np.arange(0,1,0.05) # Dim [M]
start_lvl = np.random.rand(500).reshape(-1,1)*49 # Dim [Nx1]
disc_lvl = np.arange(0,1000) # Dim [O]
@jit(nopython=True)
def my_func(disc_lvl, actions, start_lvl, inflow):
for i in range(0,100):
# Calculate new level at time i
new_lvl = start_lvl + inflow[i] + actions # Dim [N x M]
# For each new_level element, find closest discretized level
diff = (disc_lvl-new_lvl[:,:,np.newaxis]) # Dim [N x M x O]
idx_lvl = abs(diff).argmin(axis=2) # Dim [N x M]
return True
# function works fine without numba
success = my_func(disc_lvl, actions, start_lvl, inflow)
为什么上面的代码不运行?取出时会这样@jit(nopython=True)。是否有一个工作回合可以使以下计算与 numba 一起工作?
我尝试了带有 numpy repeats 和 expand_dims 的变体,以及明确定义 jit 函数的输入类型但没有成功。
HUX布斯
翻过高山走不出你
相关分类