真正矢量化函数的 numpy 数组在 python

Question

我有以下函数，它接受两个 1Dnumpy数组q_i和q_j，进行一些计算(包括取它们差值的范数)并返回一个numpy数组：

import numpy as np
import numpy.linalg as lin
def coulomb(q_i, q_j, c1=0.8, c2=0.2, k=0.5):
"""
Parameters
----------
q_i : numpy.ndarray
q_j : numpy.ndarray
c1 : float
c2 : float
k : float
Returns
-------
numpy.ndarray
"""

q_ij = q_i - q_j
q_ij_norm = lin.norm(q_i - q_j)
f_q_i = (k * c1 * c2 / q_ij_norm ** 3) * q_ij
return f_q_i

现在我有一堆这些q数组存储在另一个numpy数组中t= [q1， q2， q3， ...， qn]，我想评估t内所有唯一q_i和q_j对的函数coulomb(即 for (q1，q2)， (q1，q3)， ...， (q1，qn)， (q2，q3)， ...， (q2，qn)， ...(q_n-1，qn))。

有没有办法对此计算进行矢量化(我的意思是真正矢量化它以提高性能，因为np.vectorize只是引擎盖下的for循环)？

我当前的解决方案是嵌套的for循环，这在性能方面远非最佳：

for i, _ in enumerate(t):
for j, _ in enumerate(t[i+1:]):
f = coulomb(t[i], t[j])
...

Answer 1

这里 3 个可能的解决方案，最后一个，有点苛刻，但使用矢量化来计算nq 与 1。也是最快的

from itertools import combinations
import numpy as np
import numpy.linalg as lin
def coulomb(q_i, q_j, c1=0.8, c2=0.2, k=0.5):
"""
Parameters
----------
q_i : numpy.ndarray
q_j : numpy.ndarray
c1 : float
c2 : float
k : float
Returns
-------
numpy.ndarray
"""

q_ij = q_i - q_j
q_ij_norm = lin.norm(q_ij)
f_q_i = (k * c1 * c2 / q_ij_norm ** 3) * q_ij
return f_q_i    
def coulomb2(q_i, q_j, c1=0.8, c2=0.2, k=0.5):
"""
Parameters
----------
q_i : numpy.ndarray
q_j : numpy.ndarray
c1 : float
c2 : float
k : float
Returns
-------
numpy.ndarray
"""

q_ij = q_i - q_j
q_ij_norm = lin.norm(q_ij,axis=1).reshape(-1,1)
f_q_i = (k * c1 * c2 / q_ij_norm ** 3) * q_ij
return f_q_i    

q = np.random.randn(500,10)
from itertools import combinations
from time import time

t1= time()
v = []
for i in range(q.shape[0]):
for j in range(i+1,q.shape[0]):

v.append([coulomb(q[i], q[j])])
t2= time()
combs = combinations(range(len(q)), 2)
vv =[]
for i,j in combs:
vv.append([coulomb(q[i], q[j])])
t3 = time()
vvv = []
for i in  range(q.shape[0]):
vvv += list(coulomb2(q[i], q[i+1:]))
t4 = time()
print(t2-t1)
print(t3-t2)
print(t4-t3)
#0.9133327007293701
#1.0843684673309326
#0.04461050033569336
``

Answer 2

当您想要矢量化这些类型的 numpy 问题时，需要在内存和速度之间进行权衡，

numpy 数组上的循环很慢，但没有太多的内存要求。如果要矢量化，则必须复制，从而创建多余的内存并将其传递给numpy函数。

矢量化函数的一种方法是

import numpy as np
import numpy.linalg as lin
def coulomb(q_i, q_j, c1=0.8, c2=0.2, k=0.5):
"""
Parameters
----------
q_i : numpy.ndarray
q_j : numpy.ndarray
c1 : float
c2 : float
k : float
Returns
-------
numpy.ndarray
"""

q_ij = q_i[np.newaxis,:, :] - q_j[:,np.newaxis, :] #broadcasting, therefore creating more data
q_ij_norm = lin.norm(q_ij, axis=2) # this can be zero when qi = qj
f_q_i = (k * c1 * c2 / (q_ij_norm ** 3)[:,:,np.newaxis]) * q_ij #broadcasting again
return np.nan_to_num(f_q_i) # this returns a 3D array with shape (i,j,dim_of_q). Here Nan's will be replaced with 0's
t = np.random.rand(500,3)
f_q = coulomb(t, t)
f_q(1,2,:) #will return `coulomb` between q1 and q2