如何在Python中为多组e0参数实现多核并行计算？

问题说明 #

现有量子化学计算代码，当前固定e0=15，需要针对e0={7,10,15,20,28}多组参数完成计算，利用多核处理器将每组e0的计算任务分配到独立核心执行。

实现方案 #

采用concurrent.futures.ProcessPoolExecutor实现进程级并行——这类计算属于CPU密集型任务，进程池可以绕过Python的GIL限制，充分利用多核资源。核心思路是将单组e0的计算逻辑封装为独立函数，再通过进程池批量调度任务。

关键修改点 #

• 将原代码的计算逻辑封装为calculate_e0函数，接收e0作为参数，返回该参数下的计算结果EnB。
• 消除全局变量依赖，将einf作为函数内固定值（保持原代码的einf=15），避免进程间变量干扰。
• 使用进程池提交所有e0任务，异步获取并整理结果。

修改后的完整代码 #

import sympy as sp
import scipy as sc
import numpy as np
from concurrent.futures import ProcessPoolExecutor

def Psi(r, n):
    return 2 * np.exp(-r/n) * np.sqrt(sc.special.factorial(n)/sc.special.factorial(-1+n)) * sc.special.hyp1f1(1-n, 2, 2*r/n) / n**2

def PsiSymb(n):
    r = sp.symbols('r')
    y1 = 2 * sp.exp(-r/n) * np.sqrt(sc.special.factorial(n)/sc.special.factorial(-1+n)) / n**2
    y2 = sp.simplify(sp.functions.special.hyper.hyper([1-n], [2], 2*r/n))
    return y1 * y2

def LaplacianPsi(n):
    r = sp.symbols('r')
    ydiff = 2/r * PsiSymb(n).diff(r) + PsiSymb(n).diff(r, 2)
    return sp.lambdify(r, ydiff, "numpy")

def k(n1, n2):
    yint = sc.integrate.quad(lambda r: -0.5 * Psi(r, n2) * LaplacianPsi(n1)(r) * r**2, 0, np.inf)
    return yint[0]

def p(n1, n2, e0, einf):
    potC = sc.integrate.quad(lambda r: Psi(r, n2) * (-1/r) * Psi(r, n1) * (r**2), 0, np.inf)
    potB1 = sc.integrate.quad(lambda r: Psi(r, n2) * (1/einf - 1/e0) * ((einf/e0)**(3/5)) * (-e0/(2*r)) * (np.exp(-r*2.23)) * Psi(r, n1) * (r**2), 0, np.inf)
    potB2 = sc.integrate.quad(lambda r: Psi(r, n2) * (1/einf - 1/e0) * ((einf/e0)**(3/5)) * (-e0/(2*r)) * (np.exp(-r*2.4)) * Psi(r, n1) * (r**2), 0, np.inf)
    return potC[0] + potB1[0] + potB2[0]

def en(n1, n2, e0, einf):
    return k(n1, n2) + p(n1, n2, e0, einf)

def calculate_e0(e0_value):
    einf = 15
    nmax = 3
    EnM = [[0]*nmax for _ in range(nmax)]
    for n1 in range(nmax):
        for n2 in range(nmax):
            EnM[n2][n1] = en(n1+1, n2+1, e0_value, einf)
    EnEig = sc.linalg.eigvalsh(EnM)
    return e0_value, min(EnEig)

if __name__ == "__main__":
    e0_list = [7, 10, 15, 20, 28]
    # 进程池大小默认匹配CPU核心数，也可手动指定max_workers参数
    with ProcessPoolExecutor() as executor:
        results = executor.map(calculate_e0, e0_list)
    
    # 格式化输出结果
    for e0_val, enb in results:
        print(f"e0={e0_val} 时，EnB={enb:.6f}")

代码说明 #

1. 函数封装：calculate_e0函数负责单组e0的完整计算流程，从矩阵构建到特征值求解，最终返回e0值和对应的结果，保证每个任务独立可执行。
2. 进程池调度：ProcessPoolExecutor自动管理进程创建和任务分配，executor.map会将e0_list中的每个元素传入calculate_e0并并行执行，无需手动管理进程生命周期。
3. 结果处理：遍历返回的结果，按e0值格式化输出对应的计算结果，方便后续对比分析。

内容的提问来源于stack exchange，提问作者Mam Mam