[MRG] fix artificial cell gid assignment for non-cell-specific drives

doc/whats_new.rst

@@ -53,8 +53,9 @@ Bug
 
 - Fix bug where :func:`~hnn_core.read_spikes` wasn't returning a :class:`~hnn_core.CellResponse` instance with updated spike types, by `Ryan Thorpe`_ in `#382 <https://github.com/jonescompneurolab/hnn-core/pull/382>`_
 
-- `~hnn_core.Dipole.times` and `~hnn_core.Cell_response.times` now reflect the actual integration points instead of the intended
-   times, by `Mainak Jas`_ in `#397 <https://github.com/jonescompneurolab/hnn-core/pull/397>`_
+- :attr:`~hnn_core.Dipole.times` and :attr:`~hnn_core.Cell_response.times` now reflect the actual integration points instead of the intended times, by `Mainak Jas`_ in `#397 <https://github.com/jonescompneurolab/hnn-core/pull/397>`_
+
+- Fix overlapping non-cell-specific drive gid assignment over different ranks in `~hnn_core.MPIBackend`, by `Ryan Thorpe`_ and `Mainak Jas`_ in `#399 <https://github.com/jonescompneurolab/hnn-core/pull/399>`_
 
 API
 ~~~

examples/howto/plot_simulate_mpi_backend.py

@@ -47,7 +47,7 @@
 net.add_bursty_drive(
     'bursty', tstart=50., burst_rate=10, burst_std=20., numspikes=2,
     spike_isi=10, n_drive_cells=10, location='distal', weights_ampa=weights_ampa,
-    seedcore=4)
+    seedcore=8)
 
 ###############################################################################
 # Finally, to simulate we use the

examples/workflows/plot_simulate_beta.py

@@ -129,7 +129,7 @@ def add_beta_drives(net, beta_start):
         'beta_dist', tstart=beta_start, tstart_std=0., tstop=beta_start + 50.,
         burst_rate=1., burst_std=10., numspikes=2, spike_isi=10, n_drive_cells=10,
         location='distal', weights_ampa=weights_ampa_d1,
-        synaptic_delays=syn_delays_d1, seedcore=2)
+        synaptic_delays=syn_delays_d1, seedcore=20)
 
     # Proximal Drive
     weights_ampa_p1 = {'L2_basket': 0.00004, 'L2_pyramidal': 0.00002,
@@ -141,7 +141,7 @@ def add_beta_drives(net, beta_start):
         'beta_prox', tstart=beta_start, tstart_std=0., tstop=beta_start + 50.,
         burst_rate=1., burst_std=20., numspikes=2, spike_isi=10, n_drive_cells=10,
         location='proximal', weights_ampa=weights_ampa_p1,
-        synaptic_delays=syn_delays_p1, seedcore=8)
+        synaptic_delays=syn_delays_p1, seedcore=20)
 
     return net
 
@@ -219,6 +219,7 @@ def add_beta_drives(net, beta_start):
 axes[1].set_title('Beta + ERP Spike Raster')
 net_erp.cell_response.plot_spikes_raster(ax=axes[2], show=False)
 axes[2].set_title('ERP Spike Raster')
+plt.show()
 
 ###############################################################################
 # References

hnn_core/network.py

@@ -436,7 +436,7 @@ def _update_cells(self):
                     cell.gid = self.gid_ranges[cell_type][cell_idx]
                     cell.pos = pos
                     cells.append(cell)
