Как векторизовать функцию двух матриц в numpy?

Question

Скажем, у меня есть двоичная (смежная) матрица A измерений n x n и другая матрица U измерений n x 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

Это займет довольно много времени, даже если 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)