python - 使用二维数组作为条件

Question

我正在尝试在 if 循环中使用 2D 数组作为布尔值，似乎无法正确处理，这是一个很大的计算（400,600），所以我试图避免使用双 if 语句。这就是代码的样子。

    if obstacle == 1:

        f1_coll  =f3  
        f2_coll  =f4  
        f3_coll  =f1  
        f4_coll  =f2  

    else:
        f1_coll=tau*(f1 +f2 -f3 +f4 )
        f2_coll =tau*(f1 +f2 +f3 -f4 )
        f3_coll =tau*(-f1 +f2 +f3 +f4 )
        f4_coll =tau*(f1 -f2 +f3 +f4 )

所有变量都具有相同的大小（400,600），除了 tau 是一个常数。有蚂蚁功能可以帮助我吗？

标签： pythonmultidimensional-arrayscientific-computing

Answer 1

如果我正确理解您的问题，我认为以下内容可以解决问题：

import numpy as np 

# fake variables for the example
obstacle = np.array([[1,0],[1,1]])
f1 = np.array([[1,1],[1,1]])
f2 = 2*f1
f3 = 3*f1
f4 = 4*f2

tau = 0.5

#here starts the proposed solution
f1_coll = obstacle  * f3 + (obstacle==0) * tau *(f1 +f2 -f3 +f4)