将函数应用于多维NumPy数组的二维切片

时间:2018-09-25 18:27:16

标签: python arrays numpy

我有6D numpy数组,我有一个可以应用于2D数组的自定义函数,输出是形状相同的2D数组。我想遍历存储在索引3和4中的所有2D数组,应用自定义函数,并以原始6D数组的形状以及如上所述使用自定义函数处理的数据返回6D数组。我目前正在使用嵌套循环,但是过程非常缓慢。下面是我当前正在执行的操作(维度2和5的大小为1,因此无需遍历它们)。

orig_array = np.reshape(np.zeros(2**4), (2, 2, 1, 2, 2, 1))
def custom_fun(array_2d):
    return array_2d + 1
out_array = np.copy(orig_array)
for ind1 in range(orig_array.shape[0]):
    for ind2 in range(orig_array.shape[1]):
        out_array[ind1, ind2, 0, :, :, 0] = custom_fun(out_array[ind1, ind2, 0, :, :, 0])

我想知道是否有一种方法可以矢量化此过程,以加快计算速度。 我搜索了NumPy文档并找到了函数:np.apply_over_axes,但是它似乎并没有执行我想要的操作。

0 个答案:

没有答案