cost_function.py

#!/usr/bin/env python3

# A script that takes a segmented graph in the gt format and performs the 
# non-local cost function calculation where each edge will have a cost
# function assigned to it. 
# Inputs (to be changed within the script):
    # Orientation
    # node_id_param_file location and name
    # gt graph name and location
    # shapefiles names and location
# Output: a gt graph with cost functions and an xml graph with cost functions

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

def calculate_radii_and_rossby(start, end, inc, segment, e_overestim, handlers):
    
    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)
        date = datetime.date(1950, 1, 1) \
            + datetime.timedelta(current_eddy['date_index'])
        year = date.year
        
        # calculate the location in the shapefile
        location = util_eddies.comp_ishape(handlers[year],
                                           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[year]["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[year]["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}

###############################
# 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

orientation = input("Enter orientation (anti or cyclo): ") 

# Load the graph_tool file:

t0 = time.perf_counter()

print('Loading gt file...')
g = graph_tool.Graph()
g.load('segmented.gt')
t1 = time.perf_counter()
print(f'Loading done: {g}')
print(t1 - t0, "s")

g.set_fast_edge_removal()
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 dictionarys for the shapefiles and the ishape_last files:

shp_tr_dir = input('Directory containing %Y/SHPC_(anti|cyclo): ')
handlers = {}
    
for year in range(1993,2019):
    shp_dir = path.join(shp_tr_dir, f'{year}/SHPC_{orientation}')
    handlers[year] = util_eddies.open_shpc(shp_dir)

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

# We iterate on the years because the segments are non sequential in
# time. For each year, we are going to process only the segments that
# begin in that year.
for yr in range(1993, 2019):
    print(f'Current year: {yr}') 
    
    # iterate on the vertices 
    for n in g.vertices():
        # Get the segment and the number of days
        segment = g.vp.segment[n]      
        
        # calculate the indexes and dates
        first = report_graph.node_to_date_eddy(segment[0], e_overestim)
        
        # get year
        first_date = datetime.date(1950, 1, 1) \
            + datetime.timedelta(first['date_index'])
        first_year = first_date.year
        
        # chech if we are in the current year
        
        if (first_year == yr):
            num_of_days = len(segment)
            
            # start processing
            
            last = report_graph.node_to_date_eddy(segment[-1], e_overestim)
            last_date = datetime.date(1950, 1, 1) \
                + datetime.timedelta(last['date_index'])
            last_year = last_date.year
            
            # calculate the location in the shapefile
            first_loc = util_eddies.comp_ishape(handlers[first_year],
                                                first['date_index'],
                                                first['eddy_index'])
            last_loc = util_eddies.comp_ishape(handlers[last_year],
                                               last['date_index'],
                                               last['eddy_index'])
            
            # grab the centers
            
            first_pos = handlers[first_year]["readers"]["extremum"]\
                .shape(first_loc).points[0]
            last_pos = handlers[last_year]["readers"]["extremum"]\
                .shape(last_loc).points[0]
            
            ##### STORE POSITIONS IN THE VPS ######
            g.vp.pos_first[n] = first_pos # [deg, deg]
            g.vp.pos_last[n] = last_pos # [deg, deg]
        
            # if the segments are longer than the # of days over which to avg
        
            if (num_of_days > num_of_days_to_avg):
                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)
                
                # 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
                else:
                    # there is division by zero, average rossby is undefinied
                    pass
                    
                # 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)
                    
                # 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:
                    # there is division by zero, average rossby is undefinied
                    pass
            # else, the number of eddies in a segment is lower than the # of 
            # days over which to average, the values will be the same except
            # for the positions
            else:
                res = calculate_radii_and_rossby(0, num_of_days, 1, segment,
                                                 e_overestim, handlers)
                
                # 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
                else:
                    # there is division by zero, average rossby is undefinied
                    pass
                    
                # 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
            
###############################
# Calculate the cost function #
###############################

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

    # calculate Delta_cent: numbers used for conversion obtained from:
        # https://stackoverflow.com/questions/1253499/simple-calculations-for-working-with-lat-lon-and-km-distance
    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    


################################
# Writing to an GT Graph File #
################################

#g.save("segments_cyclo_gl_cf.xml")
#print("Done writing xml.")
g.save(f'segmented_{orientation}_cf.gt')
print("Done writing gt.")
g.save(f'segmented_{orientation}_cf.xml')
print("Done writing xml.")
#g.save(f'{orientation}_segmented_cf.gv', fmt = "dot")
#print("Done writing dot file.")

print('All done')