ALLEN INSTITUTE

# Source code for bmtk.builder.network_builder

    from .network_adaptors.dm_network import DenseNetwork

    """
    NetworkBuilder is just a slim container for the a NetworkAdaptor class object which does all the work of storing/saving
    the network. Should make it easier to
     1) Create/update the backend network-builder strategies on the fly by swapping out different network-adaptors
     2) Change the public interface/documentation without having to mess with the adaptor, see BioNetBuilder,
        PointNetBuilder, etc.
    """

    class NetworkBuilder(object):
        """The NetworkBuilder class is used for building and saving a brain network/circuit. By default it will save to
        SONATA format for running network simulations using BioNet, PointNet, PopNet or FilterNet bmtk modules, however it
        can be generalized to build any time of network for any time of simulation.

        Building the network:
            For the general use-case building a network consists of 4 steps, each with a corresponding method.

            1. Initialize the network::

                net = NetworkBuilder("network_name")

            2. Create nodes (ie cells) using the **add_nodes()** method::

                net.add_nodes(N=80, model_type='Biophysical', ei='exc')
                net.add_nodes(N=20, model_type='IntFire', ei='inh')
                ...

            3. Create Connection rules between different subsets of nodes using **add_edges()** method::

                net.add_edges(source={'ei': 'exc'}, target={'ei': 'inh'},
                              connection_rule=my_conn_func, synaptic_model='e2i')
                ...

            4. Finally **build** the network and **save** the files::

                net.build()
                net.save(output_dir='network_path')

            See the bmtk documentation, or the method doc-strings for more advanced functionality

        Network Accessor methods:
            **nodes()**

            Will return a iterable of Node objects for each node created. The Node objects can be used like dictionaries to
            fetch their properties. By default returns all nodes in the network, but you can filter out a given subset by
            passing in property/values pairs::

                for node in net.nodes(model_type='Biophysical', ei='exc'):
                    assert(node['ei'] == 'exc')
                    ...

            **edges()**

            Like the nodes() methods, but insteads returns a list of Edge type objects. It too can be filtered by an edge
            property::

                for edge in net.edges(synaptic_model='exp2'):
                    ...

            One can also pass in a list of source (or target) to filter out only those edges which belong to a specific
            subset of cells::

                for edge in net.edges(target_nodes=net.nodes(ei='exc')):
                ...

        Network Properties:
            * **name** - name of the network
            * **nnodes** - number of nodes (cells) in the network.
            * **nedges** - number of edges. Will be zero if build() method hasn't been called
        """

        def __init__(self, name, adaptor_cls=DenseNetwork, **network_props):
            self.adaptor = adaptor_cls(name, **network_props)

        @property
        def name(self):
            """Get the name (string) of this network."""
            return self.adaptor.name

        @property
        def nodes_built(self):
            """Returns True if nodes has been instantiated for this network."""
            return self.adaptor.nodes_built

        @property
        def edges_built(self):
            """Returns True if the connectivity matrix has been instantiated for this network."""
            return self.adaptor.edges_built

        @property
        def nnodes(self):
            """Returns the number of nodes for this network."""
            return self.adaptor.nnodes

        @property
        def nedges(self):
            """Returns the total number of edges for this network."""
            return self.adaptor.nedges

        def add_nodes(self, N=1, **properties):
            """Used to add nodes (eg cells) to a network. User should specify the number of Nodes (N) and can use any
            properties/attributes they require to define the nodes. By default, all individual cells will be assigned a
            unique 'node_id' to identify each node in the network and a 'node_type_id' to identify each group of nodes.

            If a property is a singular value then said property will be shared by all the nodes in the group. If a value
            is a list of length N then each property will be uniquely assigned to each node. In the below example a group
            of 100 nodes is created, all share the same 'model_type' parameter but the pos_x values will be different for
            each node::

                net.add_nodes(N=100, pos_x=np.random.rand(100), model_type='intfire1', ...)

            You can use a tuple to store property values (in which the SONATA hdf5 will save it as a dataset with multiple
            columns). For example to have one property 'positions' which keeps track of the x/y/z coordinates of each cell::

                net.add_nodes(N=100, positions=[(rand(), rand(), rand()) for _ in range(100)], ...)

            :param N: number of nodes in this group
            :param properties: Individual and group properties of given nodes
            """
            self.adaptor.add_nodes(N=N, **properties)

        def nodes(self, **properties):
            """Returns an iterator of Node (glorified dictionary) objects, filtered by parameters.

            To get all nodes on a network::

                for node in net.nodes():
                    ...

            To only get those nodes with properties that match a given list of parameter values::

                for nod in net.nodes(param1=value1, param2=value2, ...):
                    ...

            :param properties: key-value pair of node attributes to filter returned nodes
            :return: An iterator of Node objects
            """
            return self.adaptor.nodes(**properties)

        def add_edges(self, source=None, target=None, connection_rule=1, connection_params=None, iterator='one_to_one',
                      **edge_type_properties):
            """Used to create the connectivity matrix between subsets of nodes. The actually connections will not be
            created until the build() method is called, using the 'connection_rule.

            Node Selection:
                To specify what subset of nodes will be used for the pre- and post-synaptic one can use a dictionary to
                filter the nodes. In the following all inh nodes will be used for the pre-synaptic neurons, but only exc
                fast-spiking neurons will be used in the post-synaptic neurons (If target or source is not specified all
                neurons will be used)::

                    net.add_edges(source={'ei': 'inh'}, target={'ei': 'exc', 'etype': 'fast-spiking'},
                                  dynamic_params='i2e.json', synaptic_model='alpha', ...)

                In the above code there is one connection between each source/target pair of nodes, but to create a
                multi-graph with N connections between each pair use 'connection_rule' parameter with an integer value::

                    net.add_edges(
                        source={'ei': 'inh'},
                        target={'ei': 'exc', 'etype': 'fast-spiking'},
                        connection_rule=N,
                        ...
                    )

            Connection rules:
                Usually the 'connection_rule' parameter will be the name of a function that takes in source-node and
                target-node object (which can be treated like dictionaries, and returns the number of connections (ie
                synapses, 0 or None if no synapses should exists) between the source and target cell::

                    def my_conn_fnc(source_node, target_node):
                        src_pos = source_node['position']
                        trg_pos = target_node['position']
                        ...
                        return N_syns

                    net.add_edges(source={'ei': 'exc'}, target={'ei': 'inh'}, connection_rule=my_conn_fnc, **opt_edge_attrs)

                If the connection_rule function requires addition arguments use the 'connection_params' option::

                    def my_conn_fnc(source_node, target_node, min_edges, max_edges)
                        ...

                    net.add_edges(connection_rule=my_conn_fnc, connection_params={'min_edges': 0, 'max_edges': 20}, ...)

                Sometimes it may be more efficient or even a requirement that multiple connections are created at the same
                time. For example a post-synaptic neuron may only be targeted by a limited number of sources which couldn't
                be done by the previous connection_rule function. But by setting property 'iterator' to value 'all_to_one'
                the connection_rule function now takes in as a value a list of N source neurons, a single target, and should
                return a list of size N::

                    def bulk_conn_fnc(sources, target):
                        syn_list = np.zeros(len(sources))
                        for source in sources:
                            ....
                        return syn_list

                    net.add_edges(connection_rule=bulk_conn_fnc, iterator='all_to_one', ...)

                There is also a 'one_to_all' iterator option that will pair each source node with a list of all available
                target nodes.

            Edge Properties:
                Normally the properties used when creating a given type of edge will be shared by all the individual
                connections. To create unique values for each edge, the add_edges() method returns a ConnectionMap object::

                    def set_syn_weight_by_dist(source, target):
                        src_pos, trg_pos = source['position'], target['position']
                        ....
                        return syn_weight

                    cm = net.add_edges(connection_rule=my_conn_fnc, model_template='Exp2Syn', ...)
                                    delay=2.0)
                    cm.add_properties('syn_weight', rule=set_syn_weight_by_dist)
                    cm.add_properties('delay', rule=lambda *_: np.random.rand(0.01, 0.50))

                In this case the 'model_template' property has a value for all connections of this given type of edge. The
                'syn_weight' and 'delay' properties will (most likely) be unique values. See ConnectionMap documentation for
                more info.

            :param source: A dictionary or list of Node objects (see nodes() method). Used to filter out pre-synaptic
                subset of nodes.
            :param target: A dictionary or list of Node objects). Used to filter out post-synaptic subset of nodes
            :param connection_rule: Integer or a function that returns integer(s). Rule to determine number of connections
                between each source and target node
            :param connection_params: A dictionary, used when the 'connection_rule' is a function that requires additional
                argments
            :param iterator: 'one_to_one', 'all_to_one', 'one_to_all'. When 'connection_rule' is a function this sets
                how the subsets of source/target nodes are passed in. By default (one-to-one) the connection_rule is
                called for every source/target pair. 'all-to-one' will pass in a list of all possible source nodes for
                each target, and 'one-to-all' will pass in a list of all possible targets for each source.
            :param edge_type_properties: properties/attributes of the given edge type
            :return: A ConnectionMap object
            """
            return self.adaptor.add_edges(
                source=source,
                target=target,
                connection_rule=connection_rule,
                connection_params=connection_params,
                iterator=iterator,
                **edge_type_properties
            )

        def edges(self, target_nodes=None, source_nodes=None, target_network=None, source_network=None, **properties):
            """Returns a list of dictionary-like Edge objects, given filter parameters.

            To get all edges from a network::

                edges = net.edges()

            To specify the target and/or source node-set::

                edges = net.edges(target_nodes=net.nodes(type='biophysical'), source_nodes=net.nodes(ei='i'))

            To only get edges with a given edge_property::

              edges = net.edges(weight=100, syn_type='AMPA_Exc2Exc')

            :param target_nodes: gid, list of gid, dict or node-pool. Set of target nodes for a given edge.
            :param source_nodes: gid, list of gid, dict or node-pool. Set of source nodes for a given edge.
            :param target_network: name of network containing target nodes.
            :param source_network: name of network containing source nodes.
            :param properties: edge-properties used to filter out only certain edges.
            :return: list of bmtk.builder.edge.Edge properties.
            """
            return self.adaptor.edges(
                target_nodes=target_nodes,
                source_nodes=source_nodes,
                target_network=target_network,
                source_network=source_network,
                **properties
            )

        def build(self, force=False):
            """Builds nodes and edges.

            :param force: set true to force complete rebuilding of nodes and edges, if nodes() or save_nodes() has been
                called before then forcing a rebuild may change gids of each node.
            """
            self.adaptor.build(force=force)

        def save(self, output_dir='.', force_overwrite=True, compression='gzip'):
            """Used to save the network files in the appropriate (eg SONATA) format into the output_dir directory. The file
            names will be automatically generated based on the network names.

            To have more control over the output and file names use the **save_nodes()** and **save_edges()** methods.

            :param output_dir: string, directory where network files will be generated. Default, current working directory.
            :param force_overwrite: Overwrites existing network files.
            :param compression: Compression algorithm used to save hdf5 files. 'gzip' (default), 'lzf', 'none', or None.
                you can also specify an integer (1-9) to specify the level of gzip compression.
            """
            self.adaptor.save(output_dir=output_dir, force_overwrite=force_overwrite, compression=compression)

        def save_nodes(self, nodes_file_name=None, node_types_file_name=None, output_dir='.', force_overwrite=True, compression='gzip'):
            """Save the instantiated nodes in SONATA format files.

            :param nodes_file_name: file-name of hdf5 nodes file. By default will use <network.name>_nodes.h5.
            :param node_types_file_name: file-name of the csv node-types file. By default will use
                <network.name>_node_types.csv
            :param output_dir: Directory where network files will be generated. Default, current working directory.
            :param force_overwrite: Overwrites existing network files.
            :param compression: Compression algorithm used to save hdf5 files. 'gzip' (default), 'lzf', 'none', or None.
                you can also specify an integer (1-9) to specify the level of gzip compression.
            """
            self.adaptor.save_nodes(
                nodes_file_name=nodes_file_name,
                node_types_file_name=node_types_file_name,
                output_dir=output_dir,
                force_overwrite=force_overwrite,
                compression=compression,
            )

        def save_edges(self, edges_file_name=None, edge_types_file_name=None, output_dir='.', src_network=None,
                       trg_network=None, name=None, force_build=True, force_overwrite=False, compression='gzip'):
            """Save the instantiated edges in SONATA format files.

            :param edges_file_name: file-name of hdf5 edges file. By default will use <src_network>_<trg_network>_edges.h5.
            :param edge_types_file_name: file-name of csv edge-types file. By default will use
                <src_network>_<trg_network>_edges.h5.
            :param output_dir: Directory where network files will be generated. Default, current working directory.
            :param src_network: Name of the source-node populations.
            :param trg_network: Name of the target-node populations.
            :param name: Name of edge populations, eg /edges/<name> in edges.h5 file.
            :param force_build: Force to (re)build the connection matrix if it hasn't already been built.
            :param force_overwrite: Overwrites existing network files.
            :param compression: Compression algorithm used to save hdf5 files. 'gzip' (default), 'lzf', 'none', or None.
                you can also specify an integer (1-9) to specify the level of gzip compression.
            """
            self.adaptor.save_edges(
                edges_file_name=edges_file_name,
                edge_types_file_name=edge_types_file_name,
                output_dir=output_dir,
                src_network=src_network,
                trg_network=trg_network,
                name=name,
                force_build=force_build,
                force_overwrite=force_overwrite,
                compression=compression,
            )

        def import_nodes(self, nodes_file_name, node_types_file_name):
            """Import nodes from an existing sonata file

            :param nodes_file_name:
            :param node_types_file_name:
            """
            self.adaptor.import_nodes(
                nodes_file_name=nodes_file_name,
                node_types_file_name=node_types_file_name
            )

        def clear(self):
            """Resets the network removing the nodes and edges created."""
            self.adaptor.clear()

        def __getattr__(self, item):
            """Catch-all for adaptor attributes"""
            return getattr(self.adaptor, item)

    class BioNetBuilder(NetworkBuilder):
        """An instance of NetworkBuilder specifically designed for building BioNet (NEURON) complainant networks

        """
        def __init__(self, name, adaptor_cls=DenseNetwork, **network_props):
            super(BioNetBuilder, self).__init__(name=name, adaptor_cls=adaptor_cls, **network_props)

        def add_nodes(self, N=1, model_type=None, model_template=None, model_processing=None, dynamics_params=None,
                      morphology=None, x=None, y=None, z=None, rotation_angle_xaxis=None, rotation_angle_yaxis=None,
                      rotation_angle_zaxis=None, **properties):
            """Add a set of N NEURON type nodes/cells to the network

            :param N: number of nodes (cells) in this group
            :param model_type: the type of node, 'biophysical', 'virtual', 'single_compartment', or 'point_neuron'
            :param model_template: Template or class used to instantiate all instances of each cell, eg. a NeuroML or
                NEURON Hoc Template file. Should contain a prefix of the format type ('nml:<Cell.nml>' or
                'ctdb:Biophys1.hoc'
            :param model_processing: string, pre-defined or user function applied to instantiated cell property/morphology,
                if using Allen Cell-Type Database biophysical models use 'aibs_perisomatic' or 'aibs_allactive'.
            :param dynamics_params: A file-name or a dictionary of cell-dynamics parameter values used for generating the
                cell. If using models from Allen Cell-Type Database then passing in the name of <model>_fit.json file will
                apply the same model params to all N cells. If you want unique params for each cell pass in a dictionary
                where each key has an associated list of size N::

                    dynamics_params={'g_bar': [0.5, 0.2, ...], 'na_bar': [1.653e-5, 9.235e-6, ...]}

            :param morphology: string, name of morphology swc file used for creating model_type=biophysical cells
            :param x: list/array of size N floats for each cell's x-axis soma position
            :param y: list/array of size N floats for each cell's y-axis soma position
            :param z: list/array of size N floats for each cell's z-axis soma position
            :param rotation_angle_xaxis: x-axis euler rotation angle around the soma in radians
            :param rotation_angle_yaxis: y-axis euler rotation angle around the soma in radians
            :param rotation_angle_zaxis: z-axis euler rotation angle around the soma in radians
            :param properties: Individual and group properties of given nodes
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'self', 'N']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            super(BioNetBuilder, self).add_nodes(N=N, **properties)

        def add_edges(self, source=None, target=None, connection_rule=1, connection_params=None, iterator='one_to_one',
                      model_template=None, dynamics_params=None, syn_weight=None, delay=None, target_sections=None,
                      distance_range=None, **properties):
            """Add rules for creating edges between a subset of source and target cells.

            :param source: A dictionary or list of Node objects (see nodes() method). Used to filter out pre-synaptic
                subset of nodes.
            :param target: A dictionary or list of Node objects). Used to filter out post-synaptic subset of nodes
            :param connection_rule: Integer or a function that returns integer(s). Rule to determine number of connections
                between each source and target node
            :param connection_params: A dictionary, used when the 'connection_rule' is a function that requires additional
                argments
            :param iterator: 'one_to_one', 'all_to_one', 'one_to_all'. When 'connection_rule' is a function this sets
                how the subsets of source/target nodes are passed in. By default (one-to-one) the connection_rule is
                called for every source/target pair. 'all-to-one' will pass in a list of all possible source nodes for
                each target, and 'one-to-all' will pass in a list of all possible targets for each source.
            :param model_template: A predefined or user generated NEURON function name for generating connection. (default:
                exp2syn).
            :param dynamics_params: A json file path or a dictionary of edge/synapse dynamics parameter values used for
                generating the connection, according to the model_template function. If passing in the name or path of a
                json file the parameter values will be applied to a edges.
            :param syn_weight: float or list of N floats, weights applied to each edge/synpase, value dependent on model
            :param delay: float or list of N floats, delay in milliseconds for cell-cell events
            :param target_sections: list of strings, non-SONATA compliant, locations of where to place synpases when the
                target cell is model_type=biophysical. One or more of the following values: somatic, basal, apical, axon.
                If specified post-synaptic location will be randomly chosen from given sections (+ distance_range value).
                To assign specific locations use afferent_section_id/afferent_section_pos along with the ConnectionMap
                instance returned.
            :param distance_range: A numeric range of floats [beg, end]. non SONATA compliant. Start and end arc-length
                distance from SOMA where the post-synaptic location will be placed on a model_type=biophysical cell. Along
                with target_sections parameter can be used to randomly but strategically connection cells. To assign
                specific locations use afferent_section_id/afferent_section_pos along with the returned ConnectionMap.
            :param properties: properties/attributes of the given edge type
            :return: A ConnectionMap object
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'source', 'target', 'connection_rule', 'connection_params', 'iterator']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            return super(BioNetBuilder, self).add_edges(
                source=source,
                target=target,
                connection_rule=connection_rule,
                connection_params=connection_params,
                iterator=iterator,
                **properties
            )

    class PointNetBuilder(NetworkBuilder):
        """An instance of NetworkBuilder specifically designed for building PointNet (NEST) complainant networks

        """
        def __init__(self, name, adaptor_cls=DenseNetwork, **network_props):
            super(PointNetBuilder, self).__init__(name=name, adaptor_cls=adaptor_cls, **network_props)

        def add_nodes(self, N=1, model_type='point_neuron', model_template=None, dynamics_params=None, x=None, y=None,
                      z=None, **properties):
            """Add a set of N NEST model nodes/cells to the network.

            :param N:
            :param model_type: The type of node, should be value 'point_neuron'
            :param model_template: Name of NEST neuron model to used, should have a 'nest:' prefix; nest:iaf_psc_alpha,
                nest:glif_asc
            :param dynamics_params: A file-name or a dictionary of cell-dynamics parameter values used for generating the
                cell, with parameters matching that of the parameters used to initialize NEST model_template cell. If you
                want unique params for each cell pass in a dictionary where each key has an associated list of size N::

                    dynamics_params={'v_init': [-90, -85, ...], 'v_reset': [-20, -20, ...]}

            :param x: list/array of size N floats for each cell's x-axis position
            :param y: list/array of size N floats for each cell's y-axis position
            :param z: list/array of size N floats for each cell's z-axis position
            :param properties: Individual and group properties of given nodes
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'self', 'N']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            super(PointNetBuilder, self).add_nodes(N=N, **properties)

        def add_edges(self, source=None, target=None, connection_rule=1, connection_params=None, iterator='one_to_one',
                      model_template=None, dynamics_params=None, syn_weight=None, delay=None, **properties):
            """Add rules for creating edges between a subset of source and target cells.

            :param source: A dictionary or list of Node objects (see nodes() method). Used to filter out pre-synaptic
                subset of nodes.
            :param target: A dictionary or list of Node objects). Used to filter out post-synaptic subset of nodes
            :param connection_rule: Integer or a function that returns integer(s). Rule to determine number of connections
                between each source and target node
            :param connection_params: A dictionary, used when the 'connection_rule' is a function that requires additional
                argments
            :param iterator: 'one_to_one', 'all_to_one', 'one_to_all'. When 'connection_rule' is a function this sets
                how the subsets of source/target nodes are passed in. By default (one-to-one) the connection_rule is
                called for every source/target pair. 'all-to-one' will pass in a list of all possible source nodes for
                each target, and 'one-to-all' will pass in a list of all possible targets for each source.
            :param model_template: A predefined or user generated NEURON function name for generating connection. (default:
                exp2syn).
            :param dynamics_params: A json file path or a dictionary of edge/synapse dynamics parameter values used for
                generating the connection, according to the model_template function. If passing in the name or path of a
                json file the parameter values will be applied to a edges.
            :param syn_weight: float or list of N floats, weights applied to each edge/synpase, value dependent on model
            :param delay: float or list of N floats, delay in milliseconds for cell-cell events
            :param properties: properties/attributes of the given edge type
            :return: A ConnectionMap object
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'source', 'target', 'connection_rule', 'connection_params', 'iterator']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            return super(PointNetBuilder, self).add_edges(
                source=source,
                target=target,
                connection_rule=connection_rule,
                connection_params=connection_params,
                iterator=iterator,
                **properties
            )

    class PopNetBuilder(NetworkBuilder):
        """An instance of NetworkBuilder specifically designed for building PopNet (DiPDE) complainant networks

        """
        def __init__(self, name, adaptor_cls=DenseNetwork, **network_props):
            super(PopNetBuilder, self).__init__(name=name, adaptor_cls=adaptor_cls, **network_props)

        def add_nodes(self, model_type='population', model_template=None, dynamics_params=None, **properties):
            """Add a DiPDE cell population

            :param model_type: The type of node, should be value 'population'
            :param model_template: Should be either 'dipde:Internal' or 'dipde:External' (or 'virtual')
            :param dynamics_params: path to json file, or dictionary, of parameters for dipde population
            :param properties: Individual and group properties of given nodes
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'self']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            super(PopNetBuilder, self).add_nodes(N=1, **properties)

        def add_edges(self, source=None, target=None, nsyns=1, syn_weight=None, delay=None, dynamics_params=None,
                      **properties):
            """

            :param source: A dictionary or list of efferent dipde Populations
            :param target: A dictionary or list of afferent dipde Populations
            :param nsyns: Number of connections between the populations
            :param syn_weight: connection weight between the populations
            :param delay: connection delay between the populations
            :param dynamics_params: Additional parameters used for initializing connection
            :param properties: properties/attributes of the given edge type
            :return: A ConnectionMap object
            """
            for arg_var, arg_val in locals().items():
                if arg_var in ['properties', 'source', 'target']:
                    continue
                elif arg_val is not None:
                    properties[arg_var] = arg_val

            return super(PopNetBuilder, self).add_edges(
                source=source,
                target=target,
                **properties
            )