Interface.py 24 KB
Newer Older
1 2 3 4 5 6 7
#
#
import numpy as np
import os
import pickle
from netCDF4 import Dataset
from mpi4py import MPI
8
import gc
9 10
#
import sys
11 12 13 14 15
from inspect import currentframe, getframeinfo
#
localdir=os.path.dirname(getframeinfo(currentframe()).filename)
sys.path.append(localdir+'/F90subroutines')
F90=localdir+'/F90subroutines'
16
if MPI.COMM_WORLD.Get_rank() == 0 :
17
    err=os.system("cd "+F90+"; make all")
18 19 20 21 22 23 24 25
    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
26
import NcHelp as NH
27
#
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
maxpercent = config.getint("OverAll", "maxpercent", fallback=2)
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
class HydroOverlap :
#
Anthony Schrapffer's avatar
Anthony Schrapffer committed
84
    def __init__(self, nbpt, nbvmax, sub_pts, sub_index_in, sub_area_in, sub_lon_in, sub_lat_in, part, modelgrid, hydrodata) :
85 86 87
        #
        # Reshape stuff so that it fits into arrays
        #
88
        sub_index = np.zeros((nbpt,nbvmax,2), dtype=np.int32, order='F')
89 90 91 92 93 94 95 96
        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]) :
97
                sub_index[ib,ip,:] = sub_index_in[ib][:,ip]
98
        #
99
        part.landsendtohalo(np.array(sub_area), order='F')
Anthony Schrapffer's avatar
Anthony Schrapffer committed
100
        #
Anthony's avatar
Anthony committed
101 102 103 104 105 106 107 108 109 110 111 112
        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
        #
        #
113 114 115 116 117
        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')
118 119
        orog_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        floodp_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
120 121 122 123 124 125 126 127 128
        #
        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][:])
129 130
            orog_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.orog[ib][:])
            floodp_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.floodplains[ib][:])
131
        #
Anthony's avatar
Anthony committed
132 133 134 135
        del hydrodata.trip; del hydrodata.basins; del hydrodata.topoind
        del hydrodata.fac; del hydrodata.disto; del hydrodata.orog;
        del hydrodata.floodplains
        #
136 137 138 139 140 141
        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, \
142
            self.orog_bx, self.floodp_bx, \
143
            self.lon_bx, self.lat_bx, self.lshead_bx = \
144
                    routing_interface.gethydrogrid(ijdimmax, maxpercent, sub_pts, sub_index, sub_area, \
145 146
                    hydrodata.basinsmax, hydrodata.topoindmin, sub_lon, sub_lat, trip_tmp, basins_tmp, topoind_tmp, fac_tmp,\
                        hierarchy_tmp, orog_tmp, floodp_tmp)
147
        #
Anthony's avatar
Anthony committed
148 149 150
        del trip_tmp; del basins_tmp; del topoind_tmp
        del fac_tmp; del hierarchy_tmp; del orog_tmp
        del floodp_tmp
151
        #
152
        self.nwbas = nbvmax
153 154 155 156
        # 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
        #
157 158 159 160
        return
#
#
#
161
class HydroSuper :
162
    def __init__(self, nbvmax, hydrodata, hydrooverlap, nbasmax, part, comm) :
163 164 165
        #
        # Keep largest possible number of HTUs
        #
166
        self.nbasmax = nbasmax
167
        self.nbhtuext = nbvmax
168 169 170 171
        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
172
        nb_htu = nbvmax
173
        nbv = nbvmax
174 175 176 177
        #
        # Call findbasins
        #
        nb_basin, basin_inbxid, basin_outlet, basin_outtp, self.basin_sz, basin_bxout, basin_bbout, self.basin_pts, basin_lshead, coast_pts = \
178 179
                    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
180 181
                                                 basin_bx = hydrooverlap.basin_bx, fac_bx = hydrooverlap.fac_bx, \
                                                 hierarchy_bx = hydrooverlap.hierarchy_bx, \
182 183 184 185 186 187
                                                 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))
188
        # Set the number of inflows per basin. For the moment twice the maximum number of basins.
189
        self.inflowmax = max(10, self.nwbas*2)
190
        print("Maximum number of basin created : {0}".format(self.nwbas))
191 192 193 194
        ijdim=[]
        for i in range(self.nbpt) :
            ijdim.append(max(hydrooverlap.area_bx[i,:,:].shape))
        self.ijdimmax = max(ijdim)
