HydroGrid.py 9.6 KB
Newer Older
1 2 3 4 5 6 7 8
#
# Reads ans manages the model grid
#
import numpy as np
from netCDF4 import Dataset
# Doc : https://spacetelescope.github.io/spherical_geometry/api/spherical_geometry.polygon.SphericalPolygon.html
from spherical_geometry import polygon
import RPPtools as RPP
9
import os 
10 11 12 13
#
import configparser
config=configparser.ConfigParser()
config.read("run.def")
14
EarthRadius=config.getfloat("OverAll", "EarthRadius", fallback=6370000.0)
15 16 17 18 19 20 21
#
import getargs
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
#
def getbox(ncdf, corners) :
22 23 24 25 26 27
    # Add a few points to tbe box to make sure to cover everything
    halo_pts=2
    ii=[np.argmin(abs(ncdf.variables["nav_lon"][0,:]-np.min(corners[0][:]))),
        np.argmin(abs(ncdf.variables["nav_lon"][0,:]-np.max(corners[0][:])))]
    jj=[np.argmin(abs(ncdf.variables["nav_lat"][:,0]-np.min(corners[1][:]))),
        np.argmin(abs(ncdf.variables["nav_lat"][:,0]-np.max(corners[1][:])))]
28

29
    return min(ii)-halo_pts, max(ii)+halo_pts, min(jj)-halo_pts, max(jj)+halo_pts
30 31 32
#
# Get the corners for a regular lat/lon grid
#
33
def centers(lon, lat) :
34
    jjm,iim = lon.shape
35 36
    #
    index = [[j,i] for i in range(iim) for j in range(jjm)]
37 38
    centersll = [[lon[j,i], lat[j,i]] for j,i in index]
    return index, centersll
39

40 41 42
def corners(subcentersll, hdlon, hdlat) :
    cornersll = [RPP.boxit(np.array([lo, la]), hdlon, hdlat, 2) for lo,la in subcentersll]
    Llats = [[0.0000001 if p[1] == 0 else np.sign(p[1])*89.99 if np.abs(p[1]) == 90 else p[1] for p in boxll] for boxll in cornersll]
43
    Llons = [[0.0000001 if p[0] == 0 else np.sign(p[0])*89.99 if np.abs(p[0]) == 90 else p[0] for p in boxll] for boxll in cornersll]
44 45
    cornerspoly = [ polygon.SphericalPolygon.from_lonlat(lons, lats, center=cent) for lons, lats, cent in zip(Llons, Llats, subcentersll)]
    #radiusll = [RPP.maxradius(np.array(cent), np.array(lons), np.array(lats)) for cent, lons, lats in zip(subcentersll,Llons,Llats)]
46
    #
47
    return cornersll, cornerspoly
48
#
49
def gather(x, index, default = 0) :
50 51
    y=[]
    for ia in index :
52 53 54 55
        if (ia[:,0] == [-1, -1]).all():
            y.append([default])
        else:
            y.append(list(x[ia[0,i],ia[1,i]] for i in range(ia.shape[1]) ))
56 57 58 59 60 61 62 63 64
    return y
#
def getattrcontaining(nc, varname, substr) :
    att=[]
    for s in nc.variables[varname].ncattrs() :
        if s.lower().find(substr) >= 0 :
            att.append(nc.variables[varname].getncattr(s))
    return att
#
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
def calctopoindex(nf, istr, iend, jstr, jend, index, hydrogrid) :
    # The routing code in Fortran (i=1, j=2)
    # 8 --- 1 --- 2
    # |     |     |
    # 7 --- X --- 3
    # |     |     |
    # 6 --- 5 --- 4
    #
    inc = np.array([[0,1,1,1,0,-1,-1,-1],[-1,-1,0,1,1,1,0,-1]])
    elevation_precision=0.1
    #
    jlen,ilen = nf.variables["orog"].shape
    missing = nf.variables["orog"].missing_value
    land = np.where(nf.variables["orog"][jstr:jend,istr:iend] < missing)
    topoind=np.zeros((jend-jstr,iend-istr))
    topoind[:,:] = np.nan
    #
    for j,i in zip(land[:][0],land[:][1]) :
        trip = int(nf.variables["trip"][jstr+j,istr+i]-1)
        if trip < 8 :
            inf=(istr+i+inc[0,trip])%ilen
            jnf=(jstr+j+inc[1,trip])%jlen
            dz=max(elevation_precision,nf.variables["orog"][jstr+j,istr+i]-nf.variables["orog"][jnf,inf])
            dl=nf.variables["disto"][jstr+j,istr+i]-nf.variables["disto"][jnf,inf]
        else :
