python - 如何缩短这个二维数组代码？

Question

我有以下代码片段，效果很好。

P=im1.copy()
for i in range(P.shape[0]):
    for j in range(P.shape[1]):
        if (n_dens.data[i][j]==-5000 or T_k.data[i][j]==-5000):
            P.data[i][j]=-5000
        else :
            P.data[i][j]=n_dens.data[i][j]*T_k.data[i][j]

其中 P 是一个二维数组。

我想知道如何将其缩减为以下内容:

P.data=n_dens.data*T_k.data
P.data=[foo-2.5005*10**7 if n_dens.data==-5000 or T_k.data==-5000 else foo for foo in P.data]

对于上面的试用，我收到以下错误:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

如何纠正错误？或者有其他方法可以减少它吗？

Answer 1

n_dens.data==-5000 生成真/假值数组，而不是单个值。所以，if 无法处理它。不过你已经很接近这个想法了。您可以在 numpy 中使用逻辑索引。

逻辑运算符也不能在 python 中重载。因此，numpy 不会按照您希望的方式处理它们。所以，你必须做类似的事情

index = np.logical_or(n_dens.data ==-5000, T_k.data==-5000)
P.data[index] = -5000

类似地，if-else 的第二个分支的 P.data[np.logic_not(index)] = n_dens.data * T.data 。