Interface.py 36.2 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 28
    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
#
import configparser
29
config = configparser.ConfigParser()
30
config.read("run.def")
31 32 33
gendoc = config.get("OverAll", "Documentation", fallback='false')
nbxmax = config.getint("OverAll", "nbxmax", fallback=63)
largest_pos = config.getint("OverAll", "ROUTING_RIVERS", fallback=50)
34 35
#
undef_int = 999999999.9
36 37
# Order of magnitude for the area precision in m^2.
prec = 100.0
38 39 40
#
# Print the documentation for the FORTRAN interface
#
41
if gendoc.lower() == "true" :
42 43 44 45 46 47 48 49 50 51 52 53 54
    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")
55 56 57 58
    docwrapper.write(routing_interface.finish_truncate.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.killbas.__doc__)

59 60 61 62 63 64 65
    docwrapper.close
#
# Functions to access the interfaces
#
#
# initatmgrid : Initialises the grid.f90 module and passes the description of the atmospheric grid.
#
66
def initatmgrid(rank, nbcore, nbpt, modelgrid) :
67 68 69
    print("INITATMGRID corners", np.array(modelgrid.polyll).shape)
    print("INITATMGRID area", np.array(modelgrid.area).shape)
    print("INITATMGRID neighbours", np.array(modelgrid.neighbours).shape)
70 71 72 73 74 75 76 77
    routing_interface.initatmgrid(rank, nbcore, modelgrid.polyll, modelgrid.coordll, modelgrid.area, modelgrid.contfrac, modelgrid.neighbours)
    return
#
#
#
def closeinterface(comm) :
    comm.Barrier()
    routing_interface.closeinterface()
78 79 80
    return
#
#
81
#
82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 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 179 180 181
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"
        lon[:,:] = longitude[:,:]
    #
    # 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"
        lat[:] = latitude[:,:]
    #
    # 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
#
#
#
182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210
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')
    #
    maxdiff_sorted = prec*prec
    iter_count = 0
    #
    while iter_count < part.size*3 and maxdiff_sorted > prec :
        fetch_out[:,:] = 0.0
        outflow_uparea = routing_interface.finalfetch(part.landcorelist, routing_area, basin_count, route_togrid, \
                                                      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]
        # 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))
        old_sorted[:] = sorted_outareas[0:largest_pos]
        iter_count += 1
211

212 213
    #
    fetch_error = np.sum(np.abs(fetch_out[part.landcorelist,:]-fetch_in[part.landcorelist,:]), axis=1)\
Anthony's avatar
Anthony committed
214
                                                    / np.ma.sum(routing_area[part.landcorelist,:], axis=1)
215
    if np.max(fetch_error) > prec :
216
        print("Rank :"+str(part.rank)+" Too large fetch error (fraction of greid area) : ", fetch_error)
217

218 219 220 221 222 223
    print("Total fetch error in fraction of grid box : ", np.sum(fetch_error))
    #
    return fetch_out
#
#
#
224 225
class HydroOverlap :
#
Anthony Schrapffer's avatar
Anthony Schrapffer committed
226
    def __init__(self, nbpt, nbvmax, sub_pts, sub_index_in, sub_area_in, sub_lon_in, sub_lat_in, part, modelgrid, hydrodata) :
227 228 229
        #
        # Reshape stuff so that it fits into arrays
        #
230
        sub_index = np.zeros((nbpt,nbvmax,2), dtype=np.int32, order='F')
231 232 233 234 235 236 237 238
        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]) :
239
                sub_index[ib,ip,:] = sub_index_in[ib][:,ip]
240
        #
241
        part.landsendtohalo(np.array(sub_area), order='F')
Anthony Schrapffer's avatar
Anthony Schrapffer committed
242
        #
Anthony's avatar
Anthony committed
243 244 245 246 247 248 249 250 251 252 253 254
        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
        #
        #
255 256 257 258 259
        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')
260 261
        orog_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        floodp_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
262 263 264 265 266 267 268 269 270
        #
        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][:])