195 196 197
        #
        # Call Globalize
        #
198 199 200 201
        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
202
        #
203
        self.basin_count, self.basin_notrun, self.basin_area, self.basin_cg, self.basin_hierarchy, \
204 205
            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,\
206 207
            self.basin_flowdir, self.basin_lshead, self.basin_beta_fp, self.outflow_grid, self.outflow_basin,\
            self.nbcoastal, self.coastal_basin = \
208 209
                    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, \
210 211
                                                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
212 213
                                                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, \
214
                                                basin_bbout = basin_bbout, lshead = basin_lshead, coast_pts = coast_pts, nwbas = self.nwbas)
215
        #
216 217 218
        # 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
219 220 221 222 223 224 225 226 227 228 229
        #
        # Correct all variables over the halo before the Linkup
        # except : outflow_grid (corrected after linkup), basin_cg and basin_outcoor (not needed and not valid with part.landsendtohalo)
        for variable in [self.basin_count, self.basin_notrun, self.basin_area, self.basin_hierarchy,\
              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_type,self.basin_flowdir,\
               self.basin_lshead, self.outflow_basin,self.nbcoastal,self.coastal_basin]:
            part.landsendtohalo(variable, order='F')
            comm.Barrier()
        #
        #
230 231
        return
    #
232 233 234 235 236
    def linkup(self, hydrodata) :
        #
        # Call the linkup routine in routing_reg.
        #
        print("Invented basins =", hydrodata.basinsmax)
237 238 239 240 241
        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))
242
        self.nbxmax_in = self.inflow_number.shape[1]
243 244 245
        #
        #
        #
246 247
        return
    #
248

POLCHER Jan's avatar
POLCHER Jan committed
249
    def fetch(self, part) :
250
        #
251
        fetch_basin = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
252
        #
253
        self.basin_area = routing_interface.areanorm(self.basin_count, self.basin_area, self.outflow_grid)
254
        partial_sum = np.zeros(self.basin_area.shape, dtype=np.float32, order='F')
255
        #
256
        old_sorted = np.zeros(largest_pos, dtype=np.float32, order='F')
257 258 259 260
        #
        maxdiff_sorted = prec*prec
        iter_count = 0
        #
261 262
        fhalolist = [i+1 for i in range(self.nbpt) if i not in part.landcorelist]
        #
263
        while iter_count < part.size*3 and maxdiff_sorted > prec :
264
            fetch_basin[:,:] = 0.0
265
            outflow_uparea = routing_interface.fetch(fhalolist, part.landcorelist, self.basin_count, self.basin_area, self.basin_id, self.basin_hierarchy, \
266
                                                         self.basin_fac, self.outflow_grid, self.outflow_basin, partial_sum, fetch_basin)
267 268
            partial_sum = np.copy(fetch_basin)
            part.landsendtohalo(partial_sum, order='F')
269 270
            partial_sum = part.zerocore(partial_sum, order='F')
            #
271 272 273 274 275 276
            # 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.
277
            l = min(sorted_outareas.shape[0],largest_pos)
278 279 280
            if part.size == 1 :
                maxdiff_sorted = 0.0
            else :
281
                maxdiff_sorted = np.max(np.abs(sorted_outareas[0:l]-old_sorted[0:l]))
282
                old_sorted[:l] = sorted_outareas[0:largest_pos]
283
            iter_count += 1
284

285 286
        self.fetch_basin = np.copy(fetch_basin)
        #
287
        # Upstream area of the smalest river we call largest rivers.
288
        #
289
        self.largest_rivarea = sorted_outareas[l-1]
290 291 292
        #
        #
        #
293 294 295
        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)
296
        return
297 298 299

    def check_fetch(self):

POLCHER Jan's avatar
POLCHER Jan committed
300 301
        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)
302

303
        return
304 305 306

    def check_routing(self):

POLCHER Jan's avatar
POLCHER Jan committed
307 308
        routing_interface.checkrouting(nbpt = self.nbpt, nwbas = self.nwbas, outflow_grid = self.outflow_grid, outflow_basin = self.outflow_basin, \
                                       basin_count = self.basin_count)
309

310
        return
311

312 313 314 315 316 317 318 319 320
    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
321
        hg = np.copy(self.outflow_grid)
322 323 324 325 326
        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
327 328
        # Work in a copy to avoid error
        hg1 = np.copy(hg)
329
        for a in outflows_out:
Anthony's avatar
Anthony committed
330
          hg1[hg == a] = 0
331 332
        for ind in range(self.nbpt):
          hg1[hg == nbpt_glo[ind,0]] = nbpt_loc[ind,0]
Anthony's avatar
Anthony committed
333 334
        hg = hg1
        del hg1
335 336 337 338 339 340 341 342 343 344 345 346 347
        # 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
348
        nbxmax_tmp = self.inflow_grid.shape[2]
349 350 351 352 353
        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


354
    def killbas(self, tokill, totakeover, numops):
355
        ops = tokill.shape[1]
356 357 358 359
        #
        nbxmax_tmp = self.inflow_grid.shape[2]
        #
        routing_interface.killbas(nbpt = self.nbpt, inflowmax = nbxmax_tmp, \
360 361
                nbasmax = self.nbasmax, nwbas = self.nwbas, ops = ops, tokill = tokill,\
                totakeover = totakeover, numops = numops, basin_count = self.basin_count,\
362 363
                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, \
364
                basin_cg = self.basin_cg, basin_topoind = self.basin_topoind, basin_beta_fp = self.basin_beta_fp, fetch_basin = self.fetch_basin,\
365 366 367
                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)
368 369 370 371 372 373 374
    #
    def get_floodcri(self):
        nbxmax_tmp = self.inflow_grid.shape[2]
        self.floodcri = routing_interface.get_floodcri(nbpt = self.nbpt,\
             nwbas=self.nwbas, inflowmax = nbxmax_tmp, inflow_number = self.inflow_number,\
             inflow_basin = self.inflow_basin, inflow_grid = self.inflow_grid, basin_count = self.basin_count,\
             basin_floodp = self.basin_floodp, basin_orog_min = self.basin_orog_min)
375

376
    
377
    #
378 379 380 381 382 383 384 385 386 387
    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)
388

389 390 391 392 393
    #
    def dumpnetcdf(self, filename, globalgrid, procgrid, part) :
        #
        NCFillValue=1.0e20
        vtyp=np.float64
394
        inflow_size = 100
395 396 397 398 399 400 401 402 403
        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))
404
            outnf.createDimension('htuext', self.basin_id.shape[1])
405 406
            outnf.createDimension('htu', self.inflow_number.shape[1])
            outnf.createDimension('in',inflow_size )
407 408 409 410
            outnf.createDimension('bnd', nbcorners)
        else :
            outnf = None
        #
411 412
        NH.addcoordinates(outnf, globalgrid, procgrid, part, vtyp, NCFillValue, nbcorners, cornerind)
        NH.addenvironment(outnf, procgrid, part, vtyp, NCFillValue, self.nbpt)
413
        #
414 415 416
        # Variables
        # Once gathered on root_proc we transform them into float64. Careful, Integer variables do not have NaN !
        #
417 418 419 420 421 422
        # 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
423
        # contfrac
424 425 426 427 428
        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)
429 430
        #
        #self.basin_count
431
        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "basin_count", "HTU count", "-", self.basin_count, vtyp)
432
        #
433 434 435 436 437 438
        # 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)
        #
439 440 441 442 443 444
        # 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)
        #
445 446 447 448 449 450
        # 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)
451
        #
452 453 454
        # 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)
455
        #
456 457
        # type
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_type", "type", "-", self.basin_type, vtyp)
458
        #
459 460 461
        # flowdir
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_flowdir", "flowdir", "-", self.basin_flowdir, vtyp)
        #
462
        #
463
        #self.outflow_grid
464 465 466 467 468 469
        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)
470 471
        #
        #self.outflow_basin
472 473 474 475 476 477 478 479
        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
480 481
        #
        # Inflow Basin
482
        self.add_variable(outnf, procgrid, NCFillValue, part, ('in','htu','y','x'), "HTUinbas", "Inflow basin", "-", self.inflow_basin[:,:,:inflow_size], vtyp)
483 484 485
        #
        # Save the fetch of each basin
        #
486
        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "fetch_basin", "Fetch contributing to each HTU", "m^2", self.fetch_basin, vtyp)
487 488 489 490 491 492 493
        #
        # Close file
        #
        if part.rank == 0 :
            outnf.close()
        #
        return
494 495 496
#
#
#