multidimensional-array - 如何使用矢量化代码求解许多超定线性方程组？

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 个不同的线性方程组超定系统，可能使用一些矢量化技术或其他特殊方法？

谢谢

Answer 1

您可以通过使用 stack of matrices 获得一些速度。 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 倍左右，
对于那个问题大小。是否值得麻烦取决于
问题。