I have this piece of code:

other = np.random.rand((m,n,o))
prev = np.random.rand((m,n,o,m,n,o))
mu = np.zeros((m,n,o,m,n,o))
for c in range(m):
   for i in range(n):
      for j in range(o):
         mu[c,i,j,c,i,j] = other[c,i,j]*prev[c,i,j,c,i,j]

And I'd like to simplify it using einsum notation (possibly saving time by skipping the for loops in python). However after a few tries I'm eventually not sure how to approach the problem. My current try is:

np.einsum('cijklm,cij->cijklm', prev, other)

It does not achieves the same result as the "for-loop" piece of code.

CodePudding user response：

With shapes (2,3,4), I get:

In [52]: mu.shape
Out[52]: (2, 3, 4, 2, 3, 4)

This einsum expression complains that dimensions are repeated in the output:

In [53]: np.einsum('cijcij,cij->cijcij', prev, other).shape
Traceback (most recent call last):
  File "<ipython-input-53-92862a0865a2>", line 1, in <module>
    np.einsum('cijcij,cij->cijcij', prev, other).shape
  File "<__array_function__ internals>", line 180, in einsum
  File "/usr/local/lib/python3.8/dist-packages/numpy/core/einsumfunc.py", line 1359, in einsum
    return c_einsum(*operands, **kwargs)
ValueError: einstein sum subscripts string includes output subscript 'c' multiple times

Without the repeat:

In [55]: x=np.einsum('cijcij,cij->cij', prev, other)
In [56]: x.shape
Out[56]: (2, 3, 4)

Nonzero values match:

In [57]: np.allclose(mu[np.nonzero(mu)].ravel(), x.ravel())
Out[57]: True

Or by extracting the diagonals from mu:

In [59]: I,J,K = np.ix_(np.arange(2),np.arange(3),np.arange(4))
In [60]: mu[I,J,K,I,J,K].shape
Out[60]: (2, 3, 4)
In [61]: np.allclose(mu[I,J,K,I,J,K],x)
Out[61]: True

Your einsum satisfies the same 'diagonals' test:

In [68]: y=np.einsum('cijklm,cij->cijklm', prev, other)
In [69]: y.shape
Out[69]: (2, 3, 4, 2, 3, 4)
In [70]: np.allclose(y[I,J,K,I,J,K],x)
Out[70]: True

So the mu values are also present in y, but distributed in a different way. But the arrays are too big to readily view and compare.

OK, each y[i,j,k] is the same, and equal to x. In mu most of these values are 0, with only selected diagonals being nonzero.

While einsum can generate the same nonzero values, it cannot distribute them in the same 3d diagonals way as your loop.

Changing your mu calculation to produce a 3d array:

In [76]: nu = np.zeros((m,n,o))
    ...: for c in range(m):
    ...:    for i in range(n):
    ...:       for j in range(o):
    ...:          nu[c,i,j] = other[c,i,j]*prev[c,i,j,c,i,j]
    ...: 
In [77]: np.allclose(nu,x)
Out[77]: True

CodePudding user response：

It is not possible to get this result using np.einsum() only, but you can try this:

import numpy as np
from numpy.lib.stride_tricks import as_strided

m, n, o = 2, 3, 5
np.random.seed(0)
other = np.random.rand(*(m, n, o))
prev = np.random.rand(*(m, n, o, m, n, o))
mu = np.zeros((m, n, o, m, n, o))

mu_view = as_strided(mu,
                     shape=(m, n, o),
                     strides=[sum(mu.strides[i::3]) for i in range(3)]
                     )
np.einsum('cijcij,cij->cij', prev, other, out=mu_view)

The array mu should be then the same as the one produced by the code using nested loops in the question.