cost_function.py

#!/usr/bin/env python3

"""A script that takes the graph of segments without cost functions
and computes the non-local cost functions attributed to edges.

Input:

- "node_id_param.json", expected to be in the current directory

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

- shapefiles, specified as command line arguments

Output: the graph of segments with cost functions,
'segments_cost_functions.gt'

"""

import graph_tool
import time
import json
import math
from os import path
import shapefile
import datetime
from numpy import loadtxt
import report_graph
import util_eddies
import sys
import bisect

def calculate_radii_and_rossby(start, end, inc, segment, e_overestim, handlers,
                               array_d_init):

    radii = 0 #[m]
    rossby = 0 #[1/s]

    days_modifier = 0

    for i in range(start, end, inc):
        current_eddy = report_graph.node_to_date_eddy(segment[i], e_overestim)
        i_SHPC = get_SHPC(array_d_init, current_eddy['date_index'])

        # calculate the location in the shapefile
        location = util_eddies.comp_ishape(handlers[i_SHPC],
                                           current_eddy['date_index'],
                                           current_eddy['eddy_index'])

        # now that we have the location in the shapefiles, we need to
        # get the radius and the rossby number
        shapeRec = handlers[i_SHPC]["readers"]["extremum"].shapeRecord(location)
        # rossby:
        lat_in_deg = shapeRec.shape.points[0][1]
        #[deg]
        f = 2*2*math.pi/(24*3600)*math.sin(math.radians(lat_in_deg)) # [1/s]

        V_max = shapeRec.record[4] #[m/s]
        R_Vmax = handlers[i_SHPC]["readers"]["max_speed_contour"]\
            .record(location)['r_eq_area'] * 1000 #[m]

        if (V_max < 100):
            # calculate Ro and Delta_Ro
            Ro = V_max / (f * R_Vmax) #[]
        else:
            Ro = 0
            days_modifier += 1

        ####### RADII #######
        radii += R_Vmax # [m]
        ####### ROSSBY ######
        rossby += Ro # []

    return {"radii": radii, "rossby": rossby, "days_modifier": days_modifier}

def get_SHPC(array_d_ini, date_index):
    i_SHPC = bisect.bisect(array_d_init, date_index)
    assert i_SHPC >= 1
    return i_SHPC - 1

t0 = time.perf_counter()
timings = open("timings.txt", "w")

if len(sys.argv) == 1:
    sys.exit("Required arguments: SHPC-directory [SHPC-directory] ...")

# Grab e_overestim:
with open("node_id_param.json") as f: node_id_param = json.load(f)
# assign attributes to the whole graphs
e_overestim = node_id_param["e_overestim"]

################################################
# set some values needed for the cost function #
################################################

delta_cent_mean = 3.8481 # [km]
delta_cent_std = 8.0388

delta_ro_mean = -0.0025965 # []
delta_ro_std = 5.2168

delta_r_mean = -0.0094709 * 1000 #[m]
delta_r_std = 8.6953 * 1000

# Load the graph_tool file:
print('Loading graph...')
g = graph_tool.Graph()

try:
    g.load('segments.gt')
