Keep trajectory lists inside dictionaries

Keep trajectory lists inside dictionaries, indexed by trajectory
number, instead of lists. This way, we can keep the trajectory numbers
when we work on part of the graph, for example in a component.

Trajectories/Analysis/draw_segments.py

@@ -70,7 +70,7 @@ if args.t:
         ]
     )
 
-    for i, segment_list in enumerate(traj_segm):
+    for i, segment_list in traj_segm.items():
         color = next(colors)
 
         for segment in segment_list:

Trajectories/Analysis/extract_traj.py

@@ -12,5 +12,4 @@ import sys
 with open(sys.argv[1]) as my_file:
     traj = json.load(my_file)
 
-traj_index = int(sys.argv[2])
-print(traj[traj_index])
+print(traj[sys.argv[2]])

Trajectories/Analysis/filter_traj.py

@@ -23,12 +23,12 @@ with open(sys.argv[2]) as f:
 last_node = util_eddies.date_eddy_to_node(
     last_date, expanded_traj["e_overestim"], expanded_traj["e_overestim"]
 )
-filt_traj = []
+filt_traj = {}
 
-for t in expanded_traj["traj"]:
+for k, t in expanded_traj["traj"].items():
     ip = bisect.bisect(t, last_node)
     if ip != 0:
-        filt_traj.append(t[:ip])
+        filt_traj[k] = t[:ip]
 
 with open(sys.argv[3], "w") as f:
     json.dump(

Trajectories/Analysis/plot_traj.py

@@ -168,7 +168,7 @@ if __name__ == "__main__":
     random.seed(0)
 
     if args.first_node:
-        for traj in expanded_traj["traj"]:
+        for traj in expanded_traj["traj"].values():
             if traj[0] == args.first_node:
                 plot_single_traj(
                     traj,
@@ -187,7 +187,7 @@ if __name__ == "__main__":
             sys.exit("No trajectory found with this first node")
     elif args.traj_index:
         plot_single_traj(
-            expanded_traj["traj"][args.traj_index],
+            expanded_traj["traj"][str(args.traj_index)],
             expanded_traj["e_overestim"],
             SHPC,
             expanded_traj["orientation"],
@@ -221,7 +221,7 @@ if __name__ == "__main__":
             arrows = False
 
         if args.min_duration == 1:
-            for traj in expanded_traj["traj"]:
+            for traj in expanded_traj["traj"].values():
                 if args.alt:
                     color = next(colors)
                     color_initial = color
@@ -241,11 +241,12 @@ if __name__ == "__main__":
         else:
             # args.min_duration > 1
             n_long_traj = 0
+            traj_list = expanded_traj["traj"].values()
             duration_array = get_duration(
-                expanded_traj["traj"], expanded_traj["e_overestim"]
+                traj_list, expanded_traj["e_overestim"]
             )
 
-            for traj, duration in zip(expanded_traj["traj"], duration_array):
+            for traj, duration in zip(traj_list, duration_array):
                 if duration >= args.min_duration:
                     if args.alt:
                         color = next(colors)

Trajectories/Analysis/search_traj.py

@@ -13,10 +13,8 @@ if len(sys.argv) != 3: sys.exit("Required argument: JSON-FILE SEGMENT")
 with open(sys.argv[1]) as f_obj:
     traj_segm = json.load(f_obj)
 
-segment = int(sys.argv[2])
-
-for i, traj in enumerate(traj_segm):
-    if segment in traj:
+for i, traj in traj_segm.items():
+    if sys.argv[2] in traj:
         print("Trajectory index:", i)
         print(f"{traj=}")
         break

Trajectories/create_traj_lists.py

@@ -28,7 +28,6 @@ def create_traj_list(
         else:
             traj_vert_ind[traj_prop[n]] = [n]
 
-    traj_vert_ind = [traj_vert_ind[t] for t in sorted(traj_vert_ind)]
     # sorted(traj_vert_ind) instead of range(len(traj_vert_ind))
     # because create_traj_list can be called for a component of the
     # graph. An element of traj_vert_ind is a trajectory, represented
@@ -39,15 +38,15 @@ def create_traj_list(
     # We have trajectories as lists of vertex indices. Now construct
     # corresponding lists of segments and lists of instantaneous eddies:
 
-    traj_segm = []
+    traj_segm = {}
     # An element of traj_segm is a trajectory, represented by a list of
     # segments.
 
-    expanded_traj = []
+    expanded_traj = {}
     # An element of expanded_traj is a trajectory, represented by a list
     # of instantaneous eddies.
 
-    for t in traj_vert_ind:
+    for k, t in traj_vert_ind.items():
         # t is a trajectory as a list of vertex indices.
         list_segm = []  # list of segments in trajectory t
         list_eddies_traj = []  # list of eddies in trajectory t
@@ -60,8 +59,8 @@ def create_traj_list(
 
             list_eddies_traj.extend(list_eddies_segm)
 
-        traj_segm.append(list_segm)
-        expanded_traj.append(list_eddies_traj)
+        traj_segm[k] = list_segm
+        expanded_traj[k] = list_eddies_traj
 
     # Save: