Interface.py 24.7 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 11
#
import sys
#
12 13 14 15
import F90tools as F90T
#
F90T.adddirtopath("F90subroutines")
F90T.compilef90("routing_interface")
16
import routing_interface
17
import NcHelp as NH
18
#
19 20
import getargs
config = getargs.SetupConfig()
21 22
gendoc = config.get("OverAll", "Documentation", fallback='false')
nbxmax = config.getint("OverAll", "nbxmax", fallback=63)
23
largest_pos = config.getint("OverAll", "ROUTING_RIVERS", fallback=200)
24
maxpercent = config.getint("OverAll", "maxpercent", fallback=2)
25 26
#
undef_int = 999999999.9
27 28
# Order of magnitude for the area precision in m^2.
prec = 100.0
29 30 31
#
# Print the documentation for the FORTRAN interface
#
32
if gendoc.lower() == "true" :
33 34 35 36 37 38 39 40 41 42 43 44 45
    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")
46 47 48 49
    docwrapper.write(routing_interface.finish_truncate.__doc__)
    docwrapper.write("====================================================================\n")
    docwrapper.write(routing_interface.killbas.__doc__)

50 51 52 53 54 55 56
    docwrapper.close
#
# Functions to access the interfaces
#
#
# initatmgrid : Initialises the grid.f90 module and passes the description of the atmospheric grid.
#
57
def initatmgrid(rank, nbcore, nbpt, modelgrid) :
58 59 60
    print("INITATMGRID corners", np.array(modelgrid.polyll).shape)
    print("INITATMGRID area", np.array(modelgrid.area).shape)
    print("INITATMGRID neighbours", np.array(modelgrid.neighbours).shape)
61 62 63 64 65 66 67 68
    routing_interface.initatmgrid(rank, nbcore, modelgrid.polyll, modelgrid.coordll, modelgrid.area, modelgrid.contfrac, modelgrid.neighbours)
    return
#
#
#
def closeinterface(comm) :
    comm.Barrier()
    routing_interface.closeinterface()
69 70 71
    return
#
#
72
#
73 74
class HydroOverlap :
#
Anthony Schrapffer's avatar
Anthony Schrapffer committed
75
    def __init__(self, nbpt, nbvmax, sub_pts, sub_index_in, sub_area_in, sub_lon_in, sub_lat_in, part, modelgrid, hydrodata) :
76 77 78
        #
        # Reshape stuff so that it fits into arrays
        #
79
        sub_index = np.zeros((nbpt,nbvmax,2), dtype=np.int32, order='F')
80 81 82 83 84 85 86 87
        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]) :
88
                sub_index[ib,ip,:] = sub_index_in[ib][:,ip]
89
        #
90
        part.landsendtohalo(np.array(sub_area), order='F')
Anthony Schrapffer's avatar
Anthony Schrapffer committed
91
        #
Anthony's avatar
Anthony committed
92 93
        ijdim=[]
        for ib in range(nbpt) :
94 95
            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))
Anthony's avatar
Anthony committed
96 97 98 99 100 101 102 103 104
        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
        #
        #
105 106 107
        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')
108 109
        rlen_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        rdz_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
110 111
        fac_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        hierarchy_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
112 113
        orog_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
        floodp_tmp = np.zeros((nbpt,nbvmax), dtype=np.float32, order='F')
114 115 116 117 118 119 120
        #
        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][:])
121 122
            rlen_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.rlen[ib][:])
            rdz_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.rdz[ib][:])
123 124
            fac_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.fac[ib][:])
            hierarchy_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.disto[ib][:])
125 126
            orog_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.orog[ib][:])
            floodp_tmp[ib,0:sub_pts[ib]] = np.asarray(hydrodata.floodplains[ib][:])
127
        #
Anthony's avatar
Anthony committed
128 129
        del hydrodata.trip; del hydrodata.basins; del hydrodata.topoind
        del hydrodata.fac; del hydrodata.disto; del hydrodata.orog;
130
        del hydrodata.floodplains; del hydrodata.rlen; hydrodata.rdz
Anthony's avatar
Anthony committed
131
        #
132 133 134 135 136
        trip_tmp[np.isnan(trip_tmp)] = undef_int
        basins_tmp[np.isnan(trip_tmp)] = undef_int
        #
        # Go to the call of the FORTRAN interface
        #
137 138 139 140 141 142 143
        self.nbi, self.nbj, self.area_bx, self.trip_bx, self.basin_bx, self.topoind_bx, self.rlen_bx, self.rdz_bx, \
            self.rweight_bx, self.fac_bx, self.hierarchy_bx, self.orog_bx, self.floodp_bx, self.lon_bx, self.lat_bx, \
            self.lshead_bx = routing_interface.gethydrogrid(ijdimmax, maxpercent, sub_pts, sub_index, sub_area, \
                                                            hydrodata.basinsmax, hydrodata.topoindmin, hydrodata.meanlength, \
                                                            hydrodata.meandz, sub_lon, sub_lat, trip_tmp, \
                                                            basins_tmp, topoind_tmp, rlen_tmp, rdz_tmp, fac_tmp,\
                                                            hierarchy_tmp, orog_tmp, floodp_tmp)
