我可以将求解方程组(用np. linalgl .solve)推广到三维数组(仍然有两个未知数)吗? &

Question

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

假设我有一个数组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*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数组中的值求解一个线性方程组(有两个变量)，得到u和v的(1,2)数组，然后将它们堆叠在一起，得到一个总数组(2,2)，并将该数组连接到我的原始x。我现在得到一个final_x数组大小为(n+ 2,2)，与原始x相比。

我的问题是:假设我开始的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