Polish

gt_matlab_trajs.py

 #!/usr/bin/env python3
 
-# A script that loads a segmented cf graph in the gt format, iterates on all
-# of the edges and cost functions and generates the trajectories as done in
-# MATLAB. Segments (nodes) have a new vertex property called "traj" that holds
-# the trajectory ID that that segment belongs to. The output is either an
-# expanded JSON file and/or an edgelist file with all of the trajectories.
-# Inputs:
-    # name and the location of the gt segmented cf graph
-# Output:
-    # trajectories.json that contains all of the trajectories
-
-import graph_tool as gt
+"""A script that loads a segmented cf graph in the gt format, iterates
+on all of the edges and cost functions and generates the trajectories
+as done in MATLAB. Segments (nodes) have a new vertex property called
+"traj" that holds the trajectory ID that that segment belongs to. The
+output is either an expanded JSON file and/or an edgelist file with
+all of the trajectories.
+
+Inputs:
+the gt segmented cf graph
+Output:
+trajectories.json that contains all of the trajectories
+
+"""
+
+import graph_tool
 from os import path
 import json
 import csv
-import graph_tool.util as ut
+from graph_tool import util
 import sys
 
-g = gt.Graph()
-g.load(sys.argv[1])
-traj_mark = g.new_vp('int')
-g.vp['traj'] = traj_mark
+g = graph_tool.load_graph(sys.argv[1])
+g.vp['traj'] = g.new_vp('int')
 
-# assign the new vertex property
+# Assign the new vertex property:
 for node in g.vertices():
     g.vp.traj[node] = g.vp.name[node]
 
-
-##################
-# MAIN ALGORITHM #
-##################
-
-# iterating on edges:
-
 print('Algorithm starting...')
 
 for edge in g.edges():
@@ -40,8 +35,8 @@ for edge in g.edges():
     src = edge.source()
     trg = edge.target()
 
-    # source is a split
     if src.out_degree() > 1:
+        # Source is a split
         if current_cf <= min({g.ep.cost_function[e] for e in src.out_edges()}):
             if trg.in_degree() > 1:
                 if current_cf <= min({g.ep.cost_function[e]
@@ -51,7 +46,7 @@ for edge in g.edges():
             else:
                 g.vp.traj[trg] = g.vp.traj[src]
     else:
-        # source is a continuation or root with one out
+        # Source is a continuation or root with one out
         if trg.in_degree() > 1:
             if current_cf <= min({g.ep.cost_function[e]
                                   for e in trg.in_edges()}):
@@ -71,14 +66,15 @@ for node in g.vertices():
     else:
         trajectories[g.vp.traj[node]].append(g.vp.name[node])
 
-# setup a new trajectory that holds the expanded segments
+# Setup a new trajectory that holds the expanded segments:
+
 expanded_trajectories = {}
 
 for trajectory in trajectories.values():
     # Expand trajectory of current traj key = []
     key = trajectory[0]
     for val in trajectory:
-        node = ut.find_vertex(g, g.vp.name, val)
+        node = util.find_vertex(g, g.vp.name, val)
 
         if len(node) != 1:
             print('Something is wrong.')
@@ -92,7 +88,6 @@ for trajectory in trajectories.values():
             else:
                 expanded_trajectories[key].append(s)
 
-# write out the json file
 with open("trajectories.json", "w") as outfile:
     json.dump(expanded_trajectories, outfile, indent = 4)