diff --git a/Trajectories/cost_function.py b/Trajectories/cost_function.py
index 9f31bdca58d855558aa3914cc26322df56c571d1..bb15b146b665384f78a07bf137e8a02da1440b9b 100755
--- a/Trajectories/cost_function.py
+++ b/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()