ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.gids

    import numpy as np
    from collections import namedtuple

    PopulationID = namedtuple('PopulationID', 'node_id population')

    class GidPool(object):
        def __init__(self):
            # map from pool-id --> gid
            self._accumulated_offset = 0
            self._pool_offsets = {}

            # map from gid --> pop, node_id
            self._offsets = np.array([0], dtype=np.uint64)
            self._offset2pool_map = {}

        def add_pool(self, name, n_nodes):
            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)

        def get_pool_id(self, gid):
            offset_indx = np.searchsorted(self._offsets, gid, 'right')
            node_id = gid - self._offsets[offset_indx-1]
            pool_name = self._offset2pool_map[offset_indx]
            return PopulationID(int(node_id), pool_name)