Question

我对python完全不熟悉。我正在尝试做一个非常简单的事情，评估一个非浮动的函数，它将浮点数作为2D网格上的输入。下面的代码正是我想要的，但它很慢，因为循环加倍。

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()

但在考虑加快速度之前，分析会很好。我很确定在你的情况下，函数本身就是瓶颈。因此要么用编译器语言重写它，要么对它进行矢量化，要么使用所有核心。

在python中向量化一个简单的函数：避免双循环

2 个答案: