python - 如何应用依赖于当前数组元素的函数并将结果返回到当前数组元素

Question

数据格式：
我有一个形状为 (5,2,2) 的 3d Numpy 数组，基本上是 5 个 2x2 矩阵。

Import numpy as np

def mathFunc(arrSum):
    return arrSum*2

#creating 2x2 matrices
a = np.zeros((2, 2))
b = np.ones((2, 2))
c = np.random.randint(2,3,size=(2,2))
d = np.random.randint(3,4,size=(2,2))
e = np.random.randint(4,5,size=(2,2))

#Combining 2x2 matrices to form 3D matrix
arr = np.stack((a, b, c, d, e)) # you can add more 2D arrays here if needed

print("""ORIGINAL ARRAY
---------------
{0}
---------------""".format(arr))

目标
我想检查一个 2x2 数组（以下称为 3dObject）的总和是否超过某个值，如果没有，则跳过某些计算。需要进行的计算如下：

next_3dObject = next_3dObject - function(current_3dObject, next_3dObject)

我可以使用嵌套的 for 循环来完成此操作，但这会很慢。以下是当前代码。

filt_cond = 5 #filter condition (i.e. to small of a sum)
# Using Fancy indexing and slicing
for i in range (len(arr)):
    
        new_basis = np.random.randint(1,4,size=(2,2))

        '''need to replace mathFunc() with function call that can take current_3dObject and next_3dObject 
        do calculation and return value to next_3dObject, this need to happen for every 3dObject after 
        the current_3dObject(which is where the condition was first found to be True). After 3dObjects have
        been updated it will check the next 3dObject for sum condition'''

        arr[i+1:][arr[i].sum() >filt_cond] *= new_basis  #mathFunc(i) 
        print("""Array iteration {0}
        ---------------
        {1}
        ---------------""".format(i,arr))

1. 是否可以对这段代码进行矢量化？
2. 是否可以将操作应用于 Numpy 数组的某个部分，然后将这些结果反馈回数组并再次应用操作但从不同的元素开始？

标签： pythonperformancenumpyvectorization