GraphHydro.py

#
import sys
import numpy as np
from netCDF4 import Dataset
import datetime
#
import getargs
config = getargs.SetupConfig()
# Maximum error in the distance of the station in km
MaxDistErr = config.getfloat("OverAll", "MaxDistErr", fallback=25.0)
# Maximum error in the upstream area in %
MaxUpstrErr = config.getfloat("OverAll", "MaxUpstrErr", fallback=10.0)
#
log_master, log_world = getargs.getLogger(__name__)
INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug
#
import F90tools as F90T
#
F90T.adddirtopath("F90subroutines")
F90T.compilef90("routing_interface")
import routing_interface
import NcHelp as NH
import RPPtools as RPP
#
NCFillValue=1.0e20
#
# Function to compute quantiles with weights
#
def weighted_quantile(values, quantiles, sample_weight=None, 
                      values_sorted=False, old_style=False):
    """ Very close to numpy.percentile, but supports weights.
    NOTE: quantiles should be in [0, 1]!
    :param values: numpy.array with data
    :param quantiles: array-like with many quantiles needed
    :param sample_weight: array-like of the same length as `array`
    :param values_sorted: bool, if True, then will avoid sorting of
        initial array
    :param old_style: if True, will correct output to be consistent
        with numpy.percentile.
    :return: numpy.array with computed quantiles.
    """
    values = np.array(values)
    quantiles = np.array(quantiles)
    if sample_weight is None:
        sample_weight = np.ones(len(values))
    sample_weight = np.array(sample_weight)
    assert np.all(quantiles >= 0) and np.all(quantiles <= 1), \
        'quantiles should be in [0, 1]'

    if not values_sorted:
        sorter = np.argsort(values)
        values = values[sorter]
        sample_weight = sample_weight[sorter]

    weighted_quantiles = np.cumsum(sample_weight) - 0.5 * sample_weight
    if old_style:
        # To be convenient with numpy.percentile
        weighted_quantiles -= weighted_quantiles[0]
        weighted_quantiles /= weighted_quantiles[-1]
    else:
        weighted_quantiles /= np.sum(sample_weight)
    return np.interp(quantiles, weighted_quantiles, values)
