Есть ли хороший способ оптимизировать скорость этого кода 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

Как сказано в комментариях, большие части кода должны быть переписаны для достижения максимальной производительности.

Я только изменил интеграцию Симпсона и немного изменил ответ @HYRY.Это ускоряет вычисление с 26,15 до 1,76 с (15x) по предоставленным вами данным испытаний.При замене np.einsums простыми циклами это должно закончиться менее чем за секунду.(Приблизительно 0,4 с из улучшенной интеграции, 24 с из k_one_two_third(x))

Чтобы получить производительность с помощью Numba , прочитайте .Последняя версия Numba (0.39), пакет Intel SVML и такие вещи, как 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

Answer 2

Я профилировал код и обнаружил, что 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