numpy.linalg.lstsq 的广播问题

Broadcasting issues with numpy.linalg.lstsq

我正在研究一些图像分析算法,并尝试使用 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 维数组,我不确定如何告诉它它需要在前两个维度上广播。

重塑 x 以具有形状 (2, K),其中成对的像素值在列中。调用lstsq,然后恢复结果的形状。

例如,这里是Ax

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])