使用numpy的矩阵左堆叠数组的除法

Question

我正在开发一个程序来解决python中的Bloch（或更精确地说是Bloch McConnell）方程。所以要解决的方程是：

其中 A 是 NxN 矩阵， A + 是其伪逆，而 M0 和 B 是大小为N的向量。

特别的是，我想同时求解多个偏移量的方程（因此可以求解多个矩阵 A ）。新的尺寸是：

• A：MxNxN
• b：Nx1
• M0：MxNx1

程序的常规版本（在M尺寸的第一个维度上使用循环）工作正常，但我在“并行版本”中停留在某一点。

此刻我的代码如下：

def bmcsim(A, b, M0, timestep):
    ex = myexpm(A*timestep) # returns stacked array of size MxNxN
    M = np.zeros_like(M0)
    for n in range(ex.shape[0]):
        A_tmp = A[n,:,:]
        A_b = np.linalg.lstsq(A_tmp ,b, rcond=None)[0]
        M[n,:,:] = np.abs(np.real(np.dot(ex[n,:,:], M0[n,:,:] + A_b) - A_b))      
    return M

，我想摆脱那个for n in range(ex.shape[0])循环。不幸的是，np.linalg.lstsq不适用于堆叠数组，对吗？ myexpm中的np.apply_along_axis用于另一个问题：

def myexpm(A):
    vals,vects = np.linalg.eig(A)
    tmp = np.einsum('ijk,ikl->ijl', vects, np.apply_along_axis(np.diag, -1, np.exp(vals)))
    return np.einsum('ijk,ikl->ijl', tmp, np.linalg.inv(vects))

但是，这仅适用于一维输入数据。我可以在np.linalg.lstsq中使用类似的东西吗？我猜np.dot中的bmcsim将被np.einsum替换为myexpm，还是有更好的方法？

感谢您的帮助！

更新： 我刚刚意识到我可以用np.linalg.lstsq(A,b)替换np.linalg.solve(A.T.dot(A), A.T.dot(b))并设法摆脱这种循环：

def bmcsim2(A, b, M0, timestep):
    ex = myexpm(A*timestep)
    b_stack = np.repeat(b[np.newaxis, :, :], offsets.size, axis=0)
    tmp_left = np.einsum('kji,ikl->ijl', np.transpose(A), A)
    tmp_right = np.einsum('kji,ikl->ijl', np.transpose(A), b_stack)
    A_b_stack = np.linalg.solve(tmp_left , tmp_right )

    return np.abs(np.real(np.einsum('ijk,ikl->ijl',ex, M0+A_b_stack ) - A_b_stack ))

这大约快了3倍，但仍然有些复杂。我希望有更好（更短/更容易）的方法，甚至可能更快？！