使用 Dask 并行运行 PyBaMM 电池模拟

Question

我正在使用PyBaMM 包对电池进行建模，我想使用 Dask 并行运行多个模拟。下面的例子是我尝试使用dask.delayed。 Dask 方法比非 Dask 方法慢。在这个例子中是否有更好的方法来使用 Dask？我应该设置一个 Dask Client() 来并行运行模拟吗？我在我的本地机器上运行这个例子，但我最终想在集群上运行一个类似的例子。

在 8 核 MacBook Pro 上运行示例所经过的时间如下所示。注释掉 main() 中的相应部分，以便在有或没有 Dask 的情况下运行。

Example	Elapsed time
No Dask	8.02 seconds
Dask	8.74 seconds

import matplotlib.pyplot as plt
import pybamm
import time
import dask

def generate_plots(discharge, t, capacity, current, voltage):

    def styleplot(ax):
        ax.legend(loc='best')
        ax.grid(color='0.9')
        ax.set_frame_on(False)
        ax.tick_params(color='0.9')

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(t[i], current[i], label=f'{discharge[i]} A')
    ax.set_xlabel('Time [s]')
    ax.set_ylabel('Current [A]')
    styleplot(ax)

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(t[i], voltage[i], label=f'{discharge[i]} A')
    ax.set_xlabel('Time [s]')
    ax.set_ylabel('Terminal voltage [V]')
    styleplot(ax)

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(capacity[i], voltage[i], label=f'{discharge[i]} A')
    ax.set_xlabel('Discharge capacity [Ah]')
    ax.set_ylabel('Terminal voltage [V]')
    styleplot(ax)

    plt.show()

def run_simulation(dis, t_eval):

    model = pybamm.lithium_ion.SPMe()

    param = model.default_parameter_values
    param['Current function [A]'] = '[input]'

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

    return sim.solution

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

    discharge = [4, 3.5, 3, 2.5, 2, 1.8, 1.5, 1]  # discharge currents [A]
    t_eval = [0, 4000]                            # evaluation time [s]

    # No Dask
    # ------------------------------------------------------------------------

    label = 'no Dask'

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

    # Dask
    # ------------------------------------------------------------------------

    # label = 'Dask'

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

    # sols = dask.compute(*lazy_sols)

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

    t = []
    capacity = []
    current = []
    voltage = []

    for sol in sols:
        t.append(sol['Time [s]'].entries)
        capacity.append(sol['Discharge capacity [A.h]'].entries)
        current.append(sol['Current [A]'].entries)
        voltage.append(sol["Terminal voltage [V]"].entries)

    toc = time.perf_counter()
    print(f'Elapsed time ({label}) = {toc - tic:.2f} s')

    generate_plots(discharge, t, capacity, current, voltage)

if __name__ == '__main__':
    main()

Answer 1

根据@rrpelgrim 的建议，我实现了一个Client() 对象，它似乎通过使用分布式调度程序改进了我的示例代码的并行执行。修改后的示例如下所示。您可以通过注释掉 main() 中的相应部分来比较使用和不使用 Dask 所经过的时间。表中给出了使用 8 核 CPU 所经过的时间。

Example	Elapsed time
No Dask	8.57 seconds
Dask	3.83 seconds

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

def create_plots(discharge, t, capacity, current, voltage):

    def styleplot(ax, xlabel, ylabel):
        ax.legend(loc='best')
        ax.grid(color='0.9')
        ax.set_frame_on(False)
        ax.tick_params(color='0.9')
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(t[i], current[i], label=f'{discharge[i]} A')
    styleplot(ax, xlabel='Time [s]', ylabel='Current [A]')

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(t[i], voltage[i], label=f'{discharge[i]} A')
    styleplot(ax, xlabel='Time [s]', ylabel='Terminal voltage [V]')

    _, ax = plt.subplots(tight_layout=True)
    for i in range(len(discharge)):
        ax.plot(capacity[i], voltage[i], label=f'{discharge[i]} A')
    styleplot(ax, xlabel='Discharge capacity [Ah]', ylabel='Terminal voltage [V]')

    plt.show()

def run_simulation(dis, t_eval):

    model = pybamm.lithium_ion.SPMe()

    param = model.default_parameter_values
    param['Current function [A]'] = '[input]'

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

    return sim.solution

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

    discharge = [4, 3.5, 3, 2.5, 2, 1.8, 1.5, 1]  # discharge currents [A]
    t_eval = [0, 4000]                            # evaluation time [s]

    # No Dask
    # ------------------------------------------------------------------------

    # label = 'no Dask'

    # sols = []
    # for dis in discharge:
    #     sol = run_simulation(dis, t_eval)
    #     sols.append(sol)

    # Dask
    # ------------------------------------------------------------------------

    label = 'Dask'

    lazy_sols = client.map(run_simulation, discharge, t_eval=t_eval)
    sols = client.gather(lazy_sols)

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

    t = []
    capacity = []
    current = []
    voltage = []

    for sol in sols:
        t.append(sol['Time [s]'].entries)
        capacity.append(sol['Discharge capacity [A.h]'].entries)
        current.append(sol['Current [A]'].entries)
        voltage.append(sol["Terminal voltage [V]"].entries)

    toc = time.perf_counter()
    print(f'Elapsed time ({label}) = {toc - tic:.2f} s')

    create_plots(discharge, t, capacity, current, voltage)

if __name__ == '__main__':
    client = Client()
    print(client)
    main(client)