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

Question

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

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

例子	经过时间
没有黄昏	8.02 秒
达斯克	8.74 秒

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()

标签： pythondask

Answer 1

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

例子	经过时间
没有黄昏	8.57 秒
达斯克	3.83 秒

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)