向量化python中的简单函数：避免double for循环

Question

我是python的新手。 我正在尝试做一个非常简单的事情，评估将浮点数作为2D网格输入的非平凡函数。 以下代码完全符合我的要求，但是由于double for循环，它很慢。

import numpy as np
from galpy.potential import RazorThinExponentialDiskPotential


R = np.logspace(0., 2., 10)
z=R

#initialize with default values for this example
potfunc=RazorThinExponentialDiskPotential()

pot=np.zeros((R.size, z.size))


for i in range(0, R.size):
    for j in range(0, z.size):
        pot[i,j]=potfunc(R[i],z[j])

最后，阵列电位器包含了我想要的所有信息，但是现在我想提高效率。 我知道纯python很慢，特别是在循环（如IDL）上，所以我检查了np.vectorize，但这只是幕后的python循环。 问题是potfunc似乎不接受数组，而只是普通的标量。

如何优化这个简单的程序？

提前谢谢了。

Answer 1

做到这一点的标准方法是使用meshgrid ：

  r,z=np.meshgrid(R,Z)
  pot=potfunc(r,z) 

您必须避免在numpy数组上循环，否则将失去所有矢量化效率。

Answer 2

如果您无法手动向量化函数（也许可以将Razor ..类子类化并重写函数），则可以使用multiprocessing 。 除了我简单的辅助函数，您还可以使用自己喜欢的函数：

from multiprocessing import pool
import numpy as np
def worker(x):
    ai,bj = x
    return ai + bj

def run_pool():
    a = np.linspace(0,10,10)
    b = np.logspace(0,10,len(a))
    vec = [(a[i],b[j]) for i in range(len(a)) for j in range(len(b))]
    p = pool.Pool(processes=4) # as many cores as you have
    print(p.map(worker,vec))
    p.close()
    p.join()

run_pool()

但是在考虑加快处理速度之前，进行概要分析将是不错的选择。 我很确定，就您而言，函数本身就是瓶颈。 因此，您可以使用编译器语言对其进行重写，对其进行矢量化处理，或者可以使用所有内核。