Numpy: Multiple Outer Products
arrays, numpy, python
Solution
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.
Problem
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!