如何在numpy中向量化两个矩阵的函数？

Question

说，我有一个尺寸为A x n的二进制（邻接）矩阵n和另一个尺寸为U x {{1 }}。我使用以下代码来计算所需的新矩阵。

n

即使l和import numpy as np from numpy import linalg as LA new_U = np.zeros_like(U) for idx, a in np.ndenumerate(A): diff = U[idx[0], :] - U[idx[1], :] if a == 1.0: new_U[idx[0], :] += 2 * diff elif a == 0.0: norm_diff = LA.norm(U[idx[0], :] - U[idx[1], :]) new_U[idx[0], :] += -2 * diff * np.exp(-norm_diff**2) return new_U 很小，这也要花费大量时间。是否有更好的方法重写（矢量化）此代码以减少运行时间？

编辑1：示例输入和输出。

n

编辑2：以数学符号表示，我正在尝试计算以下内容：

其中l，A = np.array([[0,1,0], [1,0,1], [0,1,0]], dtype='float64') U = np.array([[2,3], [4,5], [6,7]], dtype='float64') new_U = np.array([[-4.,-4.], [0,0],[4,4]], dtype='float64') 和u_ik = U[i, k]。另外，u_jk = U[j, k]对应于代码中的u_i = U[i, :]。

Answer 1

利用broadcasting和np.einsum来减少总和-

# Get pair-wise differences between rows for all rows in a vectorized manner
Ud = U[:,None,:]-U

# Compute norm L1 values with those differences
L = LA.norm(Ud,axis=2)

# Compute 2 * diff values for all rows and mask it with ==0 condition
# and sum along axis=1 to simulate the accumulating behaviour
p1 = np.einsum('ijk,ij->ik',2*Ud,A==1.0)

# Similarly, compute for ==1 condition and finally sum those two parts
p2 = np.einsum('ijk,ij,ij->ik',-2*Ud,np.exp(-L**2),A==0.0)
out = p1+p2

或者，使用einsum来计算平方范数，并使用它们来获得p2-

Lsq = np.einsum('ijk,ijk->ij',Ud,Ud)
p2 = np.einsum('ijk,ij,ij->ik',-2*Ud,np.exp(-Lsq),A==0.0)