加速 numpy 计算

Question

我正在编写一个 Python 函数来在高维矩阵空间上集成一个向量场。 A ，形状(n, m) ，是一个矩阵，其时间导数在其每个分量A[i, j]中都是线性的。 我们可以将导数的所有系数收集到一个 4D 数组C中，使得C[i, j, k, l]是A[i, j]的导数中A[k, l]的系数。 在这种情况下， A的导数由dA[i, j] == (C[i, j] * A).sum() 。 因此计算是正确的

dA = np.array([[ (Cij * A).sum() for Cij in Ci ] for Ci in C ])

幸运的是， C可以表示为sparse.COO对象，因此上述内容只需要 O(nm) 次乘法。 但是两个 for 循环仍然很慢。 感谢有用的评论，我将其改进为

dA = (C * A).sum(axis=3).sum(axis=2)

利用广播显着加速。 有人能走得更快吗？

Answer 1

您可以使用np.einsum进一步加快速度，因为您无需进行任何中间计算。 或者至少你可以做(C * A).sum(axis=(2,3))来删除一个中间步骤。

import numpy as np
A = np.full((12,12), 2)
C = np.full((12,12,3,2), 1).T
dA = (C * A).sum(axis=3).sum(axis=2)
print(np.einsum('abkl,ijkl->ij', A[None, None], C) == dA)
print((C * A).sum(axis=(2,3)) == dA)

输出：

[[ True  True  True]
 [ True  True  True]]
[[ True  True  True]
 [ True  True  True]]

老实说，我并不完全理解你的数学问题，而且我对einsum也不是很好。 也就是说，您应该仔细检查算法和测试用例是否正确:)

编辑：添加.sum(axis=(2,3))方法