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!
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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With