.. _tutorial_cosimulation: How to use MPI to generate coupled simulations at different scales ================================================================== In this tutorial we are going to discuss how to use Arbor and MPI to drive simulations at different scales or levels of abstraction. This *requires* an MPI-enabled build of Arbor. This allows you to build compound simulations that focus effort on detailed models (in Arbor) where needed and use more abstract models to embed the more complex structures. In general, throughout this tutorial, we are assuming you are comfortable with the basics of Arbor (cells, recipes, and networks), Python package management, and MPI, as wells as installing software on your system. All source code for all intermediate steps can be in the directory `python/example/cosim `__ of the Arbor source tree. Setup ----- Since we need an MPI-capable build, we need to install from source, which is simple, but might require some guidance. We demonstrate it using the ``uv`` package manager. Before starting, you should install MPI using your system's method, e.g. using ``brew install`` (MacOS) or ``apt install`` (Ubuntu), CMake, and a C++ compiler. If these are missing, you will encounter errors during the build step. Begin by setting up the environment .. code-block:: uv venv source .venv/bin/activate uv pip install numpy mpi4py polars scikit-build-core matplotlib seaborn scipy pybind11-stubgen then install Arbor .. code-block:: CMAKE_ARGS="-DPYTHON_EXECUTABLE=`which python3` -DARB_VECTORIZE=ON -DARB_ARCH=native -DARB_WITH_MPI=ON" uv pip install --no-cache-dir --no-binary arbor --no-build-isolation arbor This will install a vectorised (making better use of modern CPUs) and MPI-enabled version of Arbor. You can test it using a file .. literalinclude:: ../../python/example/cosim/step-01.py :language: python which should print .. code-block:: python { 'mpi': True, 'mpi4py': True, 'gpu': None, 'vectorize': True, # ... more fields elided } Overview -------- We will create two models, one a spiking neural network in Arbor, the second a simple neural mass model, and then connect them over an interface. Both models will be basic, to emphasise the scaffolding, but can either be iterate upon or swapped for more interesting alternatives without touching the framework. Then, we will build the scaffolding itself, based on MPI, which is largely independent of the models. Finally, with all parts in hand, we will show how to wired them in to a coherent global model. Creating an Arbor model ----------------------- For co-simulation, we need a network model in Arbor and some disjoint model that receives and produces spikes in exchange. We will, for sake of the demonstration, use a simple ring of leaky-integrate-and-fire cells, but the concepts translate to cable cells and more complex networks straightforwardly. Save this script in ``ring.py`` for later use. We import the usual libraries .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: :lines: 1-2 and create a ``recipe`` with a configurable number of ``lif`` cells .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: :lines: 5-17 that connect in a ring .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: :lines: 19-21 At time ``0.1 ms`` we inject a single spike into the first cell .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: :lines: 23-30 that should propagate indefinitely. To confirm, we set up the simulation and run it for ``10ms``, printing out spikes at the end: .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: :lines: 33-39 and receive .. code-block:: 0.1ms gid= 0 0.6ms gid= 1 1.1ms gid= 2 1.6ms gid= 3 2.1ms gid= 0 2.6ms gid= 1 3.1ms gid= 2 3.6ms gid= 3 4.1ms gid= 0 4.6ms gid= 1 5.1ms gid= 2 5.6ms gid= 3 6.1ms gid= 0 6.6ms gid= 1 7.1ms gid= 2 7.6ms gid= 3 8.1ms gid= 0 8.6ms gid= 1 9.1ms gid= 2 9.6ms gid= 3 For reference, the full file: .. literalinclude:: ../../python/example/cosim/ring.py :language: python :linenos: :lineno-match: Creating an simple Neural Mass Model (NMM) ------------------------------------------ Next, as the co-simulation partner, we will implement a simple two-population model, the Wilson–Cowan [#f1]_ equations .. math:: \tau_\mathrm{e}\frac{dE}{dt} = -E + f(w_\mathrm{ee} E + w_\mathrm{ei} + P)\\ \tau_\mathrm{i}\frac{dI}{dt} = -I + f(w_\mathrm{ii} I + w_\mathrm{ie} + Q) with :math:`E` and :math:`I` being the activation rates of the excitatory and inhibitory populations, :math:`\tau_\mathrm{x}` decay rates, the coupling terms between populations, .. math:: w_\mathrm{ee} = 12\\ w_\mathrm{ei} = 10\\ w_\mathrm{ie} = 15\\ w_\mathrm{ii} = 0 and :math:`P = 1.25` and :math:`Q = 0` external inputs to the respective populations. The activation function is chosen as .. math:: f(x) = \frac{1}{1 + \exp(-a(x - m))} with gain :math:`a = 1.4` and threshold :math:`m = 4`. For this excercise, we implement the model from scratch, using scipy, however in real-world scenarios, one can use tools like TVB specialising in NMM simulation. We write a simple simulation .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 1-3 by noting the parameters .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 6-18 and the sigmoid activation function .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 21-22 Next, the model is defined as the temporal derivative .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 25-31 and the step function wrapping the SciPy integrator .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 34-36 For testing, we simulate for 100ms and plot .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: :lines: 39-64 receiving .. figure:: ../images/co-sim-01.svg :width: 600 :align: center We called this file ``wilson_cowan.py`` and will use it later. .. literalinclude:: ../../python/example/cosim/wilson_cowan.py :language: python :linenos: :lineno-match: MPI Setup --------- Now, we need to create a bulk-synchronous (i.e. MPI) programm in which one set of ranks drives Arbor and the other our custom model. Up front, from here on out you will need ``mpi4py`` and ``MPI`` installed. Also, the script needs to be run via ``mpirun``, i.e. .. code-block:: mpirun -n 4 python mpi.py We store this in ``mpi.py`` for later use of the communicators. First, we grab ``mpi4py`` and the ``WORLD`` communicator .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :lines: 1-3 Then, we need to determine which ranks (in ``WORLD``!) belong to which group (dubbed ``color`` in MPI parlance) in the MPI documentation, here we assume all ranks *but* the first will run Arbor. Likewise, we need to select a rank in ``WORLD`` as the leader in the respective group (the group's rank ``0``). .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :linenos: :lineno-match: :lines: 5-12 Then we cleave the ``WORLD`` in twain .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :linenos: :lineno-match: :lines: 14 Now, ``group`` is an MPI communicator similar to ``world``, but it collects ranks providing the same ``color`` value, that is, the actual object returned in ``group`` depends on the rank (or color). Thus, there's two values for rank: one in ``world`` and one in ``group``. Additionally two ranks in different groups cannot communicate using ``group``, only ``world`` which is still valid. For example, a sum reduction over ``group`` will provide the sum over (world) rank zero to (world) rank zero and the sum over all other ranks to (world) rank one, since we elected it as the group's leader. Next, we create a bridge between the halves (the first argument is for tweaking the rank ids within the new communicator and the last one is a tag for avoiding collisions. You can ignore both for now) .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :linenos: :lineno-match: :lines: 15 the communicator ``inter`` is a different beast than ``world`` and ``group``; called an intercommunicator. It joins the two groups. It behaves especially interestingly in conjunction with collective operations, however, it behaves in *exactly* the right way we need to enable conjoined simulations. There are other ways to create an intercommunicator, in particular ones that do not require a single programm and splitting of communicators. However, the presented approach is the most portable one. Then, we print our identity .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :linenos: :lineno-match: :lines: 17-23 This concludes setting up MPI functionality. Next, we will incrementally add functionality for actually running the simulation in tandem. The full file is reproduced below .. literalinclude:: ../../python/example/cosim/mpi.py :language: python :linenos: :lineno-match: Running two simulations in tandem --------------------------------- The core idea here is to write a single script that looks at its MPI rank and either runs the Wilson-Cowan NMM or the Arbor ring model. The next section will then iterate and show how to utilise the MPI intercommunicator from the last section to pass messsages between the models. We now import the custom Wilson-Cowan model and the Arbor ring model along with the libraries .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 1-6 and from ``mpi.py`` we fetch the communicators .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 7 Now, we will start running the two simulations in parallel using two nested timesteps. The smaller step sizes will drive the internal integration and must divided into the larger step size. The second, larger, step will be used to exchange data between the simulations, we call it the *epoch*. As long as the larger step is an integer multiple of both integration steps, the latter can be different. We will run the simulations for ``100ms`` .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 13-16 First, the neural mass model simulation is largely lifted from the previous example code, except the nested time loop .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 20-38 for the Arbor simulation, we will use the ``group`` communicator to create a context, which is then passed to the simulation as an additional argument .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 39-45 For all intents and purposes ``group`` behaves like ``world`` for Arbor, it just has a different number of ranks. Finally, we print out the ring's spikes: .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 47-51 This gives us *parallel* simulations, but not yet *co-simulation*, as the two models cannot exchange information. The full file is reproduced below .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: Joining the two simulations --------------------------- How, we need to transfer information from one model to the other. Copy the previous step into a new file ``cosim.py``. From here on out, we will assert that the simulation is run on exactly two ranks. This not required and asymmetric rank counts are absolutely possible and extremely useful since Arbor models are (usually) computationally intensive. Thus, our CLI invocation looks like this .. code-block:: mpirun -n 2 python simulations.py We are going to use the code we have built during the last section and edit the two parts of the ``if-else`` construct deciding which simulation to run. We will be starting with the ``else`` (Arbor) part, as it is a bit simpler. If we strip out all recording and error handling it looks like this: .. code-block:: python if world.rank == 0: # first rank, running W-C y = np.array([0.2, 0.1]) ps = Params() t = 0 while t < T: # all other ranks, running Arbor epoch = 0 while epoch < dt_com: y = step(0, dt_nmm, y, ps) t += dt_nmm epoch += dt_nmm else: # rank != 0 rec = recipe(n_cell=4) ctx = A.context(mpi=group) # `group` connects all ranks except 0 sim = A.simulation(rec, context=ctx) sim.run(T * U.ms, dt_arb * U.ms) We set up the co-simulation by importing our previous work and set the constants .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 2-16 Enabling co-simulation for Arbor ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ For the Arbor model, we need to tweak the ``recipe`` a bit, which is best done by inheritance .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 19-22 we also define a weight for the connections incoming from the neural mass model. These are defined by adding a new callback to the recipe .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :emphasize-lines: 7 :lines: 24-31 The delay for these connections is identical to the epoch time, in real models this not needed, but the minimum delay over all Arbor connections (internal and external) will be used to compute the epoch. Here, we encounter *the* central issue in this kind of simulation: As the NMM has only the notion of aggregate *populations*, how do we define where a connection comes from? Worse: we get a rate from the NMM, but Arbor handles discrete spikes exclusively. Without further information, like postulating a probability distribution, this is intractable. Thus, this must be defined as part of the overall model. One option is to assign a cell count to each population and assume distributions for the spiking activity, then drawing random numbers to generate spikes from the distributions according to the activity rates. For now, we will set the external weights to zero and return to the problem later. To keep things simple for this demonstration, we allocate a single virtual ``gid`` per remote population, i.e. assume both the excitatory and inhibitory population are represented by exactly one cell (line highlighted above). Next, we modify the simulation setup in Arbor using our newly derived ``corecipe`` to the end of the file. .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 101-107 The addition of the ``inter`` intercommunicator kicks off Arbor's internal co-simulation protocols 1. Arbor will find the epoch by scanning both the internal and external connections for the minimum value over the delays. Half of the minimum delay is used as the epoch length. 2. Arbor will periodically -- whenever an epoch is complete -- send a control message. 3. If it receives a goahead reply to the control message, the simulation will proceed to the next epoch. In this case, Arbor waits for an exchange of spikes for the last epoch next. All ranks of the NMM model **must** call ``gather_spikes`` and pass in a list of the spikes produced on their rank during the last epoch. In exchange each such ranks receives all list of all spikes occurring on the Arbor ranks during the last epoch. 4. Once the simulation is complete, a done signal is sent. 5. Internal errors will result in an abort message specifying the reason. Enabling co-simulation for Wilson-Cowan ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ With Arbor set up, we need to set up the connection infrastructure from the NMM side. Arbor has this already built-in, so this part is a bit more verbose (which is why we have put it second). Arbor, as outlined above, will progress its simulation and send messages periodically. Thus, we need to handle these messages on the ranks driving the NMM. Arbor functions as the controlling entity here, in the NMM simulation we will loop forever until we receive a message signalling termination. This loop replaces the outer loop from the previous section, i.e. .. literalinclude:: ../../python/example/cosim/simulations.py :language: python :linenos: :lineno-match: :lines: 30-38 is replaced by the code in the remainder of this section .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 45-47 In case of failure, we abort and print the reason. If Arbor sends the done signal we terminate the loop since the final step has been executed. .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 48-53 Finally, we handle the epoch message which requests proceeding into the next epoch. First, we send the list of spikes from this rank using the intercommunicator. For now, we'll leave it empty. .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 54-56 In return, we received the list of spikes from all Arbor ranks. Next, we will convert this into two rates by counting the spikes by their origin. We are going to assign ``gid=0..11`` as connected to the excitatory population, ``gid=12..15``, and leave the rest unconnected. Thus, we allocate two arrays of zeros such that we have one entry for each NMM timestep in the epoch. .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 58-60 Then, we sort the spikes into the bins .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 62-72 Now, we can proceed with the integration of the neural mass model; we choose here to add the rates we computed in the binning step directly to the variable ``E`` and ``I`` for that timestep .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 74-85 Then, we increment all timings and the bin index, record the models values, and proceed to the next timestep until we have completed the epoch. .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: :lines: 86-88 Unknown message types cause an abort. Results ^^^^^^^ After the simulation, we plot the recorded variables ``E`` and ``I`` and see this .. figure:: ../images/co-sim-02.svg :width: 600 :align: center compared to the pure Wilson-Cown model from before .. figure:: ../images/co-sim-01.svg :width: 600 :align: center we clearly see the impact of the spiking rate of from the ring model. You can find this script in ``cosim.py``: .. literalinclude:: ../../python/example/cosim/cosim.py :language: python :linenos: :lineno-match: Summary ------- We have seen how to an extremely simple model in Arbor can interact with a neural mass model over an MPI-mediated bridge by looking at each component in turn. We have also implemented a trivial neural mass model from scratch. It is important to note that both models -- Arbor and NMM -- are individually as expressive as they would be on their own. However, when joined into a single model and enriched with additional dynamics -- the connections in between and the conversions from rate to discrete spikes and back -- we are able to model more complex scenarios at reasonable computational cost. We have also seen that the adding co-simulation to Arbor is quite efficient in term of code: adding external connections and passing the intercommunicator. The parts that require care -- translating from and to spikes and modeling the virtual cells from populations -- also require the bulk of the code. This also means that adding co-simulation an existing Arbor model into this framework is straightforward. What now? --------- From here, you can a variety of ways towards more realistic models: - Generate spikes *for* Arbor from rates using an assumed distribution and random numbers. - Refine the rate calculation, e.g., by using a smoothing filter. - Use another model than Wilson-Cowan. - Use TVB or similar tools to simulate the NMM. - Iterate on the Arbor network model by - using cable cells - increasing network complexity - add plasticity We invite you to read `our paper `_ on the subject; you can find the `accompanying source code `_ in our contrib section. All source code for all intermediate steps can be in the directory `python/example/cosim `_ of the Arbor source tree. More information is found in the documentation on :ref:`interconnectivity `. For the usage of MPI in Python see `MPI4Py `_ and for MPI in general the `MPI Forum `_ References ---------- .. [#f1] Wilson, H. R., & Cowan, J. D. (1972). Excitatory and inhibitory interactions in localized populations of model neurons.