功能上的小变化。面具。Python

Question

我的职能如下：

def delta(r, dr):
   res = np.zeros(r.shape)
   mask1 = (r >= 0.5*dr) & (r <= 1.5*dr)
   res[mask1] = (5-3*np.abs(r[mask1])/dr \
    - np.sqrt(-3*(1-np.abs(r[mask1])/dr)**2+1))/(6*dr)
   mask2 = np.logical_not(mask1) & (r <= 0.5*dr)
   res[mask2] = (1+np.sqrt(-3*(r[mask2]/dr)**2+1))/(3*dr)
   return res

其中，r是大小为(shape[0],shape[1])的numpy.array，而dr是单个值。我希望修改这个函数，使dr也成为一个与r大小相同的数组，并且对于r的每个值，从dr获取类似的值。

例如，r[0,0]与dr[0,0]一起使用，r[0,1]与dr[0,1]一起使用等等。有什么想法吗？

Answer 1

您可以将2D掩码与输入数组相乘，这实际上是掩蔽，从而执行计算，从而生成一个2D数组，而不是像目前这样使用布尔索引的一维数组。唯一的区别是将值设置到输出数组中，为此您需要屏蔽要设置的数组和将从中选择值的2D计算数组。

执行过程会是这样的-

# Initialize output array   
res = np.zeros(r.shape)

# Get mask1 and compute values for all elements and use the mask to set only
# TRUE positions with the computed values
mask1 = (r >= 0.5*dr) & (r <= 1.5*dr)
V1 = (5-3*np.abs(r*mask1)/dr - np.sqrt(-3*(1-np.abs(r*mask1)/dr)**2+1))/(6*dr)
res[mask1] = V1[mask1]

# Similarly for mask2 and the computations with that mask
mask2 = np.logical_not(mask1) & (r <= 0.5*dr)
V2 = (1+np.sqrt(-3*(r*mask2/dr)**2+1))/(3*dr)
res[mask2] = V2[mask2]