在numpy中解决大量小方程组

Question

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

for i in range(n):
    x[i,:] = np.linalg.solve(A[i,:,:],b[i,:])

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

Answer 1

对于那些现在回来阅读这个问题的人，我想我会节省其他人的时间并提到 numpy 现在使用广播来处理这个问题。

因此，在 numpy 1.8.0 及更高版本中，以下可用于求解 N 个线性方程。

x = np.linalg.solve(A,b)

Answer 2

我想回答自己有点失礼，但这是我目前拥有的 Fortran 解决方案，即其他解决方案在速度和简洁性方面都在有效竞争。

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

（之所以.T需要 s，是因为 f2py 中的设计选择不佳，如果.T省略 s，这会导致不必要的复制、低效的内存访问模式以及看起来不自然的 fortran 索引。）

这也避免了 setup.py 等。我对 fortran 没有什么可挑剔的（只要不涉及字符串），但我希望 numpy 可能有一些简短而优雅的东西可以做同样的事情。