Interface.py 39.8 KB
Newer Older
1 2 3 4 5 6 7 8
#
#
import numpy as np
import os
import pickle
from netCDF4 import Dataset
import RPPtools as RPP
from mpi4py import MPI
9
import gc
10 11
#
import sys
12 13 14 15 16
from inspect import currentframe, getframeinfo
#
localdir=os.path.dirname(getframeinfo(currentframe()).filename)
sys.path.append(localdir+'/F90subroutines')
F90=localdir+'/F90subroutines'
17
if MPI.COMM_WORLD.Get_rank() == 0 :
18
    err=os.system("cd "+F90+"; make all")
19 20 21 22 23 24 25 26 27
    if err != 0 :
        print("Compilation error in the FORTRAN interface")
        sys.exit()
else :
    print("Not compiling on other cores")
MPI.COMM_WORLD.Barrier()
#
import routing_interface
#
28 29
import getargs
config = getargs.SetupConfig()
30 31
gendoc = config.get("OverAll", "Documentation", fallback='false')
nbxmax = config.getint("OverAll", "nbxmax", fallback=63)
32
largest_pos = config.getint("OverAll", "ROUTING_RIVERS", fallback=200)
33 34
#
undef_int = 999999999.9
35 36
# Order of magnitude for the area precision in m^2.
prec = 100.0
37 38 39
#
# Print the documentation for the FORTRAN interface
#
40
if gendoc.lower() == "true" :
41 42 43 44 45 46 47 48 49 50 51 52 53
    docwrapper = open('DocumentationInterface', 'w')
    docwrapper.write(routing_interface.initatmgrid.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.gethydrogrid.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.findbasins.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.globalize.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.linkup.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.fetch.__doc__)
    docwrapper.write("====================================================================\n")
54 55 56 57
    docwrapper.write(routing_interface.finish_truncate.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.killbas.__doc__)

58 59 60 61 62 63 64
    docwrapper.close
#
# Functions to access the interfaces
#
#
# initatmgrid : Initialises the grid.f90 module and passes the description of the atmospheric grid.
#
65
def initatmgrid(rank, nbcore, nbpt, modelgrid) :
66 67 68
    print("INITATMGRID corners", np.array(modelgrid.polyll).shape)
    print("INITATMGRID area", np.array(modelgrid.area).shape)
    print("INITATMGRID neighbours", np.array(modelgrid.neighbours).shape)
69 70 71 72 73 74 75 76
    routing_interface.initatmgrid(rank, nbcore, modelgrid.polyll, modelgrid.coordll, modelgrid.area, modelgrid.contfrac, modelgrid.neighbours)
    return
#
#
#
def closeinterface(comm) :
    comm.Barrier()
    routing_interface.closeinterface()
77 78 79
    return
#
#
80
#
81 82 83 84 85 86 87 88 89
def addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind) :
    #
    # Longitude
    longitude = part.gather(procgrid.lon_full)
    if part.rank == 0 :
        lon=outnf.createVariable("lon", vtyp, ('y','x'), fill_value=NCFillValue)
        lon.units="grid box centre degrees east"
        lon.title="Longitude"
        lon.axis="X"
90
        lon[:,:] = globalgrid.lonmat[:,:]
91 92 93 94 95 96 97 98 99
    #
    # Latitude
    latitude = part.gather(procgrid.lat_full)
    if part.rank == 0 :
        lat=outnf.createVariable("lat", vtyp, ('y','x'), fill_value=NCFillValue)
        lat.units="grid box centre degrees north"
        lat.standard_name="grid latitude"
        lat.title="Latitude"
        lat.axis="Y"
100
        lat[:,:] = globalgrid.latmat[:,:]
101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178
    #
    # Bounds of grid box
    #
    llonpoly=np.zeros((nbcorners,procgrid.nbland))
    llatpoly=np.zeros((nbcorners,procgrid.nbland))
    for i in range(procgrid.nbland) :
        llonpoly[:,i] = [procgrid.polyll[i][ic][0] for ic in cornerind]
        llatpoly[:,i] = [procgrid.polyll[i][ic][1] for ic in cornerind]
        lon_bnd = procgrid.landscatter(llonpoly)
        lat_bnd = procgrid.landscatter(llatpoly)
    if part.rank == 0 :
        lonbnd=outnf.createVariable("lon_bnd", vtyp, ('bnd','y','x'), fill_value=NCFillValue)
        lonbnd.units="grid box corners degrees east"
        lonbnd.title="Longitude of Corners"
        latbnd=outnf.createVariable("lat_bnd", vtyp, ('bnd','y','x'), fill_value=NCFillValue)
        latbnd.units="grid box corners degrees north"
        latbnd.title="Latitude of Corners"
    else :
        lonbnd= np.zeros((1,1,1))
        latbnd= np.zeros((1,1,1))
    lonbnd[:,:,:] = part.gather(lon_bnd[:,:,:])
    latbnd[:,:,:] = part.gather(lat_bnd[:,:,:])
    #
    # Land sea mask
    #
    if part.rank == 0 :
        land=outnf.createVariable("land", vtyp, ('y','x'), fill_value=NCFillValue)
        land.units="Land Sea mask"
        land.standard_name="landsea mask"
        land.title="Land"
        land[:,:] = globalgrid.land[:,:]
    # Area
    areas = procgrid.landscatter(np.array(procgrid.area, dtype=np.float64))
    areas[np.isnan(areas)] = NCFillValue
    if part.rank == 0 :
        area=outnf.createVariable("area", vtyp, ('y','x'), fill_value=NCFillValue)
        area.units="m^2"
        area.standard_name="grid area"
        area.title="Area"
    else :
        area = np.zeros((1,1))
    area[:,:] = part.gather(areas[:,:])
    #
    return
