加速numpy计算

Question

我正在编写一个Python函数来整合高维矩阵空间上的向量场。A，shape (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可以表示为一个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))方法