sfoi_utils.py

import json
from netCDF4 import Dataset
import netCDF4
import numpy as np
from pathlib import Path
import datetime
import copy

def calcG(m,n,junctions,basin_dict,reachlist='reach_ids_all'):
    #define G matrix. tiny tweaks from Confluence version: just removed the class objects
    
    G=np.zeros((m,n))

    for junction in junctions:
        row=junction['row_num']
        upcols=list()
        for upflow in junction['upflows']:
            try:
                 kup=basin_dict[reachlist].index(str(upflow))
                 upcols.append(kup)
            except: 
                print('did not find reach:',upflow)
                print('... in junction',junction)

        downcols=list()
        for downflow in junction['downflows']:
            try:
                 kdn=basin_dict[reachlist].index(str(downflow))
                 downcols.append(kdn)
            except:
                 print('did not find reach',downflow)
                 print('... in junction',junction)

        for upcol in upcols:
            G[row,upcol]=1
        for downcol in downcols:
            G[row,downcol]=-1

    return G

def get_basin_data(basin_json,index_to_run):
    """Extract reach identifiers and return dictionary.
    
    Dictionary is organized with a key of reach identifier and a value of
    SoS file as a Path object.
    """
    #index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX"))
    #index = 0

    if index_to_run == -235:
        index=int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX"))
    else:
        index=index_to_run
        print('Running offline, with index = ',index)


    with open(basin_json) as json_file:
        data = json.load(json_file)


    # ~~Error Handling~~
    # there is an issue where running on one basin causes an index error here
    # here we will check to see if the index we are looking for exists
    # this SHOULD allways exist if there is more than one reach
    # There should be a more elegant way to check the number of sets, 
    # but the data structure changes when only one set is written out.

    try:
        test_index = data[index]
        return {
            #"basin_id" : int(data[index]["basin_id"]),
            "basin_id" : data[index]["basin_id"], #hope it's ok not to have basin ids always integers?
            "reach_ids" : data[index]["reach_id"],
            "sos" : data[index]["sos"],
            "sword": data[index]["sword"]
        }
    except:
                return {
            #"basin_id" : int(data["basin_id"]),
            "basin_id" : data["basin_id"], #hope it's ok not to have basin ids always integers?
            "reach_ids" : data["reach_id"],
            "sos" : data["sos"],
            "sword": data["sword"]
        }

def CreateJunctionList(basin_dict,sword_dict,list_order='default'):
     # create list of junctions
     junctions=list()

     junctions_valid=True
    
     if list_order=='default':
         all_reachids=basin_dict['reach_ids_all']
     elif list_order=='reorder':
         all_reachids=basin_dict['reach_ids_all_reorder']

     for reach in all_reachids:
         reach=np.int64(reach)
         
         k=np.argwhere(sword_dict['reach_id'] == reach)
         k=k[0,0]

         # extract reach dictionary for reach k
         sword_data_reach=pull_sword_attributes_for_reach(sword_dict,k) 

         #1 try adding the upstream junction
         junction_up=dict()    
         junction_up['originating_reach_id']=reach

         #1.1 add the reaches upstream of this junction
         junction_up['upflows']=list()
         for i in range(sword_data_reach['n_rch_up']):
             junction_up['upflows'].append(sword_data_reach['rch_id_up'][i] )

         #1.2 for one of the reaches upstream of this junction, add all their downstream reaches         
         if len(junction_up['upflows'])>0:
             junction_up['downflows']=list()

             # sometimes sword says there are upstream reaches and there actually isnt
             # in these cases skip the reach and raise a warning
             if not any(junction_up['upflows']):
                warnings.warn(f'Upstream reaches not found for reach {reach}')
                junctions_valid=False
                continue

             kup=np.argwhere(sword_dict['reach_id'] == junction_up['upflows'][0])
             kup=kup[0,0]
             sword_data_reach_up=pull_sword_attributes_for_reach(sword_dict,kup)
             for j in range(sword_data_reach_up['n_rch_down']):
                 junction_up['downflows'].append(sword_data_reach_up['rch_id_dn'][j] )

             AlreadyExists,AllReachesInReachFile=ChecksPriorToAddingJunction(junctions,basin_dict,junction_up)

             if not AlreadyExists and AllReachesInReachFile:
                 junctions.append(junction_up)

         #2 try adding the downstream junction
         junction_dn=dict()
         junction_dn['originating_reach_id']=reach #just adding this for bookkeeping/debugging purposes

         #2.1 add the reaches downstream of this junction
         junction_dn['downflows']=list()
         for i in range(sword_data_reach['n_rch_down']):
             junction_dn['downflows'].append(sword_data_reach['rch_id_dn'][i] )            

         #2.2 for one of the reaches downstream of the junction, add all their upstream reaches
         if len(junction_dn['downflows'])>0:
             junction_dn['upflows']=list()

             # sometimes sword says there are downstream reaches and there actually isnt
             # in these cases skip the reach and raise a warning
             if not any(junction_dn['downflows']):
                warnings.warn(f'Downstream reaches not found for reach {reach}')
                junctions_valid=False
                continue

             kdn=np.argwhere(sword_dict['reach_id'] == junction_dn['downflows'][0])
             kdn=kdn[0,0]
             sword_data_reach_dn=pull_sword_attributes_for_reach(sword_dict,kdn)
             for j in range(sword_data_reach_dn['n_rch_up']):
                 junction_dn['upflows'].append(sword_data_reach_dn['rch_id_up'][j] )

             AlreadyExists,AllReachesInReachFile=ChecksPriorToAddingJunction(junctions,basin_dict,junction_dn)

             if junction_dn['upflows']==[0]:
                 print('gotcha!')
                 print('junction=',junction_dn)
                 print('kdn=',kdn)
                 print('junction down=',junction_dn['downflows'][0])
                 print(sword_dict['reach_id'][kdn])

             if not AlreadyExists and AllReachesInReachFile:
             #if not AlreadyExists:
                 junctions.append(junction_dn) 
                
     return junctions