#
# Add environment to netCDF file
#
def addenvironment(outnf, procgrid, part, vtyp, NCFillValue, nbpt) :
    #
    nbpt_proc = np.arange(1,nbpt+1, dtype=vtyp)
    proc = np.full(nbpt, part.rank, dtype=vtyp)
    # Environment
    # nbpt_proc
    subpt = procgrid.landscatter(nbpt_proc[:], order='F')
    subpt = subpt.astype(vtyp, copy=False)
    subpt[np.isnan(subpt)] = NCFillValue
    if part.rank == 0 :
        subptgrid = outnf.createVariable("nbpt_proc", vtyp, ('y','x'), fill_value=NCFillValue)
        subptgrid.title = "gridpoint reference inside each proc"
        subptgrid.units = "-"
    else :
        subptgrid = np.zeros((1,1))
    subptgrid[:,:] = part.gather(subpt)
    #
    # rank
    procnum = procgrid.landscatter(proc[:], order='F')
    procnum = procnum.astype(vtyp, copy=False)
    procnum[np.isnan(procnum)] = NCFillValue
    if part.rank == 0 :
        procn = outnf.createVariable("proc_num", vtyp, ('y','x'), fill_value=NCFillValue)
        procn.title = "rank"
        procn.units = "-"
    else :
        procn = np.zeros((1,1))
    procn[:,:] = part.gather(procnum)
    #
    return
