Merge branch 'issues/CSD-316-nx-loader' into develop-gh

probinet/input/loader.py

@@ -11,6 +11,8 @@
 import networkx as nx
 import numpy as np
 import pandas as pd
+import csv
+
 
 from .preprocessing import (
     create_adjacency_tensor_from_graph_list,
@@ -20,6 +22,50 @@
 from ..models.classes import GraphData
 
 
+def build_adjacency_from_networkx(
+    network: nx.Graph,
+    weight_list: list[str],
+    file_name: Optional[PathLike] = None,
+) -> GraphData:
+    """
+    Import networkx graph and convert it to the GraphData object
+    
+    Parameters
+    ----------
+    networkx
+        networkx graph that will be converted to GraphData object
+    weight_list
+        list of names of weights user would like to use from networkx graph
+    file_name
+        name of csv file (and path) created from networkx graph (used to create GraphData object)
+        e.g. /path/to/file/file_name.csv
+    Returns
+    -------
+    GraphData
+        GraphData object created from networkx graph
+    """
+    attribute_names = {
+        key for _, _, data in network.edges(data=True) for key in data
+    }
+    for w in weight_list:
+        assert w in attribute_names, f"{w} is not an attribute"
+
+    if not file_name or Path(file_name).suffix != ".csv":
+        file_name = Path.cwd() / "edge_list.csv"
+        logging.DEBUG("File will be stored at %s" % file_name)
+
+    # Save edges to a CSV file
+    with open(file_name, "w", newline="", encoding="utf-8") as edge_file:
+        writer = csv.writer(edge_file, delimiter=" ")
+        # Write header
+        writer.writerow(["source", "target"] + weight_list)  # Get edge keys.
+        # Write edge data
+        for source, target, attrs in network.edges(data=True):
+            writer.writerow([source, target] + [attrs[a] for a in weight_list])
+
+    return build_adjacency_from_file(file_name)
+
+
 def build_adjacency_from_file(
     path_to_file: PathLike,
     ego: str = "source",

tests/test_input.py

@@ -0,0 +1,28 @@
+"""
+Test cases for the input module.
+"""
+
+import unittest
+import networkx as nx
+from probinet.input.loader import build_adjacency_from_networkx
+from probinet.models.classes import GraphData
+from .fixtures import BaseTest
+
+
+class TestInput(BaseTest):
+
+    def test_build_adjacency_from_networkx(self):
+        G = nx.erdos_renyi_graph(n=14, p=0.5, seed=7, directed=False)
+        # Add weights
+        for i, edge in enumerate(G.edges()):
+            G.edges[edge]["weight"] = 3
+            G.edges[edge]["weight2"] = i
+
+        g_data = build_adjacency_from_networkx(G, weight_list=["weight", "weight2"],file_name=self.folder+"edges.csv")
+
+        self.assertIsInstance(g_data, GraphData)
+        # XXX: test didn't work; weights not added to graph using build_adjacency_from_file
+        # for n, w_name in enumerate(["weight","weight2"]):
+        #     for n1,n2,weight in g_data.graph_list[n].edges(data=w_name):
+        #         w = G.edges[(n1,n2)][w_name]
+        #         self.assertEqual(w, weight)