numpy.linalg.lstsq 的广播问题

Question

我正在研究一些图像分析算法，并尝试使用 numpy 进行最小二乘拟合。为了说明我想要做什么，我生成了一个非常简单的测试用例：

A = np.zeros((2, 2))
A[0, 0] = 1
A[0, 1] = 3
A[1, 0] = 4
A[0, 1] = 5

这是我在 Ax = b 类型方程中的简单 A 矩阵。

现在，在这个测试案例中，我的图像基本上是一个简单的 2x2 图像，并且在每个点我都有 2 个测量值。所以，就我而言，我模拟如下：

x = np.array([[13, 24], [13, 24], [13, 24], [13, 24]])
x = x.reshape((2,2,2))

现在，这代表我的 3 维图像，其中我有 2x2 网格和与每个像素关联的两个值。我可以逐个像素地解决这个问题；

np.linalg.lstsq(A, x[0, 0, :]) # fit at pixel (0, 0)

或

np.linalg.lstsq(A, x[0, 1, :]) # fit at pixel (0, 1)

但是，一旦我尝试类似：

np.linalg.lstsq(A, x) # fit at all pixels together

它抱怨 x 是 3 维数组，我不知道如何告诉它它需要在前两个维度上广播。

Answer 1

将x 重塑为形状 (2, K)，其中包含列中的像素值对。调用lstsq，然后恢复结果的形状。

例如这里是A和x：

In [86]: A

Out[86]: 
array([[ 1.,  5.],
       [ 4.,  0.]])

In [87]: x
Out[87]: 
array([[[0, 1],
        [2, 3]],

       [[4, 5],
        [6, 7]]])

将x 重塑为y，列中包含像素值：

In [88]: y = x.reshape(-1,2).T

In [89]: y
Out[89]: 
array([[0, 2, 4, 6],
       [1, 3, 5, 7]])

y 是lstsq 需要的形状：

In [90]: result = np.linalg.lstsq(A, y)

In [91]: result[0]
Out[91]: 
array([[ 0.25,  0.75,  1.25,  1.75],
       [-0.05,  0.25,  0.55,  0.85]])

获取解决方案，并恢复形状以匹配x：

In [92]: sol = result[0].T.reshape(x.shape)

In [93]: sol
Out[93]: 
array([[[ 0.25, -0.05],
        [ 0.75,  0.25]],

       [[ 1.25,  0.55],
        [ 1.75,  0.85]]])

通过单独解决几个像素来检查解决方案是否符合我们的预期：

In [94]: sol[0,0]
Out[94]: array([ 0.25, -0.05])

In [95]: np.linalg.lstsq(A, x[0,0])[0]
Out[95]: array([ 0.25, -0.05])

In [96]: sol[1,0]
Out[96]: array([ 1.25,  0.55])

In [97]: np.linalg.lstsq(A, x[1,0])[0]
Out[97]: array([ 1.25,  0.55])