提高python dblquad和多处理的速度

时间:2019-06-01 08:06:52

标签: python numpy scipy multiprocessing ray

这是我正在运行的示例代码,它生成维度矩阵(size x size)。矩阵被发送到FFT(经过多次迭代),其范数是必需的结果。由于这是一次测试,因此将size = 256和迭代次数(zaxis)设置为3。目前每个矩阵需要1-2分钟的处理时间。

实际的生产运行需要:512 x 512, 1024 x 1024(或更多)的矩阵,每个矩阵约有25次迭代,我想知道是否可以加快此python脚本的速度。

  • 简而言之 ,我生成了一个复杂矩阵=>在循环中逐个元素地分配非零值=>发送至FFT =>计算规范=>将规范保存在数组中。效果很好!

  • 繁重的工作在下面的代码中执行,其中非零值计算为val。在此,分别针对实部和虚部计算2D积分。理想情况下,我应该能够在多个内核上执行。 (*尽管我认为如果可以将不同的非零矩阵元素的分配完全卸载到多个内核,这将非常有效。我没有多处理经验。系统规格: 1700X AMD,8核,运行Python3,Win 10的32GB RAM;或者Ubuntu系统也提供12核,64 GB的RAM)

    if (r < (dim/2) ):
        c1 = fy(a,r,x,y)
        c2 = r*complexAmp(a,b,x,y)
        re = I_real ( r ) # double integral, real
        im = I_imag ( r ) # double integral, imaginary
        val=c1 * c2 *(re[0]+1j*im[0]) 

所以,我的问题。有没有什么好方法可以提高此类操作的速度(希望我可以学习更多有关python高效编程的信息)。目前,我正在检查raymultiprocessing

以下是完整输入脚本。输出显示在底部。

    import time
    import math
    import cmath as cm
    import numpy as np
    import matplotlib.pyplot as plt
    from scipy.integrate import dblquad

    #-----------------------------------------------------------
    # DEFINE FUNCTIONS

    def fy(a,b,x,y):
        return (a*b**3+(x/2.5)+y)/50

    def complexAmp(a,b,x,y):
        return ( (cm.exp(-1j*x*y*cm.sqrt(a+b)))/ cm.sqrt( a ) ) *b

    def wrap(r,  rho, phi):
        return cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 \
                     + r**2))/cm.sqrt(rho**2 + r**2)

    def wrap_real(r,  rho, phi):
        res = cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 +\
                    r**2))/cm.sqrt(rho**2 + r**2)
        return res.real

    def wrap_imag(r,  rho, phi):
        res = cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 + \
                    r**2))/cm.sqrt(rho**2 + r**2)
        return res.imag

    rMax = 5

    def I_real (value ):
        return dblquad(lambda rho, phi: wrap_real (value,  rho, phi) \
                       , 0, rMax, lambda x: 0, lambda x: 2*math.pi)

    def I_imag (value ):
        return dblquad(lambda rho, phi: wrap_imag (value,  rho, phi) ,\
                       0, rMax, lambda x: 0, lambda x: 2*math.pi)
    #-----------------------------------------------------------

    # TEST INTEGRATION
    print("\n-----------COMPLEX INTEGRATION RESULT ----------")
    print (I_real ( 6 ), I_imag ( 6 ))
    print("--------------------------------------------------")

    # parameters governing size and step of grid
    size=256
    depth=10
    step=1.75    # step of grid
    n2 = 1.45
    theta = math.asin(1.00025/n2)

    # complex matrix to keep data
    inp = np.zeros((size,size) , dtype=complex )

    zaxis = np.arange(-60, -10, 20)
    result = np.zeros(zaxis.shape[0])
    n2 = 1.454
    theta = math.asin(1.00025/n2) # update theta
    dim = 16000
    # The main program -----------------------------------------------

    for z in range(zaxis.shape[0]):
        print ("In the loop {0}".format(z))
        start = time.time()
        for i in range(inp.shape[0]):
            for j in range(inp.shape[1]):
                x = step*(i-(size/2))
                y = step*(j-(size/2))
                r = x**2 + y**2
                #print(i, (i-(size/2)),j, (j-(size/2)) )

                b = r*(  math.sin( 1.00025 /n2)) *math.sqrt(2*r**2)
                a = 250/abs(zaxis[z]-r)
                rMax =  abs(zaxis[z])*math.tan(theta)

                val=0
                if (r < (dim/2) ):
                    c1 = fy(a,r,x,y)
                    c2 = r*complexAmp(a,b,x,y)
                    re = I_real ( r ) # double integral, real
                    im = I_imag ( r ) # double integral, imaginary
                    val=c1 * c2 *(re[0]+1j*im[0])

                inp [i][j] = val # substitue the value to matrix

        end = time.time()
        print("Time taken : {0} sec \n" . format( round(end - start,7 )))

        b = np.fft.fft2(inp)
        result [z] = np.linalg.norm(b)

        # PLOTTING  - in actual run the plot is saved to disk
        title=('z={0}, val={1}'.format(z, zaxis[z]))
        fig, axs = plt.subplots(2, 2,figsize=(10,6))
        fig.suptitle(title)
        axs[0, 0].imshow(inp.real )
        axs[0, 0].set_title('inp.real')
        axs[0, 1].imshow(inp.imag)
        axs[0, 1].set_title('inp.imag')
        axs[1, 0].imshow(b.real)
        axs[1, 0].set_title('FFT real')
        axs[1, 1].imshow(b.imag)
        axs[1, 1].set_title('FFT imag')
        plt.show()

    #----------------------------------------------------------

    plt.figure(2) - in actual run the plot is saved to disk
    ax1 = plt.axes()
    plt.title('magnitude over z')
    plt.plot(zaxis, result, '.-', label="norm of FFT")
    plt.ylabel('magnitude')
    plt.xlabel('z')
    plt.grid(True)
    ax1.minorticks_on()
    plt.show()


