帮我消除Python中的for循环

Question

这里一定有更快的方法来做这个。

虽然这里的内容比较多，但其实很简单，可以慢慢理清楚。

下面是相关的Python代码（来自scipy库）

for i in arange(len(wav)):
    result[i] = sum(laser_flux * exp(-(wav[i] - laser_wav)**2) )

这里有一堆数组。

• result -- 长度为（wav）的数组
• laser_flux -- 长度为（laser）的数组
• wav -- 长度为（wav）的数组
• laser_wav -- 长度为（laser）的数组

没错，在指数运算里，我是对标量值和laser_wav数组之间的差值进行逐个平方。

一切都按预期工作（包括速度很慢），如果你能帮我去掉这个for循环，我会非常感激！

Answer 1

如果你在性能方面遇到问题，可能可以考虑重写代码，以便利用多个处理器核心（如果你的设备有的话）。

from multiprocessing import Pool
p = Pool(5) # about the number of cores you have

def f(i):
    delta = wav[i] - laser_wav
    return sum(laser_flux * exp(-delta*delta) )

result = p.map(f, arange(len(wav)) )

Answer 2

我刚开始学习Python，所以这可能不是在Python中最优的做法，但我会用同样的方法来处理Perl、Scheme等其他语言。

def func(x):
    delta = x - laser_wav
    return sum(laser_flux * exp(-delta * delta))
result = map(func, wav)

Answer 3

你可能需要使用Numpy数组（如果还没用的话）来存储你的数据。这样，你就可以利用数组广播的特性，使用np.newaxis。对于wav中的每一个值，你需要计算它和laser_wav中每一个值之间的差。这意味着你需要一个二维数组，其中一个维度是wav的维度，另一个维度是laser的维度。

在下面的例子中，我将第一个索引作为laser索引，第二个索引作为wav索引。用一些示例数据，这样就变成了：

import numpy as np

LASER_LEN  = 5
WAV_LEN    = 10
laser_flux = np.arange(LASER_LEN)
wav        = np.arange(WAV_LEN)
laser_wav  = np.array(LASER_LEN)

# Tile wav into LASER_LEN rows and tile laser_wav into WAV_LEN columns
diff    = wav[np.newaxis,:] - laser_wav[:,np.newaxis]
exp_arg = -diff ** 2
sum_arg = laser_flux[:,np.newaxis] * np.exp(exp_arg)

# Now, the resulting array sum_arg should be of size (LASER_LEN,WAV_LEN)
# Since your original sum was along each element of laser_flux/laser_wav, 
# you'll need to sum along the first axis.
result = np.sum(sum_arg, axis=0)

当然，你也可以把这个简化成一个单独的语句：

result = np.sum(laser_flux[:,np.newaxis] * 
                np.exp(-(wav[np.newaxis,:]-laser_wav[:,np.newaxis])**2),axis=0)

补充：

正如问题评论中提到的，你可以利用线性代数乘法定义中固有的“乘法求和”特性。这样就变成了：

result = np.dot(laser_flux, 
    np.exp(-(wav[np.newaxis,:] - laser_wav[:,np.newaxis])**2))