1

I would like to obtain a tensordot of two arrays with the same shape with index-dependent weight applied, without use of explicit loop. For example,

import numpy as np
    
A=np.array([1,2,3])
B=np.array([-2,6,9])
C=np.zeros((3,3))
for i in range(3):
     for j in range(3):
          C[i,j]=A[i]*B[j]*(np.exp(i-j)if i>j else 0)

Can an array similar to C be obtained with a built-in tool (e.g., with some options for tensordot)?

user16308
  • 43
  • 5

2 Answers2

1

Here's a vectorized solution:

N = 3
C = np.tril(A[:, None] * B * np.exp(np.arange(N)[:, None] - np.arange(N)), k=-1)

Output:

>>> C
array([[ -2.        ,   0.        ,   0.        ],
       [-10.87312731,  12.        ,   0.        ],
       [-44.33433659,  48.92907291,  27.        ]])
0

With np.einsum inconsistently slightly faster for some larger inputs than broadcasting, slower for others.

import numpy as np
    
A=np.array([1,2,3])
B=np.array([-2,6,9])

np.einsum('ij,i,j->ij', np.tril(np.exp(np.subtract.outer(A,A)), -1), A, B)

Output

array([[  0.        ,   0.        ,   0.        ],
       [-10.87312731,   0.        ,   0.        ],
       [-44.33433659,  48.92907291,   0.        ]])
Michael Szczesny
  • 4,911
  • 5
  • 15
  • 32