向量化 python 代码以提高性能答案

【问题标题】：Vectorize python code for improved performance向量化 python 代码以提高性能
【发布时间】：2017-10-13 05:28:51
【问题描述】：

我正在用 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 秒。

我想我做错了什么，但不知道是什么。

【问题讨论】：

你调用的函数calc_total_energy和你定义的calc_energy有区别吗
要对其进行矢量化，您需要避免在数组中的每个项目上应用该函数，因为这具有 Python 的所有开销。相反，您可以批量执行计算，例如Energy1 = cte2*((pii*diametres**2/4)*t) 会给你一个 Energy1 值的数组。
您的函数calc_energy 将按原样接受数组diametres 作为输入而无需修改。这会给你正确的输出吗？我无法理智地检查它是否给出了正确的值。
@roganjosh，是的，我刚刚检查了它calc_energy(diametres) - 将返回将产生相同总和的数组。我认为您应该将其发布为答案
对不起，我离开电脑了。该解决方案有效。它现在计算不到 0.1 秒而不是 5 秒！出色的表现

标签： python python-3.x numpy optimization scientific-computing

【解决方案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

【讨论】：