Skip to content
Snippets Groups Projects
  • Lionel GUEZ's avatar
    7360b14c
    Do not loop on years · 7360b14c
    Lionel GUEZ authored
    Since we no longer load all the contents of shapefiles for a given
    year in advance, we no longer need to process the segments year by
    year. Motivation: this will allow us to get free of the imposed
    arborescence of SHPC by year.
    7360b14c
    History
    Do not loop on years
    Lionel GUEZ authored
    Since we no longer load all the contents of shapefiles for a given
    year in advance, we no longer need to process the segments year by
    year. Motivation: this will allow us to get free of the imposed
    arborescence of SHPC by year.
cost_function.py 10.18 KiB
#!/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

# 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

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