如何将此函数上的多线程用于值的np.meshgrid

Question

以下代码生成指定间隔的r和E值的numpy 2D列表。

r = np.linspace(3, 14, 10)
E = np.linspace(0.05, 0.75, 10)
r, E = np.meshgrid(r, E)

然后，我使用以下嵌套循环为每个r和E区间值生成函数电离伽马的输出。

for ridx in trange(len(r)):
z = []
for cidx in range(len(r[ridx])):
z.append(ionisationGamma(r[ridx][cidx], E[ridx][cidx]))
Z.append(z)
Z = np.array(Z)

这个循环给了我一个2D numpy数组Z，这是我的输出，我正在将其用于3D图。它的问题是：生成所有这些间隔的输出大约需要6个小时，因为np.meshgrid导致有很多值。我刚刚在Python中发现了多线程，并想知道如何使用它来实现这一点。任何帮助都将不胜感激。

参见以下ionisationGamma代码

def ionisationGamma(r, E):
I = complex(0.1, 1.0)
a_soft = 1.0
omega = 0.057
beta = 0.0
dt = 0.1
steps = 10000
Nintervals = 60
N = 3000
xmin = float(-300)
xmax = -xmin
x = [0.0]*N
dx = (xmax - xmin) / (N - 1)
L = dx * N
dk = 2 * M_PI / L
propagator = None
in_, out_, psi0 = None, None, None
in_ = [complex(0.,0.)] * N
psi0 = [complex(0.,0.)] * N
out_ = [[complex(0.,0.)]*N for i in range(steps+1)]
overlap = exp(-r) * (1 + r + (1 / 3) * pow(r, 2))
normC = 1 / (sqrt(2 * (1 + overlap)))
gammai = 0.5
qi = 0.0 + (r / 2)
pi = 0.0
gammai1 = 0.5
gammai2 = 0.5
qi1 = 0.0 - (r / 2)
qi2 = 0.0 + (r / 2)
pi1 = 0.0
pi2 = 0.0

# split initial wavepacket
for i in range(N):
x[i] = xmin + i * dx
out_[0][i] = (normC) * ((pow(gammai1 / M_PI, 1. / 4.) * exp(complex(-(gammai1 / 2.) * pow(x[i] - qi1, 2.), pi1 * (x[i] - qi1)))) + (pow(gammai2 / M_PI, 1. / 4.) * exp(complex(-(gammai2 / 2.) * pow(x[i] - qi2, 2.), pi2 * (x[i] - qi2)))))
in_[i] = (normC) * ((pow(gammai1 / M_PI, 1. / 4.) * exp(complex(-(gammai1 / 2.) * pow(x[i] - qi1, 2.), pi1 * (x[i] - qi1)))) + (pow(gammai2 / M_PI, 1. / 4.) * exp(complex(-(gammai2 / 2.) * pow(x[i] - qi2, 2.), pi2 * (x[i] - qi2)))))
psi0[i] = in_[i]


for l in range(1, steps+1):
for i in range(N):
propagator = exp(complex(0, -potential(x[i], omega, beta, a_soft, r, E, dt, l) * dt / 2.))
in_[i] = propagator * in_[i];
in_ = np.fft.fft(in_, N)
for i in range(N):
k = dk * float(i if i < N / 2 else i - N)
propagator = exp(complex(0, -dt * pow(k, 2) / (2.)))
in_[i] = propagator * in_[i]
in_ = np.fft.ifft(in_, N)
for i in range(N):
propagator = exp(complex(0, -potential(x[i], omega, beta, a_soft, r, E, dt, l) * dt / 2.))
in_[i] = propagator * in_[i]
out_[l][i] = in_[i]

initialGammaCentre = 0.0
finalGammaCentre = 0.0
for i in range(500, 2500 +1):
initialGammaCentre += pow(abs(out_[0][i]), 2) * dx
finalGammaCentre += pow(abs(out_[steps][i]), 2) * dx

ionisationGamma = finalGammaCentre / initialGammaCentre

return ionisationGamma
def potential(x, omega, beta, a_soft, r, E, dt, l):
V = (-1. / sqrt((x - (r / 2)) * (x - (r / 2)) + a_soft * a_soft)) + ((-1. / sqrt((x + (r / 2)) * (x + (r / 2)) + a_soft * a_soft))) + E * x
return V

Answer 1

由于问题是关于如何使用多处理，因此以下代码将起作用：

import multiprocessing as mp

if __name__ == '__main__':
with mp.Pool(processes=16) as pool:
Z = pool.starmap(ionisationGamma, arguments)
Z = np.array(Z)

论点在哪里：

arguments = list()
for ridx in range(len(r)):
for cidx in range(len(r[ridx])):
arguments.append((r[ridx][cidx], E[ridx][cidx]))

我使用starmap而不是map，因为您有多个要解包的参数。这将使用ionisationGamma函数划分可在多个核心上迭代的参数，并对最终结果进行排序。

然而，我确实觉得有必要说，主要的解决方案并不是真正的多处理，而是原始的函数代码。在ionisationGamma中，您多次使用慢速pythonfor循环。如果您能够将这些操作向量化，那么您的代码将受益匪浅。

第二个观察结果是，你正在单独使用其中的许多循环，如果你能把一个大函数分成多个小函数，那就太好了。然后你可以对每个函数单独计时，并加快那些太慢的函数。