Parallel Processing with SharedMemory

[TomTranter]

Summary Doing many 1D models in parallel is a good way to solve my ECM problem and also parametric sweeps.

Motivation The problem when you have hundreds or thousands of tasks becomes sending the data back and forth, especially if you step the solutions a little bit and make adjustments on the fly as I am doing. I think this can be addressed by using shared memory. I found a package that just works on linux (or maybe POSIX systems) which seems to give great performance.

The below script demonstrates stepping 1300 SPMs with slightly different applied currents using shared arrays and a processpoolexecutor. There is also a line commented out that evaluates the terminal voltage. This slows things down a lot and highlights that evals are in need of vectorization

Additional context

import pybamm
import numpy as np
import matplotlib.pyplot as plt
from concurrent.futures import ProcessPoolExecutor
import SharedArray
import time as time_module

plt.close('all')
pybamm.set_logging_level("WARNING")

def current_function(t):
    return pybamm.InputParameter("Current")

if __name__ == '__main__':
    # load model
    model = pybamm.lithium_ion.SPMe()

    # create geometry
    geometry = model.default_geometry

    # load parameter values and process model and geometry
    param = model.default_parameter_values

    param.update({"Current function [A]": current_function,})
    param.update({"Current": "[input]"}, check_already_exists=False)
    param.process_model(model)
    param.process_geometry(geometry)

    # set mesh
    mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)

    # discretise model
    disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
    disc.process_model(model)
    inputs = {"Current": 0.67}
    # solve model
    t_eval = np.linspace(0, 1, 2)
    solver = pybamm.CasadiSolver()
    sol_init = solver.solve(model, t_eval, inputs=inputs)
    i_app = 1.0

    Nspm = 1300
    Nsteps = 10
    try:
        shm_y = SharedArray.create('shm://y', [sol_init.y.shape[0], Nspm], dtype=float)
        shm_t = SharedArray.create('shm://t', [Nspm], dtype=float)
        shm_i_app = SharedArray.create('shm://i_app', [Nspm], dtype=float)
        shm_V = SharedArray.create('shm://V', [Nsteps, Nspm], dtype=float)
    except:
        SharedArray.delete('shm://y')
        SharedArray.delete('shm://t')
        SharedArray.delete('shm://i_app')
        SharedArray.delete('shm://V')
        shm_y = SharedArray.create('shm://y', [sol_init.y.shape[0], Nspm], dtype=float)
        shm_t = SharedArray.create('shm://t', [Nspm], dtype=float)
        shm_i_app = SharedArray.create('shm://i_app', [Nspm], dtype=float)
        shm_V = SharedArray.create('shm://V', [Nsteps, Nspm], dtype=float)

    for i in range(Nspm):
        shm_y[:, i] = sol_init.y[:, -1]
        shm_t[i] = 0.0
        shm_i_app[i] = i_app*(1+(i+1)/Nspm)
    #f = ex.submit(add_one, shm_key)

    # step model
    dt = 1
    time = 0
    end_time = dt*Nsteps
    step_solver = model.default_solver
    step_solution = None

    def shm_step(args):
        ind, tstep = args
        shm_y = SharedArray.attach('shm://y')
        shm_t = SharedArray.attach('shm://t')
        shm_i_app = SharedArray.attach('shm://i_app')
        shm_V = SharedArray.attach('shm://V')
        sol_init.y[:, -1] = shm_y[:, ind]
        sol_init.t[-1] = shm_t[ind]
        inputs = {"Current": shm_i_app[ind]}
        step_solution = step_solver.step(sol_init, model, dt=dt, npts=2,
                                         inputs=inputs, save=False)
        shm_y[:, ind] = step_solution.y[:, -1]
        shm_t[ind] = step_solution.t[-1]
#        shm_V[tstep, ind] = step_solution['Terminal voltage [V]'](step_solution.t[-1])

    step = 0
    while time < end_time:
        print('Time', time)
        st = time_module.time()
        ex = ProcessPoolExecutor()
        ex.map(shm_step, zip(np.arange(Nspm, dtype=int),
                             np.ones(Nspm, dtype=int)*step))
        ex.shutdown(wait=True)
        print(shm_t[:10])
        print(shm_y[0, :10])
        time += dt
        step += 1
        print('Stepping time', np.around(time_module.time()-st, 2), 's')

    plt.figure()
    plt.plot(shm_V)

    SharedArray.delete('shm://y')
    SharedArray.delete('shm://t')
    SharedArray.delete('shm://i_app')
    SharedArray.delete('shm://V')

[valentinsulzer]

@TomTranter, it would be cool to make something like this a supported feature in core pybamm, for parameter sweeps.

I'm thinking we could pass in a list of input dictionaries to the solver, and a number of cores to use (n_cores=-1 to use all available cores), then the solver automatically calls self._integrate, looping over all the input dictionaries in the list in parallel using n_cores, and returns a list of solutions

[TomTranter]

Yes happy to discuss on Monday and work on it next week

[tlestang]

Just having a look into this, presumably the model should be setup and discontinuities calculated for each input dictionaries as well?

[valentinsulzer]

Model setup should be independent of the input values (i.e. you can reuse a setup model even if the input values change). Not sure about the discontinuities, possible that they would need to be recomputed if input parameters appear in the timescale or in the values specified by the discontinuities.

[tlestang]

base_solver.set_up takes the inputs as a keyword argument though? So I suppose in some cases the setup step depends on the inputs? The inputs dict is used in the set_up method when calculating initial conditions

# Process initial conditions
initial_conditions = process(
    model.concatenated_initial_conditions,
    "initial_conditions",
    use_jacobian=False,
)[0]
init_eval = InitialConditions(initial_conditions, model)
# ....
model.y0 = init_eval(inputs)

But I'm having a difficult time understanding if the return value of init_eval does indeed depend on the content of inputs.

This inputs dict is also passed to Symbol.evaluate when evaluating time and length scales, but since these are scalars this has no effect on the return value.

[valentinsulzer]

init_eval is a function which takes the inputs and returns the initial conditions given those inputs. So it should depend on the content of inputs in general, but only the inputs that appear in the initial conditions (if any).

It only gets created once, in setup, and then gets called in _set_initial_conditions whenever the model is solved.

I think the line model.y0 = init_eval(inputs) in setup could possibly be deleted, since the same line appears in _set_initial_conditions but I'm not sure. There is also some complexity introduced by the fact that you can give empty inputs to model.solve when using a casadi solver, in which case it will return a ProcessedSymbolicVariable which can be evaluated with any inputs

[valentinsulzer]

Can this be closed following #1261 ?

[tlestang]

I think yes as the functionality described by @TomTranter was integrated to PyBaMM in #1261 ?