Add argument to alternate colors

Trajectories/Analysis/plot_traj.py

@@ -82,6 +82,7 @@ if __name__ == "__main__":
     import json
     import argparse
     import sys
+    import itertools
 
     import matplotlib.pyplot as plt
     import cartopy.crs as ccrs
@@ -117,6 +118,11 @@ if __name__ == "__main__":
         help="only plot the trajectory with given index",
         type=int,
     )
+    parser.add_argument(
+        "--alt",
+        help="alternate colors between trajectories",
+        action="store_true",
+    )
     args = parser.parse_args()
 
     with open(args.expanded_traj) as f:
@@ -161,8 +167,31 @@ if __name__ == "__main__":
             color_initial="green",
         )
     else:
+        if args.alt:
+            colors = itertools.cycle(
+                [
+                    "blue",
+                    "orange",
+                    "green",
+                    "red",
+                    "purple",
+                    "brown",
+                    "pink",
+                    "gray",
+                    "olive",
+                    "cyan",
+                ]
+            )
+        else:
+            color = "red"
+            color_initial = "green"
+
         if args.min_duration == 1:
             for traj in expanded_traj["traj"]:
+                if args.alt:
+                    color = next(colors)
+                    color_initial = color
+
                 plot_single_traj(
                     traj,
                     expanded_traj["e_overestim"],
@@ -171,8 +200,8 @@ if __name__ == "__main__":
                     ax,
                     src_crs,
                     args.annotate,
-                    color="red",
-                    color_initial="green",
+                    color,
+                    color_initial,
                 )
         else:
             # args.min_duration > 1
@@ -183,6 +212,10 @@ if __name__ == "__main__":
 
             for traj, duration in zip(expanded_traj["traj"], duration_array):
                 if duration >= args.min_duration:
+                    if args.alt:
+                        color = next(colors)
+                        color_initial = color
+
                     n_long_traj += 1
                     plot_single_traj(
                         traj,
@@ -192,8 +225,8 @@ if __name__ == "__main__":
                         ax,
                         src_crs,
                         args.annotate,
-                        color="red",
-                        color_initial="green",
+                        color,
+                        color_initial,
                     )
 
             print(