Как сделать мой код более эффективным с помощью Cython? - PullRequest
Новая
       20

Как сделать мой код более эффективным с помощью Cython?

0 голосов
/ 22 декабря 2019

Я написал следующий код на Python. Если я использую этот код, это займет 1 минуту 20 секунд: 

def calc_cy(searchX,searchY,searchZ,dx,dz,r_vector,p_true,a_factor,sim):

calc1 = np.full((len(searchX),len(searchY),len(searchZ)),np.nan)
z_t = np.abs(np.ceil(r_vector[:,2]/dz).squeeze())

for ix in searchX:
    for iy in searchY:
        for iz in searchZ:
            dt=[]
            ir = (np.floor((((ix-r_vector[:,0])**2+(iy-r_vector[:,1])**2))**0.5)) 
            for ip in range(p_true.size):
                dt.append(sim[int(iz)][int(ir[ip]),int(z_t[ip])])

            dt_sim = dt-min(dt)
            dt_true=p_true-min(p_true[0])
            calc1[int(ix),int(iy),int(iz)]=np.linalg.norm(a_factor*(dt_sim-dt_true))
return calc1

Если я использую этот код с Cython, это займет 1 минуту 5 секунд:

код установки: 

from distutils.core import setup
from Cython.Build import cythonize
import numpy

setup(ext_modules = cythonize("calc_cy.pyx"), include_dirs = [numpy.get_include()])

Код с использованием Cython: 

import numpy as np
cimport numpy as np
import math

DTYPE = np.double
ctypedef np.double_t DTYPE_t

def calc_cy(np.ndarray [DTYPE_t,ndim = 3] sim,np.ndarray [DTYPE_t,ndim = 1] searchX,np.ndarray [DTYPE_t,ndim = 1] searchY,np.ndarray [DTYPE_t,ndim = 1] searchZ,np.ndarray [DTYPE_t,ndim = 2] r_vector,np.ndarray [DTYPE_t,ndim = 1] p_true,np.ndarray [DTYPE_t,ndim = 1] a_factor):

    cdef np.ndarray[DTYPE_t,ndim = 3] calc1 = np.zeros((len(searchX),len(searchY),len(searchZ)),dtype = DTYPE)
    cdef np.ndarray[DTYPE_t,ndim = 1] dt2 = np.zeros(len(p_true),dtype = DTYPE)
    cdef np.ndarray[DTYPE_t,ndim = 1] dt_sim = np.zeros(len(dt2),dtype = DTYPE)
    cdef np.ndarray[DTYPE_t,ndim = 1] dt_picks = np.zeros(len(p_true),dtype = DTYPE)

    cdef int ir
    cdef int k
    cdef int m

    for ix in range(len(searchX)):
        for iy in range(len(searchY)):
            for iz in range(len(searchZ)):
                for ip in range(len(p_true)):
                    ir= int(np.floor((((searchX[ix]-r_vector[ip,0])**2+(searchY[iy]-r_vector[ip,1])**2))**0.5))
                    k = int(searchZ[iz])
                    m = int(math.fabs(math.ceil(r_vector[ip,2])))
                    dt2[int(ip)] = sim[k,ir,m]
                dt_sim = dt2- min(dt2)
                dt_true = p_true - min(p_true)
                calc1[ix,iy,iz] = np.linalg.norm(a_factor*(dt_sim-dt_true))
    return calc1

Как я могу улучшить свой код и сделать его более эффективным? Спасибо!

1 Ответ

0 голосов
/ 22 декабря 2019

Глядя на ваш код, значение ir не зависит от iz, поэтому вы можете переместить его за пределы цикла iz. Значение dt_true зависит только от p_true, поэтому его нужно вычислять только один раз . Все это в совокупности устранит ненужные пересчеты и ускорит код. (Это называется движение инвариантного кода цикла .)

Кроме того, вы можете преобразовать вычисление dt в понимание списка. Это должно ускорить версию Python. Не уверен, что это поможет версии Cython.

Это сделает код похожим на это: 

def calc_cy(searchX,searchY,searchZ,dx,dz,r_vector,p_true,a_factor,sim):
    calc1 = np.full((len(searchX),len(searchY),len(searchZ)),np.nan)
    z_t = np.abs(np.ceil(r_vector[:,2]/dz).squeeze())
    dt_true = p_true - min(p_true[0])
    for ix in searchX:
        for iy in searchY:
            ir = np.floor(((ix-r_vector[:,0])**2+(iy-r_vector[:,1])**2)**0.5)
            for iz in searchZ:
                dt = [sim[int(iz)][int(ir[ip]),int(z_t[ip])] for ip in range(p_true.size)]
                dt_sim = dt-min(dt)
                calc1[int(ix),int(iy),int(iz)]=np.linalg.norm(a_factor*(dt_sim-dt_true))
    return calc1

Предполагая, что у вас многоядерный процессор, вы можете генерировать кортеживсе возможные комбинации searchX и searchY. Затем, используя multiprocessing.Pool или concurrent.futures.ProcessPoolExecutor, вы можете разделить вычисления внутри цикла searchY на все ваши ядра, поскольку они независимы. Если у вас есть n ядер, это должно сократить время работы примерно в n .

...