Numpy: several external products

Question

a common problem

Suppose I have ndarray v forms (nrow,ncols,3) . I want to compute an ndarray outer_array form (nrow,ncols,3,3) containing all the outer products of the vectors of form (3) for each index (nrow,ncol) . Of course, this is the problem for which numpy.einsum exists. Now I tried:

 outer_array = numpy.einsum("xyi,xyj->xyij",v,v.conjugate()) 

Now I'm not sure if this will work: despite the fact that outer_array has the expected shape, the elements of the matrix of external products do not match the expected.

I think this is due to the choice of labels in the einsum expression: the product is supposed to be summed over x and y , because the indices are repeated, but since I reuse them in the output expression, the result of the sum is transmitted somehow.

On the other hand, if I write:

 outer_array = numpy.einsum("xyi,uvj->...ij",v,v.conjugate()) 

numpy will calculate all possible combinations of external products for each pair (x,y) and (u,v) , as a result we get an array of the form (ncols,nrow,ncols,nrow,3,3) , where the diagonals (u,v) = (x,y) will contain the desired result.

Exact question

How to select the first two indexes in einsum notation to get an array, where in each index x,y I get the external product of the vector v without it, without resorting to the second expression?

Edit apparently this form also works:

 outer_array = numpy.einsum("...i,...j->...ij",v,v.conjugate()) 

I can only admire how useful the numpy broadcast is!

Answer 1

The key to working with 'xyi,xyj->xyij' is that xy repeated in the output line.

Let us use a simpler array:

 x = np.arange(6).reshape(3,2) np.einsum.einsum('ij->j',x) # array([6, 9]) # sums on the 1st dimension of 'x' 

Now for the external product on this x :

 In [20]: x[:,:,None]*x[:,None,:] # shape (3,2,2) Out[20]: array([[[ 0, 0], [ 0, 1]], [[ 4, 6], [ 6, 9]], [[16, 20], [20, 25]]]) 

This is an example of simple broadcasting (i.e. adding dimensions and expanding them)

In "...i,...j->...ij" , ... functions more as a placeholder for existing but anonymous dimensions.

Equivalent with einsum :

 np.einsum('ij,ik->ijk',x,x) 

(I really have to do a calculation that is not symmetrical in the last two dimensions).

I developed pure Python, similar to einsum . The focus is on parsing the argument string and how it creates input for the iter object. It is available on github: https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py. You can experiment with it. It has a debug flag to show intermediate steps.

With my einsum with debug output:

 In [23]: einsum_py.myeinsum('ij,ik->ijk',x,x,debug=True) # ... some parsing information [('ij', [105, 106], 'NONE'), ('ik', [105, 107], 'NONE')] ('ijk', [105, 106, 107], 'NONE') iter labels: [105, 106, 107],'ijk' op_axes [[0, 1, -1], [0, -1, 1], [0, 1, 2]] 

op_axes is the key argument that is used when creating iter , an object that iterates over the axes of the input and output arrays. Iterates over the 1st axis of all arrays. The 2nd axis is 1 for the 1st operation and exit, and newaxis (-1) for the 2nd operation.

With ellipsis :

 In [24]: einsum_py.myeinsum('...j,...k->...jk',x,x,debug=True) ... iter labels: [0, 106, 107],'0jk' op_axes [[0, 1, -1], [0, -1, 1], [0, 1, 2]] 

This generates the same op_axes and therefore the same calculations.