Parallel simulation of multiple neuron groups inside brainpy.DSRunner

[CloudyDory]

I have written a biological neural network that contains several groups of neurons and synapses in Brainpy. However, during simulation, it seems that the code cannot fully utilize all CPU cores. The code looks like this:

import brainpy as bp
import brainpy.math as bm

class Network(bp.DynSysGroup):
    def __init__(self):
        super().__init__()
        self.neu1 = bp.neurons.HH(size=10, **param1)
        self.neu2 = bp.neurons.HH(size=10, **param2)
        self.neu3 = bp.neurons.HH(size=10, **param3)
        ......
        self.neuron = [self.neu1, self.neu2, self.neu3, ...]

        self.syn1 = bp.synapses.AMPA(self.neu1, self.neu2, ...)
        self.syn2 = bp.synapses.AMPA(self.neu1, self.neu3, ...)
        self.syn3 = bp.synapses.GABAa(self.neu2, self.neu3, ...)
        ......
        self.synapse = [self.syn1, self.syn2, self.syn3, ...]

    def update(self, inputs):
        for neu in self.neuron:
            neu.update(inputs)
        
        for syn in self.synapse:
            syn.update()

if __name__ == '__main__':
    dt = 0.01  # ms
    bm.set_dt(dt)
            
    network = Network()
    I_input = bp.inputs.section_input(...)
    
    runner = bp.DSRunner(
        network, 
        monitors = ['neu1.V', 'neu2.V', ...], 
        jit = True,
        dt = dt,
        progress_bar = True
    )
    runner.run(inputs=I_input) 

I think updating each neuron group and synapse group in python for-loop is slow, In theory, the neuron and synapse groups can be run in parallel. I have tried modified the update function as this using the Python threading module:

def update(self, inputs):
        threads = []
        for neu in self.neuron:
            threads.append(threading.Thread(target=neu.update, args=(inputs,)))
            threads[-1].start()
        for thread in threads:
            thread.join()
        
        threads = []
        for syn in self.synapse:
            threads.append(threading.Thread(target=syn.update))
            threads[-1].start()
        for thread in threads:
            thread.join()

or this:

def update_parallel(self, neuron):
    neuron.update()

def update(self, inputs):
    bp.running.cpu_unordered_parallel(self.update_parallel, [[neu] for neu in self.neuron], num_process=min(len(self.neuron),16))
    
    for syn in self.synapse:
        syn.update()

But from the task manager (in Ubuntu 22.04), I find that the process is still running on a single core. Is there a way to run multiple neuron groups in a network in parallel in Brainpy?

Thanks!

[chaoming0625]

While, parallel simulations somehow have many things to say. Maybe we need a tutorial about how to run one BrainPy model parallelly.

bp.running.cpu_unordered_parallel is the parallelization of multiple BrainPy models on different processor cores.

A tutorial will come out, maybe next month?

[CloudyDory]

Hi, thank you for the reply. I understand that implementing parallel simulation of one BrainPy model is nontrivial. Could you let me know which function (or module) I should use? I can look into it by myself. Or does it require a major revision of the update() function or simulation process?

Thanks!

[chaoming0625]

Dear @CloudyDory , We are still working on the examples for parallel simulations. Tutorials are on the way. We will release these supports in the next week.