6

我有大量的小型线性方程组,我想使用 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 中。

4

4 回答 4

4

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

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

x = np.linalg.solve(A,b)
于 2015-12-01T18:42:44.203 回答
3

我想回答自己有点失礼,但这是我目前拥有的 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由于 f2py 中的设计选择不佳,因此需要 s,如果 s 被忽略,这都会导致不必要的复制、低效的内存访问模式以及看起来不自然的 fortran 索引.T。)

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

于 2012-11-17T17:50:19.680 回答
3

我认为您认为显式循环是一个问题是错误的。通常它只是值得优化的最里面的循环,我认为这在这里成立。例如,我们可以衡量代码的开销与实际计算的成本:

import numpy as np

n = 10**6
A = np.random.random(size=(n, 3, 3))
b = np.random.random(size=(n, 3))
x = b*0

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

np.linalg.pseudosolve = lambda a,b: b

def g():
    for i in xrange(n):
        x[i,:] = np.linalg.pseudosolve(A[i,:,:],b[i,:])

这给了我

In [66]: time f()
CPU times: user 54.83 s, sys: 0.12 s, total: 54.94 s
Wall time: 55.62 s

In [67]: time g()
CPU times: user 5.37 s, sys: 0.01 s, total: 5.38 s
Wall time: 5.40 s

IOW,除了实际解决您的问题之外,它只花费 10% 的时间做任何事情。现在,我完全可以相信,np.linalg.solve相对于你可以从 Fortran 中获得的东西,它本身对你来说太慢了,所以你想做别的事情。在小问题上可能尤其如此,想想看:IIRC 我曾经发现手动展开某些小解决方案更快,尽管那是前一段时间。

但就其本身而言,在第一个索引上使用显式循环会使整体解决方案变得非常缓慢是不正确的。如果np.linalg.solve足够快,循环不会在这里改变太多。

于 2012-11-18T15:21:01.187 回答
0

我认为您可以一次完成,使用由对角线周围的 3x3 块组成的 (3x100000,3x100000) 矩阵。

未测试:

b_new = np.vstack([ b[i,:] for i in range(len(i)) ])
x_new = np.zeros(shape=(3x10000,3) )

A_new = np.zeros(shape=(3x10000,3x10000) )
n,m = A.shape
for i in range(n):
   A_new[3*i:3*(i+1),3*i:3*(i+1)] = A[i,:,:]

x = np.linalg.solve(A_new,b_new)
于 2012-11-17T16:02:55.913 回答