diff --git a/draw_graph_script.py b/draw_graph_script.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb1950d53b5cd2bcab97aba66825d752d5e3c3c3
--- /dev/null
+++ b/draw_graph_script.py
@@ -0,0 +1,40 @@
+# coding = utf-8
+
+import argparse
+import networkx as nx
+import pandas as pd
+import joblib,json
+
+from lib.draw import draw
+
+parser = argparse.ArgumentParser()
+parser.add_argument("input_file",help="edgelist format (sep = \",\" )")
+parser.add_argument("--encoder-file",help="LabelEncoder instance that allows to obtain a label for each node")
+parser.add_argument("--country",help="if country node",action="store_true")
+parser.add_argument("-w",action="store_true")
+parser.add_argument("output_file")
+
+args = parser.parse_args()
+
+if args.w:
+ df = pd.read_csv(args.input_file,header=None,names="source target weight".split())
+ G = nx.from_pandas_edgelist(df,edge_attr="weight")
+else:
+ df = pd.read_csv(args.input_file, header=None, names="source target weight".split())
+ G = nx.from_pandas_edgelist(df, edge_attr="weight")
+
+encoder = None
+labels_dict = {}
+
+if args.encoder_file:
+ encoder = joblib.load(args.encoder_file)
+ if args.country:
+ iso2_name = json.load(open("data/ISO3166-1.alpha2.json.txt"))
+ for node in list(G.nodes()):
+ if args.country:
+ labels_dict[node] = iso2_name[encoder.inverse_transform([node])[0]]
+ else:
+ labels_dict[node] = encoder.inverse_transform([node])[0]
+
+fig, ax = draw(G,labels_dict)
+fig.savefig("test.pdf")
\ No newline at end of file
diff --git a/lib/draw.py b/lib/draw.py
index a205c2db33c3b7dd3c0e8abcb15682c9e60b8577..de68555c4074e43872e33c0a5d536651e6d041b9 100644
--- a/lib/draw.py
+++ b/lib/draw.py
@@ -89,3 +89,4 @@ def draw(G, labels_dict={}, iteration_force_atlase=2000, figsize=(40, 20), font_
if weight:
plt.colorbar(edges)
plt.axis("off")
+ return fig, ax
diff --git a/nx2gt.py b/nx2gt.py
deleted file mode 100644
index 13b7250ac860b44181f87dc40ebadb5d2eb925a7..0000000000000000000000000000000000000000
--- a/nx2gt.py
+++ /dev/null
@@ -1,119 +0,0 @@
-# code from http://bbengfort.github.io/snippets/2016/06/23/graph-tool-from-networkx.html
-
-import networkx as nx
-import graph_tool as gt
-
-
-def get_prop_type(value, key=None):
- """
- Performs typing and value conversion for the graph_tool PropertyMap class.
- If a key is provided, it also ensures the key is in a format that can be
- used with the PropertyMap. Returns a tuple, (type name, value, key)
- """
-
- # Deal with the value
- if isinstance(value, bool):
- tname = 'bool'
-
- elif isinstance(value, int):
- tname = 'float'
- value = float(value)
-
- elif isinstance(value, float):
- tname = 'float'
-
- elif isinstance(value, dict):
- tname = 'object'
-
- else:
- tname = 'string'
- value = str(value)
-
- return tname, value, key
-
-
-def nx2gt(nxG):
- """
- Converts a networkx graph to a graph-tool graph.
- """
- # Phase 0: Create a directed or undirected graph-tool Graph
- gtG = gt.Graph(directed=nxG.is_directed())
-
- # Add the Graph properties as "internal properties"
- for key, value in nxG.graph.items():
- # Convert the value and key into a type for graph-tool
- tname, value, key = get_prop_type(value, key)
-
- prop = gtG.new_graph_property(tname) # Create the PropertyMap
- gtG.graph_properties[key] = prop # Set the PropertyMap
- gtG.graph_properties[key] = value # Set the actual value
-
- # Phase 1: Add the vertex and edge property maps
- # Go through all nodes and edges and add seen properties
- # Add the node properties first
- nprops = set() # cache keys to only add properties once
- for node, data in nxG.nodes(data=True):
-
- # Go through all the properties if not seen and add them.
- for key, val in data.items():
- if key in nprops:
- continue # Skip properties already added
-
- # Convert the value and key into a type for graph-tool
- tname, _, key = get_prop_type(val, key)
-
- prop = gtG.new_vertex_property(tname) # Create the PropertyMap
- gtG.vertex_properties[key] = prop # Set the PropertyMap
-
- # Add the key to the already seen properties
- nprops.add(key)
-
- # Also add the node id: in NetworkX a node can be any hashable type, but
- # in graph-tool node are defined as indices. So we capture any strings
- # in a special PropertyMap called 'id' -- modify as needed!
- gtG.vertex_properties['id'] = gtG.new_vertex_property('string')
-
- # Add the edge properties second
- eprops = set() # cache keys to only add properties once
- for src, dst, data in nxG.edges(data=True):
-
- # Go through all the edge properties if not seen and add them.
- for key, val in data.items():
- if key in eprops:
- continue # Skip properties already added
-
- # Convert the value and key into a type for graph-tool
- tname, _, key = get_prop_type(val, key)
-
- prop = gtG.new_edge_property(tname) # Create the PropertyMap
- gtG.edge_properties[key] = prop # Set the PropertyMap
-
- # Add the key to the already seen properties
- eprops.add(key)
-
- # Phase 2: Actually add all the nodes and vertices with their properties
- # Add the nodes
- vertices = {} # vertex mapping for tracking edges later
- for node, data in nxG.nodes(data=True):
-
- # Create the vertex and annotate for our edges later
- v = gtG.add_vertex()
- vertices[node] = v
-
- # Set the vertex properties, not forgetting the id property
- data['id'] = str(node)
- for key, value in data.items():
- gtG.vp[key][v] = value # vp is short for vertex_properties
-
- # Add the edges
- for src, dst, data in nxG.edges(data=True):
-
- # Look up the vertex structs from our vertices mapping and add edge.
- e = gtG.add_edge(vertices[src], vertices[dst])
-
- # Add the edge properties
- for key, value in data.items():
- gtG.ep[key][e] = value # ep is short for edge_properties
-
- # Done, finally!
- return gtG
\ No newline at end of file