def extract_sword(sword_dir,basin_dict=None,fname=None):
    """Extracts and stores SWORD data in sword_dict.

    Parameters
    ----------
    """
    if not basin_dict and not fname:
        print('oops must have one or the other')
        return None
    
    if basin_dict:
        swordfile=sword_dir.joinpath(basin_dict['sword'])
    else:
        swordfile=sword_dir.joinpath(fname)

    sword_dataset=Dataset(swordfile)

    sword_dict={} #organized by field rather than by reaches

    # grab sizes of the data
    dimfields=['orbits','num_domains','num_reaches']
    for field in dimfields:
        sword_dict[field]=sword_dataset['reaches'].dimensions[field].size    

    # grab data    
    reachfields=['reach_id','facc','n_rch_up','n_rch_down','rch_id_up','rch_id_dn','swot_obs','swot_orbits','x','y','width']
    for field in reachfields:
        sword_dict[field]=sword_dataset['reaches/' + field][:]

    sword_dataset.close()
    
    return sword_dict

def get_all_sword_reach_in_basin(basin_dict,sword_dict):
    
    Verbose=False

    # find all those that match the basin id 
    BasinLevel=len(str(basin_dict['basin_id']))

    # basin_reach_list_all includes all reaches in SWORD that match the current basin id
    basin_reach_list_all=[]
    for reachid in sword_dict['reach_id']:
        reachidstr=str(reachid)
        if reachidstr[0:BasinLevel] == str(basin_dict['basin_id']):
            basin_reach_list_all.append(reachid)

    if Verbose:
        print('There are a total of',len(basin_reach_list_all),'reaches in SWORD for this basin')

    # create reach_ids_all list
    nadd=0
    basin_dict['reach_ids_all']=[]
    for reachid in basin_reach_list_all:
        # if str(reachid) not in basin_dict['reach_ids']:
        #     nadd+=1
        basin_dict['reach_ids_all'].append(str(reachid))

    #if Verbose:
    #   print('Total of ',nadd, 'reaches in SWORD that were not in basin json')

    return basin_dict