271 272
            orog_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.orog[ib][:])
            floodp_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.floodplains[ib][:])
273
        #
Anthony's avatar
Anthony committed
274 275 276 277
        del hydrodata.trip; del hydrodata.basins; del hydrodata.topoind
        del hydrodata.fac; del hydrodata.disto; del hydrodata.orog;
        del hydrodata.floodplains
        #
278 279 280 281 282 283
        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, \
284
            self.orog_bx, self.floodp_bx, \
285
            self.lon_bx, self.lat_bx, self.lshead_bx = \
286
                    routing_interface.gethydrogrid(ijdimmax, sub_pts, sub_index, sub_area, \
287 288
                    hydrodata.basinsmax, hydrodata.topoindmin, sub_lon, sub_lat, trip_tmp, basins_tmp, topoind_tmp, fac_tmp,\
                        hierarchy_tmp, orog_tmp, floodp_tmp)
289
        #
Anthony's avatar
Anthony committed
290 291 292
        del trip_tmp; del basins_tmp; del topoind_tmp
        del fac_tmp; del hierarchy_tmp; del orog_tmp
        del floodp_tmp
293
        #
294
        self.nwbas = nbvmax
295 296 297 298
        # 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
        #
299 300 301 302
        return
#
#
#
303
class HydroSuper :
304
    def __init__(self, nbvmax, hydrodata, hydrooverlap, nbasmax, part) :
305 306 307
        #
        # Keep largest possible number of HTUs
        #
308
        self.nbasmax = nbasmax
309
        self.nbhtuext = nbvmax
310 311 312 313
        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
314
        nb_htu = nbvmax
315
        nbv = nbvmax
316 317 318 319
        #
        # Call findbasins
        #
        nb_basin, basin_inbxid, basin_outlet, basin_outtp, self.basin_sz, basin_bxout, basin_bbout, self.basin_pts, basin_lshead, coast_pts = \
320 321
                    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
322 323
                                                 basin_bx = hydrooverlap.basin_bx, fac_bx = hydrooverlap.fac_bx, \
                                                 hierarchy_bx = hydrooverlap.hierarchy_bx, \
324 325 326 327 328 329
                                                 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))
330 331
        # Set the number of inflows per basin. For the moment twice the maximum number of basins.
        self.inflowmax = self.nwbas*2
332
        print("Maximum number of basin created : {0}".format(self.nwbas))
333 334 335 336
        ijdim=[]
        for i in range(self.nbpt) :
            ijdim.append(max(hydrooverlap.area_bx[i,:,:].shape))
        self.ijdimmax = max(ijdim)
337 338 339
        #
        # Call Globalize
        #
340 341 342 343
        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
344
        #
345 346
        self.basin_count, self.basin_notrun, self.basin_area, self.basin_cg, self.basin_hierarchy, \
            self.basin_orog, self.basin_floodp, self.basin_fac, self.basin_topoind, \
347
            self.basin_id, self.basin_outcoor, self.basin_type, self.basin_flowdir, \
348
            self.basin_lshead, self.outflow_grid, self.outflow_basin, self.nbcoastal, self.coastal_basin = \
349 350
                    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, \
351 352
                                                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
353 354
                                                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, \
355
                                                basin_bbout = basin_bbout, lshead = basin_lshead, coast_pts = coast_pts, nwbas = self.nwbas)
356

357 358 359
        # 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
360 361
        return
    #
362 363 364 365 366
    def linkup(self, hydrodata) :
        #
        # Call the linkup routine in routing_reg.
        #
        print("Invented basins =", hydrodata.basinsmax)
367 368 369 370 371
        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))
372
        self.nbxmax_in = self.inflow_number.shape[1]
373 374 375
        #
        #
        #
376 377
        return
    #
378

POLCHER Jan's avatar
POLCHER Jan committed
379
    def fetch(self, part) :
380
        #
381
        fetch_basin = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
382
        #
383
        self.basin_area = routing_interface.areanorm(self.basin_count, self.basin_area, self.outflow_grid)
