加速Numpy数组上的循环

Question

在我代码中，我让for循环遍历多维numpy数组，并使用在每次迭代中获得的子数组执行一些操作。它看起来像这样

for sub in Arr:
  #do stuff using sub

现在使用sub完成的工作是完全矢量化的，所以它应该是有效的。另一方面，这个循环迭代大约~10^5次，这是一个瓶颈。我不太愿意这样做，因为do stuff using sub使用广播，切片，智能索引技巧，这些技巧用普通的C编写会很繁琐。我也欢迎关于如何处理广播，切片，智能索引的想法和建议，当把计算卸载到C时。

Answer 1

桑，你可以看看scipy.weave。您可以使用scipy.weave.blitz透明地将表达式转换为C++代码并运行它。它将自动处理切片并消除时间，但您声称for循环的主体不会创建时间，因此您的里程可能会有所不同。

但是，如果您想用更有效的方式替换整个for循环，那么您可以使用scipy.inline。缺点是您必须编写C++代码。这应该不会太难，因为您可以使用Blitz++语法，它非常接近于numpy数组表达式。直接支持切片，但不支持广播。

有两种变通方法：

numpy

• 将使用numpy-C应用编程接口和多维迭代器。他们透明地处理广播。但是，您调用的是Numpy运行时，因此可能会有一些开销。另一种选择，也可能是更简单的选择，是使用通常的矩阵表示法进行广播。广播操作可以写成向量为1的外积。好消息是Blitz++实际上不会在内存中创建这种临时广播数组，它会弄清楚如何将其包装到一个等效的循环中。

• 对于第二种选择，请查看http://www.oonumerics.org/blitz/docs/blitz_3.html#SEC88中的索引占位符。只要你的矩阵少于11个维度，你就没问题。此链接显示如何使用它们来形成外部产品http://www.oonumerics.org/blitz/docs/blitz_3.html#SEC99 (搜索外部产品以转到文档的相关部分)。

Answer 2

如果你不能“向量化”整个操作，循环确实是瓶颈，那么我强烈推荐使用Cython。我最近一直在尝试它，它使用起来很简单，并且和numpy有一个很好的接口。对于像langevin积分器这样的东西，我看到在numpy中一个像样的实现有115倍的加速。请参阅此处的文档：

http://docs.cython.org/src/tutorial/numpy.html

我还建议您查看下面的paper

只需输入输入数组和循环计数器，您可能会看到令人满意的加速，但如果您想充分利用cython的潜力，那么您将不得不对等效的广播进行硬编码。

Answer 3

除了使用Cython之外，您还可以用Fortran编写瓶颈部分。然后使用f2py将其编译成Python .pyd文件。