有没有什么好方法可以优化这个python代码的速度？

Question

我有一段代码,它基本上评估了一些数值表达式,并使用它来集成某些范围的值.当前的一段代码在大约内部运行8.6 s,但我只是使用模拟值,而我的实际数组要大得多.特别是,我的实际大小freq_c= (3800, 101)和大小number_bin = (3800, 100),使得下面的代码真的效率低下,因为实际数组的总执行时间将接近9分钟.代码的一部分是非常缓慢的评估k_one_third和k_two_third我已经使用numexpr.evaluate("..")过的代码,它将代码加速了大约10-20%.但是,我已经避免numexpr在下面,所以任何人都可以运行它而无需安装包.还有其他方法可以提高此代码的速度吗？一些因素的改进也足够好.请注意for loop,由于内存问题,几乎不可避免,因为数组非常大,我通过循环一次操作每个轴.我也想知道numba jit这里是否可以进行优化.

import numpy as np
import scipy 
from scipy.integrate import simps as simps
import time

def k_one_third(x):
    return (2.*np.exp(-x**2)/x**(1/3) + 4./x**(1/6)*np.exp(-x)/(1+x**(1/3)))**2

def k_two_third(x):
    return (np.exp(-x**2)/x**(2/3) + 2.*x**(5/2)*np.exp(-x)/(6.+x**3))**2

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.zeros((np.size(freq_c[:,0]), np.size(theta), np.size(frequency)))
    for i in xrange(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)
        bessel_eta = np.einsum('jl, ijl->ijl',t_g_bessel_factor, k_one_third(eta))
        bessel_eta += k_two_third(eta)
        eta = None
        integrand = np.multiply(bessel_eta, b_n_omega_theta_gamma, out= bessel_eta)
        final[:,:, i] = simps(integrand, gamma)
        integrand = None
    return final

frequency = np.linspace(1, 100, 100)
theta = np.linspace(1, 3, 100)
gamma = np.linspace(2, 200, 101)
freq_c = np.random.randint(1, 200, size=(50, 101))
number_bin = np.random.randint(1, 100, size=(50, 100))
time1 = time.time()
spectra = spectrum(freq_c, number_bin, frequency, gamma, theta)
print(time.time()-time1)

Answer 1

我评测的代码,发现k_one_third()和k_two_third()缓慢.这两个函数中有一些重复的计算.

通过将两个函数合并为一个函数并用它进行装饰@numba.jit(parallel=True),我获得了4倍的加速.

@jit(parallel=True)
def k_one_two_third(x):
    x0 = x ** (1/3)
    x1 = np.exp(-x ** 2)
    x2 = np.exp(-x)
    one = (2*x1/x0 + 4*x2/(x**(1/6)*(x0 + 1)))**2
    two = (2*x**(5/2)*x2/(x**3 + 6) + x1/x**(2/3))**2
    return one, two