【问题标题】：Python: improving performance of a for loop with function call inside that depends on loop index onlyPython：通过仅依赖循环索引的内部函数调用来提高 for 循环的性能
【发布时间】：2019-01-12 19:46:32
【问题描述】：

我对在 python 中像 Fortran 一样编写代码感到内疚。我正在重写我自己在 Fortran 中编写的长代码的许多部分，因为我想显着扩展代码，并且在 python 中扩展以进行概念验证工作要容易得多。但是，如果可以加快速度，我将简单地使用 python，我实际上对一遍又一遍地启动程序并不感兴趣。一旦一个想法被证明是可行的，我就会转向下一个问题。这就是为什么我希望在 python 中工作。不幸的是，现在，用 python 编写的运行需要几周的时间。即使在以下 for 循环上加速一个数量级，也将使其成为一个可行的测试平台。

针对 R 语言提出了类似的问题

Improving loop performance with function call inside

但令人惊讶的是，我没有看到一个用于 python 的。这个类似，但是for循环中的函数有依赖关系，我的没有

Improve performance of a for loop in Python (possibly with numpy or numba)

一个巨大的瓶颈是单个for循环

简单代码

import numpy as np

part  = 3 # a random index of an array, fixed here for example purposes
nmol = 1000
energy = np.zeros((nmol),dtype=np.float_)

for i in range(nmol):
    energy[i] = np.where( part != i,function(part,i),0.0) # if i = part, energy = 0.0

加速函数本身是另一个独立的问题。必须有一种方法可以使用 numpy 或其他方法同时运行所有调用

例如，我们可以说

def function(i,j):
    for k in range(100000): # this loop is simply to make the time about a second or 2
        ener = (i + j) * (i * j) # entirely arbitrary and not my real problem
    return ener

实际上，我的函数调用了几个依赖于部分和“i”的函数。

完整的工作示例是：

import numpy as np
import time as time

def function(i,j):
    for k in range(10000): # this loop is simply to make the time about a second or 2
        ener = (i + j) * (i * j) # entirely arbitrary and not my real problem
    return ener

part  = 3 # a random index of an array, fixed here for example purposes
nmol = 1000
energy = np.zeros((nmol),dtype=np.float_)

start = time.time()
for i in range(nmol):
    energy[i] = np.where( part != i,function(part,i),0.0) # if i = part, energy = 0.0

end = time.time()
print('time: ', end-start)

我使用的是 Python 3.6 版。在循环中，索引“i”必须不与索引“part”交互。

【问题讨论】：

您对多处理持开放态度吗？如果是，您可以参考docs.python.org/3.4/library/…。很难看出如何提高循环的迭代速度。我不确定你是否可以使用列表理解，因为能量是一个 numpy 对象
energy 只是在每个元素中保存浮点标量......它可以是列表或数组或速度所需的任何东西:)。我愿意接受任何能让我的代码更快的事情！
for 循环在您的 function(i,j) 中是多余的。你知道的，对吧？
使用 numba。或者，如果这太严格，请使用 cython。另外，julia 生态越来越强大，加上（编译后的）julia 代码可以直接调用到 python 代码中。
您不想调用的函数是否包含任何非 numpy 代码？否则，使用 Numba 解决这个问题真的很容易，正如 @J_H 已经提到的那样。

标签： python python-3.x performance vectorization

【解决方案1】：

我很惊讶我能找到答案——我不是故意回答我自己的问题……我只是自己考虑了一段时间。

而不是从索引 0 到 nmol 的 for 循环，创建一个从 0 到 nmol 的整数数组。只需通过传递整数数组来调用函数。因此，数组输入接收数组输出。我修改了函数，使其不需要常量“部分”

这个矢量化解决方案比 for 循环快 ~27 倍，这给了我所需的数量级。

随着大小为 nmol 的数组长度变大，速度增益会增加，反之亦然。

import numpy as np
import time as time

def function(i):
    for k in range(10000):
        ener = (i + part) + (i * part) # entirely arbitrary and not my real problem
    return ener

part  = 3 # a random index of an array, fixed here for example purposes
nmol = 1000

start = time.time()

part_list = np.arange(0,nmol,1)
part_list = np.delete(part_list,part)  # remove the self index

energy =function(part_list) # calls the function in a vectorized form.

end = time.time()
time2 = end-start

【讨论】：

【解决方案2】：

根据您的矢量化解决方案，您可以使用 pythran 获得额外的加速

原码：

import numpy as np

def function(i, part):
    for k in range(10000):
        ener = (i + part) + (i * part)
    return ener

和相关的基准：

python -m timeit -s 'import numpy as np; part  = 3; nmol = 1000; part_list = np.arange(0,nmol,1); part_list = np.delete(part_list, part); from a import function' 'function(part_list, part)'
10 loops, best of 3: 37.3 msec per loop

然后添加pythran export 评论

import numpy as np

#pythran export function(int64[], int64)
def function(i, part):
    for k in range(10000):
        ener = (i + part) + (i * part)
    return ener

并使用以下命令编译模块：

pythran a.py

提供额外的提升：

python -m timeit -s 'import numpy as np; part  = 3; nmol = 1000; part_list = np.arange(0,nmol,1); part_list = np.delete(part_list, part); from a import function' 'function(part_list, part)'
1000000 loops, best of 3: 1.53 usec per loop

【讨论】：