4D Linear system using numpy-CodePudding

I have a linear system like A * x = b that I want so solve. The numpy solves it easily with numpy.linalg.solve.

But my A is a 4D-array of shape (m, n, m, n), b is a 2D-array of shape (m, n) and I want to find x of shape (m, n). Actually I transform A into Aexp of shape (m*n, m*n), and b into bexp of shape (m*n) to find xexp and transform back to x. My code is like

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]

But it's not nice. Is there a function like X = solve(A, B)? If not, how could I do without using loops?

How could I do with 3 index? A.shape = (m, n, p, m, n, p) and B.shape = (m, n, p)

CodePudding user response：

Does a numpy.reshape can be suitable?

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