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

Question

我有以下函数，它接受两个一维 numpy 数组 q_i 和 q_j，做一些计算（包括取它们的差异的范数）和 returns一个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] 中，我想要为 q_i 和 q_j 内部的所有唯一对评估函数 coulomb （即对于 (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

当您想向量化这些类型的 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

Answer 2

这里有 3 个可能的解决方案，最后一个，有点粗暴，但使用矢量化来计算 n q vs one。也是最快的

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

``

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

Truly vectorize function for numpy array in python

python

performance

numpy

vectorization