CODEWITHSUNDEEP
pythonarraysnumpy
tkowt
tkowt

Reputation: 347

How should I multiply 1d-array in 3d-array by 2d-matrix with numpy

I want bellow calculation.Values used in this example are variable in real case.

 A = [[1,1,1],[2,1,1]]
 B =[[[1,2,3],[2,2,2]],[[1,1,2],[1,1,1]]]
 result = apply_some_function(A,B)
 >> result = [[[6,7],[6,8]],
             [[4,5],[3,4]]]

Is their any efficient way satisfy above calculation with numpy. In fact, I don't understand clearly how to applicate numpy for multi dimensional array. So if you know helpful document to understand manipulation multi dimensional array rules, I'm very glad if you let me know.

Upvotes: 0

Views: 781

Answers (1)

unutbu
unutbu

Reputation: 879531

You could use np.tensordot or np.einsum. 

In [19]: np.tensordot(B, A, (-1, -1))
Out[19]: 
array([[[6, 7],
        [6, 8]],

       [[4, 5],
        [3, 4]]])

In [20]: np.einsum('ij,klj->kli', A, B)
Out[20]: 
array([[[6, 7],
        [6, 8]],

       [[4, 5],
        [3, 4]]])

Both of these functions compute sums of products. Thus, they are generalizations of matrix multiplication.

It is often helpful to pay attention to the shape of the arrays. If we make A, B, and result NumPy arrays:

A = np.array([[1,1,1],[2,1,1]])
B = np.array([[[1,2,3],[2,2,2]],[[1,1,2],[1,1,1]]])
result = np.array([[[6,7],[6,8]], [[4,5],[3,4]]])

then 

In [6]: A.shape
Out[6]: (2, 3)

In [7]: B.shape
Out[7]: (2, 2, 3)

In [9]: result.shape
Out[9]: (2, 2, 2)

Notice that the axis of length 3 in A and B disappears in result. That suggests the sum (of products) is being taken over the last axis of A and B.

The (-1, -1) in np.tensordot(B, A, (-1, -1)) tells np.tensordot to sum over the last axes of A and B.

Similarly, the 'ij,klj->kli' in np.einsum('ij,klj->kli', A, B) says that if A has indices i and j and if B has indices k, l and j, then the result should have indices k,l,i. Notice the j index disappears. The j index is the last index in A and B. Thus, 'ij,klj->kli' tells np.einsum to sum over the last index of A and B.

The only thing left to do is figure out the correct order of the k, l and i indices. Since each axis in result has the same length, the shape of the result gives no clue. I found the correct order by trial and error.

Upvotes: 3

Related Questions