Add the option to label a list of nodes

Overlap/plot_components.py

@@ -35,9 +35,11 @@ def is_node_type(G, n, label):
         return True
     elif label == "important":
         return G.in_degree[n] == 0 or G.out_degree[n] == 0 or G.degree[n] >= 3
-    elif isinstance(label, int):
+    elif label[0] == "date":
         return report_graph.node_to_date_eddy(n, G.graph["e_overestim"],
-                                              only_date = True) == label
+                                              only_date = True) == label[1]
+    elif label[0] == "node":
+        return n in label[1]
 
 def plot_nbunch(G, nbunch, color = '#1f78b4', label = None, ax = None):
     if ax is None: ax = plt.gca()
@@ -104,9 +106,13 @@ if __name__ == "__main__":
     parser.add_argument("edgelist", help = "path to input CSV file")
     parser.add_argument("shpc_dir", help = "directory containing SHPC, with "
                         "visible eddies at all dates")
-    parser.add_argument("-l", "--label",
-                        help = '"root", "leaf", "split", "merge", "all", '
-                        '"important" or integer for date index')
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("-l", "--label_type",
+                        choices = ["root", "leaf", "split", "merge", "all",
+                                   "important"])
+    group.add_argument("--label_date", type = int, metavar = "DATE_INDEX")
+    group.add_argument("--label_node", metavar = 'NODE', type = int,
+                       nargs = "+")
     parser.add_argument("-s", "--save", metavar = "format",
                         help = "Save file to specified format")
     parser.add_argument("-t", "--time", action = "store_true",
@@ -119,9 +125,14 @@ if __name__ == "__main__":
                        metavar = ("lllon", "lllat", "urlon", "urlat"))
     args = parser.parse_args()
 
-    if args.label not in {None, "root", "leaf", "split", "merge", "all",
-                          "important"}:
-        args.label = int(args.label)
+    if args.label_type:
+        label = args.label_type
+    elif args.label_date:
+        label = ("date", args.label_date)
+    elif args.label_node:
+        label = ("node", args.label_node)
+    else:
+        label = None
 
     plt.figure()
 
@@ -153,7 +164,7 @@ if __name__ == "__main__":
             
         dest_crs = ccrs.PlateCarree((args.window[0] + args.window[2]) / 2)
         ax = plt.axes(projection = dest_crs)
-        plot_all_components(G, args.label)
+        plot_all_components(G, label)
     elif args.node is not None:
         if args.time:
             print("Reading edge list in current directory...")
@@ -176,7 +187,7 @@ if __name__ == "__main__":
             
         dest_crs = ccrs.PlateCarree(G.nodes[args.node]["coordinates"][0])
         ax = plt.axes(projection = dest_crs)
-        plot_nbunch(G, G.nodes[args.node]["component"], label = args.label)
+        plot_nbunch(G, G.nodes[args.node]["component"], label = label)
     else:
         if args.time:
             print("Reading edge lists in current directory...")
@@ -198,7 +209,7 @@ if __name__ == "__main__":
             print("Plotting...")
             
         ax = plt.axes(projection = ccrs.PlateCarree())
-        plot_all_components(G, args.label)
+        plot_all_components(G, label)
 
     ax.coastlines()
     ax.gridlines(draw_labels = True)