python-3.x - 如何在 numpy 中向量化两个矩阵的函数？

Question

假设，我有一个维度为 nxn 的二元(邻接)矩阵 A 和另一个矩阵 U > 尺寸为 nxl。我使用以下代码来计算我需要的新矩阵。

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

即使 n 和 l 很小，这也需要相当多的时间来运行。有没有更好的方法来重写(向量化)此代码以减少运行时间？

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

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')

编辑2:用数学符号，我试图计算以下内容:

其中u_ik = U[i, k]、u_jk = U[j, k]和u_i = U[i, :].另外，(i,j)\in E 对应于代码中的a == 1.0。

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)