我可以将求解方程组（使用 np.linalg.solve）推广到三维数组（仍然有两个未知数）吗？

Question

我想知道下面的问题是否可以推广到一个三维数组，可以通过切片等方式完成：

假设我有一个大小为 (n, 2) 的数组 x。我使用数组中的值建立了一个线性方程组并求解如下：

import numpy as np
def solve(x):
        a_array = np.array([[16, -1], [-2, 1]])
        b = np.array([x[-3, 0] - 16*x[-2, 0] + 30*x[-1, 0], 1*x[-3, 0] - 2*x[-2, 0] - 0*x[-1, 0]])
        u = np.linalg.solve(a_array, b)

        c = np.array([x[-3, 1] - 16*x[-2, 1] + 30*x[-1, 1], 1*x[-3, 1] - 2*x[-2, 1] - 0*x[-1, 1]])
        v = np.linalg.solve(a_array, c)

        total_uv = np.column_stack((u, v))
        
        final_x = np.concatenate((x, total_uv))
    return final_x

本质上，我根据 x 数组中的值求解一个线性方程组（有两个变量）以获得 (1,2) 个 u 和 v 数组。然后我将它们堆叠在一起以获得 (2 ,2) 并将该数组连接到我原来的 x。与原始 x 相比，我现在得到的 final_x 数组大小为 (n+2, 2)。

我的问题是：假设我开始的 x 现在的大小为 (m, n, 2)。对于每个 (n,2) 数组（其中有 m 个），我想求解这个方程组以获得我的 total_uv 数组。然后我将每个 total_uv 数组连接到每个 (n,2) 数组（现在沿轴 = 1），所以我的最终结果是大小 (m, n+2,2)。

有没有办法在不使用 for 循环的情况下做到这一点（即遍历所有 (n, 2) 数组并进行计算、存储、然后连接）？例如仅使用切片表示法，还是更优雅的方式？

似乎通过x[:, -3, 0] + ... 对 x 数组进行切片（以获取变量 b 和 c）实际上不起作用，因为 a_array 的大小不同。

澄清 (m, n, 2) 的解：

def solve(x):
        m = x.shape[0]
        a_array = np.array([[16, -1], [-2, 1]])
        b = np.stack((x[:, -3, 0] - 16*x[:, -2, 0] + 30*x[:, -1, 0],
              1*x[:, -3, 0] - 2*x[:, -2, 0] - 0*x[:, -1, 0]), axis=1)
        u = np.linalg.solve(a_array[None, :, :], b)

        c = np.stack((x[:, -3, 0] - 16*x[:, -2, 0] + 30*x[:, -1, 0],
              1*x[:, -3, 0] - 2*x[:, -2, 0] - 0*x[:, -1, 0]), axis=1)
        v = np.linalg.solve(a_array[None, :, :], c)

        # Concatenation here:
        total_uv = np.column_stack((u, v))

        uv_reshaped = total_uv.reshape((m, 2, 2))

        final_x = np.concatenate((x, uv_reshaped), axis=1)
    return final_x

Answer 1

您可以使用切片表示法将b 和c 向量创建为二维数组：

b = np.stack((x[:, -3, 0] - 16*x[:, -2, 0] + 30*x[:, -1, 0],
              1*x[:, -3, 0] - 2*x[:, -2, 0] - 0*x[:, -1, 0]), axis=1)

然后您可以向a_array 添加一个额外的（批量）维度，以便与np.linalg.solve 一起使用：

u = np.linalg.solve(a_array[None, ...], b)

顺便说一句，由于b 和c 的系数似乎相似，因此您可以在此处执行相同操作，从而无需两次调用np.linalg.solve：

bc = np.stack((x[:, -3] - 16*x[:, -2] + 30*x[:, -1],
               1*x[:, -3] - 2*x[:, -2] - 0*x[:, -1]), axis=2)
uv = np.linalg.solve(a_array[None, None, ...], bc)

那么你有两个批次维度(n, 2)，其中第二个代表 u/v。