python - 提高 python dblquad 和多处理的速度

Question

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

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

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

• 繁重的工作在以下代码中执行，其中非零值计算为 val。这里，二维积分是分别计算实部和虚部的。理想情况下，我应该能够在多个核心上执行此操作。 (*尽管我认为如果不同的非零矩阵元素的分配可以完全卸载到多个核心，那将非常高效。我在多处理方面没有经验。系统规范: 1700X AMD，8 核，32GB RAM 运行 Python3，Win 10；或者 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 中高效编程的知识)。现在我正在检查ray和 multiprocessing 。

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

    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)

<小时/>

输出:

    -----------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]

<小时/>

Answer 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)

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