Blacken

Trajectories/cost_function.py

@@ -26,7 +26,8 @@ import graph_tool
 
 import util_eddies
 
-Omega = 2 * math.pi / 86164.
+Omega = 2 * math.pi / 86164.0
+
 
 def calculate_radii_rossby(properties):
     """Compute average on some instantaneous eddies of Rossby number and
@@ -38,24 +39,26 @@ def calculate_radii_rossby(properties):
 
     """
 
-    avg_rad = 0 # in m
+    avg_rad = 0  # in m
     avg_Rossby = 0
     n_valid_Rossby = 0
 
     for prop in properties:
-        f = 2 * Omega * math.sin(prop["pos"][1]) # in s-1
-        radius = prop["radius"] * 1000 # in m
+        f = 2 * Omega * math.sin(prop["pos"][1])  # in s-1
+        radius = prop["radius"] * 1000  # in m
 
         if abs(prop["speed"]) < 100:
             avg_Rossby += prop["speed"] / (f * radius)
             n_valid_Rossby += 1
 
-        avg_rad += radius # in m
+        avg_rad += radius  # in m
 
     avg_rad /= len(properties)
-    if n_valid_Rossby != 0: avg_Rossby /= n_valid_Rossby
+    if n_valid_Rossby != 0:
+        avg_Rossby /= n_valid_Rossby
     return avg_rad, avg_Rossby
 
