Python优化问题？

Question

最近我有个作业（别担心，我已经用C++完成了），但我很好奇如果用Python怎么做。这个问题是关于两个发光源的，它们会发出光。我就不详细讲了。

这是我写的代码（我在后面优化了一些）：

import math, array
import numpy as np
from PIL import Image

size = (800,800)
width, height = size

s1x = width * 1./8
s1y = height * 1./8
s2x = width * 7./8
s2y = height * 7./8

r,g,b = (255,255,255)
arr = np.zeros((width,height,3))
hy = math.hypot
print 'computing distances (%s by %s)'%size,
for i in xrange(width):
    if i%(width/10)==0:
        print i,    
    if i%20==0:
        print '.',
    for j in xrange(height):
        d1 = hy(i-s1x,j-s1y)
        d2 = hy(i-s2x,j-s2y)
        arr[i][j] = abs(d1-d2)
print ''

arr2 = np.zeros((width,height,3),dtype="uint8")        
for ld in [200,116,100,84,68,52,36,20,8,4,2]:
    print 'now computing image for ld = '+str(ld)
    arr2 *= 0
    arr2 += abs(arr%ld-ld/2)*(r,g,b)/(ld/2)
    print 'saving image...'
    ar2img = Image.fromarray(arr2)
    ar2img.save('ld'+str(ld).rjust(4,'0')+'.png')
    print 'saved as ld'+str(ld).rjust(4,'0')+'.png'

我已经优化了大部分代码，但在有两个for循环的部分，性能差距还是很大，我想不出用常见的数组操作来解决这个问题……欢迎大家给我建议 :D

编辑：为了回应Vlad的建议，我来详细说明一下问题： 有两个光源，每个光源发出的光都是一个正弦波： E1 = E0*sin(omega1*time+phi01) E2 = E0*sin(omega2*time+phi02) 为了简单起见，我们假设omega1=omega2=omega=2*PI/T，phi01=phi02=phi0。 假设x1是从第一个光源到平面上某点的距离，那么在这个点的光强度是 Ep1 = E0*sin(omega*time - 2*PI*x1/lambda + phi0) 其中 lambda = 光速 * T（振荡周期）。 考虑到两个光源在平面上，公式变成 Ep = 2*E0*cos(PI*(x2-x1)/lambda)sin(omegatime - PI*(x2-x1)/lambda + phi0) 从中我们可以看出，当 (x2-x1)/lambda = (2*k) * PI/2 时光强度达到最大，而当 (x2-x1)/lambda = (2*k+1) * PI/2 时光强度达到最小，中间的值则是变化的，其中k是一个整数。

在某个时刻，给定光源的坐标，以及已知的lambda和E0，我们需要写个程序来绘制光的样子。 在我看来，我已经尽可能地优化了这个问题……

Answer 1

列表推导式比循环要快得多。例如，你可以用下面的方式来写：

for j in xrange(height):
        d1 = hy(i-s1x,j-s1y)
        d2 = hy(i-s2x,j-s2y)
        arr[i][j] = abs(d1-d2)

而不是这样：

arr[i] = [abs(hy(i-s1x,j-s1y) - hy(i-s2x,j-s2y)) for j in xrange(height)]

另一方面，如果你真的想要“优化”代码，那你可能需要用C语言重新实现这个算法，然后使用SWIG之类的工具从Python中调用它。

Answer 2

如果你使用数组操作而不是循环，速度会快很多很多。对我来说，生成图像现在是最耗时的部分。与其用两个 i,j 循环，不如用我下面这个：

I,J = np.mgrid[0:width,0:height]
D1 = np.hypot(I - s1x, J - s1y)
D2 = np.hypot(I - s2x, J - s2y)

arr = np.abs(D1-D2)
# triplicate into 3 layers
arr = np.array((arr, arr, arr)).transpose(1,2,0)
# .. continue program

你未来要记住的基本点是：这不是在优化；在numpy中使用数组形式其实就是在按照它应该被使用的方式使用它。随着经验的积累，你未来的项目应该不再绕弯子去用python循环，数组形式应该是最自然的选择。

我们在这里做的事情其实很简单。我们用 numpy.hypot 替代了 math.hypot。像所有这样的numpy函数一样，它接受ndarray作为参数，并且正好做我们想要的事情。

Answer 3

干涉图样真有意思，对吧？

首先，这个程序在我笔记本上运行其实只需要十二秒半，所以影响不大。

不过，我们来看看能不能通过numpy数组操作来优化一下第一部分，好吗？基本上你想要的是：

arr[i][j] = abs(hypot(i-s1x,j-s1y) - hypot(i-s2x,j-s2y))

对于所有的 i 和 j。

因为numpy有一个可以对数组使用的 hypot 函数，我们就用这个。我们的第一个挑战是创建一个合适大小的数组，让每个元素都等于 i，再创建一个每个元素都等于 j 的数组。不过这并不难；其实下面的回答提到了一个我之前不知道的很棒的 numpy.mgrid，它正好可以做到这一点：

array_i,array_j = np.mgrid[0:width,0:height]

还有一个小问题，就是要把你的 (width, height) 大小的数组变成 (width, height, 3)，这样才能和你的图像生成语句兼容，不过这也很简单：

arr = (arr * np.ones((3,1,1))).transpose(1,2,0)

然后我们把这个放进你的程序里，让数组操作来完成任务：

import math, array
import numpy as np
from PIL import Image

size = (800,800)
width, height = size

s1x = width * 1./8
s1y = height * 1./8
s2x = width * 7./8
s2y = height * 7./8

r,g,b = (255,255,255)

array_i,array_j = np.mgrid[0:width,0:height]

arr = np.abs(np.hypot(array_i-s1x, array_j-s1y) -
             np.hypot(array_i-s2x, array_j-s2y))

arr = (arr * np.ones((3,1,1))).transpose(1,2,0)

arr2 = np.zeros((width,height,3),dtype="uint8")
for ld in [200,116,100,84,68,52,36,20,8,4,2]:
    print 'now computing image for ld = '+str(ld)
    # Rest as before

新的运行时间是……8.2秒。所以你大约节省了四秒时间。不过，现在几乎所有的时间都花在图像生成上了，所以也许你可以通过只生成需要的图像来进一步优化。