cost_function.py

#!/usr/bin/env python3

"""This script takes the graph of segments without cost functions and
computes the non-local cost functions applied to edges.

Input:

- "e_overestim.txt", expected to be in the current directory

- the graph of segments without cost functions, "segments.gt" or
  "segments.graphml", expected to be in the current directory

- the SHPC

Output: the graph of segments with cost functions,
segments_cost_functions.gt or segments_cost_functions.graphml. The
inst_eddies property of vertices is not modified by this script.

"""

import graph_tool
import time
import math
import report_graph
import util_eddies
import argparse

Omega = 2 * math.pi / 86164.

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.

    """

    avg_rad = 0 # in m
    avg_Rossby = 0
    n_valid_Rossby = 0

    for prop in properties:
        f = 2 * Omega * math.sin(math.radians(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):
    """node_list is a list of node 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.

    """

    properties = []

    for n in node_list:
        date_index, eddy_index = report_graph.node_to_date_eddy(n, e_overestim)
        i_slice = SHPC.get_slice(date_index)
        ishape = SHPC.comp_ishape(date_index, eddy_index, i_slice)
        shapeRec = SHPC.get_reader(i_slice, "extremum").shapeRecord(ishape)
        prop = {"pos": shapeRec.shape.points[0], "speed": shapeRec.record.speed}
        prop["radius"] = SHPC.get_reader(i_slice, "max_speed_contour")\
            .record(ishape).r_eq_area

        if prop["radius"] < 0:
            prop["radius"] = SHPC.get_reader(i_slice, "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("--graphml", action = "store_true",
                    help = "save to graphml format")
args = parser.parse_args()
with open("e_overestim.txt") as f: e_overestim = int(f.read())

# Set some values needed for the cost function:
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_std = 8.6953e3

# Load the graph_tool file:

print('Loading graph...')

try:
    g = graph_tool.load_graph('segments.gt')
except FileNotFoundError:
    g = graph_tool.load_graph('segments.graphml')

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

g.vp['pos_first'] = g.new_vp('vector<double>')
g.vp['pos_last'] = g.new_vp('vector<double>')
g.vp['first_av_rad'] = g.new_vp('float')
g.vp['first_av_ros'] = g.new_vp('float')
g.vp['last_av_rad'] = g.new_vp('float')
g.vp['last_av_ros'] = g.new_vp('float')
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
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], e_overestim,
                                  SHPC)
        g.vp.first_av_rad[n], g.vp.first_av_ros[n] \
            = calculate_radii_rossby(properties)
        g.vp.pos_first[n] = properties[0]["pos"] # in degrees

    if n.out_degree() != 0:
        # Define properties for end of the segment:

        len_seg = len(g.vp.inst_eddies[n])

        if n.in_degree() == 0 or len_seg > n_days_avg:
            # We have to read more from the shapefiles and redefine
            # properties.

            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:],
                                          e_overestim, SHPC)
            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:],
                                   e_overestim, SHPC)

            g.vp.last_av_rad[n], g.vp.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.
            g.vp.last_av_rad[n] = g.vp.first_av_rad[n]
            g.vp.last_av_ros[n] = g.vp.first_av_ros[n]

        g.vp.pos_last[n] = properties[- 1]["pos"] # in degrees

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()

    lat_for_conv = (g.vp.pos_last[source_node][1] +
                    g.vp.pos_first[target_node][1]) / 2
    # (latitude needed for conversion of degrees to kilometers)
    lat_for_conv = math.radians(lat_for_conv) # need to convert to radians

    # Because of the wrapping issue (360° wrapping incorrectly to 0°),
    # we check for that here:
    lon_diff = abs(g.vp.pos_last[source_node][0]
                   - g.vp.pos_first[target_node][0])
    if lon_diff > 300: lon_diff = 360 - lon_diff

    Delta_Cent = math.sqrt((lon_diff * 111.32 * math.cos(lat_for_conv))**2
                           + ((g.vp.pos_last[source_node][1]
                               - g.vp.pos_first[target_node][1]) * 110.574)**2)
    first_term = ((Delta_Cent - delta_cent_mean)/delta_cent_std)**2

    # Rossbies:
    if (g.vp.first_av_ros[target_node] and g.vp.last_av_ros[source_node]):
        Delta_Ro = g.vp.last_av_ros[source_node] \
            - g.vp.first_av_ros[target_node]
    else:
        # At least one of the rossbies is invalid.
        # Delta_Ro = delta_ro_mean
        Delta_Ro = 0

    second_term = ((Delta_Ro - delta_ro_mean)/delta_ro_std)**2

    # R_Vmax 1 and 2 already exist, just get the delta
    Delta_R_Vmax = g.vp.last_av_rad[source_node] \
        - g.vp.first_av_rad[target_node]

    third_term = ((Delta_R_Vmax - delta_r_mean)/delta_r_std)**2

    # Calculate the cost function and assign as weight to the edge:
    g.ep.cost_function[edge] = math.sqrt(first_term + second_term + third_term)

t1 = time.perf_counter()
timings.write(f"iterating on edges: {t1 - t0:.0f} s\n")
t0 = t1
print("Saving...")

if args.graphml:
    g.save('segments_cost_functions.graphml')
else:
    g.save('segments_cost_functions.gt')

print('All done')
t1 = time.perf_counter()
timings.write(f"saving: {t1 - t0:.0f} s\n")
timings.close()