POLCHER Jan's avatar
POLCHER Jan committed
90 91 92 93 94 95
            if trip > 97 :
                # Going into the ocean which is at elevation 0
                dz = max(elevation_precision,nf.variables["orog"][jstr+j,istr+i])
            else :
                # Endorehic basin so we use thte lowest dz possible.
                dz = elevation_precision
96
            #
POLCHER Jan's avatar
POLCHER Jan committed
97 98 99 100 101 102
            if nf.variables["disto"][jstr+j,istr+i] > 0 :
                dl = nf.variables["disto"][jstr+j,istr+i]
            else :
                # Some other reasonable assumption
                rx,ry = hydrogrid.resolution(i,j)
                dl = (rx+ry)*0.5*0.5
103 104 105 106 107 108 109
        topoind[j,i]=np.sqrt(dl**3./(dz*1.0e6))
        #
    minitopo = np.nanmin(topoind)
    if minitopo <= np.finfo(float).eps :
        ERROR("Topoind close to zero encoutered.")
        sys.exit()
    #
110
    return gather(topoind, index, default = 10), "Topographic index which serves for the residence time", minitopo
111
#
112
class HydroGrid :
113
    def __init__(self, lolacorners, wfile) :
114
        #
115 116 117
        self.source=config.get("OverAll", "HydroFile")
        INFO("Opening in HydroGrid : "+self.source)
        self.ncfile=Dataset(self.source,'r')
118
        istr, iend, jstr, jend = getbox(self.ncfile, lolacorners)
119 120 121
        self.hdlon = np.mean(np.abs(np.diff(self.ncfile.variables["nav_lon"][0,:])))
        self.hdlat = np.mean(np.abs(np.diff(self.ncfile.variables["nav_lat"][:,0])))
        self.hd = 2*max(self.hdlon,self.hdlat)
122 123 124 125 126 127
        # In case the box is on the edge of the routing.nc file
        if istr<0:
            istr = 0  
        if iend >=self.ncfile.variables["nav_lon"][0,:].shape[0]: 
            iend = self.ncfile.variables["nav_lon"][0,:].shape[0]
        self.box=[istr, iend, jstr, jend]   
128 129 130 131 132
        self.lon=np.copy(self.ncfile.variables["nav_lon"][jstr:jend,istr:iend])
        self.lat=np.copy(self.ncfile.variables["nav_lat"][jstr:jend,istr:iend])
        self.jjm,self.iim = self.lon.shape
        DEBUG("# Range Lon :"+str(np.min(self.lon))+" -- "+str(np.max(self.lon)))
        DEBUG("# Range Lat :"+str(np.min(self.lat))+" -- "+str(np.max(self.lat)))
133
        #
134 135
        if wfile is None or not os.path.exists(wfile):
            self.index, self.centers = centers(self.lon, self.lat) 
136 137

    def select(self, c, r) :
138 139 140 141
        indices = [i for i in range(len(self.centers)) if (RPP.loladist(np.array(c),np.array(self.centers[i])) <= r+2*self.hd)]
        subcenters = [self.centers[i] for i in indices]
        polyll, polylist = corners(subcenters, self.hdlon, self.hdlat)
        return indices, polylist
142

