def graph_to_geojson()

in libs/solaris/vector/graph.py [0:0]

• 42 lines of code
• 20 McCabe index (conditional complexity)

def graph_to_geojson(G, output_path, encoding='utf-8', overwrite=False,
                     verbose=False):
    """
    Save graph to two geojsons: one containing nodes, the other edges.
    Arguments
    ---------
    G : :class:`networkx.MultiDiGraph`
        A graph object to save to geojson files.
    output_path : str
        Path to save the geojsons to. ``'_nodes.geojson'`` and
        ``'_edges.geojson'`` will be appended to ``output_path`` (after
        stripping the extension).
    encoding : str, optional
        The character encoding for the saved files.
    overwrite : bool, optional
        Should files at ``output_path`` be overwritten? Defaults to no
        (``False``).
    verbose : bool, optional
        Switch to print relevant values.  Defaults to no (``False``).

    Notes
    -----
    This function is based on ``osmnx.save_load.save_graph_shapefile``, with
    tweaks to make it work with our graph objects. It will save two geojsons:
    a file containing all of the nodes and a file containing all of the edges.
    When writing to geojson, must convert the coordinate reference system
    (crs) to string if it's a dict, otherwise no crs will be appended to the
    geojson.

    Returns
    -------
    None
    """

    # convert directed graph G to an undirected graph for saving as a shapefile
    G_to_save = G.copy().to_undirected()
    # create GeoDataFrame containing all of the nodes
    nodes, data = zip(*G_to_save.nodes(data=True))
    gdf_nodes = gpd.GeoDataFrame(list(data), index=nodes)

    # get coordinate reference system
    g_crs = G_to_save.graph['crs']
    if type(g_crs) == dict:
        # convert from dict
        g_crs = rio.crs.CRS.from_dict(g_crs)
    gdf_nodes.crs = g_crs
    if verbose:
        print("crs:", g_crs)

    gdf_nodes['geometry'] = gdf_nodes.apply(
        lambda row: Point(row['x'], row['y']), axis=1
        )
    gdf_nodes = gdf_nodes.drop(['x', 'y'], axis=1)
    # gdf_nodes['node_idx'] = gdf_nodes['node_idx'].astype(np.int32)

    # # make everything but geometry column a string
    # for col in [c for c in gdf_nodes.columns if not c == 'geometry']:
    #    gdf_nodes[col] = gdf_nodes[col].fillna('').map(make_str)

    # create GeoDataFrame containing all of the edges
    edges = []
    for u, v, key, data in G_to_save.edges(keys=True, data=True):
        edge = {'key': key}
        for attr_key in data:
            edge[attr_key] = data[attr_key]
        if 'geometry' not in data:
            point_u = Point((G_to_save.nodes[u]['x'], G_to_save.nodes[u]['y']))
            point_v = Point((G_to_save.nodes[v]['x'], G_to_save.nodes[v]['y']))
            edge['geometry'] = LineString([point_u, point_v])
        edges.append(edge)

    gdf_edges = gpd.GeoDataFrame(edges)
    gdf_edges.crs = g_crs

    for col in [c for c in gdf_nodes.columns if c != 'geometry']:
        gdf_nodes[col] = gdf_nodes[col].fillna('').apply(str)
    for col in [c for c in gdf_edges.columns if c != 'geometry']:
        gdf_edges[col] = gdf_edges[col].fillna('').apply(str)

    # make directory structure
    if not os.path.exists(os.path.split(output_path)[0]):
        os.makedirs(os.path.split(output_path)[0])

    edges_path = os.path.splitext(output_path)[0] + '_edges.geojson'
    nodes_path = os.path.splitext(output_path)[0] + '_nodes.geojson'
    if overwrite:
        if os.path.exists(edges_path):
            os.remove(edges_path)
        if os.path.exists(nodes_path):
            os.remove(nodes_path)

    gdf_edges.to_file(edges_path, encoding=encoding, driver='GeoJSON')
    gdf_nodes.to_file(nodes_path, encoding=encoding, driver='GeoJSON')