144
        #
Anthony's avatar
Anthony committed
145 146
        del trip_tmp; del basins_tmp; del topoind_tmp
        del fac_tmp; del hierarchy_tmp; del orog_tmp
147
        del floodp_tmp; del rlen_tmp; del rdz_tmp
148
        #
149
        self.nwbas = nbvmax
150 151 152 153
        # 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
        #
154 155 156 157
        return
#
#
#
158
class HydroSuper :
159
    def __init__(self, nbvmax, hydrodata, hydrooverlap, nbasmax, part, comm) :
160 161 162
        #
        # Keep largest possible number of HTUs
        #
163
        self.nbasmax = nbasmax
164
        self.nbhtuext = nbvmax
165 166 167 168
        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
169
        nb_htu = nbvmax
170
        nbv = nbvmax
171 172 173
        #
        # Call findbasins
        #
174 175 176 177 178 179 180 181 182
        nb_basin, basin_inbxid, basin_outlet, basin_outtp, self.basin_sz, basin_bxout, basin_bbout, self.basin_pts, \
            basin_lshead, coast_pts = \
                routing_interface.findbasins(nbpt = self.nbpt, nb_htu = self.nbhtuext, nbv = nbv, nbi = hydrooverlap.nbi, \
                                             nbj = hydrooverlap.nbj, trip_bx = hydrooverlap.trip_bx, \
                                             basin_bx = hydrooverlap.basin_bx, fac_bx = hydrooverlap.fac_bx, \
                                             hierarchy_bx = hydrooverlap.hierarchy_bx, topoind_bx = hydrooverlap.topoind_bx, \
                                             rlen_bx = hydrooverlap.rlen_bx, rdz_bx = hydrooverlap.rdz_bx, \
                                             rweight_bx = hydrooverlap.rweight_bx, lshead_bx = hydrooverlap.lshead_bx, \
                                             lontmp = hydrooverlap.lon_bx, lattmp = hydrooverlap.lat_bx)
183 184 185 186
        #
        # Adjust nwbas to the maximum found over the domain
        #
        self.nwbas = part.domainmax(np.max(nb_basin))
187
        # Set the number of inflows per basin. For the moment twice the maximum number of basins.
188
        self.inflowmax = max(10, self.nwbas*2)
189
        print("Maximum number of basin created : {0}".format(self.nwbas))
190 191 192 193
        ijdim=[]
        for i in range(self.nbpt) :
            ijdim.append(max(hydrooverlap.area_bx[i,:,:].shape))
        self.ijdimmax = max(ijdim)
194 195 196
        #
        # Call Globalize
        #
197 198 199 200
        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
201
        #
202
        self.basin_count, self.basin_notrun, self.basin_area, self.basin_cg, self.basin_hierarchy, \
203
            self.basin_orog_mean, self.basin_orog_min, self.basin_orog_max, self.basin_floodp, \
204 205 206
            self.basin_fac, self.basin_topoind, self.basin_rlen, self.basin_rdz, self.basin_id, \
            self.basin_outcoor, self.basin_type, self.basin_flowdir, self.basin_lshead, self.basin_beta_fp, \
            self.outflow_grid, self.outflow_basin, self.nbcoastal, self.coastal_basin = \
207 208
                    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, \
209
                                                hierarchy_bx = hydrooverlap.hierarchy_bx, orog_bx = hydrooverlap.orog_bx, floodp_bx =  hydrooverlap.floodp_bx,\
210 211 212
                                                fac_bx = hydrooverlap.fac_bx, topoind_bx = hydrooverlap.topoind_bx, \
                                                rlen_bx = hydrooverlap.rlen_bx, rdz_bx = hydrooverlap.rdz_bx, rweight_bx = hydrooverlap.rweight_bx, \
                                                min_topoind = hydrodata.topoindmin, \
POLCHER Jan's avatar
POLCHER Jan committed
213 214
                                                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, \
215
                                                basin_bbout = basin_bbout, lshead = basin_lshead, coast_pts = coast_pts, nwbas = self.nwbas)
216
        #
217 218 219
        # 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
220 221 222 223 224 225 226 227 228 229 230
        #
        # 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()
        #
        #
231 232
        return
    #
233 234 235 236 237
    def linkup(self, hydrodata) :
        #
        # Call the linkup routine in routing_reg.
        #
        print("Invented basins =", hydrodata.basinsmax)
238 239 240 241 242
        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))
243
        self.nbxmax_in = self.inflow_number.shape[1]
244 245 246
        #
        #
        #
247 248
        return
    #
249

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

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

    def check_fetch(self):

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

304
        return
305 306 307

    def check_routing(self):

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

311
        return
312

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


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

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

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