Return networkx graph

src/snkit/network.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pandas
 import shapely.errors
+import networkx as nx
 
 from geopandas import GeoDataFrame
 from shapely.geometry import Point, MultiPoint, LineString, GeometryCollection, shape, mapping
@@ -603,3 +604,43 @@ def set_precision(geom, precision):
     geom_mapping = mapping(geom)
     geom_mapping['coordinates'] = np.round(np.array(geom_mapping['coordinates']), precision)
     return shape(geom_mapping)
+
+
+def to_networkx(network):
+    """Return a networkx graph
+    """
+    # init graph
+    G = nx.Graph()
+    # get nodes from network data
+    G.add_nodes_from(network.nodes.id.to_list())
+    # get edges from network data
+    edges_as_list = [(network.edges.loc[i].from_id, 
+                      network.edges.loc[i].to_id) 
+                      for i in network.edges.index]
+    # add edges to graph
+    G.add_weighted_edges_from(edges_as_list)
+    return G
+
+
+def connected_components(network,add_id=True,id_col='graph_id'):
+    ''' Get connected components within network and id to each individual graph
+    '''
+    # define as nx graph
+    G = to_networkx(network)
+    # get connected_components using nx
+    connected_parts = sorted(nx.connected_components(G), key = len, reverse=True)
+    # add unique id to each graph
+    if add_id is True:
+        count = 1
+        network.edges[id_col] = 0
+        for part in connected_parts:
+            network.edges.loc[ (network.edges.from_id.isin(list(part))) | \
+                               (network.edges.to_id.isin(list(part))), id_col ] = count
+            # adjust counter
+            count = count + 1
+        return Network(
+            nodes=network.nodes,
+            edges=network.edges
+        )
+    else:
+        return connected_parts