避免在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是个好主意。注意，您还可以使用简单的阵列广播 (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快--但我不确定.

编辑:用于稀疏矩阵

同样的基本想法应该适用，但是也许使用有效的索引本身而不是使用布尔掩蔽数组会更好.例如，类似：

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

Answer 2

这是与DilithiumMatrix相同的答案，但使用numpy接受的逻辑函数来生成掩码。

import numpy as np
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

Answer 3

假设您的两个数组(a和b)不是很大，您可以这样做：

import itertools
a = numpy.array([1,2,3])
b = numpy.array([4,5,6])
c = numpy.sum(list(itertools.product(a, b), 1)
def func(x, y):
    return x*y
numpy.vectorize(func)(c, 10)

请注意，对于大型数组，这是行不通的--您将在n**2中使用c元素，这意味着即使对于看起来很小的数组对，您也将使用大量的内存。对于每个有100000个元素的2个数组，所需的总内存将在74 GB的范围内。