384
        partial_sum = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
385
        #
386
        old_sorted = np.zeros(largest_pos, dtype=np.float32, order='F')
387 388 389 390 391
        #
        maxdiff_sorted = prec*prec
        iter_count = 0
        #
        while iter_count < part.size*3 and maxdiff_sorted > prec :
392
            fetch_basin[:,:] = 0.0
393
            outflow_uparea = routing_interface.fetch(part.landcorelist, self.basin_count, self.basin_area, self.basin_id, self.basin_hierarchy, \
394
                                                         self.basin_fac, self.outflow_grid, self.outflow_basin, partial_sum, fetch_basin)
395 396
            partial_sum = np.copy(fetch_basin)
            part.landsendtohalo(partial_sum, order='F')
397 398
            partial_sum = part.zerocore(partial_sum, order='F')
            #
399 400 401 402 403 404 405 406 407
            # 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.
            if part.size == 1 :
                maxdiff_sorted = 0.0
            else :
408 409 410
                l = min(sorted_outareas.shape[0],largest_pos)
                maxdiff_sorted = np.max(np.abs(sorted_outareas[0:largest_pos]-old_sorted[0:l]))
            old_sorted[:l] = sorted_outareas[0:largest_pos]
411
            iter_count += 1
412

413 414
        self.fetch_basin = np.copy(fetch_basin)
        #
415
        # Upstream area of the smalest river we call largest rivers.
416
        #
417
        self.largest_rivarea = sorted_outareas[l-1]
418 419 420
        #
        #
        #
421 422 423
        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)
424
        return
425 426 427

    def check_fetch(self):

POLCHER Jan's avatar
POLCHER Jan committed
428 429
        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)
430

431
        return
432 433 434

    def check_routing(self):

POLCHER Jan's avatar
POLCHER Jan committed
435 436
        routing_interface.checkrouting(nbpt = self.nbpt, nwbas = self.nwbas, outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, \
                                       basin_count = self.basin_count)
437

438
        return
439 440 441 442

    def check_inflows(self):
        nbxmax_tmp = self.inflow_grid.shape[2]
        routing_interface.check_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)
443
    #
444
    #
445
    def killbas(self, tokill, totakeover, numops):
446
        ops = tokill.shape[1]
447 448 449 450
        #
        nbxmax_tmp = self.inflow_grid.shape[2]
        #
        routing_interface.killbas(nbpt = self.nbpt, inflowmax = nbxmax_tmp, \
451 452 453 454 455 456 457
                nbasmax = self.nbasmax, nwbas = self.nwbas, ops = ops, tokill = tokill,\
                totakeover = totakeover, numops = numops, basin_count = self.basin_count,\
                basin_area = self.basin_area, basin_orog = self.basin_orog, basin_floodp = self.basin_floodp, \
                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)
458

459
    #
460 461 462 463 464 465 466 467 468 469
    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)
470

471 472 473 474 475
    #
    def dumpnetcdf(self, filename, globalgrid, procgrid, part) :
        #
        NCFillValue=1.0e20
        vtyp=np.float64
476
        inflow_size = 100
477 478 479 480 481 482 483 484 485 486
        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))
            outnf.createDimension('htuext', self.nbhtuext)
487 488
            outnf.createDimension('htu', self.inflow_number.shape[1])
            outnf.createDimension('in',inflow_size )
489 490 491 492 493 494
            outnf.createDimension('bnd', nbcorners)
        else :
            outnf = None
        #
        addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
495
        #
496 497 498
        # Variables
        # Once gathered on root_proc we transform them into float64. Careful, Integer variables do not have NaN !
        #
499 500 501 502 503 504
        # 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
505
        # contfrac
506 507 508 509 510
        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)
511 512
        #
        #self.basin_count
513
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "basin_count", "HTU count", "-", self.basin_count, vtyp)
514
        #
515 516 517 518 519 520
        # 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)
        #
521 522 523 524 525 526
        # 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)
        #
527 528 529 530 531 532
        # 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)
533
        #
