如何缩短此二维数组代码?

时间:2018-08-28 17:10:28

标签: python arrays

我有下面的代码片段,效果很好。

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是2D数组。

我想知道如何将其缩小为以下几行:

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()

如何纠正错误?还是有另一种修整方法?

2 个答案:

答案 0 :(得分:3)

n_dens.data==-5000生成一个true / false值数组,而不是单个值。因此,if无法处理。不过,您很接近这个主意。您可以在numpy中使用逻辑索引

逻辑运算符也不能在python中重载。因此,numpy无法按照您的意愿处理它们。因此,您必须执行类似的操作

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

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

答案 1 :(得分:2)

您可以尝试以下方法:

P.data[(n_dens.data == -5000) | (T_k.data == -5000)] = -5000
cond = ~(n_dens.data == -5000) & ~(T_k.data == -5000) # 2D array of booleans
P.data[cond] = n_dens.data[cond] * T_k.data[cond]

完整的示例:

import numpy as np
from copy import deepcopy

class IMAGE:
    def __init__(self, data):
        self.data = data
        self.shape = self.data.shape

np.random.seed(0)
P, n_dens, T_k = IMAGE(np.zeros((5,5))), IMAGE(np.reshape(np.random.choice([-5000,1,2],25), (5,5))), IMAGE(3*np.ones((5,5)))
P1 = deepcopy(P)

# with loop
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]

# vectorized
P1.data[(n_dens.data == -5000) | (T_k.data == -5000)] = -5000
cond = ~(n_dens.data == -5000) & ~(T_k.data == -5000) # 2D array of booleans
P1.data[cond] = n_dens.data[cond] * T_k.data[cond]

cond
# array([[False,  True, False,  True,  True],
#   [ True, False,  True, False, False],
#   [False,  True,  True,  True,  True],
#   [False,  True,  True,  True,  True],
#   [False,  True, False, False,  True]], dtype=bool)

# with same output for both
P.data == P1.data
# array([[ True,  True,  True,  True,  True],
#   [ True,  True,  True,  True,  True],
#   [ True,  True,  True,  True,  True],
#   [ True,  True,  True,  True,  True],
#  [ True,  True,  True,  True,  True]], dtype=bool)