如何使用向量化代码解决多个超定线性方程组？

Question

我需要解决一个线性方程组Lx=b，其中x始终是一个向量（3x1数组），L是一个Nx3的数组，而b是一个Nx1的向量。N通常在4到10之间变化。我没有使用scipy.linalg.lstsq(L,b)解决这个问题的任何问题。

然而，我需要多次执行此操作（大约200x200=40000次），因为x实际上与图像中的每个像素相关联。因此，x实际上存储在PxQx3数组中，其中P和Q大约为200-300，最后一个数字“3”指的是向量x。目前，我只是通过循环遍历每一列和行，并逐一解决方程。如何高效地解决这40000个不同的过度确定的线性方程组，可能需要使用一些矢量化技术或其他特殊方法？

谢谢

- Physicist

---

你的不同的“x”有不同的“L”吗？如果是这样，那么你可能没有太多的性能改进空间。 - fang

对于不同的x，b肯定是不同的。通常情况下，L也因x而异，但是似乎使用不同的L或相同的L得到的最终结果非常相似，因此我们先假设L对所有的x都是相同的。 - Physicist

Answer 1

你可以利用numpy.linalg例程的矩阵堆叠功能来提高一些速度。这在numpy.linalg.lstsq中尚不起作用，但numpy.linalg.svd可以，因此您可以自己实现lstsq：

import numpy as np


def stacked_lstsq(L, b, rcond=1e-10):
    """
    Solve L x = b, via SVD least squares cutting of small singular values
    L is an array of shape (..., M, N) and b of shape (..., M).
    Returns x of shape (..., N)
    """
    u, s, v = np.linalg.svd(L, full_matrices=False)
    s_max = s.max(axis=-1, keepdims=True)
    s_min = rcond*s_max
    inv_s = np.zeros_like(s)
    inv_s[s >= s_min] = 1/s[s>=s_min]
    x = np.einsum('...ji,...j->...i', v,
                  inv_s * np.einsum('...ji,...j->...i', u, b.conj()))
    return np.conj(x, x)


def slow_lstsq(L, b):
    return np.array([np.linalg.lstsq(L[k], b[k])[0]
                     for k in range(L.shape[0])])    


def test_it():
    b = np.random.rand(1234, 3)
    L = np.random.rand(1234, 3, 6)

    x = stacked_lstsq(L, b)
    x2 = slow_lstsq(L, b)

    # Check
    print(x.shape, x2.shape)
    diff = abs(x - x2).max()
    print("difference: ", diff)
    assert diff < 1e-13


test_it()

一些时间表明，对于该问题的规模，堆叠版本大约快了6倍。是否值得麻烦取决于问题本身。

- pv.