如何使用Dask并行运行PyBaMM电池模拟 解决本地运行慢于串行的问题

核心问题原因 #

你当前的Dask代码没有实际启用并行：未显式初始化dask.distributed.Client时，dask.compute默认使用单线程同步调度器，只是给串行执行额外增加了延迟包装的开销，所以耗时比原生串行更长。

优化步骤 #

1. 显式初始化Dask客户端启用并行 #

修改Dask执行区段的代码，初始化本地客户端，会自动匹配当前机器的CPU核心数调度并行任务：

# 导入Client
from dask.distributed import Client
import matplotlib.pyplot as plt
import pybamm
import time
import dask

# 其他函数保持不变

def main():
    tic = time.perf_counter()

    discharge = [4, 3.5, 3, 2.5, 2, 1.8, 1.5, 1]
    t_eval = [0, 4000]

    # Dask 并行版本
    label = 'Dask'
    # 初始化本地客户端，集群部署时只需修改此处为调度器地址即可
    client = Client()

    lazy_sols = []
    for dis in discharge:
        sol = dask.delayed(run_simulation)(dis, t_eval)
        lazy_sols.append(sol)

    sols = dask.compute(*lazy_sols)
    # 任务结束后可关闭客户端
    client.close()

    # 后续处理逻辑不变

修改后在8核机器上运行耗时会降到1~2秒区间，远快于串行执行的8秒。后续要部署到集群时，只需将Client()改为Client("tcp://集群调度器地址:端口")，其余业务代码完全不用修改。

2. 减少重复计算进一步提速 #

你当前run_simulation函数中，每次调用都会重复构建SPMe模型、加载默认参数，这部分逻辑是所有模拟任务的公共逻辑，可以提取到循环外只执行一次，既降低单任务计算量，也减少Dask的数据传输开销：

# 公共逻辑提前执行，所有任务共享
model = pybamm.lithium_ion.SPMe()
param = model.default_parameter_values
param['Current function [A]'] = '[input]'

def run_simulation(dis, t_eval):
    sim = pybamm.Simulation(model, parameter_values=param)
    sim.solve(t_eval, inputs={'Current function [A]': dis})
    return sim.solution

3. 大规模任务优化（可选） #

如果后续需要同时跑数百上千组模拟，可以将任务拆分为小批量提交，避免Dask调度大量小任务的开销占比过高，进一步提升运行效率。

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