Cupitor
Cupitor

Reputation: 11637

Element-wise effecient multiplication of arrays of matrices

Suppose array_1 and array_2 are two arrays of matrices of the same sizes. Is there any vectorised way of multiplying element-wise, the elements of these two arrays(which their elements' multiplication is well defined)?

The dummy code:

def mat_multiply(array_1,array_2):
    size=np.shape(array_1)[0]

    result=np.array([])
    for i in range(size):
        result=np.append(result,np.dot(array_1[i],array_2[i]),axis=0)
    return np.reshape(result,(size,2))

example input:

a=[[[1,2],[3,4]],[[1,2],[3,4]]]
b=[[1,3],[4,5]]

output:

[[  7.  15.]
 [ 14.  32.]]

Upvotes: 1

Views: 1369

Answers (2)

hpaulj
hpaulj

Reputation: 231335

Contrary to your first sentence, a and b are not the same size. But let's focus on your example.

So you want this - 2 dot products, one for each row of a and b

np.array([np.dot(x,y) for x,y in zip(a,b)])

or to avoid appending

X = np.zeros((2,2))
for i in range(2):
    X[i,...] = np.dot(a[i],b[i])

the dot product can be expressed with einsum (matrix index notation) as

[np.einsum('ij,j->i',x,y) for x,y in zip(a,b)]

so the next step is to index that first dimension:

np.einsum('kij,kj->ki',a,b)

I'm quite familiar with einsum, but it still took a bit of trial and error to figure out what you want. Now that the problem is clear I can compute it in several other ways

A, B = np.array(a), np.array(b)    
np.multiply(A,B[:,np.newaxis,:]).sum(axis=2)
(A*B[:,None,:]).sum(2)
np.dot(A,B.T)[0,...]
np.tensordot(b,a,(-1,-1))[:,0,:]

I find it helpful to work with arrays that have different sizes. For example if A were (2,3,4) and B (2,4), it would be more obvious the dot sum has to be on the last dimension.


Another numpy iteration tool is np.nditer. einsum uses this (in C). http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

it = np.nditer([A, B, None],flags=['external_loop'],
    op_axes=[[0,1,2], [0,-1,1], None])
for x,y,w in it:
    # x, y are shape (2,)
    w[...] = np.dot(x,y)
it.operands[2][...,0]

Avoiding that [...,0] step, requires a more elaborate setup.

C = np.zeros((2,2))
it = np.nditer([A, B, C],flags=['external_loop','reduce_ok'],
    op_axes=[[0,1,2], [0,-1,1], [0,1,-1]],
    op_flags=[['readonly'],['readonly'],['readwrite']])
for x,y,w in it:
    w[...] = np.dot(x,y)
    # w[...] += x*y 
print C
# array([[  7.,  15.],[ 14.,  32.]])

Upvotes: 4

Jaime
Jaime

Reputation: 67417

There's one more option that @hpaulj left out in his extensive and comprehensive list of options:

>>> a = np.array(a)
>>> b = np.array(b)
>>> from numpy.core.umath_tests import matrix_multiply
>>> matrix_multiply.signature
'(m,n),(n,p)->(m,p)'
>>> matrix_multiply(a, b[..., np.newaxis])
array([[[ 7],
        [15]],

       [[14],
        [32]]])
>>> matrix_multiply(a, b[..., np.newaxis]).shape
(2L, 2L, 1L)
>>> np.squeeze(matrix_multiply(a, b[..., np.newaxis]), axis=-1)
array([[ 7, 15],
       [14, 32]])

The nice thing about matrix_multiply is that, it being a gufunc, it will work not only with 1D arrays of matrices, but also with broadcastable arrays. As an example, if instead of multiplying the first matrix with the first vector, and the second matrix with the second vector, you wanted to compute all possible multiplications, you could simply do:

>>> a = np.arange(8).reshape(2, 2, 2) # to have different matrices
>>> np.squeeze(matrix_multiply(a[...,np.newaxis, :, :],
...                            b[..., np.newaxis]), axis=-1)
array([[[ 3, 11],
        [ 5, 23]],

       [[19, 27],
        [41, 59]]])

Upvotes: 2

Related Questions