-
Lionel GUEZ authoredLionel GUEZ authored
cost_function.py 11.22 KiB
#!/usr/bin/env python3
"""This script takes the graph of segments without cost and computes
the cost applied to edges.
Input:
-- the graph of segments without cost;
-- the SHPC.
Output: the graph of segments with cost.
The inst_eddies property of vertices is not modified by this
script. All the content of the input graph is part of the output graph
so the input file may be removed, if desired, after running this
script.
"""
import time
import math
import argparse
import graph_tool
import util_eddies
Omega = 2 * math.pi / 86164.0 # in s-1
r_Earth = 6371 # radius of the Earth, in km
def calculate_radii_rossby(properties):
"""Compute average on some instantaneous eddies of Rossby number and
radius of maximum speed contour. The required properties for each
eddy are position, radius and speed. "properties" is a list of
dictionaries. Each dictionary in the list contains the three
properties. If the speed is not defined for any eddy then the
returned value of avg_Rossby is 0.
"""
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
if abs(prop["speed"]) < 100:
avg_Rossby += prop["speed"] / (f * radius)
n_valid_Rossby += 1
avg_rad += radius # in m
avg_rad /= len(properties)
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 inter-date identification numbers for
instantaneous eddies. This function returns some properties of the
eddies, read from shapefiles: position of extremum, radius of
outermost contour or maximum speed contour, and speed. The three
properties are in a dictionary, for each eddy. So the function
returns a list of dictionaries: one dictionary for each element of
node_list.
"""
properties = []
for n in node_list:
date_index, eddy_index = util_eddies.node_to_date_eddy(n, e_overestim)
i_slice, ishape = SHPC.comp_ishape(date_index, eddy_index, 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
)
if prop["radius"] < 0:
prop["radius"] = (
SHPC.get_reader(i_slice, orientation, "outermost_contour")
.record(ishape)
.r_eq_area
)
properties.append(prop)
return properties
def search_beg(inst_eddies, max_delta, avg_fix, e_overestim):
"""Return an index ip in inst_eddies. inst_eddies is a list of
inter-date indices of instantaneous eddies. len(inst_eddies) must
be >= 1. max_delta must be >= 0. If avg_fix then inst_eddies[:ip]
is a fixed number of instantaneous eddies, if the list is long
enought. If not avg_fix then ip = bisect.bisect_right(inst_eddies,
date(inst_eddies[0]) + max_delta, key = date), although the
computation is not implemented that way.
"""
ip = min(max_delta + 1, len(inst_eddies))
if not avg_fix and ip >= 2:
d_max = (
util_eddies.node_to_date_eddy(
inst_eddies[0], e_overestim, only_date=True
)
+ max_delta
)
# {date(elem) > d_max for elem in inst_eddies[ip:]}
while (
ip >= 2
and util_eddies.node_to_date_eddy(
inst_eddies[ip - 1], e_overestim, only_date=True
)
> d_max
):
ip -= 1
# {date(elem) <= d_max for elem in inst_eddies[:ip] and
# date(elem) > d_max for elem in inst_eddies[ip:]}
# {1 <= ip <= min(max_delta + 1, len(inst_eddies))}
return ip
def search_end(inst_eddies, max_delta, avg_fix, e_overestim):
"""Return an index ip in inst_eddies. inst_eddies is a list of
inter-date indices of instantaneous eddies. len(inst_eddies) must
be >= 1. max_delta must be >= 0. If avg_fix then inst_eddies[ip:]
is a fixed number of instantaneous eddies, if the list is long
enought. If not avg_fix then ip = bisect.bisect_left(inst_eddies,
date(inst_eddies[-1]) - max_delta, key = date), although the
computation is not implemented that way.
"""
i_max = len(inst_eddies) - 1
ip = max(i_max - max_delta, 0)
if not avg_fix and ip < i_max:
d_min = (
util_eddies.node_to_date_eddy(
inst_eddies[-1], e_overestim, only_date=True
)
- max_delta
)
# {date(elem) < d_min for elem in inst_eddies[:ip]}
while (
ip < i_max
and util_eddies.node_to_date_eddy(
inst_eddies[ip], e_overestim, only_date=True
)
< d_min
):
ip += 1
# {date(elem) < d_min for elem in inst_eddies[:ip] and
# date(elem) >= d_min for elem in inst_eddies[ip:]}
# {max(len(inst_eddies) - max_delta - 1, 0) <= ip <= len(inst_eddies) - 1}
return ip
t0 = time.perf_counter()
timings = open("timings_cost.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(
"--avg_fix",
help="average over a fixed number of instantaneous eddies",
action="store_true",
)
args = parser.parse_args()
# Set some values needed for the cost function:
delta_cent_mean = 3.8481 # in km
delta_cent_std = 8.0388 # in km
delta_ro_mean = -0.0025965
delta_ro_std = 5.2168
delta_r_mean = -9.4709 # in m
delta_r_std = 8.6953e3 # in m
# Load the graph_tool file:
print("Loading graph...")
g = graph_tool.load_graph(args.input_segments)
print("Loading done...")
print("Input graph:")
print("Number of vertices:", g.num_vertices())
print("Number of edges:", g.num_edges())
print("Internal properties:")
g.list_properties()
t1 = time.perf_counter()
timings.write(f"loading: {t1 - t0:.0f} s\n")
t0 = t1
# It is useful to save the orientation to the output graph of this
# script for further processing of the output graph by other scripts:
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")
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")
SHPC = util_eddies.SHPC_class(args.SHPC_dir, args.orientation)
max_delta = 6
# maximum distance in number of eddies, over which we average, must be >= 0
print("Iterating on vertices...")
for n in g.vertices():
if n.in_degree() != 0:
# Define properties for beginning of the segment:
ip_beg = search_beg(
g.vp.inst_eddies[n], max_delta, args.avg_fix, g.gp.e_overestim
)
properties = node_to_prop(
g.vp.inst_eddies[n][:ip_beg],
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
else:
ip_beg = 0
if n.out_degree() != 0:
# Define properties for end of the segment:
if ip_beg < len(g.vp.inst_eddies[n]):
# We have to read more from the shapefiles and redefine
# properties.
ip_end = search_end(
g.vp.inst_eddies[n], max_delta, args.avg_fix, g.gp.e_overestim
)
if ip_beg <= ip_end:
# We cannot use part of properties from the beginning
# of the segment.
properties = node_to_prop(
g.vp.inst_eddies[n][ip_end:],
g.gp.e_overestim,
SHPC,
args.orientation,
)
else:
# assertion: ip_end < ip_beg < len(g.vp.inst_eddies[n])
# We can use part of the properties from the beginning
# of the segment.
properties = properties[ip_end:] + node_to_prop(
g.vp.inst_eddies[n][ip_beg:],
g.gp.e_overestim,
SHPC,
args.orientation,
)
last_av_rad[n], last_av_ros[n] = calculate_radii_rossby(properties)
else:
# The number of eddies in the segment is lower than or
# equal to the number of days over which to average. The
# values for the end of the segment will be the same as
# for the begining, except for the position.
last_av_rad[n] = first_av_rad[n]
last_av_ros[n] = first_av_ros[n]
pos_last[n] = properties[-1]["pos"] # in rad
t1 = time.perf_counter()
timings.write(f"iterating on vertices: {t1 - t0:.0f} s\n")
t0 = t1
print("Iterating on edges...")
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
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
Delta_Cent = r_Earth * math.sqrt(
(lon_diff * math.cos(latitude)) ** 2
+ (pos_last[source_node][1] - pos_first[target_node][1]) ** 2
)
# Rossby numbers:
if first_av_ros[target_node] != 0 and last_av_ros[source_node] != 0:
Delta_Ro = first_av_ros[target_node] - last_av_ros[source_node]
else:
# At least one of the Rossby numbers (computed by
# calculate_radii_rossby) is invalid.
Delta_Ro = 0
# R_Vmax 1 and 2 already exist, just get the delta
Delta_R_Vmax = first_av_rad[target_node] - last_av_rad[source_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
)
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")
t1 = time.perf_counter()
timings.write(f"saving: {t1 - t0:.0f} s\n")
timings.close()