Lets say, I have bunch of matrices As and vectors bs.
As = array([[[1, 7], [3, 8]],
[[2, 1], [5, 9]],
[[7, 2], [8, 3]]])
bs = array([[8, 0], [8, 8], [7, 3]])
When I do np.inner(As, bs), I get:
array([[[ 8, 64, 28], [ 24, 88, 45]],
[[ 16, 24, 17], [ 40, 112, 62]],
[[ 56, 72, 55], [ 64, 88, 65]]])
But I do not need all inner products. What I want is, to calculate each matrix with each vector once.
I can do something like this:
np.array(map(lambda (a, b): np.inner(a, b), zip(As, bs)))
Then I get the expected matrix:
array([[ 8, 24], [ 24, 112], [ 55, 65]])
Now I do not want to use zip, map etc. because I need this operation > 10**6 time (for image processing, exactly for GMM).
Is there any way to use numpy, scipy etc. that can do this for me? (fast and efficent)
You can use np.einsum
-
np.einsum('ijk,ik->ij',As, bs)
Explanation
With np.array(map(lambda (a, b): np.inner(a, b), zip(As, bs)))
, we are selecting the first element off As
as a
and off bs
as b
and doing inner product. Thus, we are doing :
In [19]: np.inner(As[0],bs[0])
Out[19]: array([ 8, 24])
In [20]: np.inner(As[1],bs[1])
Out[20]: array([ 24, 112])
In [21]: np.inner(As[2],bs[2])
Out[21]: array([55, 65])
Think of it as a loop, we iterate 3 times, corresponding to the length of first axis of As
, which is same as for bs
. Thus, looking at the lambda
expression, at each iteration, we have a = As[0] & b = bs[0]
, a = As[1] & b = bs[1]
and so on.
As
and bs
being 3D
and 2D
, let's represent them as iterators imagining the inner-product in our minds. Thus, at iteration, we would have a : j,k
and b : m
. With that inner product between a
and b
, we would lose the second axis of a
and first of b
. Thus, we need to align k
with m
. Thus, we could assume b
to have the same iterator as k
. Referencing back from a
to As
and b
to bs
, in essence, we would lose the third axis from As
and second from bs
with the inner product/sum-reduction. That iterating along the first axis for As
and bs
signifies that we need to keep those aligned under these sum-reductions.
Let's summarize.
We have the iterators involved for the input arrays like so -
As : i x j x k
bs : i x k
Steps involved in the intended operation :
- Keep the first axis of
As
aligned with first of bs
.
- Lose the third axis of
As
with sum-reduction against second of bs
.
Thus, we would be left with the iterators i,j
for the output.
np.einsum
is a pretty efficient implementation and is specially handy when we need to keep one or more of the axes of the input arrays aligned against each other.
For more info on einsum
, I would suggest following the docs link supplied earlier and also this Q&A
could be helpful!