#
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202
# Add time constants and possibly other prameters to the graph file
#
def addparameters(outnf, part, tcst, vtyp, NCFillValue) :
    if part.rank == 0 :
        outnf.createDimension('dimpara', 1)
        stream = outnf.createVariable("StreamTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        stream.title = "Time constant for the stream reservoir"
        stream.units = "day/m"
        stream[:] = tcst.stream_tcst
        fast = outnf.createVariable("FastTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        fast.title = "Time constant for the fast reservoir"
        fast.units = "day/m"
        fast[:] = tcst.fast_tcst
        slow = outnf.createVariable("SlowTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        slow.title = "Time constant for the slow reservoir"
        slow.units = "day/m"
        slow[:] = tcst.slow_tcst
        flood = outnf.createVariable("FloodTimeCst", vtyp, ('dimpara'), fill_value=NCFillValue)
        flood.title = "Time constant for the flood reservoir"
        flood.units = "day/m"
        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 = "-"
203
        swamp[:] = tcst.swamp_cst
204 205
    return
#
206 207
#
#
208 209 210 211 212
def finalfetch(part, routing_area, basin_count, route_togrid, route_tobasin, fetch_in) :
    #
    fetch_out = np.zeros(routing_area.shape, dtype=np.float32, order='F')
    partial_sum = np.zeros(routing_area.shape, dtype=np.float32, order='F')
    old_sorted = np.zeros(largest_pos, dtype=np.float32, order='F')
213
    #
Anthony Schrapffer's avatar
Anthony Schrapffer committed
214 215 216
    nbpt = routing_area.shape[0]
    fhalolist = [i+1 for i in range(nbpt) if i not in part.landcorelist]
    in_core = np.array([1 if i in part.landcorelist else 0 for i in range(nbpt)])
217 218 219 220 221 222
    #
    maxdiff_sorted = prec*prec
    iter_count = 0
    #
    while iter_count < part.size*3 and maxdiff_sorted > prec :
        fetch_out[:,:] = 0.0
Anthony Schrapffer's avatar
Anthony Schrapffer committed
223
        outflow_uparea = routing_interface.finalfetch(in_core,fhalolist, part.landcorelist, routing_area, basin_count, route_togrid, \
224 225 226 227 228 229 230 231 232 233
                                                      route_tobasin, partial_sum, fetch_out)
        partial_sum = np.copy(fetch_out)
        part.landsendtohalo(partial_sum, order='F')
        partial_sum = part.zerocore(partial_sum, order='F')
        #
        # Find area the largest basins we need to have right.
        #
        xtmp = np.hstack(part.comm.allgather(outflow_uparea[np.where(outflow_uparea > 0.0)]))
        # Precision in m^2 of the upstream areas when sorting.
        sorted_outareas = (np.unique(np.rint(np.array(xtmp)/prec))*prec)[::-1]
234 235 236 237
        if sorted_outareas.shape[0]<largest_pos:
           s = sorted_outareas[:]
           sorted_outareas = np.zeros(largest_pos, dtype=np.float32, order='F')
           sorted_outareas[:s.shape[0]] = s[:]
238 239 240 241 242
        # If mono-proc no need to iterate as fetch produces the full result.
        if part.size == 1 :
            maxdiff_sorted = 0.0
        else :
            maxdiff_sorted = np.max(np.abs(sorted_outareas[0:largest_pos]-old_sorted))
243
            old_sorted[:] = sorted_outareas[0:largest_pos]
244
        iter_count += 1
245

246 247
    #
    fetch_error = np.sum(np.abs(fetch_out[part.landcorelist,:]-fetch_in[part.landcorelist,:]), axis=1)\
Anthony's avatar
Anthony committed
248
                                                    / np.ma.sum(routing_area[part.landcorelist,:], axis=1)
249
    if np.max(fetch_error) > prec :
250
        print("Rank :"+str(part.rank)+" Too large fetch error (fraction of greid area) : ", fetch_error)
251

252 253 254 255 256 257
    print("Total fetch error in fraction of grid box : ", np.sum(fetch_error))
    #
    return fetch_out
#
#
#
258 259
class HydroOverlap :
#
Anthony Schrapffer's avatar
Anthony Schrapffer committed
260
    def __init__(self, nbpt, nbvmax, sub_pts, sub_index_in, sub_area_in, sub_lon_in, sub_lat_in, part, modelgrid, hydrodata) :
261 262 263
        #
        # Reshape stuff so that it fits into arrays
        #
264
        sub_index = np.zeros((nbpt,nbvmax,2), dtype=np.int32, order='F')
265 266 267 268 269 270 271 272
        sub_area = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        sub_lon = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        sub_lat = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        for ib in range(nbpt) :
            sub_area[ib,0:sub_pts[ib]] = sub_area_in[ib][:]
            sub_lon[ib,0:sub_pts[ib]] = sub_lon_in[ib][:]
            sub_lat[ib,0:sub_pts[ib]] = sub_lat_in[ib][:]
            for ip in range(sub_pts[ib]) :
273
                sub_index[ib,ip,:] = sub_index_in[ib][:,ip]
274
        #
275
        part.landsendtohalo(np.array(sub_area), order='F')
Anthony Schrapffer's avatar
Anthony Schrapffer committed
276
        #
Anthony's avatar
Anthony committed
277 278 279 280 281 282 283 284 285 286 287 288
        ijdim=[]
        for ib in range(nbpt) :
            ijdim.append(max(np.max(sub_index[ib,:sub_pts[ib],0])-np.min(sub_index[ib,:sub_pts[ib],0])+1,np.max(sub_index[ib,:sub_pts[ib],1])-np.min(sub_index[ib,:sub_pts[ib],1])+1))
        ijdimmax = max(ijdim)
        #
        print("GETHYDROGRID : nbpt = {0}".format(nbpt))
        print("GETHYDROGRID : nbvmax = {0}".format(nbvmax))
        print("GETHYDROGRID : ijdimmax = {0}".format(ijdimmax))
        #
        del sub_area_in; del sub_lon_in; del sub_lat_in; del sub_index_in
        #
        #
289 290 291 292 293
        trip_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        basins_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        topoind_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        fac_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        hierarchy_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
294 295
        orog_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        floodp_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
296 297 298 299 300 301 302 303 304
        #
        trip_tmp[:,:] = np.nan
        basins_tmp[:,:] = np.nan
        for ib in range(nbpt) :
            trip_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.trip[ib][:])
            basins_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.basins[ib][:])
            topoind_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.topoind[ib][:])
            fac_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.fac[ib][:])
            hierarchy_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.disto[ib][:])
305 306
            orog_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.orog[ib][:])
            floodp_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.floodplains[ib][:])
307
        #
Anthony's avatar
Anthony committed
308 309 310 311
        del hydrodata.trip; del hydrodata.basins; del hydrodata.topoind
        del hydrodata.fac; del hydrodata.disto; del hydrodata.orog;
        del hydrodata.floodplains
        #
312 313 314 315 316 317
        trip_tmp[np.isnan(trip_tmp)] = undef_int
        basins_tmp[np.isnan(trip_tmp)] = undef_int
        #
        # Go to the call of the FORTRAN interface
        #
        self.nbi, self.nbj, self.area_bx, self.trip_bx, self.basin_bx, self.topoind_bx, self.fac_bx, self.hierarchy_bx, \
318
            self.orog_bx, self.floodp_bx, \
319
            self.lon_bx, self.lat_bx, self.lshead_bx = \
320
                    routing_interface.gethydrogrid(ijdimmax, sub_pts, sub_index, sub_area, \
321 322
                    hydrodata.basinsmax, hydrodata.topoindmin, sub_lon, sub_lat, trip_tmp, basins_tmp, topoind_tmp, fac_tmp,\
                        hierarchy_tmp, orog_tmp, floodp_tmp)
