scipy.linalg.solve的向量化

问题描述

我有许多较小的矩阵，例如5 x 5，A = numpy.random.rand(5,5,7,77)，其中的一侧为y = numpy.random.rand(5)。我想解决所有7 x 77问题A_{ij} x = b，以使结果x的形状为5,77。我可以简单地遍历它们，

from scipy.linalg import solve
import numpy

A = numpy.random.rand(5,77)
b = numpy.random.rand(5)

x = []
for i in range(A.shape[2]):
    x.append([])
    for j in range(A.shape[3]):
        x[-1].append(solve(A[:,:,i,j],b))

x = numpy.array(x)
x = numpy.moveaxis(x,-1,0)

print(x.shape)

但这很慢。感觉应该有可能通过将A而不是5 x 5 x 7 x 77张量或浮点数，而是作为7 x 77浮点数数组的5 x 5矩阵进行向量化，并执行{这些阵列上的{1}}。有提示吗？

（我经常遇到这类问题，因此，如果有一个图书馆来处理这些问题，我也将很高兴听到它。）

解决方法

如果您首先重新排列尺寸，则可以使用np.linalg.solve进行操作。

import numpy as np

# Make random problem
np.random.seed(0)
a = np.random.rand(5,5,7,77)
b = np.random.rand(5)
# Put additional axes at the end
at = np.moveaxis(a,(0,1),(2,3))
# Solve
xt = np.linalg.solve(at,b[np.newaxis,np.newaxis])
# Put axes back in place
x = np.moveaxis(xt,2,0)
print(x.shape)
# (5,77)
# Test some result
print(np.allclose(a[:,:,4,36] @ x[:,36],b))
# True