Add option `--first_node`

Trajectories/plot_traj.py

@@ -66,6 +66,7 @@ def get_duration(expanded_traj):
 if __name__ == "__main__":
     import json
     import argparse
+    import sys
 
     import matplotlib.pyplot as plt
     import cartopy.crs as ccrs
@@ -84,12 +85,14 @@ if __name__ == "__main__":
         action="store_true",
         help="annotate the first point of trajectory with node number",
     )
-    parser.add_argument(
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument(
         "--min_duration",
         type=int,
         default=1,
         help="minimum duration of plotted trajectories (in time steps), >= 1",
     )
+    group.add_argument("--first_node", help = "only plot the trajectory beginning with given node", type=int)
     args = parser.parse_args()
 
     with open(args.expanded_traj) as f:
@@ -103,26 +106,25 @@ if __name__ == "__main__":
     fig, ax = plt.subplots(subplot_kw={"projection": projection})
     random.seed(0)
 
-    if args.min_duration == 1:
+    if args.first_node:
         for traj in expanded_traj["traj"]:
-            plot_single_traj(
-                traj,
-                expanded_traj["e_overestim"],
-                SHPC,
-                expanded_traj["orientation"],
-                ax,
-                src_crs,
-                args.annotate,
-                color="red",
-            )
+            if traj[0] == args.first_node:
+                plot_single_traj(
+                    traj,
+                    expanded_traj["e_overestim"],
+                    SHPC,
+                    expanded_traj["orientation"],
+                    ax,
+                    src_crs,
+                    args.annotate,
+                    color="red",
+                )
+                break
+        else:
+            sys.exit("No trajectory found with this first node")
     else:
-        # args.min_duration > 1
-        n_long_traj = 0
-        duration_array = get_duration(expanded_traj)
-
-        for traj, duration in zip(expanded_traj["traj"], duration_array):
-            if duration >= args.min_duration:
-                n_long_traj += 1
+        if args.min_duration == 1:
+            for traj in expanded_traj["traj"]:
                 plot_single_traj(
                     traj,
                     expanded_traj["e_overestim"],
@@ -133,12 +135,30 @@ if __name__ == "__main__":
                     args.annotate,
                     color="red",
                 )
-
-        print("Number of trajectories with sufficient duration:", n_long_traj)
-        ax.set_title(
-            rf"lifetime $\geq$ {args.min_duration} time steps"
-            f"\nnumber of trajectories: {n_long_traj}"
-        )
+        else:
+            # args.min_duration > 1
+            n_long_traj = 0
+            duration_array = get_duration(expanded_traj)
+
+            for traj, duration in zip(expanded_traj["traj"], duration_array):
+                if duration >= args.min_duration:
+                    n_long_traj += 1
+                    plot_single_traj(
+                        traj,
+                        expanded_traj["e_overestim"],
+                        SHPC,
+                        expanded_traj["orientation"],
+                        ax,
+                        src_crs,
+                        args.annotate,
+                        color="red",
+                    )
+
+            print("Number of trajectories with sufficient duration:", n_long_traj)
+            ax.set_title(
+                rf"lifetime $\geq$ {args.min_duration} time steps"
+                f"\nnumber of trajectories: {n_long_traj}"
+            )
 
     ax.set_title(expanded_traj["orientation"], loc="left")
     ax.add_feature(cfeature.LAND, edgecolor="black")