Calculation of the basins mask for StationDiag.

F90subroutines/Makefile

@@ -9,15 +9,18 @@ FFLAG = -fPIC
 FPY = f2py
 #PDBG =  --debug
 #
-all : routing_interface.so
+all : routing_interface.so diagst.so
 
 clean :
-	rm -f compilation.log *.o *.so *.mod routing_interface.so *~ 
+	rm -f compilation.log *.o *.so *.mod routing_interface.so diagst.so  *~ 
 
 routing_interface.so : $(obj) routing_interface.f90
 	rm -f compilation.log
 	$(FPY) -c routing_interface.f90 $(PDBG) --fcompiler=gfortran -m routing_interface $(obj) > compilation.log
 
+diagst.so : diagst.f90
+	$(FPY) -c diagst.f90 $(PDBG) --fcompiler=gfortran -m diagst > compilation.log
+
 defprec.o : defprec.f90
 	$(FC) $(FFLAG) $(FDBG) -c defprec.f90
 

Locations.py

 #
 #
 import numpy as np
+import numpy.ma as ma
 from netCDF4 import Dataset
 from netCDF4 import stringtochar, chartostring, stringtoarr
 import sys
 import codecs
+import os
+from inspect import currentframe, getframeinfo
+import sys
+
+localdir=os.path.dirname(getframeinfo(currentframe()).filename)
+sys.path.append(localdir+'/F90subroutines')
+"""
+F90=localdir+'/F90subroutines'
+err=os.system("cd "+F90+"; make all")
+if err != 0 :
+    print("Compilation error in the FORTRAN interface")
+    sys.exit()
+"""
+from diagst import *
 #
 import getargs
 config = getargs.SetupConfig()
@@ -278,9 +293,65 @@ class PlacedLocations :
         self.FMONindex = locinfo[4,:].astype(int).tolist()
         self.lonrange = [np.min(nc.variables["lon"][:,:]), np.max(nc.variables["lon"][:,:])]
         self.latrange = [np.min(nc.variables["lat"][:,:]), np.max(nc.variables["lat"][:,:])]
+        # Parameters
+        #to improve
+        basid = nc.variables["basinid"][:,:,:]
+        self.nbasmax = basid.shape[0]
+        land = nc.variables["land"][:,:]
+        self.nbpt = int(np.sum(land))
+        #
+        # Grid points 
+        self.nbpt_glo = nc.variables["nbpt_glo"][:]
+        a = np.unravel_index(ma.argsort(self.nbpt_glo, axis = None)[:self.nbpt], self.nbpt_glo.shape)
+        self.conv_land2ij = [[int(a[0][i]), int(a[1][i])] for i in range(a[0].shape[0])]
+        
+        # Variables
+        self.routetogrid = self.conv_2land(nc,"routetogrid")
+        self.routetobasin = self.conv_2land(nc,"routetobasin")
+        self.routenbintobas = self.conv_2land(nc,"routenbintobas") 
+        self.basin_area = self.conv_2land(nc,"basin_area") 
+        self.area = self.conv_2land(nc,"area") 
         nc.close()
         return
-    def mask(self, i) :
+
+    def conv_2land(self, nc, varname):
+       """
+       Convert the variable from nhtu/lat/lon to ig/ib format.
+       """
+       var = nc.variables[varname][:]
+       if len(var.shape) == 3:
+          return np.array([var[:,j,i] for j,i in self.conv_land2ij], order = "F")
+       elif len(var.shape) == 2:
+          return np.array([var[j,i] for j,i in self.conv_land2ij], order = "F")
+    # 
+    def upstmask(self, i, j, k):
+       """
+       Get the upstream area mask from the ig/ib coordinates.
+       """
+       g0 = self.nbpt_glo[int(j)-1,int(i)-1] # -1 Because F -> C
+       b0 = int(k)
+       mask = diagst.upstmask(
+                    nbpt = self.nbpt, 
+                    nbasmax = self.nbasmax,
+                    routetogrid = self.routetogrid,
+                    routetobasin = self.routetobasin,
+                    routenbbasin = self.routenbintobas,
+                    basin_area = self.basin_area,
+                    area = self.area,
+                    g0 = g0,
+                    b0 = b0,
+                    )
+       return mask
+
+    def mask(self, index):
         # Function to compute mask of area upstream of function with information in the GraphFile
-        mask=np.zeros(self.gridshape)
-        return mask
+       i = self.Fiindex[index]
+       j = self.Fjindex[index]
+       k = self.FHTUindex[index]
+       # 
+       mask_raw = self.upstmask(i,j,k)
+       mask = np.zeros(self.gridshape)
+       for g in range(self.nbpt): 
+          jj,ii = self.conv_land2ij[g]
+          mask[jj,ii] = min(mask_raw[g],1)
+       return mask