323
        #
Anthony's avatar
Anthony committed
324 325 326
        del trip_tmp; del basins_tmp; del topoind_tmp
        del fac_tmp; del hierarchy_tmp; del orog_tmp
        del floodp_tmp
327
        #
328
        self.nwbas = nbvmax
329 330 331 332
        # Clean-up these arrays so that they are easy to use in Python.
        self.lon_bx[self.lon_bx > 360.]=np.nan
        self.lat_bx[self.lat_bx > 90.]=np.nan
        #
333 334 335 336
        return
#
#
#
337
class HydroSuper :
338
    def __init__(self, nbvmax, hydrodata, hydrooverlap, nbasmax, part) :
339 340 341
        #
        # Keep largest possible number of HTUs
        #
342
        self.nbasmax = nbasmax
343
        self.nbhtuext = nbvmax
344 345 346 347
        self.nbpt = hydrooverlap.nbi.shape[0]
        #
        # nb_htu can be adjusted with self.nwbas
        # nb_htu can be lowered with a larger maxpercent (routing_reg.f90)
POLCHER Jan's avatar
POLCHER Jan committed
348
        nb_htu = nbvmax
349
        nbv = nbvmax
350 351 352 353
        #
        # Call findbasins
        #
        nb_basin, basin_inbxid, basin_outlet, basin_outtp, self.basin_sz, basin_bxout, basin_bbout, self.basin_pts, basin_lshead, coast_pts = \
354 355
                    routing_interface.findbasins(nbpt = self.nbpt, nb_htu = self.nbhtuext, nbv = nbv, nbi = hydrooverlap.nbi, \
                                                 nbj = hydrooverlap.nbj, trip_bx = hydrooverlap.trip_bx, \
POLCHER Jan's avatar
POLCHER Jan committed
356 357
                                                 basin_bx = hydrooverlap.basin_bx, fac_bx = hydrooverlap.fac_bx, \
                                                 hierarchy_bx = hydrooverlap.hierarchy_bx, \
358 359 360 361 362 363
                                                 topoind_bx = hydrooverlap.topoind_bx, lshead_bx = hydrooverlap.lshead_bx, \
                                                 lontmp = hydrooverlap.lon_bx, lattmp = hydrooverlap.lat_bx)
        #
        # Adjust nwbas to the maximum found over the domain
        #
        self.nwbas = part.domainmax(np.max(nb_basin))
364
        # Set the number of inflows per basin. For the moment twice the maximum number of basins.
365
        self.inflowmax = max(10, self.nwbas*2)
366
        print("Maximum number of basin created : {0}".format(self.nwbas))
367 368 369 370
        ijdim=[]
        for i in range(self.nbpt) :
            ijdim.append(max(hydrooverlap.area_bx[i,:,:].shape))
        self.ijdimmax = max(ijdim)
371 372 373
        #
        # Call Globalize
        #
374 375 376 377
        lon_bx_tmp = hydrooverlap.lon_bx
        lon_bx_tmp[np.isnan(lon_bx_tmp)] = undef_int
        lat_bx_tmp = hydrooverlap.lat_bx
        lat_bx_tmp[np.isnan(lat_bx_tmp)] = undef_int
378
        #
379
        self.basin_count, self.basin_notrun, self.basin_area, self.basin_cg, self.basin_hierarchy, \
380 381 382
            self.basin_orog_mean, self.basin_orog_min, self.basin_orog_max, self.basin_floodp, \
            self.basin_fac, self.basin_topoind, self.basin_id, self.basin_outcoor, self.basin_type,\
            self.basin_flowdir, self.basin_lshead, self.outflow_grid, self.outflow_basin, self.nbcoastal, self.coastal_basin = \
383 384
                    routing_interface.globalize(nbpt = self.nbpt, nb_htu = self.nbhtuext, nbv = nbv, ijdimmax = self.ijdimmax, \
                                                area_bx = hydrooverlap.area_bx, lon_bx = lon_bx_tmp, lat_bx = lat_bx_tmp, trip_bx = hydrooverlap.trip_bx, \
385 386
                                                hierarchy_bx = hydrooverlap.hierarchy_bx, orog_bx = hydrooverlap.orog_bx, floodp_bx =  hydrooverlap.floodp_bx,\
                                                fac_bx = hydrooverlap.fac_bx, topoind_bx = hydrooverlap.topoind_bx, min_topoind = hydrodata.topoindmin, \
POLCHER Jan's avatar
POLCHER Jan committed
387 388
                                                nb_basin = nb_basin, basin_inbxid = basin_inbxid, basin_outlet = basin_outlet, basin_outtp = basin_outtp, \
                                                basin_sz = self.basin_sz, basin_pts = self.basin_pts, basin_bxout = basin_bxout, \
389
                                                basin_bbout = basin_bbout, lshead = basin_lshead, coast_pts = coast_pts, nwbas = self.nwbas)
390

