Einstein Notation for product of stacked matrices

Question

Background Information:

I recently started using the Einstein summation notation to express certain operations over an "image" $\mathbf{A}$ where to each pixel a square matrix is attached. That is, an array of the shape $\left[n_{row}\times n_{col}\times n \times n\right]$. Now, if I want to represent the multiplication of each "submatrix" $A_{pq}$ of size $n \times n$ by a matrix $\mathbf{B}$ of compatible size, I can express it as:

$A^{'}_{pqij} = B_{ik}A_{pqkj}$.

The problem: I have an array $\mathbf{B}$ of the same size as $\mathbf{A}$ and I want to compute a third array $\mathbf{A}^{'}$ where each pixel $\mathbf{A}_{pix}^{'}$ at the first two indices $p$ and $q$ corresponds to the matrix multiplication of matrix $\mathbf{A}_{pix}$ with $\mathbf{B}_{pix}$ in array $\mathbf{B}$ at the same indices $p$ and $q$.

In the notation of linear algebra, for each $n \times n$ matrix indicized by $p$ and $q$ I want:

$A_{pix}^{'} = B_{pix}A_{pix}$

Attempt:

I initially tried something of the form: $A_{mnloij}^{'} = B_{mnik}A_{lokj}$ but if I think this will compute the submatrix multiplication at all combinations of indices $m,l,n,o$.

The Question in Brief: Is there a way, using the Einstein summation notation, to compute the sum:

$A_{pqij}^{'} = B_{mnik}A_{lokj}$

for the combinations of indices where $(m,n)=(l,o)=\left(p,q\right)$ only?

Additional Attempt: In the worst case, I suppose that I have to first compute the expression above and then only slice the array in order to select the "diagonals" where $(m,n)=(l,o)$.

Answer 1

Numpy.einsum has a shortcut for this. You can specify the axes you want out.

Normally you line up, and sum repeated indices.

If you specify the output indices, it alines things the same way, but summs the indices that aren't in the specified output.

#Hadamard product
c = np.einsum('ij,ij->ij',a,b)

#select the diagonal of a matrix
c = np.einsum('ii->i',a)

#sum
c = np.einsum('ij->',a)

#squared length of each vector in a 2d array
c = np.einsum('ijk,ijk->ij',a,a)

I think your application can be said as:

#matrix product of each matrix in a 2d array
c = np.einsum('ijkl,ijlm->ijkm',a,b)

I find this is much simpler and more direct. I don't know if it matters, but I also think that this will be much more efficient than a direct implementation of the accepted answer.

Answer 2

Define the tensor

$${E_{ijk}} = \mathrm{1\ \ if\ i=j=k,\ 0\ otherwise}$$

$\mathbf{E}$ has some interesting properties: $$ \mathbf{E.a = Diag(a)} = \mathrm{matrix\ (\mathbf{a}\ along\ the\ diagonal)} $$ $$ \mathbf{E:A = diag(A)} = \mathrm{vector\ (diagonal\ of\ \mathbf{A})}$$ $$ \mathbf{a.E.b = E:ab} = \mathrm{Schur\ product\ (\mathbf{a \odot b})}$$

Your question is similar to the Schur product case

$${A'_{pqij} = E_{pml}\ E_{qno}\ B_{mnik}\ A_{lokj}}$$