python - 无需for循环解决广播错误，加快代码速度

Question

我可能误解了Python中广播的工作原理，但我仍然遇到错误。

scipy 提供了许多接受两个参数的“特殊函数”，特别是 eval_XX(n, x[,out]) 函数。 请参阅http://docs.scipy.org/doc/scipy/reference/special.html

我的程序使用许多正交多项式，因此我必须在不同的点评估这些多项式。让我们举一个具体的例子scipy.special.eval_hermite(n, x, out=None)。

我希望 x 参数为矩阵形状 (50, 50)。然后，我想在多个点评估该矩阵的每个条目。让我们将 n 定义为一个 numpy 数组 narr = np.arange(10) (其中我们将 numpy 导入为 np，即 将 numpy 导入为 np)。

所以，打电话

scipy.special.eval_hermite(narr, matrix)

应返回埃尔米特多项式H_0(矩阵)、H_1(矩阵)、H_2(矩阵)等。每个H_X(矩阵)的形状为 (50,50)，原始输入矩阵的形状。

然后，我想对这些值求和。所以，我打电话

matrix1 = np.sum( [scipy.eval_hermite(narr, matrix)], axis=0 )

但是我收到广播错误!

ValueError: operands could not be broadcast together with shapes (10,) (50,50)

我可以用 for 循环来解决这个问题，即

matrix2 = np.sum( [scipy.eval_hermite(i, matrix) for i in narr], axis=0)

这给了我正确的答案，以及输出matrix2.shape = (50,50)。但是使用这个 for 循环会大大减慢我的代码速度。请记住，我们正在处理矩阵条目。

有没有办法不用 for 循环来做到这一点？

Answer 1

eval_hermite 使用 x 广播 n，然后在每个点计算 H_n(x)。因此，输出形状将是 n 与 x 广播的结果。因此，如果您想实现此目的，则必须使 n 和 x 具有兼容的形状:

import scipy.special as ss
import numpy as np
matrix = np.ones([100,100]) # example
narr = np.arange(10) # example
ss.eval_hermite(narr[:,None,None], matrix).shape # => (10, 100, 100)

但请注意，这实际上可能更快:

out = np.zeros_like(matrix)
for n in narr:
    out += ss.eval_hermite(n, matrix)

在测试中，它似乎比上面的 np.sum(...) 快 5-10%。