解决无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，即import numpy as np）。

所以，打电话

scipy.special.eval_hermite(narr, matrix)

应该返回厄米多项式H_0(matrix), H_1(matrix), H_2(matrix) 等。每个H_X(matrix) 的形状为(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 广播 n 和 x，然后在每个点计算 H_n(x)。因此，输出形状将是用x 广播n 的结果。所以，如果你想完成这项工作，你必须让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%。

Answer 2

这些函数的文档很简陋，而且编译了很多代码，所以这只是基于实验：

special.eval_hermite(n, x, out=None)

n 显然是一个标量或整数数组。 x 可以是浮点数数组。

special.eval_hermite(np.ones(5,int)[:,None],np.ones(6)) 给我一个(5,6) 结果。这与我从np.ones(5,int)[:,None] * np.ones(6) 得到的形状相同。

np.ones(5,int)[:,None] 是一个 (5,1) 数组，np.ones(6) 是一个 (6,)，为此，它等效于 (1,6)。两者都可以扩展为(5,6)。

据我所知，这些special 函数中的广播规则与* 等运算符的广播规则相同。

---

由于special.eval_hermite(nar[:,None,None], x) 产生(10,50,50)，您只需将sum 应用于其轴0 即可产生(50,50)。

special.eval_hermite(nar[:,Nar,Nar], x).sum(axis=0)

就像我之前写的，同样的广播（和求和）规则适用于 hermite，就像它们适用于像 * 这样的基本操作。