def apply_sword_patches(sword_dict_in,fname='sword_patches_v216.json'):
    # this is included here to test custom patches. 
    # not run as part of normal confluence runs
    #patch_json = Path("/home/mdurand_umass_edu/dev-confluence/mnt/").joinpath('sword_patches_v216.json')
    # patch_json = Path("/Users/mtd/Analysis/SWOT/Discharge/Confluence/ohio_offline_runs/mnt/").joinpath('sword_patches_v216.json')
    patch_json = Path(".").joinpath(fname)

    print('using patches from:',patch_json)
    
    with open(patch_json) as json_file:
        patch_data = json.load(json_file)
        
    Verbose=False

    reaches_to_patch=list(patch_data['reach_data'].keys())

    if Verbose:
        print('Read in patches for:',len(reaches_to_patch))
        print('... for reaches: ',list(reaches_to_patch))

    # sword_dict_out=sword_dict_in.deepcopy()
    sword_dict_out=copy.deepcopy(sword_dict_in)

    for reachid in reaches_to_patch:
       
        try:
            k=np.argwhere(sword_dict_out['reach_id'][:]==np.int64(reachid))
            k=k[0,0]
        except: 
            if Verbose:
                print(reachid , 'is not in this domain. not patching')
            continue

        if Verbose:
           print('Patching reach:',reachid)

        for data_element in patch_data['reach_data'][reachid]:

            if data_element != 'metadata':

                if data_element == 'n_rch_up' or data_element == 'n_rch_down' or data_element == 'facc':
                    data_type='scalar'
                elif data_element == 'rch_id_up' or data_element == 'rch_id_dn':
                     data_type='vector'
                else:
                     print('unknown data type found in patch! crash imminent...')

                #if Verbose:
                   #print('  Patching data element:',data_element)
                   #print('    In the patch:',patch_data['reach_data'][reachid][data_element])
                   #if data_type == 'vector':
                   #    print('    In SWORD:',sword_dict[data_element][:,k])
                   #elif data_type == 'scalar':
                   #    print('    In SWORD:',sword_dict[data_element][k])

                # apply patch
                if data_type=='vector':
                    sword_dict_out[data_element][:,k]=patch_data['reach_data'][reachid][data_element]
                elif data_type=='scalar':
                    sword_dict_out[data_element][k]=patch_data['reach_data'][reachid][data_element]

                
                #if Verbose:
                #   if data_type=='vector':
                #       print('    In SWORD after fix:',sword_dict[data_element][:,k])
                #   elif data_type=='scalar':
                #       print('    In SWORD after fix:',sword_dict[data_element][k])

    return sword_dict_out

def pull_sword_attributes_for_reach(sword_dict,k):
     """
     Pull out needed SWORD data from the continent dataset arrays for a particular reach    
     """

     sword_data_reach={}
     # extract all single-dimension variables, including number of orbits and reach ids needed for multi-dim vars
     for key in sword_dict:
         if np.shape(sword_dict[key]) == (sword_dict['num_reaches'],):
             sword_data_reach[key]=sword_dict[key][k]

     # extract multi-dim vars
     for key in sword_dict:            
         if key == 'rch_id_up':
             sword_data_reach[key]=sword_dict[key][0:sword_data_reach['n_rch_up'],k]
         elif key == 'rch_id_dn':
             sword_data_reach[key]=sword_dict[key][0:sword_data_reach['n_rch_down'],k]
         elif key == 'swot_orbits':
             sword_data_reach[key]=sword_dict[key][0:sword_data_reach['swot_obs'],k]

     return sword_data_reach

def ChecksPriorToAddingJunction(junctions,basin_dict,junction_to_check):
    #check to see if this one already exists
     AlreadyExists=False
     for junction in junctions:
         if junction['upflows']==junction_to_check['upflows'] and junction['downflows']==junction_to_check['downflows']: 
             AlreadyExists=True    

     #check to see if all reaches we've identified area in the basin 
     AllReachesInReachFile=True
     for r in junction_to_check['upflows']:
         if str(r) not in basin_dict['reach_ids_all']:
             AllReachesInReachFile=False
     for r in junction_to_check['downflows']:
         if str(r) not in basin_dict['reach_ids_all']:
             AllReachesInReachFile=False            

     return AlreadyExists,AllReachesInReachFile

