Define `traj_vert_ind`

Preparing to define `traj_segm` and `expanded_traj` out of the loop on
vertices.

Trajectories/trajectories.py

@@ -17,7 +17,8 @@ from graph_tool import topology
 import json
 import sys
 
-def new_traj(ind_traj, traj_prop, n, g, traj_segm, expanded_traj):
+def new_traj(ind_traj, traj_prop, n, g, traj_segm, expanded_traj,
+             traj_vert_ind):
     """Assign new trajectory to vertex index n."""
 
     ind_traj += 1
@@ -25,6 +26,7 @@ def new_traj(ind_traj, traj_prop, n, g, traj_segm, expanded_traj):
     traj_segm.append([g.vertex_properties.name[n]])
     inst_eddies = list(g.vertex_properties.inst_eddies[n])
     expanded_traj.append(inst_eddies)
+    traj_vert_ind.append([n])
     return ind_traj
 
 if len(sys.argv) != 2: sys.exit("Required argument: input-graph")
@@ -53,9 +55,14 @@ else:
 
 traj_prop = g.new_vertex_property('int')
 # Trajectory index to which a segment belongs. The trajectory index is
-# also the index in lists traj_segm and expanded_traj.
+# also the index in lists traj_vert_ind, traj_segm and expanded_traj.
 
 ind_traj = - 1 # initialize index of trajectory
+
+traj_vert_ind = []
+# An element of traj_vert_ind is a trajectory, represented by a list
+# of vertex indices.
+
 traj_segm = [] # an element of traj_segm is a trajectory, represented
                # by a list of segments
 expanded_traj = [] # an element of expanded_traj is a trajectory,
@@ -74,23 +81,25 @@ for n in topology.topological_sort(g):
 
         if closest_succ[closest_pred] == n:
             # Assign to n the trajectory of closest_pred. This means
-            # updating ind_traj, traj_prop, traj_segm and
-            # expanded_traj. closest_pred is already in a trajectory,
-            # traj_prop[closest_pred]] is defined. We extend the
-            # trajectory of closest_pred forward.
+            # updating ind_traj, traj_prop, traj_vert_ind, traj_segm
+            # and expanded_traj. closest_pred is already in a
+            # trajectory, traj_prop[closest_pred]] is defined. We
+            # extend the trajectory of closest_pred forward.
             traj_prop[n] = traj_prop[closest_pred]
+            traj_vert_ind[traj_prop[n]].append(n)
             traj_segm[traj_prop[n]].append(g.vertex_properties.name[n])
             inst_eddies = list(g.vertex_properties.inst_eddies[n])
             expanded_traj[traj_prop[n]].extend(inst_eddies)
         else:
             ind_traj = new_traj(ind_traj, traj_prop, n, g, traj_segm,
-                                expanded_traj)
+                                expanded_traj, traj_vert_ind)
     else:
-        ind_traj = new_traj(ind_traj, traj_prop, n, g, traj_segm, expanded_traj)
+        ind_traj = new_traj(ind_traj, traj_prop, n, g, traj_segm,
+                            expanded_traj, traj_vert_ind)
 
     # traj_prop[n] is defined
 
-print("Number of trajectories: ", len(traj_segm))
+print("Number of trajectories: ", len(traj_vert_ind))
 
 with open("traj_segm.json", "w") as outfile:
     json.dump(traj_segm, outfile, indent = 4)