CODEWITHSUNDEEP
matrixnumpyscipy
alex
alex

Reputation: 3118

fast way to invert or dot kxnxn matrix

Is there a fast way to calculate the inverse of a kxnxn matrix using numpy (the inverse being calculated at each k-slice)? In other words, is there a way to vectorize the following code:

>>>from numpy.linalg import inv
>>>a-random(4*2*2).reshape(4,2,2)
>>>b=a.copy()
>>>for k in range(len(a)):
>>>    b[k,:,:] = inv(a[k,:,:])

Upvotes: 5

Views: 1448

Answers (1)

egpbos
egpbos

Reputation: 502

First about getting the inverse. I have looked into both np.linalg.tensorinv and np.linalg.tensorsolve.

I think unfortunately tensorinv will not give you what you want. It needs the array to be "square". This excludes what you want to do, because their definition of square is that np.prod(a[:i]) == np.prod(a[i:]) where i is 0, 1 or 2 (one of the axes of the array in general); this can be given as the third argument ind of tensorinv. This means that if you have a general array of NxN matrices of length M, you need to have e.g. (for i = 1) NxN == NxM, which is not true in general (it actually is true in your example, but it does not give the correct answer anyway).

Now, maybe something is possible with tensorsolve. This would however involve some heavy construction work on the a matrix-array before it is passed as the first argument to tensorsolve. Because we would want b to be the solution of the "matrix-array equation" a*b = 1 (where 1 is an array of identity matrices) and 1 would have the same shape as a and b, we cannot simply supply the a you defined above as the first argument to tensorsolve. Rather, it needs to be an array with shape (M,N,N,M,N,N) or (M,N,N,N,M,N) or (M,N,N,N,N,M). This is necessary, because tensorsolve would multiply with b over these last three axes and also sum over them so that the result (the second argument to the function) is again of shape (M,N,N).

Then secondly, about dot products (your title suggests that's also part of your question). This is very doable. Two options.

First: this blog post by James Hensman gives some good suggestions.

Second: I personally like using np.einsum better for clarity. E.g.:

a=np.random.random((7,2,2))
b=np.random.random((7,2,2))
np.einsum('ijk,ikl->ijl', a,b)

This will matrix-multiply all 7 "matrices" in arrays a and b. It seems to be about 2 times slower than the array-method from the blog post above, but it's still about 70 times faster than using a for loop as in your example. In fact, with larger arrays (e.g. 10000 5x5 matrices) the einsum method seems to be slightly faster (not sure why).

Hope that helps.

Upvotes: 3

Related Questions