534 535 536
        # 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)
537
        #
538 539
        # type
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_type", "type", "-", self.basin_type, vtyp)
540
        #
541 542 543
        # flowdir
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_flowdir", "flowdir", "-", self.basin_flowdir, vtyp)
        #
544
        #
545
        #self.outflow_grid
546 547 548 549 550 551
        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)
552 553
        #
        #self.outflow_basin
554 555 556 557 558 559 560 561
        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
562 563
        #
        # Inflow Basin
564
        self.add_variable(outnf, procgrid, NCFillValue, part, ('in','htu','y','x'), "HTUinbas", "Inflow basin", "-", self.inflow_basin[:,:,:inflow_size], vtyp)
565 566 567
        #
        # Save the fetch of each basin
        #
568
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "fetch_basin", "Fetch contributing to each HTU", "m^2", self.fetch_basin, vtyp)
569 570 571 572 573 574 575
        #
        # Close file
        #
        if part.rank == 0 :
            outnf.close()
        #
        return
576 577 578
#
#
#
579
class HydroGraph :
580
    def __init__(self, nbasmax, hydrosuper, part, modelgrid) :
581
        #
582
        self.nbasmax = nbasmax
583
        self.nbpt = hydrosuper.basin_count.shape[0]
584
        nwbas = hydrosuper.basin_topoind.shape[1]
585 586
        nbxmax_in = hydrosuper.inflow_grid.shape[2]
        #
587
        #
588 589 590
        self.routing_area, self.routing_orog, self.routing_floodp, self.routing_cg, self.topo_resid, 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 = \
591
                                    routing_interface.finish_truncate(nbpt = self.nbpt, inflowmax = nbxmax_in, nbasmax = nbasmax, nwbas = nwbas, \
POLCHER Jan's avatar
POLCHER Jan committed
592 593 594 595 596 597 598 599 600 601 602 603
                                                                      num_largest = hydrosuper.num_largest, gridarea = modelgrid.area, \
                                                                      cfrac = modelgrid.contfrac, basin_count = hydrosuper.basin_count, \
                                                                      basin_notrun = hydrosuper.basin_notrun, basin_area = hydrosuper.basin_area, \
                                                                      basin_orog = hydrosuper.basin_orog, basin_floodp = hydrosuper.basin_floodp, \
                                                                      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)
604

605 606
        #
        self.routing_fetch = finalfetch(part, self.routing_area, self.route_nbbasin, self.route_togrid, self.route_tobasin, self.routing_fetch)
607
        #
608
        self.num_largest = routing_interface.finalrivclass(part.landcorelist, self.route_togrid, self.route_tobasin, self.routing_fetch, \
609
                      hydrosuper.largest_rivarea)
610
        #
611 612 613
        # Inflows
        self.max_inflow = part.domainmax(np.max(hydrosuper.inflow_number))
        gingrid = part.l2glandindex( hydrosuper.inflow_grid[:,:,:self.max_inflow])
614 615 616 617
        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])
618

619
        return
620 621 622 623 624 625 626 627 628
    #
    #
    #
    def add_variable(self,outnf, procgrid, NCFillValue, part, coord, name, title, units, data, vtyp, orig_type = "float"):
        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
629
            var[np.isnan(var)] = RPP.IntFillValue
630
            var[var>=np.abs(RPP.IntFillValue)] = NCFillValue
631 632 633 634 635 636 637 638

        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)
639
    #
640
    #
641
    #
642 643 644 645 646
    def dumpnetcdf(self, filename, globalgrid, procgrid, part) :
        #
        NCFillValue=1.0e20
        vtyp=np.float64
        cornerind=[0,2,4,6]
647
        nbcorners = len(cornerind)
648 649 650 651 652 653 654 655
        #
        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('htu', self.nbasmax)
656
            outnf.createDimension('bnd', nbcorners)
657
            outnf.createDimension('inflow', self.max_inflow)
658
        else :
659
            outnf = None
660
        #
661 662
        addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
663
        #
664
        # land grid index -> to facilitate the analyses of the routing