输出:


    -----------COMPLEX INTEGRATION RESULT ----------
    (-0.0003079405888916291, 1.0879642638692853e-17) (-0.0007321233659418995, 2.5866160149768244e-17)
    --------------------------------------------------
    In the loop 0
    Time taken : 138.8842542 sec 
    [plot]

    In the loop 1
    Time taken : 134.3815458 sec 
    [plot]

    In the loop 2
    Time taken : 56.848331 sec 
    [plot]

    [plot]


1 个答案:

答案 0 :(得分:1)

使用Ray,我能够大大提高上述脚本的速度。现在以并行方式求解双积分。

下面是时间的比较。

Time in seconds
+-------+-----------------+-------------------+
| loop  | serial version  | parallel with Ray |
+-------+-----------------+-------------------+
|     0 |           138.8 |            34.391 |
|     1 |           134.3 |            34.303 |
|     2 |           56.84 |            32.647 |
+-------+-----------------+-------------------+

以下是更新的脚本。

from sys import exit
import time
import math
import cmath as cm
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import dblquad
import ray

ray.init(num_cpus=6) # initializing ray here

#-----------------------------------------------------------

# DEFINE FUNCTIONS : 

def fy(a,b,x,y):
    return (a*b**3+(x/2.5)+y)/50

def complexAmp(a,b,x,y):
    return ( (cm.exp(-1j*x*y*cm.sqrt(a+b)))/ cm.sqrt( a ) ) *b

def wrap(r,  rho, phi):
    return cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 \
                 + r**2))/cm.sqrt(rho**2 + r**2)


def wrap_real(r,  rho, phi):
    res = cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 +\
                r**2))/cm.sqrt(rho**2 + r**2)
    return res.real


def wrap_imag(r,  rho, phi):
    res = cm.cos(phi)*cm.exp(-1j*2*math.pi*cm.sqrt(rho**2 + \
                r**2))/cm.sqrt(rho**2 + r**2)
    return res.imag


rMax = 5
def I_real (value ):
    return dblquad(lambda rho, phi: wrap_real (value,  rho, phi) \
                   , 0, rMax, lambda x: 0, lambda x: 2*math.pi)

def I_imag (value ):
    return dblquad(lambda rho, phi: wrap_imag (value,  rho, phi) ,\
                   0, rMax, lambda x: 0, lambda x: 2*math.pi)

############################################################

# DEFINE RAY FUNCTIONS

@ray.remote
def I_real_mod (value ):
    out= dblquad(lambda rho, phi: wrap_real (value,  rho, phi) \
                   , 0, rMax, lambda x: 0, lambda x: 2*math.pi)
    return out[0]
#-----------------------------------------------------------
@ray.remote
def I_imag_mod (value ):
    out= dblquad(lambda rho, phi: wrap_imag (value,  rho, phi) ,\
                   0, rMax, lambda x: 0, lambda x: 2*math.pi)
    return out[0]

#-----------------------------------------------------------
@ray.remote
def compute_integral( v ):
    i1 = I_real_mod.remote( v )
    i2 = I_imag_mod.remote( v )
    result_all = ray.get([i1, i2])
    return (result_all[0]+1j*result_all[1])
#-----------------------------------------------------------

# TEST INTEGRATION

print("\n-----------COMPLEX INTEGRATION RESULT  : SERIAL ----------")
print (I_real ( 6 ), I_imag ( 6 ))

print("--------------------------------------------------")

print("\n-----------COMPLEX INTEGRATION RESULT  : PARALLEL with RAY ----------")
v1=compute_integral.remote( 6 )
print(ray.get(v1))

print("--------------------------------------------------")

#exit(0)

# parameters governing size and step of grid
size=256
depth=10
step=1.75    # step of grid
n2 = 1.45
theta = math.asin(1.00025/n2)

# complex matrix to keep data
inp = np.zeros((size,size) , dtype=complex )

zaxis = np.arange(-60, -10, 20)
result = np.zeros(zaxis.shape[0])

n2 = 1.454
theta = math.asin(1.00025/n2)
dim = 16000

# The main program -----------------------------------------------

for z in range(zaxis.shape[0]):
    print ("In the loop {0}".format(z))
    start = time.time()
    for i in range(inp.shape[0]):
        for j in range(inp.shape[1]):
            x = step*(i-(size/2))
            y = step*(j-(size/2))
            r = x**2 + y**2
            #print(i, (i-(size/2)),j, (j-(size/2)) )

            b = r*(  math.sin( 1.00025 /n2)) *math.sqrt(2*r**2)
            a = 250/abs(zaxis[z]-r)
            rMax =  abs(zaxis[z])*math.tan(theta)

            val=0

            if (r < (dim/2) ):
                c1 = fy(a,r,x,y)
                c2 = r*complexAmp(a,b,x,y)
                o1 = compute_integral.remote( r ) # using RAY decorated integral here
                val=c1 * c2 *(ray.get(o1))
            inp [i][j] = val # substitue the value to matrix

    end = time.time()
    print("Time taken : {0} sec \n" . format( round(end - start,7 )))

    b = np.fft.fft2(inp)
    result [z] = np.linalg.norm(b)

#----------------------------------------------------------