How to vectorize a function of two matrices in numpy?

Question

Say, I have a binary (adjacency) matrix A of dimensions nxn and another matrix U of dimensions nxl. I use the following piece of code to compute a new matrix that I need.

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

This takes quite a lot of time to run even when n and l are small. Is there a better way to rewrite (vectorize) this code to reduce the runtime?

Edit 1: Sample input and output.

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

Edit 2: In mathematical notation, I am trying to compute the following:

where u_ik = U[i, k],u_jk = U[j, k], and u_i = U[i, :]. Also, (i,j) \in E corresponds to a == 1.0 in the code.

Answer 1

Leveraging broadcasting and np.einsum for the sum-reductions -

# 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

Alternatively, use einsum for computing squared-norm values and using those to get p2 -

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