使用numpy的4D线性系统

Question

我有一个像A * x = b这样的线性系统，我想要这样解。numpy用numpy.linalg.solve很容易地解决了这个问题。

但是我的A是形状(m, n, m, n)的4D数组，b是形状(m, n)的2D数组，我想要找到形状(m, n)的x。实际上，我将A转换为Aexp of shape (m*n, m*n)，将b转换为shape (m*n)的bexp，以找到xexp并将其转换回x。我的代码就像

Aexp = np.zeros((m*n, m*n))
Bexp = np.zeros(m*n)
for i in range(m):
    for j in range(n):
        Bexp[i*n + j] = B[i, j]
        for r in range(m):
            for s in range(n):
                Aexp[i*n + j, r*n + s] = A[i, j, r, s]
Xexp = np.linalg.solve(Aexp, Bexp)
X = np.zeros((n, m))
for i in range(m):
    for j in range(n):
        X[i, j] = Xexp[i*n+j]

但这不太好。有像X = solve(A, B)这样的函数吗？如果没有，我怎么才能不使用循环呢？

我如何处理3个索引？A.shape = (m, n, p, m, n, p)和B.shape = (m, n, p)

Answer 1

numpy.reshape可以合适吗？

import numpy as np
m = 2
n = 3
A = np.zeros((m,n, m,n))
B = np.zeros((m,n))
Aexp = A.reshape((m*n, m*n))
Bexp = B.reshape((m*n))