Return networkx graph

src/snkit/network.py

@@ -5,6 +5,11 @@
 import numpy as np
 import pandas
 import shapely.errors
+try:
+    import networkx as nx
+    USE_NX = True
+except ImportError:
+    USE_NX = False
 
 from geopandas import GeoDataFrame
 from shapely.geometry import Point, MultiPoint, LineString, GeometryCollection, shape, mapping
@@ -603,3 +608,49 @@ 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
+    """
+    if USE_NX is True:
+        # 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 = list(zip(network.edges.from_id,network.edges.to_id))
+        # add edges to graph
+        G.add_weighted_edges_from(edges_as_list)
+        return G
+    else:
+        raise ImportError('No module named networkx')
+
+
+def get_connected_components(network):
+    ''' Get connected components within network and id to each individual graph
+    '''
+    if USE_NX is True:
+        G = to_networkx(network)
+        return sorted(nx.connected_components(G), key = len, reverse=True)
+    else:
+        raise ImportError('No module named networkx')
+
+
+def add_component_ids(network,id_col='component_id'):
+    ''' Add column of component IDs to network edge data
+    '''
+    # get connected components
+    connected_parts = get_connected_components(network)
+    # add unique id to each graph
+    count = 1
+    network.edges[id_col] = 0 # init id_col
+    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
+    return Network(
+        nodes=network.nodes,
+        edges=network.edges