draw_graph_script.py

# coding = utf-8

import argparse
import networkx as nx
import pandas as pd
import joblib
import json
import geopandas as gpd

from lib.draw import draw

parser = argparse.ArgumentParser()
parser.add_argument("input_file",help="edgelist format (sep = \",\" )")
parser.add_argument("output_file")
parser.add_argument("--country",help="if country node",action="store_true")
parser.add_argument("-w",action="store_true")

args = parser.parse_args()

G = nx.read_gexf(args.input_file)

encoder = None
labels_dict = {}
positions = {}

if args.country:
    iso2_name = json.load(open("data/ISO3166-1.alpha2.json.txt"))
    world = gpd.read_file("data/TM_WORLD_BORDERS-0/TM_WORLD_BORDERS-0.3.shp")
    world["centroid_c"] = world.centroid
    iso2_togeom = dict(world["ISO2 centroid_c".split()].values)
    positions = {k: [v.x, v.y] for k, v in iso2_togeom.items() if k in G}

for node in list(G.nodes()):
    if args.country:
        labels_dict[node] = iso2_name[node]
    else:
        labels_dict[node] = node

fig, ax = draw(G,labels_dict,positions)
if args.country:
    world.boundary.plot(ax=ax)
fig.savefig(args.output_file)