From da2dd9e29c49e34fc67eb3943f77840d8a59f1af Mon Sep 17 00:00:00 2001
From: Lionel GUEZ <guez@lmd.ens.fr>
Date: Wed, 24 Jul 2024 10:10:10 +0200
Subject: [PATCH] Add argument to alternate colors
---
Trajectories/Analysis/plot_traj.py | 41 +++++++++++++++++++++++++++---
1 file changed, 37 insertions(+), 4 deletions(-)
diff --git a/Trajectories/Analysis/plot_traj.py b/Trajectories/Analysis/plot_traj.py
index 6760e964..1ac1cdad 100755
--- a/Trajectories/Analysis/plot_traj.py
+++ b/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(
--
GitLab