665
        #
666 667 668 669
        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)
670
        #
671
        ################
672
        #
673
        # TEST: l2glandindex
674 675 676 677 678 679 680
        itarget=-1
        for il in range(procgrid.nbland) :
            lo = procgrid.lon_full[procgrid.indP[il][0],procgrid.indP[il][1]]
            la = procgrid.lat_full[procgrid.indP[il][0],procgrid.indP[il][1]]
            d=np.sqrt((lo-3.13)**2+(la-39.70)**2)
            if d < 0.05 :
                itarget = il
681

682
        if itarget >+ 0 :
POLCHER Jan's avatar
POLCHER Jan committed
683
            print(part.rank, itarget, " Before route_togrid = ", self.route_togrid[itarget,:])
684
        # Conversion
685
        grgrid = part.l2glandindex(self.route_togrid[:,:])
686
        if itarget >+ 0 :
POLCHER Jan's avatar
POLCHER Jan committed
687
            print(part.rank, itarget, " After route_togrid = ", self.route_togrid[itarget,:])
688 689
        ################
        #
690
        # The field route_togrid is with indices on the local grid. That needs to be converted to the global grid.
691 692 693 694 695 696 697 698
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "routetogrid", "Grid into which the basin flows", "-", grgrid, vtyp, "int")
        # route to basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "routetobasin", "Basin in to which the water goes", "-", self.route_tobasin[:,:], vtyp, "int")
        # basin id
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "basinid", "ID of basin", "-", self.global_basinid[:,:], vtyp, "int")
        #
        # basin area
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "basin_area", "area of basin", "m^2", self.routing_area[:,:], vtyp, "float")
699 700 701 702 703

        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "basin_orog", "Mean orography", "m", self.routing_orog[:,:], vtyp, "float")

        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu', 'y','x'), "basin_floodp", "area of floodplains", "m^2", self.routing_floodp[:,:], vtyp, "float")

704 705
        # 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")
706
        #
707 708
        # 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")
709
        #
710 711 712 713 714
        # latitude of outlet
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "outletlat", "Latitude of Outlet", "degrees north", self.route_outlet[:,:,0], vtyp, "float")
        # longitude of outlet
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "outletlon", "Longitude of Outlet", "degrees east", self.route_outlet[:,:,1], vtyp, "float")
        # type of outlet
715
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "outlettype", "Type of outlet", "code", self.route_type[:,:], vtyp, "float")
716
        #
717 718
        # topographic index
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "topoindex", "Topographic index of the retention time", "m", self.topo_resid[:,:], vtyp, "float")
719
        #
720

721 722
        # Inflow number
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "route_innum", "Number of inflow", "-", self.route_innum[:,:], vtyp, "int")
723
        #
724 725 726
        # Inflow grid
        #gingrid = part.l2glandindex(self.inflow_grid[:,:,:inflow_size])
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow', 'htu','y','x'), "route_ingrid", "Grid from which the water flows", "-", self.route_ingrid[:,:,:], vtyp, "int")
727
        #
728 729
        # Inflow basin
        self.add_variable(outnf, procgrid, NCFillValue, part, ('inflow', 'htu','y','x'), "route_inbasin", "Basin from which the water flows", "-", self.route_inbasin[:,:,:], vtyp, "int")
730

731
        #
732 733
        # Save centre of gravity of HTU
        #
734 735
        # Check if it works
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "CG_lon", "Longitude of centre of gravity of HTU", "degrees east", self.routing_cg[:,:,1], vtyp, "float")
736 737

        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "CG_lat", "Latitude of centre of gravity of HTU", "degrees north", self.routing_cg[:,:,0], vtyp, "float")
738
        #
739 740
        # Save the fetch of each basin
        #
741
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htu','y','x'), "fetch", "Fetch contributing to each HTU", "m^2", self.routing_fetch[:,:], vtyp, "float")
742
        #
743 744
        # Close the file
        if part.rank == 0:
745
            outnf.close()
746 747
        #
        return