391 392 393
        # Memory management
        del basin_bbout; del basin_lshead; del coast_pts; del basin_bxout; del self.basin_pts;
        del basin_outtp; del basin_outlet; del basin_inbxid; del nb_basin
394 395
        return
    #
396 397 398 399 400
    def linkup(self, hydrodata) :
        #
        # Call the linkup routine in routing_reg.
        #
        print("Invented basins =", hydrodata.basinsmax)
401 402 403 404 405
        self.inflow_number,self.inflow_grid,self.inflow_basin = \
            routing_interface.linkup(self.ijdimmax, self.inflowmax, self.basin_count, self.basin_area, \
                                     self.basin_id, self.basin_flowdir, self.basin_lshead, self.basin_hierarchy, \
                                     self.basin_fac, self.outflow_grid, self.outflow_basin, \
                                     self.nbcoastal, self.coastal_basin, float(hydrodata.basinsmax))
406
        self.nbxmax_in = self.inflow_number.shape[1]
407 408 409
        #
        #
        #
410 411
        return
    #
412

POLCHER Jan's avatar
POLCHER Jan committed
413
    def fetch(self, part) :
414
        #
415
        fetch_basin = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
416
        #
417
        self.basin_area = routing_interface.areanorm(self.basin_count, self.basin_area, self.outflow_grid)
418
        partial_sum = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
419
        #
420
        old_sorted = np.zeros(largest_pos, dtype=np.float32, order='F')
421 422 423 424
        #
        maxdiff_sorted = prec*prec
        iter_count = 0
        #
425 426
        fhalolist = [i+1 for i in range(self.nbpt) if i not in part.landcorelist]
        #
427
        while iter_count < part.size*3 and maxdiff_sorted > prec :
428
            fetch_basin[:,:] = 0.0
429
            outflow_uparea = routing_interface.fetch(fhalolist, part.landcorelist, self.basin_count, self.basin_area, self.basin_id, self.basin_hierarchy, \
430
                                                         self.basin_fac, self.outflow_grid, self.outflow_basin, partial_sum, fetch_basin)
431 432
            partial_sum = np.copy(fetch_basin)
            part.landsendtohalo(partial_sum, order='F')
433 434
            partial_sum = part.zerocore(partial_sum, order='F')
            #
435 436 437 438 439 440
            # Find area the largest basins need at least to have.
            #
            xtmp = np.hstack(part.comm.allgather(outflow_uparea[np.where(outflow_uparea > 0.0)]))
            # Precision in m^2 of the upstream areas when sorting.
            sorted_outareas = (np.unique(np.rint(np.array(xtmp)/prec))*prec)[::-1]
            # If mono-proc no need to iterate as fetch produces the full result.
441
            l = min(sorted_outareas.shape[0],largest_pos)
442 443 444
            if part.size == 1 :
                maxdiff_sorted = 0.0
            else :
445
                maxdiff_sorted = np.max(np.abs(sorted_outareas[0:l]-old_sorted[0:l]))
446
                old_sorted[:l] = sorted_outareas[0:largest_pos]
447
            iter_count += 1
448

449 450
        self.fetch_basin = np.copy(fetch_basin)
        #
451
        # Upstream area of the smalest river we call largest rivers.
452
        #
453
        self.largest_rivarea = sorted_outareas[l-1]
454 455 456
        #
        #
        #
457 458 459
        self.num_largest =  routing_interface.rivclassification( part.landcorelist, self.basin_count, self.outflow_grid, self.outflow_basin, \
                   self.fetch_basin, self.largest_rivarea)
        #print("Rank :"+str(part.rank)+" Area of smallest large rivers : ", self.largest_rivarea, " Nb of Large rivers on proc : ",self.num_largest)
460
        return
461 462 463

    def check_fetch(self):

POLCHER Jan's avatar
POLCHER Jan committed
464 465
        routing_interface.checkfetch(nbpt = self.nbpt, nwbas = self.nwbas, fetch_basin = self.fetch_basin, outflow_grid = self.outflow_grid, \
                                     outflow_basin = self.outflow_basin, basin_count = self.basin_count)
466

467
        return
468 469 470

    def check_routing(self):

POLCHER Jan's avatar
POLCHER Jan committed
471 472
        routing_interface.checkrouting(nbpt = self.nbpt, nwbas = self.nwbas, outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, \
                                       basin_count = self.basin_count)
473

474
        return
475

476 477 478 479 480 481 482 483 484
    def correct_outflows(self, part):
        # Global index of the proc domain
        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)
        # Halo points
        fhalo = np.array([pt+1 for pt in range(self.nbpt) if pt not in part.landcorelist], order = "F")
        #
        # Outflow grid in global index and send to halo
Anthony's avatar
Anthony committed
485
        hg = np.copy(self.outflow_grid)
486 487 488 489 490
        hg = part.l2glandindex(self.outflow_grid)
        part.landsendtohalo(hg, order='F')
        # Convert to local index
        outflows = np.unique(hg)
        outflows_out = [a for a in outflows if (a not in nbpt_glo and a>0)]