143 144 145 146 147 148 149 150 151 152 153 154 155
    def resolution(self,i,j) :
        if i+1 >= self.iim or j+1 >= self.jjm :
            inx=i-1; jnx=j-1
        else :
            inx=i+1; jnx=j+1
        dlon = np.radians(np.abs(self.lon[j,i]-self.lon[j,inx]))
        alon = np.cos(np.radians(self.lat[j,i]))**2*(np.sin(dlon/2))**2
        clon = 2*np.arctan2(np.sqrt(alon), np.sqrt(1-alon))
        dlat = np.radians(np.abs(self.lat[j,i]-self.lat[jnx,i]))
        alat = (np.sin(dlat/2))**2
        clat = 2*np.arctan2(np.sqrt(alat), np.sqrt(1-alat))
        return EarthRadius*clon, EarthRadius*clat
    
156
class HydroData :
157
    def __init__(self, nf, box, index, part, hydrogrid) :
158
        istr, iend, jstr, jend = box[:]
159 160 161
        #
        # Flow direction
        #
162
        self.trip=gather(nf.variables["trip"][jstr:jend,istr:iend].astype(np.float32), index, default = 98)
163 164
        self.tripdesc=nf.variables["trip"].long_name
        #
165 166
        # ID of basin
        #
167
        self.basins=gather(nf.variables["basins"][jstr:jend,istr:iend], index, 999)
168 169 170 171 172
        self.basinsdesc=nf.variables["basins"].long_name
        att = getattrcontaining(nf, "basins", "max")
        if len(att) > 0 :
            self.basinsmax=att[0]
        else :
173
            INFO("We need to scan full file to find the largest basin ID")
174
            # This variable seems not to be used further
175 176 177
            ##self.basinsmax = part.domainmax(np.ma.max(np.ma.masked_where(self.basins < 1.e10, self.basins)))
            ##self.basinsmax = part.domainmax(np.max(np.array(self.basins)[np.array(self.basins) < missing]))
            self.basinsmax = 999
178
        #
179
        # Distance to ocean
180
        #
181
        self.disto=gather(nf.variables["disto"][jstr:jend,istr:iend].astype(np.float32), index, default = 500)
182 183
        self.distodesc=nf.variables["disto"].long_name
        #
184 185
        # Flow accumulation
        #
186
        self.fac=gather(nf.variables["fac"][jstr:jend,istr:iend].astype(np.float32), index, default = 0)
187
        self.facdesc=nf.variables["fac"].long_name
188
        #
189 190 191
        # If Orography is present in the file then we can compute the topographic index.
        # Else the topographic index needs to be present in the file.
        #
POLCHER Jan's avatar
POLCHER Jan committed
192
        if "orog" in nf.variables.keys() and "disto" in nf.variables.keys() :
193
            self.orog = gather(nf.variables["orog"][jstr:jend,istr:iend].astype(np.float32), index, default = 0)
194
            self.orogdesc=nf.variables["orog"].long_name
195 196
            self.topoind, self.topoinddesc, mintopo = calctopoindex(nf, istr, iend, jstr, jend, index, hydrogrid)
            self.topoindmin = part.domainmin(mintopo)
197
        elif "topoind" in nf.variables.keys() and "orog" in nf.variables.keys() :
198
            self.orog = gather(np.zeros((jend-jstr,iend-istr)).astype(np.float32), index)
199
            self.topoind=gather(nf.variables["topoind"][jstr:jend,istr:iend].astype(np.float32), index, default = 10)
200 201 202 203 204 205 206
            self.topoinddesc=nf.variables["topoind"].long_name
            att = getattrcontaining(nf, "topoind", "min")
            if len(att) > 0 :
                self.topoindmin=att[0]
            else :
                INFO("We need to scan full file to find minimum topoind over domain")
                self.topoindmin=np.min(np.where(nf.variables["topoind"][:,:] <  1.e15))
POLCHER Jan's avatar
POLCHER Jan committed
207 208 209
        else :
            ERROR("Not enough information in the HydroFile in order to compute topographic residence time.")
            sys.exit()
210
        #
211 212
        #
        if "floodplains" in nf.variables.keys():
213
            self.floodplains = gather(nf.variables["floodplains"][jstr:jend,istr:iend].astype(np.float32), index, 0)
214 215 216
            self.floodplainsdesc=nf.variables["floodplains"].long_name
        else:
            self.floodplains = gather(np.zeros((jend-jstr,iend-istr)).astype(np.float32), index)