Coverage for src/abm_colony_collection/convert_to_network.py: 100%

1import networkx as nx 

2import pandas as pd 

3 

4 

5def convert_to_network(neighbors: pd.DataFrame) -> nx.Graph: 

6 """ 

7 Convert lists of neighbors to a network object. 

8 

9 Parameters 

10 ---------- 

11 neighbors 

12 Lists of neighbors for each node id. 

13 

14 Returns 

15 ------- 

16 : 

17 The network object. 

18 """ 

19 

20 nodes = list(neighbors["ID"].values) 

21 edges = [ 

22 (node_id, neighbor_id) 

23 for node_id, neighbor_ids in zip(neighbors["ID"], neighbors["NEIGHBORS"]) 

24 for neighbor_id in neighbor_ids 

25 ] 

26 

27 network = nx.Graph() 

28 network.add_nodes_from(nodes) 

29 network.add_edges_from(edges) 

30 

31 return network