4d Array Processing (используя einsum?)

Question

У меня есть матричная проблема, которую, я думаю, можно решить (вычислительно дешево) в одной строке кода, используя numpy (возможно, einsum?), Но я не могу найти решение.

Интересно, кто-нибудь может сделать какие-либо предложения, пожалуйста? Проблема заключается в следующем:

A = 4d array of shape: (5, 5, 5, 5)
M = 2d array of shape: (100, 5)

Target result, R = 2d array of shape (100, 5) where...
R[z, p] = sum for p having values of [i, j, k, l] over: A[i,j,k,l] * M[z,i] * M[z, j] * sgn[p]

where sgn[p] = -ve if p = i or j; +ve if p = k or l

В качестве примера:

# Enter two values in A array
A = np.zeros([5, 5, 5, 5])
A[2, 1, 3, 0] = 1
A[2, 2, 4, 1] = 1

M = np.zeros([100, 5])
M[None, :] = np.arange(5)  # fill with dummy data

R = SOLUTION_FUNC(A, M), computed as:

# For R[:, 0] there's only one term to add
# since 0 only appears once as an index for non-zero values in the A array.
# As 0 appears in second half of the index set, we keep the value positive
R[z, 0] = A[2, 1, 3, 0] * M[z, 2] * M[z, 1]

# For R[:, 1], we add both terms in A since 1 appears as an index in both non-zero values
# As 1 is in first half of the list the first time, we make this value negative
R[z, 1] = -1 * A[2, 1, 3, 0] * M[z, 2] * M[z, 1] + A[2, 2, 4, 1] * M[z, 2] * M[z,2]

# For R[z:, 2], there are three terms to add since 2 appears as a 
# non-zero index three times in the A matrix.
# The 2 always appears in first half of list, so we make all terms negative
R[z, 2] = -1* A[2, 1, 3, 0] * M[z, 2] * M[z, 1] + -1*A[2, 2, 4, 1] * M[z, 2] * M[z,2] + -1 * A[2, 2, 4, 1] * M[z, 2] * M[z,2]

# etc....
R[z, 3] = A[2, 1, 3, 0] * M[z, 2] * M[z, 1]
R[z, 4] = A[2, 2, 4, 1] * M[z, 2] * M[z, 2]

Если это помогает, A является относительно редким (вероятно, содержит < 20 ненулевых членов - но они могут быть в любой случайной позиции).

РЕДАКТИРОВАТЬ: петля добавлена

Я мог бы добиться этого, используя вложенные циклы, но это кажется плохим решением:

length = 5
A = np.zeros([length, length, length, length])
A[2, 1, 3, 0] = 1
A[2, 2, 4, 1] = 2

M = np.zeros([100, length])
M[None, :] = np.arange(length) + 1 # fill with dummy data

result = np.zeros_like(M, dtype=float)

# Many loop solution
non_zero_indexes = np.array(np.nonzero(A)).T

for z in range(len(M)):
    for non_zero_index_set in non_zero_indexes:
        for i in range(length):
            # The set indices are inspected in two halves: if i in first half, return negative value
            if i in non_zero_index_set[:2]:
                # If a number appears twice in the list, we need to add twice, so include a factor
                factor = -1 * list(non_zero_index_set).count(i)  # negative factor
                result[z, i] += factor * A[non_zero_index_set[0], non_zero_index_set[1], non_zero_index_set[2], non_zero_index_set[3]]   \
                                 * M[z, non_zero_index_set[0]] * M[z, non_zero_index_set[1]]
            elif i in non_zero_index_set[-2:]:
                factor = +1 * list(non_zero_index_set).count(i)  # positive factor
                result[z, i] += factor * A[non_zero_index_set[0], non_zero_index_set[1], non_zero_index_set[2], non_zero_index_set[3]]   \
                                 * M[z, non_zero_index_set[0]] * M[z, non_zero_index_set[1]]

Я пробовал многочисленные перестановки einsum, который, я думаю, мог бы решить это. например, формы np.einsum('ijkl,mi->mi') но как-то включая суммирование для всех ийкл, как объяснено выше.

Надеемся, что это имеет смысл, и любые указатели (или дальнейшие советы по использованию einsum / других маршрутов обработки массивов) были бы отличными ура.