python - 避免在 np.array 操作中使用 for 循环。 Python

Question

我有这两个数组:

import numpy as np
a = np.array([0, 10, 20])
b = np.array([20, 30, 40, 50])  

我想通过以下方式添加两者:

for i in range (len(a)):
   for j in range(len(b)):
      c = a[i] + b[j]
      d = delta(c, dr)

如您所见，每次迭代我都会得到一个值 c，我将其传递给函数 delta(请参阅帖子末尾的注释)。问题是当数组很大时，我想避免缓慢的 Python“for”循环。

我可以做的一件事是:

c = np.ravel(a(-1, 1) + b)

这要快得多。问题是现在 c 是一个数组，我将不得不再次使用 for 循环将其抛出。那么，您是否知道我如何在完全不使用 for 循环的情况下执行此操作。

注意:delta 是我按以下方式定义的函数:

def delta(r,dr):
   if r >= 0.5*dr and r <= 1.5*dr:
     delta = (5-3*abs(r)/dr-np.sqrt(-3*(1-abs(r)/dr)**2+1))/(6*dr)
   elif r <= 0.5*dr:
     delta = (1+np.sqrt(-3*(r/dr)**2+1))/(3*dr)
   else:
     delta = 0
   return delta

Answer 1

使用 ravel是个好主意。请注意，您也可以使用简单的 array broadcasting (a[:, np.newaxis] + b[np.newaxis, :])。

为您的function ，您可以对其进行很大改进，因为它仅由三个特定案例组成。最好的方法可能是对这三个部分中的每一个部分都使用掩码。

您开始于:

def delta(r,dr):
   if r >= 0.5*dr and r <= 1.5*dr:
     delta = (5-3*abs(r)/dr-np.sqrt(-3*(1-abs(r)/dr)**2+1))/(6*dr)
   elif r <= 0.5*dr:
     delta = (1+np.sqrt(-3*(r/dr)**2+1))/(3*dr)
   else:
     delta = 0

一个常见的替代方法是这样的:

def delta(r, dr):
    res = np.zeros_like(r)
    ma = (r >= 0.5*dr) & (r <= 1.5*dr)  # Create first mask
    res[ma] = (5-3*np.abs(r[ma])/dr[ma]-np.sqrt(-3*(1-np.abs(r[ma])/dr[ma])**2+1))/(6*dr[ma])
    ma = (r <= 0.5*dr)    # Create second mask
    res[ma] = (1+np.sqrt(-3*(r[ma]/dr[ma])**2+1))/(3*dr[ma])
    return res

初始化为零处理最终的 else案件。我还假设 np.abs比 abs 快--- 但我实际上不确定...

---

编辑:对于稀疏矩阵

应该应用相同的基本思想，但也许不使用 bool 掩码数组，而是使用有效索引本身会更好......例如像这样的东西:

res = scipy.sparse.coo_matrix(np.shape(r))
ma = np.where((r >= 0.5*dr) & (r <= 1.5*dr))  # Create first mask
res[ma] = ...