Truncate class for monoproc alternative

Truncate.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Sat Feb 22 10:19:21 2020
+
+@author: anthony
+"""
+
+from netCDF4 import Dataset as NetCDFFile
+import numpy as np
+import numpy.ma as ma
+import time
+import sys
+import os
+from inspect import currentframe, getframeinfo
+
+localdir=os.path.dirname(getframeinfo(currentframe()).filename)
+sys.path.append(localdir+'/F90subroutines')
+
+import routing_interface
+
+######################
+
+class truncation:
+    """
+    Indices in Fortran index : need to adapt when using numpy
+    """
+    def __init__(self, dfile, nbasmax):
+        self.nc  = NetCDFFile(dfile, "r")
+        
+        self.nbasmax = nbasmax
+        self.num_largest = 10
+        self.FillValue=np.nan
+        self.IntFillValue=999999
+        #
+        self.proc_num = self.nc.variables["proc_num"][:]
+        self.nbpt_proc = self.nc.variables["nbpt_proc"][:]
+        #
+        self.nbpt = self.nc.variables["nbpt_glo"][:]
+        self.nbpt_num = int(np.max(self.nbpt))
+        self.nj, self.ni = self.nbpt.shape
+        #
+        self.il_ji = [[k[0][0], k[1][0]] for k in [np.where(self.nbpt == i) for i in range(1,self.nbpt_num+1)]] 
+        #
+        # Get i,j variables
+        self.contfrac_e = self.nc.variables["contfrac"][:]
+        self.area_e = self.nc.variables["area"][:]
+        self.basin_count_e = self.nc.variables["basin_count"][:]
+        self.basin_notrun_e = self.nc.variables["basin_notrun"][:]
+        self.basin_area_e = self.nc.variables["basin_area"][:]
+        self.basin_cg_lon_e = self.nc.variables["CG_lon"][:]
+        self.basin_cg_lat_e = self.nc.variables["CG_lat"][:]
+        self.basin_topoind_e = self.nc.variables["basin_topoind"][:]
+        self.fetch_basin_e = self.nc.variables["fetch_basin"][:]
+        self.basin_id_e = self.nc.variables["basin_id"][:]
+        self.basin_outcoor_lon_e = self.nc.variables["outcoor_lon"][:]
+        self.basin_outcoor_lat_e = self.nc.variables["outcoor_lat"][:]
+        self.basin_type_e = self.nc.variables["basin_type"][:]
+        self.basin_flowdir_e = self.nc.variables["basin_flowdir"][:]
+        self.outflow_grid_e = self.nc.variables["HTUoutgrid"][:] # name
+        self.outflow_basin_e = self.nc.variables["HTUoutbasin"][:] # name
+        self.inflow_grid_e = self.nc.variables["HTUingrid"][:] # name
+        self.inflow_basin_e = self.nc.variables["HTUinbas"][:] # name
+        self.inflow_number_e = self.nc.variables["HTUinnum"][:]
+        #
+        # Conversion
+        #
+        self.gridarea = self.convert_toland(self.area_e).astype(np.int32)
+        self.contfrac = self.convert_toland(self.contfrac_e).astype(np.int32)
+        self.basin_count = self.convert_toland(self.basin_count_e).astype(np.int32)
+        self.basin_notrun = self.convert_toland(self.basin_notrun_e).astype(np.int32)
+        self.basin_area = self.convert_toland(self.basin_area_e).astype(np.float32)
+        self.basin_cg_lon = self.convert_toland(self.basin_cg_lon_e).astype(np.float32)
+        self.basin_cg_lat = self.convert_toland(self.basin_cg_lat_e).astype(np.float32)
+        self.basin_cg = np.zeros((self.basin_cg_lon.shape[0], self.basin_cg_lon.shape[1], 2), order = 'F').astype(np.float32)
+        self.basin_cg[:,:,0] = self.basin_cg_lat
+        self.basin_cg[:,:,1] = self.basin_cg_lon
+        
+        self.basin_topoind = self.convert_toland(self.basin_topoind_e).astype(np.float32)
+        self.fetch_basin = self.convert_toland(self.fetch_basin_e).astype(np.float32)
+        self.basin_id = self.convert_toland(self.basin_id_e).astype(np.int32)
+        self.basin_outcoor_lon = self.convert_toland(self.basin_outcoor_lon_e)
+        self.basin_outcoor_lat = self.convert_toland(self.basin_outcoor_lat_e)
+        self.basin_outcoor = np.zeros((self.basin_outcoor_lon.shape[0], self.basin_outcoor_lon.shape[1], 2), order = 'F').astype(np.float32)
+        self.basin_outcoor[:,:,0] = self.basin_outcoor_lat
+        self.basin_outcoor[:,:,1] = self.basin_outcoor_lon
+        
+        self.basin_type = self.convert_toland(self.basin_type_e).astype(np.float32)
+        self.basin_flowdir = self.convert_toland(self.basin_flowdir_e).astype(np.int32)
+        self.outflow_grid = self.convert_toland(self.outflow_grid_e).astype(np.int32)
+        self.outflow_basin = self.convert_toland(self.outflow_basin_e).astype(np.int32)
+        
+        self.inflow_grid = self.convert_toland(self.inflow_grid_e).astype(np.int32)
+        self.inflow_basin = self.convert_toland(self.inflow_basin_e).astype(np.int32)
+        self.inflow_number = self.convert_toland(self.inflow_number_e).astype(np.int32)
+        
+        
+    def convert_toland(self, array):
+        if array.ndim == 2:
+            dim = (self.nbpt_num)
+        if array.ndim == 3:
+            dim = (self.nbpt_num, array.shape[0])
+        if array.ndim == 4:
+            dim = (self.nbpt_num, array.shape[1], array.shape[1]) # (inflow, htu, lat, lon) -> nbpt, htu, inflow
+            
+        # Ne pas mettre à 0, mettre à IntFillValue
+        out = np.zeros(dim, order = 'F')
+         
+        for il in range(self.nbpt_num):
+            j,i = self.il_ji[il]
+            if array.ndim == 2:
+                out[il] = array[j,i]
+            if array.ndim == 3:
+                out[il,:] = array[:,j,i]
+            if array.ndim == 4:
+                out[il,:,:array.shape[0]] = np.swapaxes(array[:,:,j,i],0,1)
+        return out
+
+    ###############
+
+    def convert_toij(self, array):
+        if array.ndim == 1:
+            dim = (self.nj,self.ni)
+        if array.ndim == 2:
+            dim = (array.shape[1], self.nj, self.ni)
+        if array.ndim == 3:
+            dim = (array.shape[2], array.shape[1], self.nj, self.ni) # (grid, htu, inflow) -> inflow, htu, lat, lon
+            
+        # Ne pas mettre à 0, mettre à IntFillValue
+        dtype = array.dtype
+        out = np.zeros(dim, order = 'F', dtype = dtype)
+        if dtype == np.int32:
+            out[:] = self.IntFillValue
+        elif dtype == np.float32:
+            out[:] = self.FillValue
+         
+        for il in range(self.nbpt_num):
+            j,i = self.il_ji[il]
+            if array.ndim == 1:
+                out[j,i] = array[il]
+            if array.ndim == 2:
+                out[:,j,i] = array[il,:]
+            if array.ndim == 3:
+                out[:,:,j,i] = np.swapaxes(array[il,:,:], 0, 1) # (inflow, htu, lat, lon)
+        return out
+
+    ###############
+    
+    def truncate(self):
+        nwbas = self.basin_topoind.shape[1]
+        nbxmax_in = self.inflow_grid.shape[1]
+        inf = np.copy(self.inflow_number)
+        self.routing_area, 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 = \
+                                    routing_interface.truncate(nbpt = self.nbpt_num, nbxmax_in = nbxmax_in, nbasmax = self.nbasmax, nwbas = nwbas, num_largest = self.num_largest, gridarea = self.gridarea, cfrac = self.contfrac, basin_count = self.basin_count, \
+                                                               basin_notrun = self.basin_notrun, basin_area = self.basin_area, basin_cg = self.basin_cg, \
+                                                               basin_topoind = self.basin_topoind, fetch_basin = self.fetch_bis, 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)
+
+        # Adjust inflows -> add this part in fortran
+        for il in range(self.nbpt_num):
+            for ib in range(self.nbasmax):
+                self.inflow_grid[il,ib,self.inflow_number[il,ib]:] = self.IntFillValue
+                self.inflow_basin[il,ib,self.inflow_number[il,ib]:] = self.IntFillValue
+
+        self.inflow_grid[il,self.nbasmax:,:] = self.IntFillValue
+        self.inflow_basin[il,self.nbasmax:,:] = self.IntFillValue        
+        self.inflow_number[il,self.nbasmax:] = self.IntFillValue
+                
+
+    def finalrivclass(self, largest_rivarea):
+
+        A = np.where((self.route_tobasin > self.nbasmax) & (self.routing_fetch > largest_rivarea))
+        A = [[A[0][l], A[1][l]] for l in range(len(A[0]))]
+        num_largest = len(A)
+        for il, ib in A:
+            self.route_tobasin[il,ib] = self.nbasmax +3
+        
+        print('NUMBER OF RIVERS :', len(A))
+
+
+    def get_downstream(self, L, ig, ib):
+        """
+        Useful to add / remove fetch
+        """
+        L.append([ig,ib])
+        if self.outflow_grid[ig-1,ib-1]>0:
+            self.get_downstream(L, int(self.outflow_grid[ig-1, ib-1]), int(self.outflow_basin[ig-1, ib-1]))
+        else:
+            L.append(int(self.outflow_grid[ig-1, ib-1]))#,int(self.outbasin[ib-1,j,i])])                                
+                                    
+    def get_fetch(self):
+    
+        L_area = []
+        for ig, J in enumerate(self.il_ji):
+            j,i = J
+            L_area.append([self.basin_area[ig,k] for k in range(self.basin_count[ig])])
+        
+        L_fetch_calc = []
+        for ig, J in enumerate(self.il_ji):
+            j,i = J
+            L_fetch_calc.append([0 for k in range(self.basin_count[ig])])
+
+        for ig in range(1, len(self.basin_count)+1): # plus propre nbpt_num
+            for ib in range(1, self.basin_count[ig-1]+1):
+                fe = L_area[ig-1][ib-1]
+                L = []
+                self.get_downstream(L, ig, ib)
+                for igf,ibf in L[:-1]:
+                    L_fetch_calc[igf-1][ibf-1] += fe
+                
+        self.fetch_bis = np.zeros(self.fetch_basin.shape, order = 'F').astype(np.float32)
+        for ig in range(1, len(self.basin_count)+1):# plus propre nbpt_num
+            self.fetch_bis[ig-1,:self.basin_count[ig-1]] = np.array(L_fetch_calc[ig-1])
+ 
+        return
+
+    ####################################################
+    #
+    # Generation of the NetCDF output
+    # 
+
+    def addcoordinates(self, outnf) :
+        #
+        for varn in ["longitude", "latitude", "lon_bnd", "land", "area", "proc_num", "nbpt_proc", "nbpt_glo"]: 
+            ovar = self.nc.variables[varn]
+            nvar = outnf.createVariable(varn, ovar.dtype, ovar.dimensions)
+            nvar[:] = ovar[:]
+            for attrn in ovar.ncattrs():
+                attrv = ovar.getncattr(attrn)
+                nvar.setncattr(attrn,attrv)
+            outnf.sync()
+        return
+    
+    def add_variable(self, outnf, data, NCFillValue, name, vtyp, dim, title, unit, type_data):
+        dtype = data.dtype
+        data_ij = self.convert_toij(data)
+        data_ij = data_ij.astype(vtyp)
+
+        if dtype == np.int32:
+            data_ij[data_ij >= self.IntFillValue] = NCFillValue
+        elif dtype == np.float32:
+            data_ij[np.isnan(data_ij)] = NCFillValue
+
+        ncvar = outnf.createVariable(name, vtyp, dim, fill_value=NCFillValue)
+        ncvar.title = title
+        ncvar.units = unit
+        ncvar[:] = data_ij[:]
+
+    def dumpnetcdf(self, filename) :
+        #
+        NCFillValue=1.0e20
+        insize = 100
+        vtyp=np.float64
+        cornerind=[0,2,4,6]
+        nbcorners = len(cornerind)
+        #
+        outnf=NetCDFFile(filename, 'w', format='NETCDF4_CLASSIC')
+        # Dimensions
+        outnf.createDimension('x', self.ni)
+        outnf.createDimension('y', self.nj)
+        outnf.createDimension('in', insize)
+        outnf.createDimension('land', self.nbpt_num)
+        outnf.createDimension('htu', self.nbasmax)
+        outnf.createDimension('bnd', nbcorners)
+
+        self.add_variable(outnf, self.route_togrid, NCFillValue, "routetogrid", vtyp, ('htu','y','x'), "Grid into which the basin flows", "-", "int")
+        self.add_variable(outnf, self.route_tobasin, NCFillValue, "routetobasin", vtyp, ('htu','y','x'), "Basin in to which the water goes", "-", "int")
+        self.add_variable(outnf, self.global_basinid, NCFillValue, "basinid", vtyp, ('htu','y','x'), "ID of basin", "-", "int")
+        self.add_variable(outnf, self.route_nbintobas, NCFillValue, "routenbintobas", vtyp, ('y','x'), "Number of basin into current one", "-", "int")
+        self.add_variable(outnf, self.origin_nbintobas, NCFillValue, "originnbintobas", vtyp, ('y','x'), "Number of sub-grid basin into current one before truncation", "-", "int")
+        self.add_variable(outnf, self.route_outlet[:,:,0], NCFillValue, "outletlat", vtyp, ('htu','y','x'), "Latitude of Outlet","degrees north", "float")
+        self.add_variable(outnf, self.route_outlet[:,:,1], NCFillValue, "outletlon", vtyp, ('htu','y','x'), "Longitude of outlet","degrees east", "float")
+        self.add_variable(outnf, self.route_type[:,:], NCFillValue, "outlettype", vtyp, ('htu','y','x'), "Type of outlet","code", "int")
+        self.add_variable(outnf, self.topo_resid[:,:], NCFillValue, "topoindex", vtyp, ('htu','y','x'), "Topographic index of the retention time","m", "float")
+        self.add_variable(outnf, self.routing_cg[:,:,1], NCFillValue, "CG_lon", vtyp, ('htu','y','x'), "Longitude of centre of gravity of HTU","degrees east", "float")    
+        self.add_variable(outnf, self.routing_cg[:,:,0], NCFillValue, "CG_lat", vtyp, ('htu','y','x'), "Latitude of centre of gravity of HTU","degrees east", "float")
+        self.add_variable(outnf, self.routing_fetch, NCFillValue, "fetch", vtyp, ('htu','y','x'), "Fetch contributing to each HTU","m^2", "float")
+
+        self.add_variable(outnf, self.inflow_number[:,:self.nbasmax], NCFillValue, "inflow_number", vtyp, ('htu','y','x'), "Fetch contributing to each HTU","m^2", "int")
+        self.add_variable(outnf, self.inflow_grid[:,:self.nbasmax,:100], NCFillValue, "inflow_grid", vtyp, ('in', 'htu','y','x'), "Grid of HTUs flowing into the basin","-", "int")
+        self.add_variable(outnf, self.inflow_basin[:,:self.nbasmax,:100], NCFillValue, "inflow_basin", vtyp, ('in', 'htu','y','x'), "Basin of HTUs flowing into the basin","-", "int")
+        # Revoir inflows
+
+        outnf.close()
+        return