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)

然而，＆＃34;矢量化＆＃34;解决方案没有正常工作。当计时我得到第一个解决方案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

Vectorize python代码以提高性能

1 个答案: