加速numpy计算

Question

我正在编写一个Python函数来对高维矩阵空间上的向量场进行积分。A，形状(n, m)，是一个矩阵，它的时间导数在它的每个分量A[i, j]中都是线性的。我们可以将所有的导数系数集合成一个四维数组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可以表示为一个稀疏。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))方法