CODEWITHSUNDEEP
pythonnumpymatrixscipyvectorization
H. Arponen
H. Arponen

Reputation: 557

Vectorized (partial) inverse of an N*M*M tensor with numpy

I'm almost exactly in a similar situation as the asker here over a year ago: fast way to invert or dot kxnxn matrix

So I have a tensor with indices a[n,i,j] of dimensions (N,M,M) and I want to invert the M*M square matrix part for each n in N.

For example, suppose I have

In [1]:    a = np.arange(12)
           a.shape = (3,2,2)
           a

Out[1]: array([[[ 0,  1],
                  [ 2,  3]],

                  [[ 4,  5],
                  [ 6,  7]],

                  [[ 8,  9],
                  [10, 11]]])

Then a for loop inversion would go like this:

In [2]: inv_a = np.zeros([3,2,2])
        for m in xrange(0,3):
            inv_a[m] = np.linalg.inv(a[m])
        inv_a

Out[2]: array([[[-1.5,  0.5],
                  [ 1. ,  0. ]],

                  [[-3.5,  2.5],
                  [ 3. , -2. ]],

                  [[-5.5,  4.5],
                 [ 5. , -4. ]]])

This will apparently be implemented in NumPy 2.0, according to this issue on github...

I guess I need to install the dev version as seberg noted in the github issue thread, but is there another way to do this in vectorized manner right now?

Upvotes: 4

Views: 2181

Answers (1)

IanH
IanH

Reputation: 10700

Update: In NumPy 1.8 and later, the functions in numpy.linalg are generalized universal functions. Meaning that you can now do something like this:

import numpy as np
a = np.random.rand(12, 3, 3)
np.linalg.inv(a)

This will invert each 3x3 array and return the result as a 12x3x3 array. See the numpy 1.8 release notes.

Original Answer:

Since N is relatively small, how about we compute the LU decomposition manually for all the matrices at once. This ensures that the for loops involved are relatively short.

Here's how this can be done with normal NumPy syntax:

import numpy as np
from numpy.random import rand

def pylu3d(A):
    N = A.shape[1]
    for j in xrange(N-1):
        for i in xrange(j+1,N):
            #change to L
            A[:,i,j] /= A[:,j,j]
            #change to U
            A[:,i,j+1:] -= A[:,i,j:j+1] * A[:,j,j+1:]

def pylusolve(A, B):
    N = A.shape[1]
    for j in xrange(N-1):
        for i in xrange(j+1,N):
            B[:,i] -= A[:,i,j] * B[:,j]
    for j in xrange(N-1,-1,-1):
        B[:,j] /= A[:,j,j]
        for i in xrange(j):
            B[:,i] -= A[:,i,j] * B[:,j]

#usage
A = rand(1000000,3,3)
b = rand(3)
b = np.tile(b,(1000000,1))
pylu3d(A)
# A has been replaced with the LU decompositions
pylusolve(A, b)
# b has been replaced to the solutions of
# A[i] x = b[i] for each A[i] and b[i]

As I have written it, pylu3d modifies A in place to compute the LU decomposition. After replacing each NxN matrix with its LU decomposition, pylusolve can be used to solve an MxN array b representing the right hand sides of your matrix systems. It modifies b in place and does the proper back substitutions to solve the system. As it is written, this implementation does not include pivoting, so it isn't numerically stable, but it should work well enough in most cases.

Depending on how your array is arranged in memory, it is probably still a good bit faster to use Cython. Here are two Cython functions that do the same thing, but they iterate along M first. It's not vectorized, but it is relatively fast.

from numpy cimport ndarray as ar
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def lu3d(ar[double,ndim=3] A):
    cdef int n, i, j, k, N=A.shape[0], h=A.shape[1], w=A.shape[2]
    for n in xrange(N):
        for j in xrange(h-1):
            for i in xrange(j+1,h):
                #change to L
                A[n,i,j] /= A[n,j,j]
                #change to U
                for k in xrange(j+1,w):
                    A[n,i,k] -= A[n,i,j] * A[n,j,k]

@cython.boundscheck(False)
@cython.wraparound(False)
def lusolve(ar[double,ndim=3] A, ar[double,ndim=2] b):
    cdef int n, i, j, N=A.shape[0], h=A.shape[1]
    for n in xrange(N):
        for j in xrange(h-1):
            for i in xrange(j+1,h):
                b[n,i] -= A[n,i,j] * b[n,j]
        for j in xrange(h-1,-1,-1):
            b[n,j] /= A[n,j,j]
            for i in xrange(j):
                b[n,i] -= A[n,i,j] * b[n,j]

You could also try using Numba, though I couldn't get it to run as fast as Cython in this case.

Upvotes: 4

Related Questions