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

问题描述

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

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

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] - 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,2)的解：

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

        c = np.stack((x[:,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

解决方法

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

b = np.stack((x[:,-3,0] - 16*x[:,-2,0] + 30*x[:,-1,0],1*x[:,0] - 2*x[:,0] - 0*x[:,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],-3] - 2*x[:,-2] - 0*x[:,-1]),axis=2)
uv = np.linalg.solve(a_array[None,None,bc)

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