Anthony's avatar
Anthony committed
491 492
        # Work in a copy to avoid error
        hg1 = np.copy(hg)
493
        for a in outflows_out:
Anthony's avatar
Anthony committed
494 495 496 497 498
          hg1[hg == a] = 0
        for loc, glo in zip(nbpt_loc,nbpt_glo):
          hg1[hg == glo[0]] = loc[0]
        hg = hg1
        del hg1
499 500 501 502 503 504 505 506 507 508 509 510 511
        # Send Outflow basin to the halo and adapt it
        hb = np.copy(self.outflow_basin)
        part.landsendtohalo(hb, order='F')
        hb[hg <= 0] = 999999999
        for ig in range(self.nbpt):
            hb[ig,self.basin_count[ig]:] = 0
        #
        # Correct the routing graph in the halo
        routing_interface.correct_outflows(nbpt = self.nbpt, nwbas = self.nwbas, nbhalo = fhalo.shape[0], \
                    outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, \
                    basin_count = self.basin_count, hg = hg, hb = hb, halopts = fhalo)
        #
        # Correct the inflows
512
        nbxmax_tmp = self.inflow_grid.shape[2]
513 514 515 516 517
        routing_interface.correct_inflows(nbpt = self.nbpt, nwbas = self.nwbas, inflowmax = nbxmax_tmp, outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, basin_count = self.basin_count, inflow_number = self.inflow_number, inflow_grid = self.inflow_grid, inflow_basin = self.inflow_basin)

        return


518
    def killbas(self, tokill, totakeover, numops):
519
        ops = tokill.shape[1]
520 521 522 523
        #
        nbxmax_tmp = self.inflow_grid.shape[2]
        #
        routing_interface.killbas(nbpt = self.nbpt, inflowmax = nbxmax_tmp, \
524 525
                nbasmax = self.nbasmax, nwbas = self.nwbas, ops = ops, tokill = tokill,\
                totakeover = totakeover, numops = numops, basin_count = self.basin_count,\
526 527
                basin_area = self.basin_area, basin_orog_mean = self.basin_orog_mean, \
                basin_orog_min = self.basin_orog_min, basin_orog_max = self.basin_orog_max, basin_floodp = self.basin_floodp, \
528 529 530 531
                basin_cg = self.basin_cg, basin_topoind = self.basin_topoind, fetch_basin = self.fetch_basin,\
                basin_id = self.basin_id, basin_coor = self.basin_outcoor, basin_type = self.basin_type,\
                basin_flowdir = self.basin_flowdir, outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, \
                inflow_number = self.inflow_number, inflow_grid = self.inflow_grid, inflow_basin = self.inflow_basin)
532

533
    #
534 535 536 537 538 539 540 541 542 543
    def add_variable(self,outnf, procgrid, NCFillValue, part, coord, name, title, units, data, vtyp):
        var = procgrid.landscatter(data.astype(vtyp), order='F')
        var[np.isnan(var)] = 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,1,1))
        ncvar[:] = part.gather(var)
544

545 546 547 548 549
    #
    def dumpnetcdf(self, filename, globalgrid, procgrid, part) :
        #
        NCFillValue=1.0e20
        vtyp=np.float64
550
        inflow_size = 100
551 552 553 554 555 556 557 558 559
        cornerind=[0,2,4,6]
        nbcorners = len(cornerind)
        #
        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))
560
            outnf.createDimension('htuext', self.basin_id.shape[1])
561 562
            outnf.createDimension('htu', self.inflow_number.shape[1])
            outnf.createDimension('in',inflow_size )
563 564 565 566 567 568
            outnf.createDimension('bnd', nbcorners)
        else :
            outnf = None
        #
        addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
569
        #
570 571 572
        # Variables
        # Once gathered on root_proc we transform them into float64. Careful, Integer variables do not have NaN !
        #
573 574 575 576 577 578
        # nbpt_glo
        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", "Grid point Global", "-", nbpt_glo[:,0], vtyp)
        #
Anthony's avatar
Anthony committed
579
        # contfrac
580 581 582 583 584
        contfrac = np.array(procgrid.contfrac)
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "contfrac", "Land fraction", "-", np.array(procgrid.contfrac), vtyp)
        #
        # basin_id
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_id", "ID for each HTU", "-", self.basin_id, vtyp)
585 586
        #
        #self.basin_count
587
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "basin_count", "HTU count", "-", self.basin_count, vtyp)
588
        #
589 590 591 592 593 594
        # self.basin_notrun
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "basin_notrun", "Not run", "-", self.basin_notrun, vtyp)
        #
        # self.basin_area
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_area", "Basin area", "-", self.basin_area, vtyp)
        #
595 596 597 598 599 600
        # self.basin_orog
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_orog", "Basin orography", "-", self.basin_orog, vtyp)
        #
        # self.basin_floodp
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_floodp", "Basin floodplains", "-", self.basin_floodp, vtyp)
        #