-                    self.cells[cell_type] = cells
+                self.cells[cell_type] = cells
                 self.n_cells += len(cells)
 
     def add_evoked_drive(self, name, *, mu, sigma, numspikes, location,

hnn_core/network_builder.py

@@ -294,6 +294,8 @@ def __init__(self, net, trial_idx=0):
         if len(self.net.rec_arrays) > 0:
             self._expose_imem = True
 
+        self._rank = 0
+
         self._build()
 
     def _build(self):
@@ -302,10 +304,12 @@ def _build(self):
         global _CVODE, _PC
         _create_parallel_context(expose_imem=self._expose_imem)
 
+        self._rank = _get_rank()
+
         # load mechanisms needs ParallelContext for get_rank
         load_custom_mechanisms()
 
-        if _get_rank() == 0:
+        if self._rank == 0:
             print('Building the NEURON model')
 
         self._clear_last_network_objects()
@@ -337,35 +341,48 @@ def _build(self):
         if len(self.net.rec_arrays) > 0:
             self._record_extracellular()
 
-        if _get_rank() == 0:
+        if self._rank == 0:
             print('[Done]')
 
-    # this happens on EACH node
-    # creates self._gid_list for THIS node
-    def _gid_assign(self):
+    def _gid_assign(self, rank=None, n_hosts=None):
+        """Assign cell IDs to this node
+
+        Parameters
+        ----------
+        rank : int | None
+            If not None, override the rank set
+            automatically using Neuron. Used for testing.
+        n_hosts : int | None
+            If not None, override the number of hosts set
+            automatically using Neuron. Used for testing.
+        """
+        if rank is not None:
+            self._rank = rank
+        if n_hosts is None:
+            n_hosts = _get_nhosts()
 
         self.net._update_cells()  # updates net.n_cells
 
-        rank = _get_rank()
-        nhosts = _get_nhosts()
-
         # round robin assignment of gids
-        for gid in range(rank, self.net.n_cells, nhosts):
+        for gid in range(self._rank, self.net.n_cells, n_hosts):
             # set the cell gid
-            _PC.set_gid2node(gid, rank)
             self._gid_list.append(gid)
 
         # loop over all drives, then all cell types, then all artificial cells
         # only assign a "source" artificial cell to this rank if its target
         # exists in _gid_list. "Global" drives get placed on different ranks
         for drive in self.net.external_drives.values():
-            for conn in drive['conn'].values():  # all cell types
-                for src_gid, target_gid in zip(conn['src_gids'],
-                                               conn['target_gids']):
-                    if (target_gid in self._gid_list and
-                            src_gid not in self._gid_list):
-                        _PC.set_gid2node(src_gid, rank)
-                        self._gid_list.append(src_gid)
+            if drive['cell_specific']:
+                for conn in drive['conn'].values():  # all cell types
+                    for src_gid, target_gid in zip(conn['src_gids'],
+                                                   conn['target_gids']):
+                        if (target_gid in self._gid_list and
+                                src_gid not in self._gid_list):
+                            self._gid_list.append(src_gid)
+            else:
+                src_gids = list(self.net.gid_ranges[drive['name']])
+                for gid_idx in range(self._rank, len(src_gids), n_hosts):
+                    self._gid_list.append(src_gids[gid_idx])
 
         # extremely important to get the gids in the right order
         self._gid_list.sort()
@@ -380,6 +397,10 @@ def _create_cells_and_drives(self, threshold, record_vsoma=False,
         These drives are spike SOURCES but cells are also targets.
         External inputs are not targets.
         """
+
+        for gid in self._gid_list:
+            _PC.set_gid2node(gid, self._rank)
+
         # loop through ALL gids
         # have to loop over self._gid_list, since this is what we got
         # on this rank (MPI)
@@ -402,6 +423,7 @@ def _create_cells_and_drives(self, threshold, record_vsoma=False,
 
                 # this call could belong in init of a _Cell (with threshold)?
                 nrn_netcon = cell.setup_source_netcon(threshold)
+                assert cell.gid in self._gid_list
                 _PC.cell(cell.gid, nrn_netcon)
                 self._cells.append(cell)
 
@@ -588,6 +610,7 @@ def _clear_neuron_objects(self):
 
         self._gid_list = list()
         self._cells = list()
+        self._drive_cells = list()
 
     def _clear_last_network_objects(self):
         """Clears NEURON objects and saves the current Network instance"""

hnn_core/tests/test_parallel_backends.py

@@ -15,6 +15,7 @@
 from hnn_core import MPIBackend, jones_2009_model, read_params
 from hnn_core.dipole import simulate_dipole
 from hnn_core.parallel_backends import requires_mpi4py, requires_psutil
+from hnn_core.network_builder import NetworkBuilder
 
 
 def _terminate_mpibackend(event, backend):
@@ -32,6 +33,43 @@ def _terminate_mpibackend(event, backend):
         sleep(0.01)
 
 
+def test_gid_assignment():
+    """Test that gids are assigned without overlap across ranks"""
+
+    net = jones_2009_model(add_drives_from_params=False)
+    weights_ampa = {'L2_basket': 1.0, 'L2_pyramidal': 3.0,
+                    'L5_basket': 2.0, 'L5_pyramidal': 4.0}
+    syn_delays = {'L2_basket': 1.0, 'L2_pyramidal': 2.0,
+                  'L5_basket': 3.0, 'L5_pyramidal': 4.0}
+
+    net.add_bursty_drive(
+        'bursty_dist', location='distal', burst_rate=10,
+        weights_ampa=weights_ampa, synaptic_delays=syn_delays,
+        cell_specific=False, n_drive_cells=5)
+    net.add_evoked_drive(
+        'evoked_prox', mu=1.0, sigma=1.0, numspikes=1,
+        weights_ampa=weights_ampa, location='proximal',
+        synaptic_delays=syn_delays, cell_specific=True,
+        n_drive_cells='n_cells')
+    net._instantiate_drives(tstop=20, n_trials=2)
+
+    all_gids = list()
+    for type_range in net.gid_ranges.values():
+        all_gids.extend(list(type_range))
+    all_gids.sort()
+
+    n_hosts = 3
+    all_gids_instantiated = list()
+    for rank in range(n_hosts):
+        net_builder = NetworkBuilder(net)
+        net_builder._gid_list = list()
+        net_builder._gid_assign(rank=rank, n_hosts=n_hosts)
+        all_gids_instantiated.extend(net_builder._gid_list)
+    all_gids_instantiated.sort()
+    assert all_gids_instantiated == sorted(set(all_gids_instantiated))
+    assert all_gids == all_gids_instantiated
+
+
 # The purpose of this incremental mark is to avoid running the full length
 # simulation when there are failures in previous (faster) tests. When a test
 # in the sequence fails, all subsequent tests will be marked "xfailed" rather