Tip

An interactive online version of this notebook is available, which can be accessed via Open this notebook in Google Colab

Alternatively, you may download this notebook and run it offline.

Attention

You are viewing this notebook on the latest version of the documentation, where these notebooks may not be compatible with the stable release of PyBaMM since they can contain features that are not yet released. We recommend viewing these notebooks from the stable version of the documentation. To install the latest version of PyBaMM that is compatible with the latest notebooks, build PyBaMM from source.

Running many simulations in parallel using OpenMP#

[1]:

%pip install "pybamm[plot,cite]" -q    # install PyBaMM if it is not installed
import pybamm
import numpy as np
import time

Note: you may need to restart the kernel to use updated packages.

Workflows such as parameter sweeps require running many identical simulations with different input parameters. PyBaMM provides a way to run many simulations in parallel using OpenMP using the IDAKLU solver. Looking at the API docs for this solver, we see a num_threads argument that can be used to set the number of threads to use when creating the solver, like so:

[2]:

solver = pybamm.IDAKLUSolver(options={"num_threads": 2})

This option will have no effect if you try to solve a single simulation, but if use an input parameter in your model and pass in a list of values for that parameter, PyBaMM will automatically distribute the simulations across the number of threads you specify. For example,

[3]:

model = pybamm.lithium_ion.DFN()
params = model.default_parameter_values
params["Current function [A]"] = "[input]"
sim = pybamm.Simulation(model, parameter_values=params, solver=solver)
sim.solve(
    [0, 3600], inputs=[{"Current function [A]": 1}, {"Current function [A]": 0.5}]
)

[3]:

[<pybamm.solvers.solution.Solution at 0x7d2a1418dfc0>,
 <pybamm.solvers.solution.Solution at 0x7d2a1418e170>]

The previous code will then run the 2 simulations in parallel on 2 threads. Since we are solving such a small number of simulations, the overhead of parallelization will not be worth it and it is likely that this will be slower than solving in serial. Below we show an example of solving 1000 simulations in parallel, varying the number of threads used to see the effect on the runtime.

[4]:

n = 1e3
current_inputs = [
    {"Current function [A]": current} for current in np.linspace(0, 0.6, int(n))
]
num_threads_list = [1, 2, 4, 8, 16, 32]
for num_threads in reversed(num_threads_list):
    model = pybamm.lithium_ion.DFN()
    params = model.default_parameter_values
    params.update(
        {
            "Current function [A]": "[input]",
        }
    )
    solver = pybamm.IDAKLUSolver(options={"num_threads": num_threads})
    sim = pybamm.Simulation(model, solver=solver, parameter_values=params)
    start_time = time.perf_counter()
    sol = sim.solve([0, 3600], inputs=current_inputs)
    end_time = time.perf_counter()
    print(
        f"Time taken to solve 1000 DFN simulation for {num_threads} threads: {end_time - start_time:.2f} s"
    )

Time taken to solve 1000 DFN simulation for 32 threads: 2.65 s
Time taken to solve 1000 DFN simulation for 16 threads: 3.72 s
Time taken to solve 1000 DFN simulation for 8 threads: 6.07 s
Time taken to solve 1000 DFN simulation for 4 threads: 10.11 s
Time taken to solve 1000 DFN simulation for 2 threads: 17.73 s
Time taken to solve 1000 DFN simulation for 1 threads: 26.46 s

You can see that the speed-up from using more threads starts to diminish after a certain point (about 20-30 threads), and this effect will be more pronounced for smaller numbers of simulations, or if the simulations are very quick to solve (e.g. an SPM model). Below we show the same example, but this time using the SPM model.

[5]:

n = 1e3
current_inputs = [
    {"Current function [A]": current} for current in np.linspace(0, 0.6, int(n))
]
num_threads_list = [1, 2, 4, 8, 16, 32]
for num_threads in reversed(num_threads_list):
    model = pybamm.lithium_ion.SPM()
    params = model.default_parameter_values
    params.update(
        {
            "Current function [A]": "[input]",
        }
    )
    solver = pybamm.IDAKLUSolver(options={"num_threads": num_threads})
    sim = pybamm.Simulation(model, solver=solver, parameter_values=params)
    start_time = time.perf_counter()
    sol = sim.solve([0, 3600], inputs=current_inputs)
    end_time = time.perf_counter()
    print(
        f"Time taken to solve 1000 SPM simulation for {num_threads} threads: {end_time - start_time:.2f} s"
    )

Time taken to solve 1000 SPM simulation for 32 threads: 0.42 s
Time taken to solve 1000 SPM simulation for 16 threads: 0.70 s
Time taken to solve 1000 SPM simulation for 8 threads: 0.43 s
Time taken to solve 1000 SPM simulation for 4 threads: 0.53 s
Time taken to solve 1000 SPM simulation for 2 threads: 0.90 s
Time taken to solve 1000 SPM simulation for 1 threads: 0.86 s

So in this case the speed-up for using multiple threads to solve 1000 SPM simulations is much less than for the DFN simulations, and above 4 threads no speed-up is observed at all.