Python：优化循环

Posted 2023-03-11

【中文标题】Python：优化循环

【英文标题】：Python: optimising loops

【发布时间】：2013-07-17 14:25:44

【问题描述】：

我希望优化一些由两个嵌套循环组成的 python 代码。我对 numpy 不是很熟悉，但我知道它应该使我能够提高此类任务的效率。下面是我编写的测试代码，它反映了实际代码中发生的情况。目前使用 numpy 范围和迭代器比通常的 python 慢。我究竟做错了什么？这个问题的最佳解决方案是什么？

感谢您的帮助！

import numpy
import time

# setup a problem analagous to that in the real code
npoints_per_plane = 1000
nplanes = 64
naxis = 1000
npoints3d = naxis + npoints_per_plane * nplanes
npoints = naxis + npoints_per_plane
specres = 1000

# this is where the data is being mapped to
sol = dict()
sol["ems"] = numpy.zeros(npoints3d)
sol["abs"] = numpy.zeros(npoints3d)

# this would normally be non-random input data
data = dict()
data["ems"] = numpy.zeros((npoints,specres))
data["abs"] = numpy.zeros((npoints,specres))
for ip in range(npoints):
    data["ems"][ip,:] = numpy.random.random(specres)[:]
    data["abs"][ip,:] = numpy.random.random(specres)[:]
ems_mod = numpy.random.random(1)[0]
abs_mod = numpy.random.random(1)[0]
ispec = numpy.random.randint(specres)

# this the code I want to optimize

t0 = time.time()

# usual python range and iterator
for ip in range(npoints_per_plane):
    jp = naxis + ip
    for ipl in range(nplanes):
        ip3d = jp + npoints_per_plane * ipl
        sol["ems"][ip3d] = data["ems"][jp,ispec] * ems_mod
        sol["abs"][ip3d] = data["abs"][jp,ispec] * abs_mod

t1 = time.time()

# numpy ranges and iterator
ip_vals = numpy.arange(npoints_per_plane)
ipl_vals = numpy.arange(nplanes)
for ip in numpy.nditer(ip_vals):
    jp = naxis + ip
    for ipl in numpy.nditer(ipl_vals):
        ip3d = jp + npoints_per_plane * ipl
        sol["ems"][ip3d] = data["ems"][jp,ispec] * ems_mod
        sol["abs"][ip3d] = data["abs"][jp,ispec] * abs_mod


t2 = time.time()

print "plain python: %0.3f seconds" % ( t1 - t0 )
print "numpy: %0.3f seconds" % ( t2 - t1 )

编辑：仅将“jp = naxis + ip”放在第一个 for 循环中

补充说明：

我想出了如何让 numpy 快速执行内部循环，而不是外部循环：

# numpy vectorization
for ip in xrange(npoints_per_plane):
    jp = naxis + ip
    sol["ems"][jp:jp+npoints_per_plane*nplanes:npoints_per_plane] = data["ems"][jp,ispec] * ems_mod
    sol["abs"][jp:jp+npoints_per_plane*nplanes:npoints_per_plane] = data["abs"][jp,ispec] * abs_mod

以下 Joe 的解决方案展示了如何同时进行这两种操作，谢谢！

【问题讨论】：

我不熟悉numpy.range，但使用Python的range与创建n元素列表相同，而xrange是数字生成器——避免存储整个列表。

谢谢。你是对的，但是当我使用 xrange 而不是 range 时，运行时没有变化。

编写快速 numpy 代码的关键是避免通过矢量化操作（或者，更确切地说，将循环推低到快速 C 级别）循环，而不是迭代 numpy 对象是比较快的。 IOW，惯用的 numpy 代码没有很多 for 循环，而是作为一个整体作用于向量和数组。

2 个问题。 1) 数组的大小是否与您的真实应用程序相同？ 2) 你需要多快的代码运行速度？

1) 对于一个小问题，是的。对于最大的问题，它更像是 npoints_per_plane = 50000，nplanes = 64，naxis = 1000，specres = 1000。2）我希望代码以与 c 相似的数量级运行。我刚刚花了一个星期将一个 swig 接口转换为纯 python，但现在一个小问题需要两倍的时间，这就是瓶颈。

【参考方案1】：

在 numpy 中编写循环的最佳方式是不编写循环，而是使用向量化操作。例如：

c = 0
for i in range(len(a)):
    c += a[i] + b[i]

变成

c = np.sum(a + b, axis=0)

对于形状为 (100000, 100) 的 a 和 b，第一个变体需要 0.344 秒，第二个变体需要 0.062 秒。

在您的问题中提出的情况下，以下是您想要的：

sol['ems'][naxis:] = numpy.ravel(
    numpy.repeat(
        data['ems'][naxis:,ispec,numpy.newaxis] * ems_mod,
        nplanes,
        axis=1
    ),
    order='F'
)

这可以使用some tricks 进一步优化，但这会降低清晰度并且可能是过早的优化，因为：

普通蟒蛇：0.064 秒

numpy：0.002 秒

解决方案如下：

您的原始版本包含 jp = naxis + ip 仅跳过第一个 naxis 元素 [naxis:] 选择除第一个 naxis 元素之外的所有元素。您的内部循环将 data[jp,ispec] 的值重复 nplanes 次并将其写入多个位置 ip3d = jp + npoints_per_plane * ipl，这相当于一个扁平的二维数组偏移 naxis。因此，通过numpy.newaxis 将第二个维度添加到（以前的一维）data['ems'][naxis:, ispec]，这些值通过numpy.repeat 沿着这个新维度重复nplanes 次。然后通过numpy.ravel（按 Fortran 顺序，即最低轴具有最小步幅）再次展平生成的二维数组，并写入sol['ems'] 的适当子数组。如果目标数组实际上是二维的，则可以使用自动数组广播来跳过重复。

如果遇到无法避免使用循环的情况，可以使用Cython（在numpy 数组上支持efficient buffer views）。

【讨论】：

谢谢，是的，我正在阅读docs.scipy.org/doc/numpy/reference/arrays.nditer.html 以尝试确定是否可以矢量化我的循环。问题是它不是 1:1 映射，所以也许我仍然需要循环。我可以将 cython 代码放在 python 脚本中吗？不知道它是如何工作的......

实际上代码看起来确实是一对一的映射，只是它是reshaped。 Cython 代码可以通过cython.inline 内联使用。

如何向量化“for i in range(len(a)): c += a[i] + b[2*i]” 例如？这将帮助我了解该怎么做。

c = numpy.sum(a + b[::2])

好的，我现在看到了语法（[i:j:k] where i=start, j=stop, k=increment），但我不知道如何应用这个循环。例如，如何矢量化“for ipl in xrange(nplanes): sol["ems"][jp + npoints_per_plane * ipl] = data["ems"][jp,ispec] * ems_mod"？