diff --git a/gt_matlab_trajs.py b/gt_matlab_trajs.py index 196021bdd5a2de8e895fc51aacdebf6c65d1705e..83256902a7e09b5a5c1d7a3c0d7752767d49536f 100755 --- a/gt_matlab_trajs.py +++ b/gt_matlab_trajs.py @@ -22,11 +22,14 @@ from graph_tool import util, search import sys class Visitor(search.BFSVisitor): + def __init__(self): + self.trajectories = {} + def discover_vertex(self, u): - if g.vp.traj[u] not in trajectories: - trajectories[g.vp.traj[u]] = [g.vp.name[u]] + if g.vp.traj[u] not in self.trajectories: + self.trajectories[g.vp.traj[u]] = [g.vp.name[u]] else: - trajectories[g.vp.traj[u]].append(g.vp.name[u]) + self.trajectories[g.vp.traj[u]].append(g.vp.name[u]) g = graph_tool.load_graph(sys.argv[1]) g.vp['traj'] = g.new_vp('int') @@ -65,17 +68,17 @@ print('Algorithm done, saving...') # Make a dictionary of trajectories: -trajectories = {} -search.bfs_search(g, visitor = Visitor()) +visitor = Visitor() +search.bfs_search(g, visitor = visitor) with open("trajectories.json", "w") as outfile: - json.dump(trajectories, outfile, indent = 4) + json.dump(visitor.trajectories, outfile, indent = 4) # Setup a new trajectory that holds the expanded segments: expanded_trajectories = {} -for trajectory in trajectories.values(): +for trajectory in visitor.trajectories.values(): # Expand trajectory of current traj key = [] key = trajectory[0] for val in trajectory: @@ -98,4 +101,4 @@ for trajectory in trajectories.values(): with open("expanded_trajectories.json", "w") as outfile: json.dump(expanded_trajectories, outfile, indent = 4) -print(f'Done saving json: {len(trajectories)} trajectories.') +print(f'Done saving json: {len(visitor.trajectories)} trajectories.')