except FileNotFoundError:
    g.load('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} s\n")
t0 = t1

g.vp['pos_first'] = g.new_vp('object')
g.vp['pos_last'] = g.new_vp('object')
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['nl_cost_function'] = g.new_ep('float')

# Set up the list of SHPC:

handlers = [util_eddies.open_shpc(shp_dir) for shp_dir in sys.argv[1:]]

array_d_init = [handler["d_init"] for handler in handlers]
# (create the list once and for all)

# change if there is a change over the number of days to average
num_of_days_to_avg = 7

print("Iterating on vertices...")

for n in g.vertices():
    segment = g.vp.segment[n]
    num_of_days = len(segment)

    # Calculate the indexes and dates:

    first = report_graph.node_to_date_eddy(segment[0], e_overestim)
    first_SHPC = get_SHPC(array_d_init, first['date_index'])

    last = report_graph.node_to_date_eddy(segment[-1], e_overestim)
    last_SHPC = get_SHPC(array_d_init, last['date_index'])

    # calculate the location in the shapefile
    first_loc = util_eddies.comp_ishape(handlers[first_SHPC],
                                        first['date_index'],
                                        first['eddy_index'])
    last_loc = util_eddies.comp_ishape(handlers[last_SHPC],
                                       last['date_index'],
                                       last['eddy_index'])

    # grab the centers:
    first_pos = handlers[first_SHPC]["readers"]["extremum"]\
        .shape(first_loc).points[0]
    last_pos = handlers[last_SHPC]["readers"]["extremum"]\
        .shape(last_loc).points[0]

    # Store positions in the vertex properties:
    g.vp.pos_first[n] = first_pos # [deg, deg]
    g.vp.pos_last[n] = last_pos # [deg, deg]

    if (num_of_days > num_of_days_to_avg):
        # The segment is longer than the number of days over which to average

        first_radii = 0 # [m]
        last_radii = 0 # [m]

        first_rossby = 0 # []
        last_rossby = 0 # []

        # First 7 days calculation
        first_res = calculate_radii_and_rossby(0, num_of_days_to_avg,
                                               1, segment, e_overestim,
                                               handlers, array_d_init)

        # average and assign radii
        first_radii = first_res['radii'] / num_of_days_to_avg
        g.vp.first_av_rad[n] = first_radii

        # grab the days modifier
        modifier = first_res['days_modifier']

        if (num_of_days_to_avg - modifier > 0):
            # Average and assign the rossbies:
            first_rossby = first_res['rossby'] / (num_of_days_to_avg - modifier)
            g.vp.first_av_ros[n] = first_rossby

        # Last 7 days calculation
        last_res = calculate_radii_and_rossby(len(segment) - 1, len(segment) -
                                              (num_of_days_to_avg + 1), -1,
                                              segment, e_overestim, handlers,
                                              array_d_init)

        # Average and assign the last radii
        last_radii = last_res['radii'] / num_of_days_to_avg
        g.vp.last_av_rad[n] = last_radii

        # grab the days modifier
        modifier = last_res['days_modifier']

        if (num_of_days_to_avg - modifier > 0):
            # Average and assign the rossbies:
            last_rossby = last_res['rossby'] / (num_of_days_to_avg - modifier)
            g.vp.last_av_ros[n] = last_rossby
    else:
        # The number of eddies in a segment is lower than the number
        # of days over which to average. The values will be the same
        # except for the positions.
        res = calculate_radii_and_rossby(0, num_of_days, 1, segment,
                                         e_overestim, handlers, array_d_init)

        # grab the days modifier
        modifier = res['days_modifier']

        if (num_of_days - modifier > 0):
            # Average and assign the rossbies:
            rossby = res['rossby'] / (num_of_days - modifier)
            g.vp.first_av_ros[n] = rossby
            g.vp.last_av_ros[n] = rossby

        # Average and assign the radii
        radii = res['radii'] / num_of_days
        g.vp.first_av_rad[n] = radii
        g.vp.last_av_rad[n] = radii

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

for edge in g.edges():
    source_node = edge.source()
    target_node = edge.target()

    cf = -10000

    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)

    # calculate the first term
    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:
        print("At least one of the rossbies is invalid.")
        #Delta_Ro = delta_ro_mean
        Delta_Ro = 0

    # Calculate the second term
    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]

    # Calculate the third term
    third_term = ((Delta_R_Vmax - delta_r_mean)/delta_r_std) ** 2

    #############################
    # calculate the cost function
    #############################

    cf = math.sqrt(first_term + second_term + third_term)

    # assign as weight to the edge
    g.ep.nl_cost_function[edge] = cf


t1 = time.perf_counter()
timings.write(f"iterating on edges: {t1 - t0} s\n")
t0 = t1
print("Saving...")
g.save('segments_cost_functions.gt')
print('All done')
t1 = time.perf_counter()
timings.write(f"saving: {t1 - t0} s\n")
timings.close()