Use another algorithm, sorting on cost function

Rewrite following an idea of R\'emi Laxenaire: loop on edges in order of
ascending cost function.

trajectories.py

@@ -20,80 +20,92 @@ import json
 import csv
 from graph_tool import util, search
 import sys
-
-class Visitor(search.BFSVisitor):
-    def __init__(self):
-        self.traj_segm = {}
-
-    def discover_vertex(self, u):
-        if traj_prop[u] not in self.traj_segm:
-            self.traj_segm[traj_prop[u]] = [g.vp.name[u]]
-        else:
-            self.traj_segm[traj_prop[u]].append(g.vp.name[u])
+import numpy as np
 
 g = graph_tool.load_graph(sys.argv[1])
-traj_prop = g.new_vp('int')
-
-# Assign the new vertex property:
-for node in g.vertices():
-    traj_prop[node] = g.vp.name[node]
-
-print('Setting trajectory vertex property...')
-
-for edge in g.edges():
-    current_cf = g.ep.cost_function[edge]
-    source = edge.source()
-    target = edge.target()
-
-    if source.out_degree() > 1:
-        # Source is a split
-        if current_cf <= min({g.ep.cost_function[e]
-                              for e in source.out_edges()}):
-            if target.in_degree() > 1:
-                if current_cf <= min({g.ep.cost_function[e]
-                                      for e in target.out_edges()}):
-                    # Set target trajectory to source trajectory
-                    traj_prop[target] = traj_prop[source]
-            else:
-                traj_prop[target] = traj_prop[source]
-    else:
-        # source.out_degree() == 1
-        if target.in_degree() > 1:
-            if current_cf <= min({g.ep.cost_function[e]
-                                  for e in target.in_edges()}):
-                traj_prop[target] = traj_prop[source]
+traj_prop = g.new_vp('int', val = - 1)
+all_cost_functions = g.get_edges([g.ep.cost_function])
+
+# Array of indices that sort on cost function:
+ind_cf = np.argsort(all_cost_functions[:, 2])
+
+ind_traj = - 1 # initialize index of trajectory
+traj_segm = [] # trajectories as lists of segments
+expanded_traj = [] # trajectories as lists of instantaneous eddies
+
+# Loop over edges in ascending order of cost function:
+for i in ind_cf:
+    source = all_cost_functions[i,0]
+    target = all_cost_functions[i, 1]
+
+    if traj_prop[source] == - 1 and traj_prop[target] == - 1:
+        ind_traj += 1
+        traj_prop[source] = ind_traj
+        traj_prop[target] = ind_traj
+        segments = [int(g.vp.name[source]), int(g.vp.name[target])]
+        traj_segm.append(segments)
+        inst_eddies = list(g.vp.inst_eddies[source]) \
+            + list(g.vp.inst_eddies[target])
+        expanded_traj.append(inst_eddies)
+    elif traj_prop[source] == - 1:
+        # assert traj_prop[target] != - 1
+        # target is already in a trajectory. Can we extend the
+        # trajectory of target backward?
+        in_neighbors = g.get_in_neighbors(target, [traj_prop])
+
+        if np.any(in_neighbors[:, 1] == traj_prop[target]):
+            # target already has a direct predecessor in its
+            # trajectory so we create a new trajectory for source.
+            ind_traj += 1
+            traj_prop[source] = ind_traj
+            segments = [int(g.vp.name[source])]
+            traj_segm.append(segments)
+            inst_eddies = list(g.vp.inst_eddies[source])
+            expanded_traj.append(inst_eddies)
+        else:
+            traj_prop[source] = traj_prop[target]
+            segment = int(g.vp.name[source])
+            traj_segm[traj_prop[source]].append(segment)
+            inst_eddies = list(g.vp.inst_eddies[source])
+            expanded_traj[traj_prop[source]].extend(inst_eddies)
+    elif traj_prop[target] == - 1:
+        # assert traj_prop[source] != - 1
+        # source is already in a trajectory. Can we extend the
+        # trajectory of source forward?
+        out_neighbors = g.get_out_neighbors(source, [traj_prop])
+
+        if np.any(out_neighbors[:, 1] == traj_prop[source]):
+            # source already has a direct successor in its trajectory
+            # so we create a new trajectory for target.
+            ind_traj += 1
+            traj_prop[target] = ind_traj
+            segments = [int(g.vp.name[target])]
+            traj_segm.append(segments)
+            inst_eddies = list(g.vp.inst_eddies[target])
+            expanded_traj.append(inst_eddies)
         else:
             traj_prop[target] = traj_prop[source]
-
-visitor = Visitor()
-search.bfs_search(g, visitor = visitor)
-print("Number of trajectories: ", len(visitor.traj_segm))
+            segment = int(g.vp.name[target])
+            traj_segm[traj_prop[target]].append(segment)
+            inst_eddies = list(g.vp.inst_eddies[target])
+            expanded_traj[traj_prop[source]].extend(inst_eddies)
+
+# Create trajectories for isolated segments:
+for v in g.vertices():
+    if traj_prop[v] == - 1:
+        ind_traj += 1
+        traj_prop[v] = ind_traj
+        segments = [int(g.vp.name[v])]
+        traj_segm.append(segments)
+        inst_eddies = list(g.vp.inst_eddies[v])
+        expanded_traj.append(inst_eddies)
+
+print("Number of trajectories: ", len(traj_segm))
 
 with open("traj_segm.json", "w") as outfile:
-    json.dump(visitor.traj_segm, outfile, indent = 4)
+    json.dump(traj_segm, outfile, indent = 4)
 
 print("Created traj_segm.json.")
-expanded_traj = {}
-
-for trajectory in visitor.traj_segm.values():
-    # Expand trajectory of current traj key = []
-    key = trajectory[0]
-    for val in trajectory:
-        node = util.find_vertex(g, g.vp.name, val)
-        # (Maybe a loop with a break on the first match is faster. To
-        # be tested.)
-
-        if len(node) != 1:
-            print('Something is wrong.')
-
-        node = node[0]
-        seg = g.vp.inst_eddies[node]
-
-        for s in seg:
-            if key not in expanded_traj:
-                expanded_traj[key] = [s]
-            else:
-                expanded_traj[key].append(s)
 
 with open("expanded_traj.json", "w") as outfile:
     json.dump(expanded_traj, outfile, indent = 4)