在Python中创建跨平台和Python版本独立的快速循环的最佳方法是什么？

Question

我正在用Python写一个科学应用，里面有一个非常耗费处理器资源的循环。我想尽可能优化这个循环，同时又不想给最终用户带来太多麻烦，因为他们可能会把它当作一堆未编译的Python脚本来使用，操作系统包括Windows、Mac和（主要是）Ubuntu Linux。

目前这个应用是用Python写的，稍微用了一点NumPy，下面是我的代码。

1. 有没有一种解决方案可以比较快，而且不需要编译？这样似乎是保持跨平台兼容性最简单的方法。
2. 如果使用像Pyrex这样的工具，它需要编译，那有没有简单的方法可以把很多模块打包，让Python根据检测到的操作系统和Python版本来选择使用哪个模块？有没有简单的方法可以构建这些模块的集合，而不需要访问每个系统和每个版本的Python？
3. 有没有哪种方法特别适合多处理器优化？

（如果你感兴趣，这个循环是用来计算晶体内部某一点的磁场，通过把周围很多磁离子的贡献加起来，这些离子被当作小棒磁铁。基本上就是一个巨大的求和过程，涉及到这些。）

# calculate_dipole
# -------------------------
# calculate_dipole works out the dipole field at a given point within the crystal unit cell
# ---
# INPUT
# mu = position at which to calculate the dipole field
# r_i = array of atomic positions
# mom_i = corresponding array of magnetic moments
# ---
# OUTPUT
# B = the B-field at this point

def calculate_dipole(mu, r_i, mom_i):
    relative = mu - r_i
    r_unit = unit_vectors(relative)
    #4pi / mu0 (at the front of the dipole eqn)
    A = 1e-7
    #initalise dipole field
    B = zeros(3,float)

    for i in range(len(relative)):
        #work out the dipole field and add it to the estimate so far
        B += A*(3*dot(mom_i[i],r_unit[i])*r_unit[i] - mom_i[i]) / sqrt(dot(relative[i],relative[i]))**3
    return B

Answer 1

你的Python代码可以通过用生成器表达式替换循环来加快一些速度，同时使用itertools.izip并行遍历所有三个序列，避免多次查找mom_i[i]、relative[i]和r_unit[i]。

也就是说，把

B = zeros(3,float)

for i in range(len(relative)):
    #work out the dipole field and add it to the estimate so far
    B += A*(3*dot(mom_i[i],r_unit[i])*r_unit[i] - mom_i[i]) / sqrt(dot(relative[i],relative[i]))**3
return B

换成：

from itertools import izip
...
return sum((A*(3*dot(mom,ru)*ru - mom) / sqrt(dot(rel,rel))**3 
            for mom, ru, rel in izip(mom_i, r_unit, relative)),
           zeros(3,float)) 

这样写也更容易理解，因为核心公式不会到处都是[i]，看起来更清晰。

不过，我觉得这样做的速度提升可能不大，跟用像Cython这样的编译语言写整个函数比起来，效果可能微乎其微。

Answer 2

Numpy确实对数组处理做了一些本地优化，这样可以提高效率。你还可以把Numpy数组和Cython结合起来使用，这样可以让处理速度更快。

Answer 3

如果你去掉循环，改用Numpy的向量化操作，这样可以让程序运行得快很多。把你的数据放进形状为(3,N)的numpy数组，然后试试下面的代码：

import numpy as np

N = 20000
mu = np.random.random((3,1))
r_i = np.random.random((3,N))
mom_i = np.random.random((3,N))

def unit_vectors(r):
     return r / np.sqrt((r*r).sum(0))

def calculate_dipole(mu, r_i, mom_i):
    relative = mu - r_i
    r_unit = unit_vectors(relative)
    A = 1e-7

    num = A*(3*np.sum(mom_i*r_unit, 0)*r_unit - mom_i)
    den = np.sqrt(np.sum(relative*relative, 0))**3
    B = np.sum(num/den, 1)
    return B

对我来说，这样的运行速度比用for循环快大约50倍。