601 602 603 604 605 606
        # self.basin_cg
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "CG_lon", "CG lon", "-", self.basin_cg[:,:,1], vtyp)
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "CG_lat", "CG lat", "-", self.basin_cg[:,:,0], vtyp)
        #
        # self.topoind
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_topoind", "Topoindex", "-", self.basin_topoind, vtyp)
607
        #
608 609 610
        # outcoor
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "outcoor_lon", "outcoor lon", "-", self.basin_outcoor[:,:,1], vtyp)
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "outcoor_lat", "outcoor lat", "-", self.basin_outcoor[:,:,0], vtyp)
611
        #
612 613
        # type
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_type", "type", "-", self.basin_type, vtyp)
614
        #
615 616 617
        # flowdir
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_flowdir", "flowdir", "-", self.basin_flowdir, vtyp)
        #
618
        #
619
        #self.outflow_grid
620 621 622 623 624 625
        grgrid = part.l2glandindex(self.outflow_grid)
        grgrid[self.outflow_grid == 0 ] = -2 # in case it flows out of the domain, the 0 should not remain
        grgrid[self.outflow_grid == -1 ] = -1
        grgrid[self.outflow_grid == -2 ] = -2
        grgrid[self.outflow_grid == -3 ] = -3
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "HTUoutgrid", "HTU outflow grid", "-", grgrid, vtyp)
626 627
        #
        #self.outflow_basin
628 629 630 631 632 633 634 635
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "HTUoutbasin", "Outflow HTU of grid", "-", self.outflow_basin, vtyp)
        #
        # self.inflow_number
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "HTUinnum", "Inflow number", "-", self.inflow_number, vtyp)
        #
        # Inflow Grid -> convert to global
        gingrid = part.l2glandindex(self.inflow_grid[:,:,:inflow_size])
        self.add_variable(outnf, procgrid, NCFillValue, part, ('in','htu','y','x'), "HTUingrid", "Inflow grid", "-", gingrid, vtyp)
Anthony's avatar
Anthony committed
636 637
        #
        # Inflow Basin
638
        self.add_variable(outnf, procgrid, NCFillValue, part, ('in','htu','y','x'), "HTUinbas", "Inflow basin", "-", self.inflow_basin[:,:,:inflow_size], vtyp)
639 640 641
        #
        # Save the fetch of each basin
        #
642
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "fetch_basin", "Fetch contributing to each HTU", "m^2", self.fetch_basin, vtyp)
643 644 645 646 647 648 649
        #
        # Close file
        #
        if part.rank == 0 :
            outnf.close()
        #
        return
650 651 652
#
#
#
653
class HydroGraph :
654
    def __init__(self, nbasmax, hydrosuper, part, modelgrid) :
655
        #
656
        self.nbasmax = nbasmax
657
        self.nbpt = hydrosuper.basin_count.shape[0]
658
        nwbas = hydrosuper.basin_topoind.shape[1]
659 660
        nbxmax_in = hydrosuper.inflow_grid.shape[2]
        #
661
        #
662 663
        self.routing_area, self.routing_orog_mean, self.routing_orog_min,self.routing_orog_max, \
            self.routing_floodp, self.routing_cg, self.topo_resid, self.route_nbbasin,\
664 665
            self.route_togrid, self.route_tobasin, self.route_nbintobas, self.global_basinid, \
            self.route_outlet, self.route_type, self.origin_nbintobas, self.routing_fetch = \
666
                                    routing_interface.finish_truncate(nbpt = self.nbpt, inflowmax = nbxmax_in, nbasmax = nbasmax, nwbas = nwbas, \
POLCHER Jan's avatar
POLCHER Jan committed
667
                                                                      num_largest = hydrosuper.num_largest, gridarea = modelgrid.area, \
668 669
                                                                      cfrac = modelgrid.contfrac, gridcenters = np.array(modelgrid.centers), \
                                                                      basin_count = hydrosuper.basin_count, \
POLCHER Jan's avatar
POLCHER Jan committed
670
                                                                      basin_notrun = hydrosuper.basin_notrun, basin_area = hydrosuper.basin_area, \
671 672
                                                                      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, \
POLCHER Jan's avatar
POLCHER Jan committed
673 674 675 676 677 678 679 680
                                                                      basin_cg = hydrosuper.basin_cg, \
                                                                      basin_topoind = hydrosuper.basin_topoind, 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)
681

682
        #
Anthony Schrapffer's avatar
Anthony Schrapffer committed
683 684
        # This step is no more necessary
        #self.routing_fetch = finalfetch(part, self.routing_area, self.route_nbbasin, self.route_togrid, self.route_tobasin, self.routing_fetch)
685
        #
686
        self.num_largest = routing_interface.finalrivclass(part.landcorelist, self.route_togrid, self.route_tobasin, self.routing_fetch, \
687
                      hydrosuper.largest_rivarea)
688
        #
689 690 691
        # Inflows
        self.max_inflow = part.domainmax(np.max(hydrosuper.inflow_number))
        gingrid = part.l2glandindex( hydrosuper.inflow_grid[:,:,:self.max_inflow])
692 693 694 695
        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])
