ALLEN INSTITUTE

# Source code for bmtk.simulator.core.node_sets

    from .io_tools import io

    class NodeSet(object):
        def __init__(self, filter_params, network):
            self._network = network
            self._populations = []
            self._preselected_gids = None

            if isinstance(filter_params, list):
                self._preselected_gids = filter_params
            elif isinstance(filter_params, dict):
                self._filter = filter_params.copy()
                self._populations = self._find_populations()
            else:
                io.log_exception('Unknown node set params type {}'.format(type(filter_params)))

        def _find_populations(self):
            for k in ['population', 'populations']:
                if k in self._filter:
                    node_pops = []
                    for pop_name in to_list(self._filter[k]):
                        node_pops.append(self._network.get_node_population(pop_name))
                    del self._filter[k]
                    return node_pops

            return self._network.get_node_populations()

        def populations(self):
            return self._populations

        def population_names(self):
            return [p.name for p in self._populations]

        def gids(self):
            if self._preselected_gids is not None:
                for gid in self._preselected_gids:
                    yield gid
            else:
                for pop in self._populations:
                    for node in pop.filter(self._filter):
                        yield self._network.gid_pool.get_gid(name=pop.name, node_id=node.node_id)

        @property
        def node_ids(self):
            node_ids = []
            for pop in self._populations:
                for node in pop.filter(self._filter):
                    node_ids.append(node.node_id)
            return node_ids

        def nodes(self):
            return None

        def fetch_nodes(self):
            for pop in self._populations:
                for node in pop.filter(self._filter):
                    yield node

    class NodeSetAll(NodeSet):
        def __init__(self, network):
            super(NodeSetAll, self).__init__({}, network)

    def to_list(val):
        if isinstance(val, list):
            return val
        else:
            return [val]