cost_function.py

#!/usr/bin/env python3

"""This script takes a graph of segments and computes the cost applied
to edges.

Input:

-- the graph of segments;

-- the SHPC.

Output: the graph of segments with cost.

The inst_eddies property of vertices is not modified by this
script. The cost is an internal property of the output graph. This
internal property is created if it was not already in the input
graph. The cost is overwritten if it was already in the input graph.

"""

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
FACTOR = 2 * Omega * 1000  # 1000 for conversion of km to m


def avg_radii_rossby(latitude_list, speed_list, radius_list):
    """Compute average on some instantaneous eddies of Rossby number and
    radius of maximum speed contour. The required properties for each
    eddy are latitude, radius and speed. 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 latitude, speed, radius in zip(latitude_list, speed_list, radius_list):
        if abs(speed) < 100:
            avg_Rossby += speed / (FACTOR * math.sin(latitude) * radius)
            # (convert radius to m)

            n_valid_Rossby += 1

        avg_rad += radius

    avg_rad /= len(radius_list)

    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
    function returns a tuple of lists: one list for each property.

    """

    longitude_list = []
    latitude_list = []
    speed_list = []
    radius_list = []

    for n in node_list:
        i_slice, ishape = SHPC.comp_ishape_n(n, e_overestim, orientation)
        shapeRec = SHPC.get_reader(
            i_slice, orientation, "extremum"
        ).shapeRecord(ishape)
        longitude, latitude = [
            math.radians(x) for x in shapeRec.shape.points[0]
        ]
        longitude_list.append(longitude)
        latitude_list.append(latitude)
        speed_list.append(shapeRec.record.speed)
        radius = (
            SHPC.get_reader(i_slice, orientation, "max_speed_contour")
            .record(ishape)
            .r_eq_area
        )

        if radius < 0:
            radius = (
                SHPC.get_reader(i_slice, orientation, "outermost_contour")
                .record(ishape)
                .r_eq_area
            )

        radius_list.append(radius)

    return longitude_list, latitude_list, speed_list, radius_list


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
        )
        # assertion: date(inst_eddies[0]) <= d_max
        # assertion: 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

        # assertion: date(elem) <= d_max for elem in inst_eddies[:ip]
        # and date(elem) > d_max for elem in inst_eddies[ip:]

    # assertion: 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:]
    contains a fixed number of instantaneous eddies, if the list is
    long enough. 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
        )
        # assertion: date(inst_eddies[i_max]) >= d_min
        # assertion: 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

        # assertion: date(elem) < d_min for elem in inst_eddies[:ip] and
        # date(elem) >= d_min for elem in inst_eddies[ip:]

    # assertion: 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.4709e-3  # in km
delta_r_std = 8.6953  # in km

# Load the graph_tool file:

print("Loading graph...")
g = graph_tool.load_graph(
    args.input_segments,
    ignore_vp=[
        "pos_first",
        "pos_last",
        "first_av_rad",
        "first_av_ros",
        "last_av_rad",
        "last_av_ros",
    ],
    ignore_ep=["cost_function"],
)
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")
longitude_prop = g.new_vp("vector<double>")
latitude_prop = g.new_vp("vector<double>")
speed_prop = g.new_vp("vector<double>")
radius_prop = g.new_vp("vector<double>")

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)

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():
    (
        longitude_prop[n],
        latitude_prop[n],
        speed_prop[n],
        radius_prop[n],
    ) = node_to_prop(
        g.vp.inst_eddies[n], g.gp.e_overestim, SHPC, args.orientation
    )

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
        )
        first_av_rad[n], first_av_ros[n] = avg_radii_rossby(
            latitude_prop[n][:ip_beg],
            speed_prop[n][:ip_beg],
            radius_prop[n][:ip_beg],
        )
        pos_first[n] = [longitude_prop[n][0], latitude_prop[n][0]]  # 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
            )
            last_av_rad[n], last_av_ros[n] = avg_radii_rossby(
                latitude_prop[n][ip_end:],
                speed_prop[n][ip_end:],
                radius_prop[n][ip_end:],
            )
        else:
            # assertion: ip_beg == len(g.vp.inst_eddies[n])
            # assertion: len(g.vp.inst_eddies[n]) <= max_delta + 1
            # 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] = [longitude_prop[n][-1], latitude_prop[n][-1]]  # 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
        # avg_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()