def get_average_gage_q(reachid,results,priors,gauge_reach,agency='USGS'):

    # there is a bug in here somewhere!

    method='all'

    # get agency id
    gdx=np.where(priors[agency][agency+'_reach_id'][:] == np.int64(reachid) )[0]
    agencyid=str(netCDF4.chartostring( priors[agency][agency+'_id'][gdx,:] ))

    print('processing',reachid,'with agency id=',agencyid)
    
    # 1) get list of times with data
    rdx = np.where(results['reaches']['reach_id'][:] == np.int64(reachid) )    
    
    ts=list(results['reaches']['time'][rdx][0])
    
    swot_ts = datetime.datetime(2000,1,1,0,0,0)
    
    time_strs = []
    for t in ts:
        if t==-999999999999:
            time_strs.append('NODATA')
        else:
            time_strs.append( (swot_ts + datetime.timedelta(seconds=t)).strftime('%Y-%m-%dT%H:%M:%S') )
    
    # print('... there are a total of ', len(time_strs),'observed times.')
    # print(time_strs)    
    
    # 2) pull gauge time and discharge data
    reach_gauge_index = np.where(gauge_reach == np.int64(reachid)) # hm. but don't use this with priors

    # Get discharge and filter out missing values
    missing = priors[agency][agency+"_q"]._FillValue
    # gauge_discharge = priors["USGS"]["USGS_q"][reach_gauge_index].filled()[0] #this is wrong
    gauge_discharge = priors[agency][agency+"_q"][gdx].filled()[0]
    nonmissing_indexes_g = np.where(gauge_discharge != missing)
    gauge_discharge = gauge_discharge[nonmissing_indexes_g]
    #print(f"Number of gauge discharge values: {len(gauge_discharge)}.")

    # Get time and filter out missing values
    # gauge_time_ordinal = priors["USGS"]["USGS_qt"][reach_gauge_index].filled().astype(int)[0] #wrong
    gauge_time_ordinal = priors[agency][agency+"_qt"][gdx].filled().astype(int)[0]
    gauge_time_ordinal = gauge_time_ordinal[nonmissing_indexes_g]
    #print(f"Number of gauge time values: {len(gauge_time)}.")

    # Convert time from ordinal value
    gauge_time = [ datetime.datetime.fromordinal(gt).strftime("%Y%m%d") for gt in gauge_time_ordinal ]    

    if method=='true':
        # 3) get gague data and average them
        Qgages=[]
        tgages=[]
    
        for time_str in time_strs:
            if not time_str=='NODATA':        
                try: 
                    idx=gauge_time.index(time_str[0:10].replace('-',''))
                    Qgages.append(gauge_discharge[idx])
                    tgages.append(gauge_time_ordinal[idx])
                except:
                    Qgages.append(np.nan)
                    tgages.append(np.nan)
                        
    
        # Qbar=sum(Qgages)/len(Qgages)
        Qbar=np.nanmean(np.array(Qgages))
    if method =='all':
        if time_strs[0][0:10].replace('-','') in gauge_time:
            idxstart=gauge_time.index(time_strs[0][0:10].replace('-',''))
        else:
            print('could not find start index, with',time_strs[0])
            idxstart=None
        if time_strs[-1][0:10].replace('-','') in gauge_time:
            idxend=gauge_time.index(time_strs[-1][0:10].replace('-',''))
        else:
            print('could not find end index')
            idxend=None

        if idxstart is None or idxend is None:            
            Qgages=None
            tgages=None
            Qbar=None
        else:
            Qgages=gauge_discharge[idxstart:idxend]
            tgages=gauge_time_ordinal[idxstart:idxend]
            Qbar=np.nanmean(Qgages)
    
    return Qbar,Qgages,time_strs,tgages,agencyid

def get_junction_downflows(junctions):
    # make a list of the "downstream flows" of each junction
    junction_downflows=[]
    for junction in junctions:
        downflows=junction['downflows']
        if len(downflows)>1:
            print('oops, we have a distributary, junction=',junction)
        else:
            junction_downflows.append(str(downflows[0]))
    return junction_downflows

def table_iteration(indices_not_done,updf,junctions,junction_downflows,ΔD,reachids_reorder):
    idx_to_do=indices_not_done[0]
    reachid=updf.at[idx_to_do,'reachid']    
    reaches_to_add=updf.iloc[idx_to_do]['upstream reachids']

    for reach_to_add in reaches_to_add: 
        idx_new=len(updf)
        updf.at[idx_new,'reachid']=reach_to_add
        updf.at[idx_to_do,'upstream ids added']=True

        if reach_to_add not in junction_downflows:
            updf.at[idx_new,'upstream reachids']=None
            updf.at[idx_new,'upstream ids added']=True
            updf.at[idx_new,'drainage area']=None
        else:                    
            next_jdx=junction_downflows.index(reach_to_add)
            updf.at[idx_new,'upstream reachids']=[str(rid) for rid in junctions[next_jdx]['upflows']]
            updf.at[idx_new,'upstream ids added']=False
            updf.at[idx_new,'drainage area']=ΔD[reachids_reorder.index(reachid)]
            
    indices_not_done=list(updf[~updf['upstream ids added'].astype(bool)].index) 
    
    return updf,indices_not_done