python 速度有改进吗_有没有什么好方法来优化这个python代码的速度?

正如评论中所说,应该重写大部分代码以获得最佳性能。

我只修改了辛普森集成和@hyry答案一点。这加快了计算速度

22.15s

到

1.76s

(15倍),根据您提供的测试数据。用简单的循环替换np.einsums,结果应该不到一秒钟。(改进后的集成约0.4s,24秒后

k_one_two_third(x)

)

使用numba获得性能

read

. 最新的numba版本(0.39)、IntelSVML包以及fastMath=true等对您的示例产生了很大的影响。

代码

#a bit faster than HYRY's version

@nb.njit(parallel=True,fastmath=True,error_model='numpy')

def k_one_two_third(x):

one=np.empty(x.shape,dtype=x.dtype)

two=np.empty(x.shape,dtype=x.dtype)

for i in nb.prange(x.shape[0]):

for j in range(x.shape[1]):

for k in range(x.shape[2]):

x0 = x[i,j,k] ** (1/3)

x1 = np.exp(-x[i,j,k] ** 2)

x2 = np.exp(-x[i,j,k])

one[i,j,k] = (2*x1/x0 + 4*x2/(x[i,j,k]**(1/6)*(x0 + 1)))**2

two[i,j,k] = (2*x[i,j,k]**(5/2)*x2/(x[i,j,k]**3 + 6) + x1/x[i,j,k]**(2/3))**2

return one, two

#improved integration

@nb.njit(fastmath=True)

def simpson_nb(y_in,dx):

s = y[0]+y[-1]

n=y.shape[0]//2

for i in range(n-1):

s += 4.*y[i*2+1]

s += 2.*y[i*2+2]

s += 4*y[(n-1)*2+1]

return(dx/ 3.)*s

@nb.jit(fastmath=True)

def spectrum(freq_c, number_bin, frequency, gamma, theta):

theta_gamma_factor = np.einsum('i,j->ij', theta**2, gamma**2)

theta_gamma_factor += 1.

t_g_bessel_factor = 1.-1./theta_gamma_factor

number = np.concatenate((number_bin, np.zeros((number_bin.shape[0], 1), dtype=number_bin.dtype)), axis=1)

number_theta_gamma = np.einsum('jk, ik->ijk', theta_gamma_factor**2*1./gamma**3, number)

final = np.empty((np.size(frequency),np.size(freq_c[:,0]), np.size(theta)))

#assume that dx is const. on integration

#speedimprovement of the scipy.simps is about 4x

#numba version to scipy.simps(y,x) is about 60x

dx=gamma[1]-gamma[0]

for i in range(np.size(frequency)):

b_n_omega_theta_gamma = frequency[i]**2*number_theta_gamma

eta = theta_gamma_factor**(1.5)*frequency[i]/2.

eta = np.einsum('jk, ik->ijk', eta, 1./freq_c)

one,two=k_one_two_third(eta)

bessel_eta = np.einsum('jl, ijl->ijl',t_g_bessel_factor, one)

bessel_eta += two

integrand = np.multiply(bessel_eta, b_n_omega_theta_gamma, out= bessel_eta)

#reorder array

for j in range(integrand.shape[0]):

for k in range(integrand.shape[1]):

final[i,j, k] = simpson_nb(integrand[j,k,:],dx)

return final