Allow the module to be imported

Because we are going to define a function.

Trajectories/plot_traj.py

@@ -9,71 +9,74 @@ import cartopy.crs as ccrs
 import random
 import argparse
 
-parser = argparse.ArgumentParser()
-parser.add_argument("expanded_traj", help = "JSon file")
-parser.add_argument("e_overestim", help = "text file")
-parser.add_argument("SHPC", help = "directory")
-parser.add_argument("orientation", choices = ["Anticyclones", "Cyclones"])
-parser.add_argument("--save", metavar = "format",
-                    help = "Save file to specified format")
-parser.add_argument("--annotate", action = "store_true", help = "annotate the "
-                    "first point of trajectory with node number")
-parser.add_argument("--min_duration", type = int, default = 1, help = "minimum"
-                    "duration of plotted trajectories (in time steps), >= 1")
-args = parser.parse_args()
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("expanded_traj", help = "JSon file")
+    parser.add_argument("e_overestim", help = "text file")
+    parser.add_argument("SHPC", help = "directory")
+    parser.add_argument("orientation", choices = ["Anticyclones", "Cyclones"])
+    parser.add_argument("--save", metavar = "format",
+                        help = "Save file to specified format")
+    parser.add_argument("--annotate", action = "store_true", help = "annotate "
+                        "the first point of trajectory with node number")
+    parser.add_argument("--min_duration", type = int, default = 1,
+                        help = "minimumduration of plotted trajectories (in "
+                        "time steps), >= 1")
+    args = parser.parse_args()
 
-with open(args.expanded_traj) as f:
-    trajectories = json.load(f)
+    with open(args.expanded_traj) as f:
+        trajectories = json.load(f)
 
-print("Number of trajectories:", len(trajectories))
+    print("Number of trajectories:", len(trajectories))
 
-with open(args.e_overestim) as f:
-    e_overestim = int(f.read())
+    with open(args.e_overestim) as f:
+        e_overestim = int(f.read())
 
-SHPC = util_eddies.SHPC_class(args.SHPC, def_orient = args.orientation)
-src_crs = ccrs.PlateCarree()
-projection = ccrs.PlateCarree(central_longitude = 90)
-fig, ax = plt.subplots(subplot_kw = {"projection": projection})
-random.seed(0)
-n_long_traj = 0
+    SHPC = util_eddies.SHPC_class(args.SHPC, def_orient = args.orientation)
+    src_crs = ccrs.PlateCarree()
+    projection = ccrs.PlateCarree(central_longitude = 90)
+    fig, ax = plt.subplots(subplot_kw = {"projection": projection})
+    random.seed(0)
+    n_long_traj = 0
 
-for traj in trajectories:
-    if args.min_duration > 1:
-        duration = report_graph.node_to_date_eddy(traj[- 1], e_overestim,
-                                                  only_date = True) \
-            - report_graph.node_to_date_eddy(traj[0], e_overestim,
-                                             only_date = True)
-    if args.min_duration == 1 or duration >= args.min_duration:
-        if args.min_duration > 1: n_long_traj += 1
-        xy = []
+    for traj in trajectories:
+        if args.min_duration > 1:
+            duration = report_graph.node_to_date_eddy(traj[- 1], e_overestim,
+                                                      only_date = True) \
+                - report_graph.node_to_date_eddy(traj[0], e_overestim,
+                                                 only_date = True)
+        if args.min_duration == 1 or duration >= args.min_duration:
+            if args.min_duration > 1: n_long_traj += 1
+            xy = []
 
-        for node in traj:
-            date_index, eddy_index = report_graph.node_to_date_eddy(node,
-                                                                    e_overestim)
-            i_slice = SHPC.get_slice(date_index)
-            ishape = SHPC.comp_ishape(date_index, eddy_index, i_slice)
-            shape = SHPC.get_reader(i_slice, layer = "extremum").shape(ishape)
-            xy.append(shape.points[0])
+            for node in traj:
+                date_index, eddy_index \
+                    = report_graph.node_to_date_eddy(node, e_overestim)
+                i_slice = SHPC.get_slice(date_index)
+                ishape = SHPC.comp_ishape(date_index, eddy_index, i_slice)
+                shape = SHPC.get_reader(i_slice, layer = "extremum")\
+                            .shape(ishape)
+                xy.append(shape.points[0])
 
-        xy = np.array(xy)
-        ax.plot(xy[:, 0], xy[:, 1], color = "red", transform = src_crs)
-        ax.plot(xy[0, 0], xy[0, 1], marker = "s", color = "black",
-                transform = src_crs)
-        if args.annotate:
-            ax.annotate(str(traj[0]),
-                        ax.projection.transform_point(xy[0, 0], xy[0, 1],
-                                                      src_crs),
-                        xytext = (3 * random.random(), 3 * random.random()),
-                        textcoords = "offset points")
+            xy = np.array(xy)
+            ax.plot(xy[:, 0], xy[:, 1], color = "red", transform = src_crs)
+            ax.plot(xy[0, 0], xy[0, 1], marker = "s", color = "black",
+                    transform = src_crs)
+            if args.annotate:
+                ax.annotate(str(traj[0]),
+                            ax.projection.transform_point(xy[0, 0], xy[0, 1],
+                                                          src_crs),
+                            xytext = (3 * random.random(), 3 * random.random()),
+                            textcoords = "offset points")
 
-ax.coastlines()
-ax.gridlines(draw_labels = True)
+    ax.coastlines()
+    ax.gridlines(draw_labels = True)
 
-if args.min_duration > 1:
-    print("Number of trajectories with sufficient duration:", n_long_traj)
+    if args.min_duration > 1:
+        print("Number of trajectories with sufficient duration:", n_long_traj)
 
-if args.save:
-    plt.savefig(f"plot_traj.{args.save}")
-    print(f'Created "plot_traj.{args.save}".')
-else:
-    plt.show()
+    if args.save:
+        plt.savefig(f"plot_traj.{args.save}")
+        print(f'Created "plot_traj.{args.save}".')
+    else:
+        plt.show()