ALLEN INSTITUTE

# Source code for bmtk.tests.builder.test_network

    import pytest
    import numpy as np

    from bmtk.builder import NetworkBuilder

    def test_basic():
        net = NetworkBuilder('CA1')
        assert(net.name == 'CA1')
        assert(net.nnodes == 0)
        assert(net.nedges == 0)
        assert(net.nodes_built is False)
        assert(net.edges_built is False)

        assert(len(net.nodes()) == 0)
        assert(len(net.edges()) == 0)
        assert(net.nodes_built is True)
        assert(net.edges_built is True)

    def test_no_name():
        # Fail if network name is invalid
        with pytest.raises(ValueError):
            NetworkBuilder(name='')

        with pytest.raises(ValueError):
            NetworkBuilder(name=None)

    def test_add_nodes():
        # Tests that mutliple models of nodes can be added to network
        net = NetworkBuilder('V1')
        net.add_nodes(N=10, arg_list=range(10), arg_const='pop1', arg_shared='global')
        net.add_nodes(N=1, arg_list=[11], arg_const='pop2', arg_shared='global')
        net.add_nodes(N=5, arg_unique=range(12, 17), arg_const='pop3', arg_shared='global')  # diff param signature
        net.build()

        assert(net.nodes_built is True)
        assert(net.nnodes == 16)
        assert(net.nedges == 0)
        assert(len(net.nodes()) == 16)
        assert(len(net.nodes(arg_const='pop1')) == 10)
        assert(len(net.nodes(arg_const='pop2')) == 1)
        assert(len(net.nodes(arg_shared='global')) == 16)
        assert(len(net.nodes(arg_shared='invalid')) == 0)

        node_set = net.nodes(arg_list=2)
        assert(len(node_set) == 1)
        node = list(node_set)[0]
        assert(node['arg_const'] == 'pop1')
        assert(node['arg_shared'] == 'global')

        node_set = net.nodes(arg_unique=12)
        assert(len(node_set) == 1)
        node = list(node_set)[0]
        assert(node['arg_const'] == 'pop3')
        assert(node['arg_shared'] == 'global')
        assert('arg_list' not in node)

    def test_add_nodes_tuples():
        # Should be able to store tuples of values in single parameters for a given node
        net = NetworkBuilder('V1')
        net.add_nodes(N=10,
                      arg_list=range(10),
                      arg_tuples=[(r, r+1) for r in range(10)],
                      arg_const=('a', 'b'))
        net.build()

        assert(net.nodes_built is True)
        assert(net.nnodes == 10)
        for node in net.nodes():
            assert(len(node['arg_tuples']) == 2)
            assert(node['arg_tuples'][0] == node['arg_list'] and node['arg_tuples'][1] == node['arg_list']+1)
            assert(len(node['arg_const']) == 2)
            assert(node['arg_const'][0] == 'a' and node['arg_const'][1] == 'b')

    def test_add_node_ids():
        # Special case if parameters node_id and node_type_id are explicitly defined by the user
        net = NetworkBuilder('V1')
        net.add_nodes(N=3, node_id=[100, 200, 300], node_type_id=101, name=['one', 'two', 'three'])
        net.build()
        node_one = list(net.nodes(name='one'))[0]
        assert(node_one['name'] == 'one')
        assert(node_one['node_id'] == 100)
        assert(node_one.node_id == 100)
        assert(node_one['node_type_id'] == 101)

        node_three = list(net.nodes(name='three'))[0]
        assert(node_three['name'] == 'three')
        assert(node_three['node_id'] == 300)
        assert (node_one.node_id == 100)
        assert(node_three['node_type_id'] == 101)

    def test_add_node_ids_mixed():
        net = NetworkBuilder('V1')
        net.add_nodes(N=3, node_id=[0, 2, 4], vals=[0, 2, 4])
        net.add_nodes(N=3, vals=[1, 1, 1])
        net.add_nodes(nodes_ids=[6], vals=[6])
        unique_ids = set()
        for n in net.nodes():
            unique_ids.add(n.node_id)
            if n.node_id % 2 == 0:
                assert(n['vals'] == n.node_id)
            else:
                assert(n['vals'] == 1)
        assert(len(unique_ids) == 7)

    def test_add_nodes_mismatch_params():
        # Should fail if a parameter list does not equal to the number of nodes in a group
        net = NetworkBuilder('V1')
        with pytest.raises(Exception):
            net.add_nodes(N=100, list1=[100]*99)

    def test_duplicate_node_ids():
        # Check if the same node_id is being used twice
        net = NetworkBuilder('V1')
        net.add_nodes(N=1, node_id=[100])

        with pytest.raises(ValueError):
            net.add_nodes(N=1, node_id=[100])

    def test_duplicate_node_type_ids():
        # Should fail if their is a node_type_id clash between two groups
        net = NetworkBuilder('V1')
        net.add_nodes(N=2, node_type_id=0)
        with pytest.raises(Exception):
            net.add_nodes(N=2, node_type_id=0)

    def test_add_edges():
        net = NetworkBuilder('V1')
        net.add_nodes(N=10, cell_type='Scnna1', ei='e')
        net.add_nodes(N=10, cell_type='PV1', ei='i')
        net.add_nodes(N=10, cell_type='PV2', ei='i')
        net.add_edges(
            source={'ei': 'i'},
            target={'ei': 'e'},
            connection_rule=lambda s, t: 1,
            edge_arg='i2e'
        )
        net.add_edges(
            source=net.nodes(cell_type='Scnna1'),
            target=net.nodes(cell_type='PV1'),
            connection_rule=2,
            edge_arg='e2i'
        )
        net.build()
        assert(net.nedges == 200 + 200)
        assert(net.edges_built is True)

        for e in net.edges(target_nodes=net.nodes(cell_type='Scnna1')):
            assert(e['edge_arg'] == 'i2e')
            assert(e['nsyns'] == 1)

        for e in net.edges(target_nodes=net.nodes(cell_type='PV1')):
            assert(e['edge_arg'] == 'e2i')
            assert(e['nsyns'] == 2)

    def test_add_edges_custom_params():
        # Uses connection map functionality to create edges with unique parameters
        net = NetworkBuilder('V1')
        net.add_nodes(N=10, arg_list=range(10), arg_ctype='e')
        net.add_nodes(N=5, arg_list=range(10, 15), arg_ctype='i')

        cm = net.add_edges(
            source={'arg_ctype': 'e'},
            target={'arg_ctype': 'i'},
            connection_rule=2
        )
        cm.add_properties('syn_weight', rule=0.5, dtypes=float)
        cm.add_properties(
            ['src_num', 'trg_num'],
            rule=lambda s, t: [s['node_id'], t['node_id']],
            dtypes=[int, int]
        )
        net.build()

        assert(net.nedges == 2*50)
        assert(net.edges_built is True)

        for e in net.edges():
            assert(e['syn_weight'] == 0.5)
            assert(e['src_num'] == e.source_node_id)
            assert(e['trg_num'] == e.target_node_id)

    def test_cross_pop_edges():
        # Uses connection map functionality to create edges with unique parameters
        net1 = NetworkBuilder('V1')
        net1.add_nodes(N=10, arg_list=range(10), arg_ctype='e')
        net1.build()

        net2 = NetworkBuilder('V2')
        net2.add_nodes(N=5, arg_list=range(10, 15), arg_ctype='i')

        net2.add_edges(
            source={'arg_ctype': 'i'},
            target=net1.nodes(arg_ctype='e'),
            connection_rule=lambda s, t: 1,
            edge_arg='i2e'
        )
        net2.build()
        assert(net2.nedges == 50)

    if __name__ == '__main__':
        # test_basic()
        # test_no_name()
        # test_add_nodes()
        # test_add_node_ids_mixed()
        # test_add_nodes_tuples()
        # test_add_nodes_ids()
        # test_duplicate_node_ids()
        # test_build_nodes_id_clash()
        # test_add_edges()
        test_add_edges_custom_params()
        # test_cross_pop_edges()