Changes to get a more complete intermediate .nc output

Interface.py

@@ -362,11 +362,23 @@ class HydroSuper :
         return
     #
     #
+    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)
+    
     #
     def dumpnetcdf(self, filename, globalgrid, procgrid, part) :
         #
         NCFillValue=1.0e20
         vtyp=np.float64
+        inflow_size = 100
         cornerind=[0,2,4,6]
         nbcorners = len(cornerind)
         #
@@ -377,6 +389,8 @@ class HydroSuper :
             outnf.createDimension('y', globalgrid.nj)
             outnf.createDimension('land', len(procgrid.area))
             outnf.createDimension('htuext', self.nbhtuext)
+            outnf.createDimension('htu', self.inflow_number.shape[1])
+            outnf.createDimension('in',inflow_size )
             outnf.createDimension('bnd', nbcorners)
         else :
             outnf = None
@@ -387,72 +401,74 @@ class HydroSuper :
         # Variables
         # Once gathered on root_proc we transform them into float64. Careful, Integer variables do not have NaN !
         #
-        # self.basin_id
-        bid = procgrid.landscatter(self.basin_id[:,:].astype(vtyp), order='F')
-        bid[np.isnan(bid)] = NCFillValue
-        if part.rank == 0 :
-            basinid = outnf.createVariable("HTUid", vtyp, ('htuext','y','x'), fill_value=NCFillValue)
-            basinid.title = "ID for each HTU"
-            basinid.units = "-"
-        else :
-            basinid = np.zeros((1,1,1))
-        basinid[:,:,:] = part.gather(bid)
+        # 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)
+        #
+        # gridarea
+        contfrac = np.array(procgrid.contfrac)
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "contfrac", "Land fraction", "-", np.array(procgrid.contfrac), vtyp)
+        #
+        # gridarea
+        nbptarea = np.array(procgrid.area)
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "gridarea", "Grid area", "-", nbptarea, vtyp)
+        #
+        # basin_id
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_id", "ID for each HTU", "-", self.basin_id, vtyp)
         #
         #self.basin_count
-        bcnt = procgrid.landscatter(self.basin_count[:].astype(vtyp), order='F')
-        bcnt[np.isnan(bcnt)] = NCFillValue
-        if part.rank == 0 :
-            basincnt = outnf.createVariable("HTUcount", vtyp, ('y','x'), fill_value=NCFillValue)
-            basincnt.title = "HTU count"
-            basincnt.units = "-"
-        else :
-            basincnt = np.zeros((1,1))
-        basincnt[:,:] = part.gather(bcnt)
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('y','x'), "basin_count", "HTU count", "-", self.basin_count, vtyp)
         #
-        #self.basin_area
-        ba = procgrid.landscatter(self.basin_area[:,:].astype(vtyp), order='F')
-        ba[np.isnan(ba)] = NCFillValue
-        if part.rank == 0 :
-            basinarea = outnf.createVariable("HTUarea", vtyp, ('htuext','y','x'), fill_value=NCFillValue)
-            basinarea.title = "HTU area"
-            basinarea.units = "m^2"
-        else :
-            basinarea = np.zeros((1,1,1))
-        basinarea[:,:,:] = part.gather(ba)
+        # 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)
+        #
+        # 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)
+        # 
+        # 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)
+        # 
+        # type
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_type", "type", "-", self.basin_type, vtyp)
         #
+        # flowdir
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "basin_flowdir", "flowdir", "-", self.basin_flowdir, vtyp)
+        #
+        # 
         #self.outflow_grid
-        og = procgrid.landscatter(self.outflow_grid[:,:].astype(vtyp), order='F')
-        og[np.isnan(og)] = NCFillValue
-        if part.rank == 0 :
-            outgrid = outnf.createVariable("HTUoutgrid", vtyp, ('htuext','y','x'), fill_value=NCFillValue)
-            outgrid.title = "HTU outflow grid"
-            outgrid.units = "-"
-        else :
-            outgrid = np.zeros((1,1,1))
-        outgrid[:,:,:] = part.gather(og)
+        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)
         #
         #self.outflow_basin
-        ob = procgrid.landscatter(self.outflow_basin[:,:].astype(vtyp), order='F')
-        ob[np.isnan(ob)] = NCFillValue
-        if part.rank == 0 :
-            outbas = outnf.createVariable("HTUoutbasin", vtyp, ('htuext','y','x'), fill_value=NCFillValue)
-            outbas.title = "Outflow HTU of grid"
-            outbas.units = "-"
-        else :
-            outbas = np.zeros((1,1,1))
-        outbas[:,:,:] = part.gather(ob)
+        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)
+
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('in','htu','y','x'), "HTUinbas", "Inflow basin", "-", self.inflow_basin[:,:,:inflow_size], vtyp)
+
         #
         # Save the fetch of each basin
         #
-        fe = procgrid.landscatter(self.fetch_basin[:,:].astype(vtyp), order='F')
-        fe[np.isnan(fe)] = NCFillValue
-        if part.rank == 0 :
-            fetch = outnf.createVariable("fetch", vtyp, ('htuext','y','x'), fill_value=NCFillValue)
-            fetch.title = "Fetch contributing to each HTU"
-            fetch.units = "m^2"
-        else :
-            fetch = np.zeros((1,1,1))
-        fetch[:,:,:] = part.gather(fe)
+        self.add_variable(outnf, procgrid, NCFillValue, part, ('htuext','y','x'), "fetch_basin", "Fetch contributing to each HTU", "m^2", self.fetch_basin, vtyp)
         #
         # Close file
         #

Partition.py

@@ -434,16 +434,29 @@ class partition :
     # Convert local index of land points to global index
     #
     def l2glandindex(self, x) :
-        nl,nh = x.shape
-        y = np.zeros(x.shape, dtype=x.dtype)
-        if nl == self.nbland :
-            for i in range(nl) :
-                for j in range(nh) :
-                    # Land indices are in FORTRAN !!
-                    y[i,j] = self.l2glandind[x[i,j]-1]+1
-        else :
-            ERROR("The first dimension does not have the length of the number of land points")
-            sys.exit()
+        if x.ndim == 2:
+           nl,nh = x.shape
+           y = np.zeros(x.shape, dtype=x.dtype)
+           if nl == self.nbland :
+               for i in range(nl) :
+                   for j in range(nh) :
+                       # Land indices are in FORTRAN !!
+                       y[i,j] = self.l2glandind[x[i,j]-1]+1
+           else :
+               ERROR("The first dimension does not have the length of the number of land points")
+               sys.exit()
+        if x.ndim == 3:
+           nl,nh1,nh2 = x.shape
+           y = np.zeros(x.shape,dtype = x.dtype)
+           if nl == self.nbland:
+               for i in range(nl):
+                   for j in range(nh1):
+                       for k in range(nh2):
+                           # Land indices are in Fortran
+                           y[i,j,k] = self.l2glandind[x[i,j,k]-1]+1
+           else:
+               ERROR("The first dimension does not have the length of the number of land points")
+               sys.exit()
         return y
     #
     # Set to zero all points in the core