非常大的就地numpy数组操作: numba、pythran或其他？

Question

tI需要在非常大的数组(数百万个条目)上执行操作，这些数组的累积大小接近可用内存。我的理解是，当使用像a=a*3+b-c**2这样的numpy进行朴素操作时，会创建几个临时数组，从而占用更多内存。

由于我计划在内存占用的限制下工作，我担心这种简单的方法将不起作用。因此，我想从正确的方法开始我的开发。

我知道像numba或pythran这样的包可以在操作数组时帮助提高性能，但我不清楚它们是否可以自动处理就地操作，避免临时对象……？

作为一个简单的例子，这里有一个我必须在大型数组上使用的函数：

def find_bins(a, indices):
    global offset, width, nstep
    i = (a-offset) *nstep/ width 
    i = np.where(i<0,0,i)
    i = np.where(i>=nstep,nstep, i)
    indices[:] = i.astype(int)

所以混合了算术运算和对numpy函数的调用。

使用numba或pythran (或其他东西)编写这样的函数有多容易？每种情况下的利弊是什么？

谢谢你的任何提示！

ps:我知道numexpr，但我不确定它是否方便，或者是否适合比单个算术表达式更复杂的函数？

Answer 1

Pythran在设计上避免了许多临时数组。对于您所指向的简单表达式，这将是

#pythran export find_bins(float[], int[], float, float, int)
import numpy as np
def find_bins(a, indices, offset, width, nstep):
    i = (a-offset) *nstep/ width #
    i = np.where(i<0,0,i)
    i = np.where(i>=nstep,nstep, i)
    indices[:] = i.astype(int)

这既避免了临时计算，又加快了计算速度。

并不是说你应该在这里使用np.clip函数，Pythran也支持它。

Answer 2

使用numexpr。例如：

import numexpr
numexpr.evaluate("a+b*c", out=a)

这可以帮助你避免tmp变量，你可以参考High Performance Python，M.G，I.O.