python - 向量化 python 代码以提高性能

Question

我正在用 Python 编写科学代码来计算系统的能量。这是我的函数:cte1、cte2、cte3、cte4 是先前计算的常量； pii 是 np.pi(预先计算，否则会减慢循环速度)。我计算总能量的 3 个分量，然后将它们相加。

def calc_energy(diam): 
    Energy1 = cte2*((pii*diam**2/4)*t)
    Energy2 = cte4*(pii*diam)*t
    d=diam/t
    u=np.sqrt((d)**2/(1+d**2))
    cc= u**2
    E = sp.special.ellipe(cc) 
    K = sp.special.ellipk(cc) 
    Id=cte3*d*(d**2+(1-d**2)*E/u-K/u)
    Energy3 = cte*t**3*Id
    total_energy = Energy1+Energy2+Energy3
    return (total_energy,Energy1)

我的第一个想法是简单地遍历所有直径值:

start_diam, stop_diam, step_diam = 1e-10, 500e-6, 1e-9 #Diametre
diametres = np.arange(start_diam,stop_diam,step_diam)

for d in diametres:  
    res1,res2 = calc_energy(d)
    totalEnergy.append(res1)
    Energy1.append(res2)

为了加快计算速度，我决定使用numpy进行向量化，如下图:

diams = diametres.reshape(-1,1) #If not reshaped, calculations won't run
r1 = np.apply_along_axis(calc_energy,1,diams)

但是，“矢量化”解决方案无法正常工作。计时时，第一个解决方案用时 5 秒，第二个解决方案用时 18 秒。

我想我做错了，但无法弄清楚是什么。

Answer 1

使用当前的方法，您将对数组的每个元素应用一个 Python 函数，这会带来额外的开销。相反，您可以将整个数组传递给您的函数并返回一组答案。您现有的功能似乎无需任何修改即可正常工作。

import numpy as np
from scipy import special
cte = 2
cte1 = 2
cte2 = 2
cte3 = 2
cte4 = 2
pii = np.pi

t = 2

def calc_energy(diam): 
    Energy1 = cte2*((pii*diam**2/4)*t)
    Energy2 = cte4*(pii*diam)*t
    d=diam/t
    u=np.sqrt((d)**2/(1+d**2))
    cc= u**2
    E = special.ellipe(cc) 
    K = special.ellipk(cc) 
    Id=cte3*d*(d**2+(1-d**2)*E/u-K/u)
    Energy3 = cte*t**3*Id
    total_energy = Energy1+Energy2+Energy3
    return (total_energy,Energy1)

start_diam, stop_diam, step_diam = 1e-10, 500e-6, 1e-9 #Diametre
diametres = np.arange(start_diam,stop_diam,step_diam)

a = calc_energy(diametres) # Pass the whole array