Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Combining element-wise and matrix multiplication with multi-dimensional arrays in NumPy

I have two multidimensional NumPy arrays, A and B, with A.shape = (K, d, N) and B.shape = (K, N, d). I would like to perform an element-wise operation over axis 0 (K), with that operation being matrix multiplication over axes 1 and 2 (d, N and N, d). So the result should be a multidimensional array C with C.shape = (K, d, d), so that C[k] = np.dot(A[k], B[k]). A naive implementation would look like this:

C = np.vstack([np.dot(A[k], B[k])[np.newaxis, :, :] for k in xrange(K)])

but this implementation is slow. A slightly faster approach looks like this:

C = np.dot(A, B)[:, :, 0, :]

which uses the default behaviour of np.dot on multidimensional arrays, giving me an array with shape (K, d, K, d). However, this approach computes the required answer K times (each of the entries along axis 2 are the same). Asymptotically it will be slower than the first approach, but the overhead is much less. I am also aware of the following approach:

from numpy.core.umath_tests import matrix_multiply
C = matrix_multiply(A, B)

but I am not guaranteed that this function will be available. My question is thus, does NumPy provide a standard way of doing this efficiently? An answer which applies to multidimensional arrays in general would be perfect, but an answer specific to only this case would be great too.

Edit: As pointed out by @Juh_, the second approach is incorrect. The correct version is:

C = np.dot(A, B).diagonal(axis1=0, axis2=2).transpose(2, 0, 1)

but the overhead added makes it slower than the first approach, even for small matrices. The last approach is winning by a long shot on all my timing tests, for small and large matrices. I'm now strongly considering using this if no better solution crops up, even if that would mean copying the numpy.core.umath_tests library (written in C) into my project.

like image 484
gerardlouw Avatar asked Apr 25 '12 23:04

gerardlouw


People also ask

How element by element multiplication is is performed on NumPy arrays?

The NumPy library's np. multiply(x1, x2) method receives two matrices as input and executes element-wise multiplication over them before returning the resultant matrix. We must send the two matrices as input to the np. multiply() method to execute element-wise input.

How do you perform matrix multiplication on the NumPy arrays?

To multiply two matrices use the dot() function of NumPy. It takes only 2 arguments and returns the product of two matrices.

Does NumPy do element wise multiplication?

multiply() in Python. numpy. multiply() function is used when we want to compute the multiplication of two array. It returns the product of arr1 and arr2, element-wise.


2 Answers

A possible solution to your problem is:

C = np.sum(A[:,:,:,np.newaxis]*B[:,np.newaxis,:,:],axis=2)

However:

  1. it is quicker than the vstack approach only if K is much bigger than d and N
  2. their might be some memory issue: in the above solution an KxdxNxd array is allocated (i.e. all possible product paires, before summing). Actually I could not test with big K,d and N as I was going out of memory.

btw, note that:

C = np.dot(A, B)[:, :, 0, :]

does not give the correct result. It got me tricked because I first checked my method by comparing the results to those given by this np.dot command.

like image 144
Juh_ Avatar answered Sep 30 '22 17:09

Juh_


I have this same issue in my project. The best I've been able to come up with is, I think it's a little faster (maybe 10%) than using vstack:

K, d, N = A.shape
C = np.empty((K, d, d))
for k in xrange(K):
    C[k] = np.dot(A[k], B[k])

I'd love to see a better solution, I can't quite see how one would use tensordot to do this.

like image 20
Bi Rico Avatar answered Sep 30 '22 17:09

Bi Rico