#
# Add time constants and possibly other prameters to the graph file
#
def addparameters(outnf, version, part, tcst, vtyp, NCFillValue) :
    if part.rank == 0 :
        # Add global attributes to the netCDF file
        outnf.setncattr("RoutingPPVersion", version)
        outnf.setncattr("GenerationDate", datetime.datetime.utcnow().isoformat(sep=" "))
        # Add routing parameters
        outnf.createDimension('dimpara', 1)
        stream = outnf.createVariable("StreamTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        stream.title = "Time constant for the stream reservoir"
        stream.units = "s/km"
        stream[:] = tcst.stream_tcst
        fast = outnf.createVariable("FastTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        fast.title = "Time constant for the fast reservoir"
        fast.units = "s/km"
        fast[:] = tcst.fast_tcst
        slow = outnf.createVariable("SlowTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        slow.title = "Time constant for the slow reservoir"
        slow.units = "s/km"
        slow[:] = tcst.slow_tcst
        flood = outnf.createVariable("FloodTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        flood.title = "Time constant for the flood reservoir"
        flood.units = "s/km"
        flood[:] = tcst.flood_tcst
        swamp = outnf.createVariable("SwampCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        swamp.title = "Fraction of the river transport that flows to the swamps"
        swamp.units = "-"
        swamp[:] = tcst.swamp_cst
    return
#
#
#
class HydroGraph :
    def __init__(self, nbasmax, hydrosuper, part, modelgrid) :
        #
        self.nbasmax = nbasmax
        self.nbpt = hydrosuper.basin_count.shape[0]
        nwbas = hydrosuper.basin_topoind.shape[1]
        nbxmax_in = hydrosuper.inflow_grid.shape[2]
        #
        #
        self.routing_area, self.routing_orog_mean, self.routing_orog_min,self.routing_orog_max, \
            self.routing_floodp,self.routing_beta, self.routing_cg, self.topo_resid, self.topo_resid_stream, self.topo_rlen, \
            self.topo_rdz, self.route_nbbasin, self.route_togrid, self.route_tobasin, self.route_nbintobas, \
            self.global_basinid, self.route_outlet, self.route_type, self.origin_nbintobas, self.routing_fetch, self.floodcri = \
                                    routing_interface.finish_truncate(nbpt = self.nbpt, inflowmax = nbxmax_in, nbasmax = nbasmax, nwbas = nwbas, \
                                                                      num_largest = hydrosuper.num_largest, gridarea = modelgrid.area, \
                                                                      cfrac = modelgrid.contfrac, gridcenters = np.array(modelgrid.centers), \
                                                                      basin_count = hydrosuper.basin_count, \
                                                                      basin_notrun = hydrosuper.basin_notrun, basin_area = hydrosuper.basin_area, \
                                                                      basin_orog_mean = hydrosuper.basin_orog_mean, basin_orog_min = hydrosuper.basin_orog_min,\
                                                                      basin_orog_max = hydrosuper.basin_orog_max, basin_floodp = hydrosuper.basin_floodp, \
                                                                      basin_cg = hydrosuper.basin_cg, basin_topoind = hydrosuper.basin_topoind, \
                                                                      basin_topoindex_stream = hydrosuper.basin_topoindex_stream, \
                                                                      basin_rlen = hydrosuper.basin_rlen, basin_rdz = hydrosuper.basin_rdz, \
                                                                      basin_beta_fp = hydrosuper.basin_beta_fp , fetch_basin = hydrosuper.fetch_basin, \
                                                                      basin_id = hydrosuper.basin_id, \
                                                                      basin_coor = hydrosuper.basin_outcoor, basin_type = hydrosuper.basin_type, \
                                                                      basin_flowdir = hydrosuper.basin_flowdir, \
                                                                      outflow_grid = hydrosuper.outflow_grid, outflow_basin = hydrosuper.outflow_basin, \
                                                                      inflow_number = hydrosuper.inflow_number, inflow_grid = hydrosuper.inflow_grid, \
                                                                      inflow_basin = hydrosuper.inflow_basin, floodcri = hydrosuper.floodcri, \
                                                                      rfillval = NCFillValue, ifillval = RPP.IntFillValue)
        #
        self.num_largest = routing_interface.finalrivclass(part.landcorelist, self.route_togrid, self.route_tobasin, self.routing_fetch, \
                      hydrosuper.largest_rivarea)
        #
        # Inflows
        self.max_inflow = part.domainmax(np.max(hydrosuper.inflow_number))
        gingrid = part.l2glandindex( hydrosuper.inflow_grid[:,:,:self.max_inflow])
        self.route_innum, self.route_ingrid, self.route_inbasin = \
            routing_interface.finish_inflows(nbpt = self.nbpt, nwbas = nwbas, nbasmax = nbasmax, inf_max = self.max_inflow, \
                                             basin_count = hydrosuper.basin_count, inflow_number = hydrosuper.inflow_number, \
                                             inflow_grid = gingrid, inflow_basin = hydrosuper.inflow_basin[:,:,:self.max_inflow])
        #
        return
    #
    #
    #
    def position(self, modelgrid, locations) :
        npt, nbas = self.routing_fetch.shape
        if len(locations.lid) > 0 :
          for i,lid in enumerate(locations.lid) :
            # Some stations have a negative upstream area (because undefined)
            if locations.upstream[i] > 0:
                dist = RPP.distances(np.array(modelgrid.coordll)[:,0], np.array(modelgrid.coordll)[:,1], \
                                    locations.lon[i], locations.lat[i])
                # Calculation in km
                disterr = (MaxDistErr-dist/1000.)/MaxDistErr
                disterr[disterr < 0] = np.nan
                # Calculation in km^2 and result in %
                upstrerr = (MaxUpstrErr-np.abs(self.routing_fetch/1.e+6-locations.upstream[i])/locations.upstream[i]*100)/MaxUpstrErr
                upstrerr[upstrerr<0] = np.nan
                cost = np.tile(disterr,nbas).reshape((nbas,npt)).transpose()+upstrerr
                if np.count_nonzero(~np.isnan(cost)) > 0 :
                    ij, ib = np.unravel_index(np.nanargmax(cost),cost.shape)
                    userr = (self.routing_fetch[ij,ib]/1.e+6-locations.upstream[i])/locations.upstream[i]*100
                    locations.addhtupos(i, ij, ib, np.nanmax(cost), userr)
                else :
                    # Only report a missed station is there was a geographical solution on the domain.
                    if np.count_nonzero(~np.isnan(disterr)) > 0 :
                        if lid < 9999 :
                            typename="Structure"
                        else :
                            typename="GRDC"
                        #
                        INFO("Could not locate "+typename+" number : "+str(lid)+" of type "+\
                             locations.ltype[i]+" at "+locations.lname[i])
                        ij = np.nanargmax(disterr)
                        INFO(typename+" number : "+str(lid)+" best distance match [km] : "+str(dist[ij]/1000.))
                        if np.count_nonzero(~np.isnan(upstrerr)) > 0 :
                            ib, ij = np.unravel_index(np.nanargmax(upstrerr),upstrerr.shape)
                            INFO(typename+" number : "+str(lid)+" smallest upstream error [%] : "+\
                                 str(upstrerr[ib,ij]*MaxUpstrErr+MaxUpstrErr))
                        else :
                            INFO(typename+" number : "+str(lid)+" No appropriate HTU with an upstream "+\
                                 "error less than "+str(MaxUpstrErr)+" % ")
    #
    #
    #
    def add_variable(self, outnf, procgrid, NCFillValue, part, coord, name, title, units, data, vtyp, orig_type = "float", arrayorder = 'F'):
        
        var = procgrid.landscatter(data.astype(vtyp), order=arrayorder)
            
        if orig_type == "float":
            var[np.isnan(var)] = NCFillValue
        elif orig_type == "int":
            var[np.isnan(var)] = RPP.IntFillValue
            var[var==RPP.IntFillValue] = NCFillValue

        if part.rank == 0:
            ncvar = outnf.createVariable(name, vtyp, coord, fill_value=NCFillValue)
            ncvar.title = title
            ncvar.units = units
        else :
            ncvar = np.zeros((1,)*len(coord))
        ncvar[:] = part.gather(var, default=NCFillValue)
        return
    #
    #
    #
    def add_tstepdistrib(self, outnf, procgrid, NCFillValue, part, data, area, vtyp, nbbins, tcst, arrayorder = 'F') :

        
        var = procgrid.landscatter(data.astype(vtyp), order=arrayorder)
        gvar = part.gather(var, default=NCFillValue)
        larea = procgrid.landscatter(area.astype(vtyp), order=arrayorder)
        garea = part.gather(larea, default=NCFillValue)
        
        if part.rank == 0:
            gvar[gvar >= NCFillValue] = np.nan
            #
            hist, bin_edges = np.histogram(gvar[~np.isnan(gvar)], bins=nbbins, range=(0,min(10000, np.nanmax(gvar))))
            whist, bin_edges = np.histogram(gvar[~np.isnan(gvar)], weights=garea[~np.isnan(gvar)], \
                                             bins=nbbins, range=(0,min(10000, np.nanmax(gvar))))
            #
            topobins = outnf.createVariable("topobins", vtyp, ('nbbnds','nbbins'))
            topobins.title = "Topographic index"
            topobins.units = "km"
            topobins[0,:] = bin_edges[:-1]
            topobins[1,:] = bin_edges[1:]
            topohist = outnf.createVariable("topocount", vtyp, ("nbbins",), fill_value=NCFillValue)
            topohist.title = "Counts of HTU by topoindex"
            topohist.units = "counts"
            topohist[:] = hist[:]
            wtopohist = outnf.createVariable("wtopocount", vtyp, ("nbbins",), fill_value=NCFillValue)
            wtopohist.title = "Area weighted counts of HTU by topoindex"
            wtopohist.units = "counts"
            wtopohist[:] = whist[:]
            #
            thist, tbin_edges = np.histogram(gvar[~np.isnan(gvar)]*tcst.stream_tcst, bins=nbbins, range=(0,RPP.OneDay/8))
            tbins = outnf.createVariable("tstepbins", vtyp, ('nbbnds','nbbins'))
            tbins.title = "Time bins"
            tbins.units = "s"
            tbins[0,:] = tbin_edges[:-1]
            tbins[1,:] = tbin_edges[1:]
            ncvar = outnf.createVariable("stabtstepcount", vtyp, ("nbbins",), fill_value=NCFillValue)
            ncvar.title = "Stable timestep distribution"
            ncvar.units = "count"
            ncvar[:] = thist[:]
            #
            maxtstep = weighted_quantile(gvar[~np.isnan(gvar)]*tcst.stream_tcst, 0.50, sample_weight=garea[~np.isnan(gvar)])
            mtstep = outnf.createVariable("MaxTimeStep", vtyp, ('dimpara'), fill_value=NCFillValue)
            mtstep.title = "The maximum time-step possible given the distribution of HTU properties"
            mtstep.units = "s"
            mtstep[:] = maxtstep
        return
    #
    #
    #
    def dumpnetcdf(self, filename, version, globalgrid, procgrid, part, tcst, locations) :
        #
        vtyp=np.float64
        cornerind=[0,2,4,6]
        nbcorners = len(cornerind)
        nbbins = 2000
        #
        nlmax, nblocated = locations.maxlocations(self.nbasmax, part)
        #
        if part.rank == 0 :
            outnf=Dataset(filename, 'w', format='NETCDF4_CLASSIC')
            # Dimensions
            outnf.createDimension('x', globalgrid.ni)
            outnf.createDimension('y', globalgrid.nj)
            outnf.createDimension('land', len(procgrid.area))
            outnf.createDimension('z', self.nbasmax)
            outnf.createDimension('bnd', nbcorners)
            outnf.createDimension('inflow', self.max_inflow)
            outnf.createDimension('stnperhtu', nlmax)
            outnf.createDimension('nbbins', nbbins)
            outnf.createDimension('nbbnds',2)
            outnf.createDimension('locations', nblocated)
        else :
            outnf = None
        #
        NH.addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        NH.addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
        addparameters(outnf, version, part, tcst, vtyp, NCFillValue)
        #
        # land grid index -> to facilitate the analyses of the routing
        #
        nbpt_loc = np.zeros((self.nbpt,1)).astype(np.int32)
        nbpt_loc[:,0] = np.arange(1, self.nbpt+1)
        nbpt_glo = part.l2glandindex(nbpt_loc)
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), \
                          "nbpt_glo", "Land point indices on global grid", "-", nbpt_glo[:,0], vtyp)
        #
        ################
        #
        # Conversion of grid indices for route_togrid
        grgrid = part.l2glandindex(self.route_togrid[:,:])
        #
        #
        # The field route_togrid is with indices on the local grid. That needs to be converted to the global grid.
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "routetogrid", "Grid into which the basin flows", "-", grgrid, vtyp, orig_type="int")
        # route to basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "routetobasin", "Basin in to which the water goes", "-", self.route_tobasin[:,:], vtyp, orig_type="int")
        # basin id
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basinid", "ID of basin", "-", self.global_basinid[:,:], vtyp, orig_type="int")
        #
        # basin area
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_area", "area of basin", "m^2", self.routing_area[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_orog_mean", "Mean orography", "m", self.routing_orog_mean[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_orog_min", "Min orography", "m", self.routing_orog_min[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_orog_max", "Max orography", "m", self.routing_orog_max[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_floodp", "Fraction of floodplains", "-", self.routing_floodp[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "floodcri", "Height at which all basin is flooded", "mm", self.floodcri[:,:], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "basin_beta_fp", "Shape of the floodplins HTUs (concave / convex)", "-", self.routing_beta[:,:], vtyp)

        # route number into basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), \
                          "routenbintobas", "Number of basin into current one", "-", self.route_nbintobas[:], vtyp, orig_type="int")
        #
        # original number into basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ( 'y','x'), \
                          "originnbintobas", "Number of sub-grid basin into current one before truncation", "-", \
                          self.origin_nbintobas[:], vtyp, orig_type="int")
        #
        # latitude of outlet
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "outletlat", "Latitude of Outlet", "degrees north", self.route_outlet[:,:,0], vtyp)
        # longitude of outlet
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "outletlon", "Longitude of Outlet", "degrees east", self.route_outlet[:,:,1], vtyp)
        # type of outlet
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "outlettype", "Type of outlet", "code", self.route_type[:,:], vtyp)
        #
        # topographic index
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "topoindex", "Topographic index of the retention time", "km", self.topo_resid[:,:], vtyp)
        self.add_tstepdistrib(outnf, procgrid, NCFillValue, part, self.topo_resid[:,:], self.routing_area[:,:], \
                              vtyp, nbbins, tcst)
        #
        # topoindex stream
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "topoindex_stream", "Topographic index of the retention time for stream res.", "km", self.topo_resid_stream[:,:], vtyp)
        # Add another distribution ? Or replace the other ?
        #self.add_tstepdistrib(outnf, procgrid, NCFillValue, part, self.topo_resid_stream[:,:], self.routing_area[:,:], \
        #                      vtyp, nbbins, tcst)

        #
        # Geometric properties of HTU
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "riverlen", "Mean river length within HTU", "m", self.topo_rlen[:,:], vtyp)
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "riverdz", "Mean river elevation change within HTU", "m", self.topo_rdz[:,:], vtyp)
        # Inflow number
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "route_innum", "Number of inflow", "-", self.route_innum[:,:], vtyp, orig_type="int")
        #
        # Inflow grid
        #gingrid = part.l2glandindex(self.inflow_grid[:,:,:inflow_size])
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow','z','y','x'), \
                          "route_ingrid", "Grid from which the water flows", "-", \
                          self.route_ingrid[:,:,:], vtyp, orig_type="int")
        #
        # Inflow basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow','z','y','x'), \
                          "route_inbasin", "Basin from which the water flows", "-", \
                          self.route_inbasin[:,:,:], vtyp, orig_type="int")

        #
        # Save centre of gravity of HTU
        #
        # Check if it works
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "CG_lon", "Longitude of centre of gravity of HTU", "degrees east", self.routing_cg[:,:,1], vtyp)

        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "CG_lat", "Latitude of centre of gravity of HTU", "degrees north", self.routing_cg[:,:,0], vtyp)
        #
        # Save the fetch of each basin
        #
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), \
                          "fetch", "Fetch contributing to each HTU", "m^2", self.routing_fetch[:,:], vtyp)
        #
        # Build and save a map of all the structures located on the graph
        #
        locations.addtonetcdf(self.nbpt, self.nbasmax, outnf, procgrid, part, ('y','x'), ('locations',), vtyp)
        #
        # Close the file
        #
        if part.rank == 0:
            outnf.close()
        #
        return