Solve large number of small equation systems in numpy

linear-algebra, numpy, python

Solution

I guess answering yourself is a bit of a faux pas, but this is the fortran solution I have a the moment, i.e. what the other solutions are effectively competing against, both in speed and brevity.

function pixsolve(A, b) result(x)
    implicit none
    real*8    :: A(:,:,:), b(:,:), x(size(b,1),size(b,2))
    integer*4 :: i, n, m, piv(size(b,1)), err
    n = size(A,3); m = size(A,1)
    x = b
    do i = 1, n
        call dgesv(m, 1, A(:,:,i), m, piv, x(:,i), m, err)
    end do
end function

This would be compiled as:

f2py -c -m foo{,.f90} -llapack -lblas

And called from python as

x = foo.pixsolve(A.T, b.T).T

(The `.T`s are needed due to a poor design choice in f2py, which both causes unnecessary copying, inefficient memory access patterns and unnatural looking fortran indexing if the `.T`s are left out.)

This also avoids a setup.py etc. I have no bone to pick with fortran (as long as strings aren't involved), but I was hoping that numpy might have something short and elegant which could do the same thing.

Problem

I have a large number of small linear equation systems that I'd like to solve efficiently using numpy. Basically, given `A[:,:,:]` and `b[:,:]`, I wish to find `x[:,:]` given by `A[i,:,:].dot(x[i,:]) = b[i,:]`. So if I didn't care about speed, I could solve this as ``` for i in range(n): x[i,:] = np.linalg.solve(A[i,:,:],b[i,:]) ``` But since this involved explicit looping in python, and since `A` typically has a shape like `(1000000,3,3)`, such a solution would be quite slow. If numpy isn't up to this, I could do this loop in fortran (i.e. using f2py), but I'd prefer to stay in python if possible.

Original source