如何提高循环遍历 numpy 线性求解的计算效率答案

【问题标题】：How can I improve the computational efficiency of looping over numpy linear solve如何提高循环遍历 numpy 线性求解的计算效率
【发布时间】：2021-09-27 16:25:55
【问题描述】：

我有一个方程组，我需要求解大量变量。代码比下面复杂，但问题如下：

我有一个迭代器：

iterator= np.arange(2000)   # array!

以及依赖于该迭代器函数的数组：

A_11 = function_A11(iterator)
A_12 = function_A12(iterator)
A_21 = function_A21(iterator)
A_22 = function_A22(iterator)

B_1 = function_B1(iterator)
B_2 = function_B2(iterator)

X = np.zeros(2, 2000)

for i, (A11, A12, A21, A22, B1, B2) in enumerate(zip(A_11, A_12, A_21, A_22, B_1, B_2):
    A = np.array([[A11, A12],[A21,A22]]) 
    B = np.array([B1, B2])
    X[:,i] = np.linalg.solve(A,B)

该方法有效，但计算量很大，我觉得我应该能够优化它，例如通过使用 3D 数组。有人有什么建议吗？

谢谢！

提姆

【问题讨论】：

如果您必须将该函数应用于数组中的每个条目，它的速度不能超过 O(N)，其中 N = # 个元素。如果你有多个核心，你可以并行化，但是通过上下文切换，单个核心不能做得更好。
你没有迭代器，你有一个数组。
根据文档，A 可以是 (…, M, M) 和 B (…, M,)。在您的情况下，这意味着制作 (2000,2,2) 和 (2000,2) 数组

标签： python numpy loops linear-algebra

【解决方案1】：

arr= np.arange(2000)   # array!

A_11 = function_A11(arr)
A_12 = function_A12(arr)
A_21 = function_A21(arr)
A_22 = function_A22(arr)

B_1 = function_B1(arr)
B_2 = function_B2(arr)

结合

AA = np.array([[A_11,A_12],[A_21, A_22]])

我希望 AA.shape 是 (2,2,2000)。

A = AA.transpose(2,0,1)

获取solve_ivp 可以使用的 (2000,2,2)。

我还没有测试过。

【讨论】：

你也应该转置B。从文档中，如果A.shape==(2000,2,2) 那么它应该是B.shape==(2000,2)。然后 OP 可以使用X = np.linalg.solve(A,B.T).. 给 OP：如果你的数组是整数数组，你应该将它们转换为float，否则结果可能会有所不同。
有了上面的评论，效果很好！