diff --git a/trajectories.py b/trajectories.py
index d50fc6353d4e74aa973b7ce6c9b1f42bec245a37..bd932a77e1e019538ef7b45f2068b94f3f1eb67a 100755
--- a/trajectories.py
+++ b/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)