4d Array processing (using einsum?)

I have a matrix problem that I think can be solved (computationally cheaply) in one line of code using numpy (maybe einsum?), But can't get to a solution.

I wonder if anyone can make any suggestions? The problem is this:

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

      

As an example:

# 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 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]

      

If that helps, A is relatively sparse (probably has <20 non-zero terms, but they can be in any random position).

EDIT: Added a loop

I could achieve this with nested loops, but this seems like a bad solution:

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]]

      

I have tried many permutations einsum

which I think can solve this. for example forms np.einsum('ijkl,mi->mi')

, but somehow including summation for all ijkl as explained above.

Hopefully this makes sense and any pointers (or further advice using einsum / other array processing routes) will be great greetings.

+3


source to share





All Articles