Numpy: Multiple Outer Products

Question

General Problem

Suppose that I have a ndarray v of shape (nrow,ncols,3). I want to compute the ndarray outer_array of shape (nrow,ncols,3,3) containing all outer products of the vectors of shape (3) at each index (nrow,ncol). Of course, this is the the kind of problem for which numpy.einsum exists. Now, what I've tried is:

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

Now, I'm not sure that this will work: despite the fact that outer_array has the expected shape, the elements of the matrices of outer products do not correspond to what I'm expecting.

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'm reusing them in the output expression, the result of the sum is somehow broadcast.

On the other hand, if I write:

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

numpy will compute all possible combinations of outer products for each pair (x,y) and (u,v), resulting in an array of shape (ncols,nrow,ncols,nrow,3,3), where the diagonals (u,v) = (x,y) will contain the desired output.

The Precise Question

How do I choose the first two indices in the einsum notation in order to obtain an array where at each index x,y I get the outer product of vector v with itself without having to resort to the second expression?

Edit apparently, this form seems to work too:

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

I can only admire how useful numpy broadcasting is!

Answer 1

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

Let's 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 an outer 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 numpy broadcasting (i.e. adding dimensions and expanding them)

In "...i,...j->...ij", ... is functioning more as a place holder for existing, but anonymous dimensions.

The equivalent with einsum is:

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

(I should really do a calculation that isn't symmetric in the last 2 dimensions).

I've worked out a pure Python work-alike of einsum. The focus is on parsing the argument string, and how it creates the inputs for an iter object. It's available on github: https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py You are welcome to play around with it. It has a debug flag to show intermediate steps.

With my einsum with debugging 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 in creating a iter, the object that iterates over the axes of the input and output arrays. It iterates over the 1st axis of all arrays. The 2nd axis is 1 for 1st op and output, and a newaxis (-1) for the 2nd op.

With the 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 hence the same calculation.