Python Numba和多维添加-不使用numpy.newaxis？

试图在python上加速DP算法，numba似乎是合适的人选

我正在用一个一维数组减去一个二维数组，这个一维数组提供一个三维数组。然后，我沿着第三维方向使用

.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）

时会发生这种情况。 是否有一个工作来使以下计算与麻木工作

---

我尝试过使用numpy repeats&expand_dims的变体，并明确定义了jit函数的输入类型，但都没有成功。

要使其正常工作，您需要做一些更改：

添加带有

arr[：，：，None]

的维度：对于Numba，它看起来像

getitem

，因此更喜欢使用

重塑

使用

np.abs

代替内置的

abs

带有

轴的
argmin
关键字参数为。更喜欢使用循环，这是Numba设计用来优化的

修复所有这些问题后，您可以运行jitted函数：

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
        new_lvl_3d = new_lvl.reshape(*new_lvl.shape, 1)
        diff = np.abs(disc_lvl - new_lvl_3d)  # Dim [N x M x O]

        idx_lvl = np.empty(new_lvl.shape)
        for i in range(diff.shape[0]):
            for j in range(diff.shape[1]):
                idx_lvl[i, j] = diff[i, j, :].argmin()

        return True

# function works fine without numba
success = my_func(disc_lvl, actions, start_lvl, inflow)

---

在下面找到我的第一篇文章中更正的代码，您可以使用或不使用numba库的jitted模式执行（通过删除以@jit开头的行）。我观察到这个例子的速度增加了2倍

from numba import jit
import numpy as np
import datetime as dt

inflow = np.arange(1,0,-0.01)                       # Dim [T]
nbTime = np.shape(inflow)[0]
actions = np.arange(0,1,0.01)                       # 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(nbTime, disc_lvl, actions, start_lvl, inflow):
    # Initialize result 
    res = np.empty((nbTime,np.shape(start_lvl)[0],np.shape(actions)[0]))

    for t in range(0,nbTime):
        # Calculate new level at time t
        new_lvl = start_lvl + inflow[t] + actions  # Dim [N x M]      
        print(t)

        # For each new_level element, find closest discretized level
        new_lvl_3d = new_lvl.reshape(*new_lvl.shape, 1)
        diff = np.abs(disc_lvl - new_lvl_3d)  # Dim [N x M x O]

        idx_lvl = np.empty(new_lvl.shape)
        for i in range(diff.shape[0]):
            for j in range(diff.shape[1]):
                idx_lvl[i, j] = diff[i, j, :].argmin()

        res[t,:,:] = idx_lvl

    return res

# Call function and print running time
start_time = dt.datetime.now()
result = my_func(nbTime, disc_lvl, actions, start_lvl, inflow)
print('Execution time :',(dt.datetime.now() - start_time))

---

我相信这是因为

newaxis

是

None

的别名，而这种索引在numba中还不受支持。您可以尝试使用

重塑

。谢谢您的建议@Ralvi Isufaj。我已经试过了，但不幸的是它不起作用。我使用参数（数组（int64，1d，C））*参数化的边界函数（数组（float64，2d，C））无效使用了

TypingError:I已替换定义了diff的行。同样，它在没有@jit
diff=disc\u lvl-new\u lvl.reforme（np.append（np.shape（new\u lvl），1））的情况下工作得很好。
作为将来的参考，当发布错误时，请复制您收到的错误消息，这对找出错误有很大帮助：）非常感谢@Jacques Gaudin。你的解决方案很有效，帮助我解决了问题。我稍微增加了矩阵的大小，并观察到与相同函数的非jitted版本相比，计算时间减少了因子2。这对你来说合理吗？这可能是你能得到的最多了，因为你的大部分原始函数都是矢量化的。如果有的话，我会研究产生不同维度结果的操作，看看循环版本是否更快。我不认为有多少可以节省，但它可能值得一试。