696

697
        return
698 699 700
    #
    #
    #
POLCHER Jan's avatar
POLCHER Jan committed
701
    def add_variable(self, outnf, procgrid, NCFillValue, part, coord, name, title, units, data, vtyp, orig_type = "float"):
702 703 704 705 706
        var = procgrid.landscatter(data.astype(vtyp), order='F')

        if orig_type == "float":
            var[np.isnan(var)] = NCFillValue
        elif orig_type == "int":
Anthony's avatar
Anthony committed
707
            var[np.isnan(var)] = RPP.IntFillValue
708
            var[var==RPP.IntFillValue] = NCFillValue
709 710 711 712 713 714 715 716

        if part.rank == 0:
            ncvar = outnf.createVariable(name, vtyp, coord, fill_value=NCFillValue)
            ncvar.title = title
            ncvar.units = units
        else :
            ncvar = np.zeros((1,1,1))
        ncvar[:] = part.gather(var)
717
    #
718
    #
719
    #
720
    def dumpnetcdf(self, filename, globalgrid, procgrid, part, tcst) :
721 722 723 724
        #
        NCFillValue=1.0e20
        vtyp=np.float64
        cornerind=[0,2,4,6]
725
        nbcorners = len(cornerind)
726 727 728 729 730 731 732
        #
        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))
733
            outnf.createDimension('z', self.nbasmax)
734
            outnf.createDimension('bnd', nbcorners)
735
            outnf.createDimension('inflow', self.max_inflow)
736
        else :
737
            outnf = None
738
        #
739 740
        addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
741
        addparameters(outnf, part, tcst, vtyp, NCFillValue)
742
        #
743
        # land grid index -> to facilitate the analyses of the routing
744
        #
745 746 747
        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)
POLCHER Jan's avatar
POLCHER Jan committed
748
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "nbpt_glo", "Land point indices on global grid", "-", nbpt_glo[:,0], vtyp)
749
        #
750
        ################
751
        #
POLCHER Jan's avatar
POLCHER Jan committed
752
        # Conversion of grid indices for route_togrid
753
        grgrid = part.l2glandindex(self.route_togrid[:,:])
POLCHER Jan's avatar
POLCHER Jan committed
754
        #
755
        #
756
        # The field route_togrid is with indices on the local grid. That needs to be converted to the global grid.
757
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z', 'y','x'), "routetogrid", "Grid into which the basin flows", "-", grgrid, vtyp, "int")
758
        # route to basin
759
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z', 'y','x'), "routetobasin", "Basin in to which the water goes", "-", self.route_tobasin[:,:], vtyp, "int")
760
        # basin id
761
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z', 'y','x'), "basinid", "ID of basin", "-", self.global_basinid[:,:], vtyp, "int")
762 763
        #
        # basin area
764
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z', 'y','x'), "basin_area", "area of basin", "m^2", self.routing_area[:,:], vtyp, "float")
765

766 767 768 769 770
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z', 'y','x'), "basin_orog_mean", "Mean orography", "m", self.routing_orog_mean[:,:], vtyp, "float")

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

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

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

774 775
        # route number into basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "routenbintobas", "Number of basin into current one", "-", self.route_nbintobas[:], vtyp, "int")
776
        #
777 778
        # 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, "int")
779
        #
780
        # latitude of outlet
781
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "outletlat", "Latitude of Outlet", "degrees north", self.route_outlet[:,:,0], vtyp, "float")
782
        # longitude of outlet
783
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "outletlon", "Longitude of Outlet", "degrees east", self.route_outlet[:,:,1], vtyp, "float")
784
        # type of outlet
785
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "outlettype", "Type of outlet", "code", self.route_type[:,:], vtyp, "float")
786
        #
787
        # topographic index
788
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "topoindex", "Topographic index of the retention time", "m", self.topo_resid[:,:], vtyp, "float")
789
        #
790

791
        # Inflow number
792
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "route_innum", "Number of inflow", "-", self.route_innum[:,:], vtyp, "int")
793
        #
794 795
        # Inflow grid
        #gingrid = part.l2glandindex(self.inflow_grid[:,:,:inflow_size])
796
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow', 'z','y','x'), "route_ingrid", "Grid from which the water flows", "-", self.route_ingrid[:,:,:], vtyp, "int")
797
        #
798
        # Inflow basin
799
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow', 'z','y','x'), "route_inbasin", "Basin from which the water flows", "-", self.route_inbasin[:,:,:], vtyp, "int")
800

801
        #
802 803
        # Save centre of gravity of HTU
        #
804
        # Check if it works
805
        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, "float")
806

807
        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, "float")
808
        #
809 810
        # Save the fetch of each basin
        #
811
        self.add_variable(outnf, procgrid, NCFillValue, part, ('z','y','x'), "fetch", "Fetch contributing to each HTU", "m^2", self.routing_fetch[:,:], vtyp, "float")
812
        #
813 814
        # Close the file
        if part.rank == 0:
815
            outnf.close()
816 817
        #
        return