+
 def node_to_prop(node_list, e_overestim, SHPC, orientation):
     """node_list is a list of node identification numbers for
     instantaneous eddies. This function returns some properties of the
@@ -73,48 +76,64 @@ def node_to_prop(node_list, e_overestim, SHPC, orientation):
         date_index, eddy_index = util_eddies.node_to_date_eddy(n, e_overestim)
         i_slice = SHPC.get_slice(date_index)
         ishape = SHPC.comp_ishape(date_index, eddy_index, i_slice, orientation)
-        shapeRec = SHPC.get_reader(i_slice, orientation, "extremum")\
-                       .shapeRecord(ishape)
-        prop = {"pos": [math.radians(x) for x in shapeRec.shape.points[0]], "speed": shapeRec.record.speed}
-        prop["radius"] \
-            = SHPC.get_reader(i_slice, orientation, "max_speed_contour")\
-                  .record(ishape).r_eq_area
+        shapeRec = SHPC.get_reader(
+            i_slice, orientation, "extremum"
+        ).shapeRecord(ishape)
+        prop = {
+            "pos": [math.radians(x) for x in shapeRec.shape.points[0]],
+            "speed": shapeRec.record.speed,
+        }
+        prop["radius"] = (
+            SHPC.get_reader(i_slice, orientation, "max_speed_contour")
+            .record(ishape)
+            .r_eq_area
+        )
 
         if prop["radius"] < 0:
-            prop["radius"] = SHPC.get_reader(i_slice, orientation,
-                                             "outermost_contour")\
-                .record(ishape).r_eq_area
+            prop["radius"] = (
+                SHPC.get_reader(i_slice, orientation, "outermost_contour")
+                .record(ishape)
+                .r_eq_area
+            )
 
         properties.append(prop)
 
     return properties
 
+
 t0 = time.perf_counter()
 timings = open("timings.txt", "w")
 parser = argparse.ArgumentParser()
 parser.add_argument("SHPC_dir")
-parser.add_argument("orientation", choices = ["Anticyclones", "Cyclones"])
-parser.add_argument("input_segments", help = "input graph of segments without "
-                    "cost, suffix .gt (graph-tool) or .graphml")
-parser.add_argument("output_segments", help = "output graph of segments with "
-                    "cost, suffix .gt (graph-tool) or .graphml")
-parser.add_argument("--debug", help = "save properties to output file",
-                    action = "store_true")
+parser.add_argument("orientation", choices=["Anticyclones", "Cyclones"])
+parser.add_argument(
+    "input_segments",
+    help="input graph of segments without "
+    "cost, suffix .gt (graph-tool) or .graphml",
+)
+parser.add_argument(
+    "output_segments",
+    help="output graph of segments with "
+    "cost, suffix .gt (graph-tool) or .graphml",
+)
+parser.add_argument(
+    "--debug", help="save properties to output file", action="store_true"
+)
 args = parser.parse_args()
 
 # Set some values needed for the cost function:
-delta_cent_mean = 3.8481 # in km
+delta_cent_mean = 3.8481  # in km
 delta_cent_std = 8.0388
 delta_ro_mean = -0.0025965
 delta_ro_std = 5.2168
-delta_r_mean = -9.4709 # in m
+delta_r_mean = -9.4709  # in m
 delta_r_std = 8.6953e3
 
 # Load the graph_tool file:
 
-print('Loading graph...')
+print("Loading graph...")
 g = graph_tool.load_graph(args.input_segments)
-print('Loading done...')
+print("Loading done...")
 print("Input graph:")
 print("Number of vertices:", g.num_vertices())
 print("Number of edges:", g.num_edges())
@@ -129,34 +148,38 @@ t0 = t1
 g.graph_properties["orientation"] = g.new_graph_property("string")
 g.graph_properties["orientation"] = args.orientation
 
-pos_first = g.new_vp('vector<double>')
-pos_last = g.new_vp('vector<double>')
-first_av_rad = g.new_vp('float')
-first_av_ros = g.new_vp('float')
-last_av_rad = g.new_vp('float')
-last_av_ros = g.new_vp('float')
+pos_first = g.new_vp("vector<double>")
+pos_last = g.new_vp("vector<double>")
+first_av_rad = g.new_vp("float")
+first_av_ros = g.new_vp("float")
+last_av_rad = g.new_vp("float")
+last_av_ros = g.new_vp("float")
 
 if args.debug:
     # Make the properties internal to the graph:
-    g.vp['pos_first'] = pos_first
-    g.vp['pos_last'] = pos_last
-    g.vp['first_av_rad'] = first_av_rad
-    g.vp['first_av_ros'] = first_av_ros
-    g.vp['last_av_rad'] = last_av_rad
-    g.vp['last_av_ros'] = last_av_ros
-
-g.ep['cost_function'] = g.new_ep('float')
+    g.vp["pos_first"] = pos_first
+    g.vp["pos_last"] = pos_last
+    g.vp["first_av_rad"] = first_av_rad
+    g.vp["first_av_ros"] = first_av_ros
+    g.vp["last_av_rad"] = last_av_rad
+    g.vp["last_av_ros"] = last_av_ros
+
+g.ep["cost_function"] = g.new_ep("float")
 SHPC = util_eddies.SHPC_class(args.SHPC_dir, args.orientation)
-n_days_avg = 7 # number of days to average
+n_days_avg = 7  # number of days to average
 print("Iterating on vertices...")
 
 for n in g.vertices():
     if n.in_degree() != 0:
         # Define properties for beginning of the segment:
-        properties = node_to_prop(g.vp.inst_eddies[n][:n_days_avg],
-                                  g.gp.e_overestim, SHPC, args.orientation)
+        properties = node_to_prop(
+            g.vp.inst_eddies[n][:n_days_avg],
+            g.gp.e_overestim,
+            SHPC,
+            args.orientation,
+        )
         first_av_rad[n], first_av_ros[n] = calculate_radii_rossby(properties)
-        pos_first[n] = properties[0]["pos"] # in rad
+        pos_first[n] = properties[0]["pos"]  # in rad
 
     if n.out_degree() != 0:
         # Define properties for end of the segment:
@@ -170,18 +193,24 @@ for n in g.vertices():
             if n.in_degree() == 0 or len_seg > 2 * n_days_avg:
                 # We cannot use part of properties from the beginning
                 # of the segment.
-                properties = node_to_prop(g.vp.inst_eddies[n][- n_days_avg:],
-                                          g.gp.e_overestim, SHPC,
-                                          args.orientation)
+                properties = node_to_prop(
+                    g.vp.inst_eddies[n][-n_days_avg:],
+                    g.gp.e_overestim,
+                    SHPC,
+                    args.orientation,
+                )
             else:
                 # assertion: n.in_degree() != 0 and n_days_avg <
                 # len_seg < 2 * n_days_avg
 
                 # We can use part of the properties from the beginning
                 # of the segment.
-                properties = properties[len_seg - n_days_avg:] \
-                    + node_to_prop(g.vp.inst_eddies[n][n_days_avg:],
-                                   g.gp.e_overestim, SHPC, args.orientation)
+                properties = properties[len_seg - n_days_avg :] + node_to_prop(
+                    g.vp.inst_eddies[n][n_days_avg:],
+                    g.gp.e_overestim,
+                    SHPC,
+                    args.orientation,
+                )
 
             last_av_rad[n], last_av_ros[n] = calculate_radii_rossby(properties)
         else:
@@ -192,7 +221,7 @@ for n in g.vertices():
             last_av_rad[n] = first_av_rad[n]
             last_av_ros[n] = first_av_ros[n]
 
-        pos_last[n] = properties[- 1]["pos"] # in rad
+        pos_last[n] = properties[-1]["pos"]  # in rad
 
 t1 = time.perf_counter()
 timings.write(f"iterating on vertices: {t1 - t0:.0f} s\n")
@@ -203,20 +232,22 @@ for edge in g.edges():
     source_node = edge.source()
     target_node = edge.target()
 
-    latitude = (pos_last[source_node][1] +
-                             pos_first[target_node][1]) / 2
+    latitude = (pos_last[source_node][1] + pos_first[target_node][1]) / 2
 
     # Because of the wrapping issue (360° wrapping incorrectly to 0°),
     # we check for that here:
     lon_diff = abs(pos_last[source_node][0] - pos_first[target_node][0])
-    if lon_diff > math.radians(300): lon_diff = 2 * math.pi - lon_diff
+    if lon_diff > math.radians(300):
+        lon_diff = 2 * math.pi - lon_diff
 
-    Delta_Cent = math.sqrt((lon_diff * 6378.166175 * math.cos(latitude))**2
-                           + ((pos_last[source_node][1]
-                               - pos_first[target_node][1]) * 6335.423523)**2)
+    Delta_Cent = math.sqrt(
+        (lon_diff * 6378.166175 * math.cos(latitude)) ** 2
+        + ((pos_last[source_node][1] - pos_first[target_node][1]) * 6335.423523)
+        ** 2
+    )
 
     # Rossbies:
-    if (first_av_ros[target_node] and last_av_ros[source_node]):
+    if first_av_ros[target_node] and last_av_ros[source_node]:
         Delta_Ro = last_av_ros[source_node] - first_av_ros[target_node]
     else:
         # At least one of the rossbies is invalid.
@@ -226,17 +257,18 @@ for edge in g.edges():
     Delta_R_Vmax = last_av_rad[source_node] - first_av_rad[target_node]
 
     # Calculate the cost and assign to the edge:
-    g.ep.cost_function[edge] \
-        = math.sqrt(((Delta_Cent - delta_cent_mean) / delta_cent_std)**2
-                    + ((Delta_Ro - delta_ro_mean) / delta_ro_std)**2
-                    + ((Delta_R_Vmax - delta_r_mean) / delta_r_std)**2)
+    g.ep.cost_function[edge] = math.sqrt(
+        ((Delta_Cent - delta_cent_mean) / delta_cent_std) ** 2
+        + ((Delta_Ro - delta_ro_mean) / delta_ro_std) ** 2
+        + ((Delta_R_Vmax - delta_r_mean) / delta_r_std) ** 2
+    )
 
 t1 = time.perf_counter()
 timings.write(f"iterating on edges: {t1 - t0:.0f} s\n")
 t0 = t1
 print("Saving...")
 g.save(args.output_segments)
-print('All done')
+print("All done")
 t1 = time.perf_counter()
 timings.write(f"saving: {t1 - t0:.0f} s\n")
 timings.close()