我有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,但是它似乎并没有执行我想要的操作。