解决numpy中的大量小方程组 [英] Solve large number of small equation systems in numpy

问题描述

我有大量的小型线性方程组，我想使用numpy有效地求解.基本上，给定A[:,:,:]和b[:,:]，我希望找到A[i,:,:].dot(x[i,:]) = b[i,:]给定的x[:,:].因此，如果我不在乎速度，可以解决此问题

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,:])

但是由于这涉及到python中的显式循环，并且由于A通常具有类似(1000000,3,3)的形状，因此这种解决方案将非常慢.如果numpy无法做到这一点，我可以在fortran中执行此循环(即使用f2py)，但如果可能的话，我宁愿留在python中.

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.

推荐答案

我想回答自己有点虚假，但这是我目前正在讨论的fortran解决方案，即其他解决方案正在与之竞争的是什么，既快捷又简洁.

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

这将被编译为:

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

从python调用为

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

(由于在f2py中糟糕的设计选择而需要.T，这会导致不必要的复制，低效的内存访问模式以及如果忽略了.T会导致fortran索引看起来不自然.)

(The .Ts 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 .Ts are left out.)

这也避免了setup.py等.fortran我没有骨头可以选择(只要不涉及字符串)，但是我希望numpy可能简短而优雅，可以做同样的事情.

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.

这篇关于解决numpy中的大量小方程组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！