提高python dblquad和多处理的速度

使用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)

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