diff --git a/bmtk/simulator/bionet/bionetwork.py b/bmtk/simulator/bionet/bionetwork.py index 8dea6895..9c125ec3 100644 --- a/bmtk/simulator/bionet/bionetwork.py +++ b/bmtk/simulator/bionet/bionetwork.py @@ -133,7 +133,7 @@ def local_gids(self): return list(self._rank_node_gids.keys()) def add_nodes(self, node_population): - self._gid_pool.add_pool(node_population.name, node_population.n_nodes()) + self._gid_pool.add_pool(node_population.name, node_population.node_ids) super(BioNetwork, self).add_nodes(node_population) def get_virtual_cells(self, population, node_id, spike_trains, spikes_generator=None, sim=None): diff --git a/bmtk/simulator/bionet/gids.py b/bmtk/simulator/bionet/gids.py index ced964c1..e35ef9f7 100644 --- a/bmtk/simulator/bionet/gids.py +++ b/bmtk/simulator/bionet/gids.py @@ -14,14 +14,17 @@ def __init__(self): self._offsets = np.array([0], dtype=np.uint64) self._offset2pool_map = {} - def add_pool(self, name, n_nodes): + def add_pool(self, name, node_ids): + n_nodes = np.max(node_ids) + 1 offset_index = len(self._offsets) + self._offset2pool_map[offset_index] = name self._offsets = np.append(self._offsets, np.array([self._accumulated_offset + n_nodes], dtype=np.uint64)) self._pool_offsets[name] = self._accumulated_offset self._accumulated_offset += n_nodes + def get_gid(self, name, node_id): return int(self._pool_offsets[name] + node_id) diff --git a/bmtk/simulator/core/sonata_reader/network_reader.py b/bmtk/simulator/core/sonata_reader/network_reader.py index 71303088..4c3a44ef 100644 --- a/bmtk/simulator/core/sonata_reader/network_reader.py +++ b/bmtk/simulator/core/sonata_reader/network_reader.py @@ -151,6 +151,10 @@ def get_nodes(self): for node in node_group: yield node_adaptor.get_node(node) + @property + def node_ids(self): + return self._node_pop.node_ids + class SonataEdges(EdgesReader): def __init__(self, edge_population, adaptor): diff --git a/tests/simulator/bionet/test_gids.py b/tests/simulator/bionet/test_gids.py index 344f7308..8337a25c 100644 --- a/tests/simulator/bionet/test_gids.py +++ b/tests/simulator/bionet/test_gids.py @@ -1,14 +1,16 @@ import pytest +import numpy as np + from .conftest import * @pytest.mark.skipif(not nrn_installed, reason='NEURON is not installed') def test_gid_pool(): gid_map = GidPool() - gid_map.add_pool(name='p1', n_nodes=1000) - gid_map.add_pool(name='p2', n_nodes=10000) - gid_map.add_pool(name='p3', n_nodes=1) - gid_map.add_pool(name='p4', n_nodes=500) + gid_map.add_pool(name='p1', node_ids=list(range(1000))) + gid_map.add_pool(name='p2', node_ids=list(range(10000))) + gid_map.add_pool(name='p3', node_ids=[0]) + gid_map.add_pool(name='p4', node_ids=np.arange(500, dtype=int)) assert(gid_map.get_gid(name='p1', node_id=0) == 0) assert(gid_map.get_pool_id(0) == (0, 'p1'))